NX-AI/xlstm

xLSTM – 拡張された長期短期記憶

何であるか – xLSTMは古典的なLSTMを基盤とする新しい再帰型ニューラルネットワークアーキテクチャです。指数ゲート、行列メモリ、およびいくつかの正規化/安定化テクニックを導入することで、非常に大きな言語モデル（例：70億パラメータモデル）へのスケーリングを可能にしつつ、推論速度とメモリ効率を維持しています。著者らは、Transformerや状態空間モデルと同等の性能を言語モデルベンチマークで達成できることを主張しています。

主な構成要素

xLSTMBlockStack – Transformerブロックのスタックの即時置き換え可能。内部では3種類のブロック（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に例あり）。

実験リポジトリには、2つのサブブロックの相補的な強みを示す合成タスク（Parity、Multi-Query Associative Recall）が含まれています（sLSTM：状態追跡、mLSTM：記憶）。実行は以下の通りです：

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

（注：トレーニングループは最小限。早期停止やテスト評価は含まれません。）

引用コードまたは7Bモデルを使用する場合、READMEに記載の2つの論文（NeurIPS 2024 xLSTM論文、ICML 2025 xLSTM-7B論文）を引用してください。

すべての情報はリポジトリのREADMEから直接取得。外部の仮定は一切追加されていません。

NX-AI/xlstm

xLSTM – 拡張された長期短期記憶

関連