NX-AI/xlstm

Official repository of the xLSTM.

What it solves

xLSTM(Extended Long Short-Term Memory)은 기존 LSTM 아키텍처의 한계를 극복하고 Transformer 및 State Space Models(SSM)와 경쟁하도록 설계되었습니다. 대규모 언어 모델에서 학습과 추론 모두 효율적인 순환 신경망(RNN) 대안을 제공하는 것이 목표입니다.

How it works

이 아키텍처는 기존 LSTM에 두 가지 주요 향상을 도입합니다:

  1. Exponential Gating: 정규화와 안정화 기술을 사용하여 네트워크의 정보 흐름 관리를 개선합니다.
  2. Matrix Memory: 새로운 메모리 구조로, 네트워크가 복잡한 데이터 패턴을 더 잘 처리할 수 있게 합니다.

프로젝트는 두 가지 주요 구성 요소를 제공합니다: xLSTMBlockStack(일반 애플리케이션용 백본)과 xLSTMLMModel(토큰 기반 언어 모델링용 래퍼). 또한 학습 처리량과 안정성을 최적화한 xLSTMLarge 아키텍처도 포함되어 있으며, 이는 2.3조 토큰으로 7B 파라미터 모델을 학습하는 데 사용되었습니다.

Who it’s for

  • AI Researchers: 시퀀스 모델링을 위한 Transformer 대안을 찾는 연구자.
  • ML Engineers: 빠르고 효율적인 추론을 위해 순환 모델을 구현하는 엔지니어.
  • Developers: 기존 프로젝트에 xLSTM 백본을 통합하여 Transformer 블록의 대안으로 사용하고자 하는 개발자.

Highlights

  • 7B Parameter Model: 2.3T 토큰으로 학습된 대규모 순환 LLM.
  • Optimized Kernels: mlstm_kernels 패키지를 통해 Triton 및 CUDA 커널을 지원하여 NVIDIA와 AMD GPU에서 고성능 실행을 가능하게 함.
  • Flexible Architecture: sLSTM 및 mLSTM 블록을 모두 지원하며, 상태 추적과 기억 능력의 균형을 맞추기 위해 조합 가능.
  • Hardware Compatibility: Apple Metal(커뮤니티가 MLX용으로 포팅한 버전 포함) 등 플랫폼을 위한 네이티브 PyTorch 구현 제공.