xlstm: 언어 모델링에서 Transformer와 경쟁하기 위해 LSTM을 확장한 순환 신경망 아키텍처

해결하는 문제

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

작동 방식

이 아키텍처는 기존 LSTM에 대해 두 가지 주요 개선 사항을 도입합니다:

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

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

대상 사용자

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

주요 특징

• 7B Parameter Model: 2.3T 토큰으로 훈련된 대규모 순환 LLM.
• Optimized Kernels: NVIDIA 및 AMD GPU에서 고성능 실행을 위해 mlstm_kernels 패키지를 통한 Triton 및 CUDA 커널 지원.
• Flexible Architecture: sLSTM 및 mLSTM 블록을 모두 지원하며, 이를 조합하여 상태 추적 및 암기 능력을 균력할 수 있습니다.
• Hardware Compatibility: Apple Metal (MLX용 커뮤니티 포트 포함)과 같은 플랫폼을 위한 네이티브 PyTorch 구현이 사용 가능합니다.