NX-AI/xlstm

xLSTM – 擴展的長期短期記憶

是什麼 – xLSTM 是一種基於經典 LSTM 的新型循環神經網路架構。它引入了 指數門控、矩陣記憶 以及多種歸一化/穩定化技巧，使其能擴展至非常大型的語言模型（例如，一個 70 億參數的模型），同時保持推論速度快且記憶體效率高。作者聲稱它在語言建模基準測試中可與 Transformer 和狀態空間模型相媲美。

核心元件

xLSTMBlockStack – 可直接取代 Transformer 塊堆疊的元件；內部混合三種塊類型（mLSTM、sLSTM 和前饋層），可透過 dataclasses 配置。
xLSTMLMModel – 在塊堆疊之上添加詞元嵌入和語言模型頭的語言模型包裝器。
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 – 使用社群維護的 xLSTM-metal 版本（MLX）以獲得 Metal 原生實作。

模型

xLSTM-Large 7B – 在 2.3T 個 token 上訓練的 70 億參數循環語言模型；權重託管於 Hugging Face 上（https://huggingface.co/NX-AI/xLSTM-7b）。
更小的研究模型可透過 xLSTMBlockStack 和 xLSTMLMModel 類從 YAML 設定檔實例化（README 中有範例）。

實驗倉儲包含合成任務（Parity、Multi-Query Associative Recall），展示兩種子塊（sLSTM 用於狀態追蹤，mLSTM 用於記憶）的互補優勢。執行方式如下：

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

（註：訓練迴圈極簡——無早期停止或測試評估。）

引用如果您使用程式碼或 7B 模型，請引用 README 中列出的兩篇論文（NeurIPS 2024 xLSTM 論文與 ICML 2025 xLSTM-7B 論文）。

所有資訊均直接來自倉儲的 README；未添加任何外部假設。

NX-AI/xlstm

xLSTM – 擴展的長期短期記憶

相關