NX-AI/xlstm

xLSTM – 확장된 장기 단기 기억

무엇인가요 – xLSTM은 고전적인 LSTM을 기반으로 하는 새로운 순환 신경망 아키텍처입니다. 지수 게이팅, 행렬 메모리, 그리고 여러 정규화/안정화 기법을 도입하여, 매우 큰 언어 모델(예: 70억 파라미터 모델)로 확장하면서도 추론 속도와 메모리 효율성을 유지할 수 있습니다. 저자들은 이 모델이 언어 모델링 벤치마크에서 트랜스포머와 상태 공간 모델과 경쟁할 수 있다고 주장합니다.

주요 구성 요소

xLSTMBlockStack – Transformer 블록 스택의 즉시 대체 가능한 대안. 내부적으로 mLSTM, sLSTM, 그리고 피드포워드 블록을 혼합하며, dataclasses를 통해 구성 가능합니다.
xLSTMLMModel – 블록 스택 위에 토큰 임베딩과 LM 헤드를 추가하는 언어 모델 래퍼입니다.
mlstm_kernels – 커스텀 CUDA/Triton 커널("sLSTM" 커널)을 제공하는 별도 패키지로, 순환 연산을 극적으로 가속화합니다.
xLSTMLarge – 후속 논문에서 사용된 70억 파라미터 모델의 단일 파일 구현이며, 다양한 하드웨어용 구성 가능한 커널을 제공합니다.

설치

# 선택 사항: 저자들이 사용한 정확한 conda 환경 생성
conda env create -f environment_pt240cu124.yaml
conda activate xlstm

# 빠른 커널 설치 (7B 모델에 필수)
pip install mlstm_kernels

# 라이브러리 자체 설치
pip install xlstm   # 또는: git clone https://github.com/NX-AI/xlstm && pip install -e .

PyTorch ≥ 1.8 필요. GPU 가속을 위해서는 최신 NVIDIA GPU (CUDA 계산 능력 ≥ 8.0)가 필요하여 Triton 커널을 사용할 수 있으며, 그렇지 않으면 순수 PyTorch 폴백이 모든 플랫폼에서 작동합니다.

빠른 시작 (7B 추론)

import torch
from xlstm.xlstm_large.model import xLSTMLargeConfig, xLSTMLarge

cfg = xLSTMLargeConfig(
    embedding_dim=512,
    num_heads=4,
    num_blocks=6,
    vocab_size=2048,
    return_last_states=True,
    mode="inference",
    chunkwise_kernel="chunkwise--triton_xl_chunk",
    sequence_kernel="native_sequence__triton",
    step_kernel="triton",
)
model = xLSTMLarge(cfg).to("cuda")
inputs = torch.randint(0, 2048, (3, 256), device="cuda")
out = model(inputs)
print(out.shape)   # (3, 256, 2048)

노트북 (notebooks/xlstm_large/demo.ipynb)에서 동일한 워크플로우를 보여줍니다.

하드웨어 추천

NVIDIA GPU – Triton 커널로 최고 성능. RTX 3080/3090 및 이후 모델(CC 8.0+)에서 성공적으로 실행됨.
AMD GPU – Triton 커널이 여전히 작동할 수 있지만, 저자들은 네이티브 PyTorch 커널로 폴백하는 것을 권장합니다.
Apple Silicon – Metal 네이티브 구현을 위한 커뮤니티 유지 관리 xLSTM-metal (MLX)을 사용하세요.

모델

xLSTM-Large 7B – 2.3T 토큰으로 훈련된 70억 파라미터 순환형 LLM. Hugging Face에 호스팅됨 (https://huggingface.co/NX-AI/xLSTM-7b).
더 작은 연구용 모델은 xLSTMBlockStack 및 xLSTMLMModel 클래스를 통해 YAML 설정 파일에서 인스턴스화 가능 (README에 예시 있음).

실험 리포지토리에는 두 하위 블록의 보완적인 강점(예: sLSTM은 상태 추적, mLSTM은 기억)을 보여주는 합성 작업(Parity, Multi-Query Associative Recall)이 포함되어 있습니다. 실행은 다음과 같습니다:

PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm11.yaml

(참고: 학습 루프는 최소화되어 있으며, 조기 중단이나 테스트 평가가 포함되지 않음.)

인용 코드 또는 7B 모델을 사용할 경우, README에 나열된 두 논문(NeurIPS 2024 xLSTM 논문, ICML 2025 xLSTM-7B 논문)을 인용하세요.

모든 정보는 리포지토리의 README에서 직접 가져왔으며, 외부 가정은 추가되지 않았습니다.

NX-AI/xlstm

xLSTM – 확장된 장기 단기 기억

관련