论文: https://arxiv.org/abs/2405.04517

作者: 马克西米利安·贝克、科尔比尼安·珀佩尔、马库斯·施潘林、安德烈亚斯·奥尔、奥莱克桑德拉·普鲁德尼科娃、迈克尔·科普、京特·克拉姆鲍尔、约翰内斯·布兰德施泰特、塞普·霍赫赖特

关于

xLSTM 是一种基于原始 LSTM 思想的新型循环神经网络架构。通过采用指数门控机制，并结合适当的归一化和稳定化技术，以及引入新的矩阵存储单元，xLSTM 克服了原始 LSTM 的局限性，在语言建模任务上表现出与 Transformer 或状态空间模型相媲美的优异性能。

:rotating_light: 我们使用 2.3T 个 token 训练了一个 7B 参数的 xLSTM 语言模型！ :rotating_light:

我们将其优化后的 7B 参数架构称为 xLSTM Large。

最小安装步骤

首先从 environment_pt240cu124.yaml 文件创建一个 conda 环境。然后仅将模型代码（即 xlstm 模块）以包的形式安装：

若要使用 xLSTM Large 7B 模型，请先通过以下命令安装 mlstm_kernels：

pip install mlstm_kernels

接着再通过 pip 安装 xlstm 包：

pip install xlstm

或者直接从 GitHub 克隆：

git clone https://github.com/NX-AI/xlstm.git
cd xlstm
pip install -e .

系统要求

本包基于 PyTorch 开发，经测试适用于 >=1.8 版本。为确保环境稳定运行，可按如下方式安装已验证过的环境配置文件：

conda env create -n xlstm -f environment_pt240cu124.yaml
conda activate xlstm

对于 xLSTM Large 7B 模型，我们需要使用我们的 mlstm_kernels 包，它提供了针对 xLSTM 的高效计算内核。

我们在训练吞吐量和稳定性方面对 xLSTM 架构进行了优化。更新后的架构代码位于 xlstm/xlstm_large 目录中。

该模型的权重已在 Hugging Face 上发布，地址为：https://huggingface.co/NX-AI/xLSTM-7b。

如何使用 xLSTM Large 7B 及其架构

我们在 xlstm/xlstm_large/model.py 中提供了一个独立的单文件实现版本，用于 xLSTM Large 架构。此实现依赖于我们的 mlstm_kernels 包，除此之外不依赖于 NeurIPS 论文中的 xLSTM 架构实现。

为了方便快速上手，我们还提供了一个用于 xLSTM Large 架构的笔记本文件：demo.ipynb，路径为 notebooks/xlstm_large/demo.ipynb。

在该笔记本中，我们导入配置和模型类，初始化一个随机模型并执行前向传播，具体操作如下：

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

# 使用 TFLA Triton 内核配置模型
xlstm_config = 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", # xl_chunk == TFLA 内核
    sequence_kernel="native_sequence__triton",
    step_kernel="triton",
)
# 实例化模型
xlstm = xLSTMLarge(xlstm_config)
xlstm = xlstm.to("cuda")
# 创建输入
input = torch.randint(0, 2048, (3, 256)).to("cuda")
# 运行前向传播
out = xlstm(input)
out.shape[1:] == (256, 2048)

其他硬件平台的建议

我们主要在 NVIDIA GPU 上测试了该模型，不过我们的 Triton 内核也应该能在 AMD GPU 上运行。对于其他平台，比如 Apple Metal，目前建议使用原生 PyTorch 实现：

xlstm_config = xLSTMLargeConfig(
    embedding_dim=512,
    num_heads=4,
    num_blocks=6,
    vocab_size=2048,
    return_last_states=True,
    mode="inference",
    chunkwise_kernel="chunkwise--native_autograd", # 不使用 Triton 内核
    sequence_kernel="native_sequence__native", # 不使用 Triton 内核
    step_kernel="native", # 不使用 Triton 内核
)

如果您正在使用 Apple 的 MLX 生态系统，可以尝试社区驱动的 xLSTM-metal 移植项目，它提供了面向 Apple Silicon 的 MLX 原生 xLSTM 实现。

来自 xLSTM NeurIPS 论文的模型

本节介绍如何使用 xLSTM 论文中提出的模型。

如何使用我们 NeurIPS 论文中的 xLSTM 架构

对于非语言类应用或需要与其他架构集成的情况，您可以使用 xLSTMBlockStack；而对于语言建模或其他基于 token 的应用，则可以使用 xLSTMLMModel。

使用 sLSTM CUDA 内核

对于 sLSTM 的 CUDA 版本，您的设备需具备 Compute Capability >= 8.0，请参阅 https://developer.nvidia.com/cuda-gpus。如果编译过程中遇到问题，您可以尝试以下方法（感谢 @zia1138 的提示）：

export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"

对于所有涉及 PyTorch 和 CUDA 的自定义设置，请务必确保版本匹配。此外，为确保正确加载 CUDA 库，您现在可以使用 XLSTM_EXTRA_INCLUDE_PATHS 环境变量来注入不同的包含路径，例如：

export XLSTM_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'

或者在 Python 中设置：

import os
os.environ['XLSTM_EXTRA_INCLUDE_PATHS']='/usr/local/include/cuda/:/usr/include/cuda/'

如需更加快速的 sLSTM 独立内核，欢迎使用 FlashRNN 库。

xLSTM 块堆栈

xLSTMBLockStack 旨在作为现有项目中的替代骨干网络。它类似于 Transformer 块的堆栈，但使用的是 xLSTM 块：

import torch

from xlstm import (
    xLSTMBlockStack,
    xLSTMBlockStackConfig,
    mLSTMBlockConfig,
    mLSTMLayerConfig,
    sLSTMBlockConfig,
    sLSTMLayerConfig,
    FeedForwardConfig,
)

cfg = xLSTMBlockStackConfig(
    mlstm_block=mLSTMBlockConfig(
        mlstm=mLSTMLayerConfig(
            conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4
        )
    ),
    slstm_block=sLSTMBlockConfig(
        slstm=sLSTMLayerConfig(
            backend="cuda",
            num_heads=4,
            conv1d_kernel_size=4,
            bias_init="powerlaw_blockdependent",
        ),
        feedforward=FeedForwardConfig(proj_factor=1.3, act_fn="gelu"),
    ),
    context_length=256,
    num_blocks=7,
    embedding_dim=128,
    slstm_at=[1],

)

xlstm_stack = xLSTMBlockStack(cfg)

x = torch.randn(4, 256, 128).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)
y.shape == x.shape

如果你使用 YAML 字符串或文件进行配置，也可以使用 dacite 来创建配置数据类。这与上面的代码片段相同：

from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm import xLSTMBlockStack, xLSTMBlockStackConfig

xlstm_cfg = """ 
mlstm_block:
  mlstm:
    conv1d_kernel_size: 4
    qkv_proj_blocksize: 4
    num_heads: 4
slstm_block:
  slstm:
    backend: cuda
    num_heads: 4
    conv1d_kernel_size: 4
    bias_init: powerlaw_blockdependent
  feedforward:
    proj_factor: 1.3
    act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""
cfg = OmegaConf.create(xlstm_cfg)
cfg = from_dict(data_class=xLSTMBlockStackConfig, data=OmegaConf.to_container(cfg), config=DaciteConfig(strict=True))
xlstm_stack = xLSTMBlockStack(cfg)

x = torch.randn(4, 256, 128).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)
y.shape == x.shape

xLSTM 语言模型

xLSTMLMModel 是一个围绕 xLSTMBlockStack 的封装器，添加了词嵌入和语言模型头。

from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm import xLSTMLMModel, xLSTMLMModelConfig

xlstm_cfg = """ 
vocab_size: 50304
mlstm_block:
  mlstm:
    conv1d_kernel_size: 4
    qkv_proj_blocksize: 4
    num_heads: 4
slstm_block:
  slstm:
    backend: cuda
    num_heads: 4
    conv1d_kernel_size: 4
    bias_init: powerlaw_blockdependent
  feedforward:
    proj_factor: 1.3
    act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""
cfg = OmegaConf.create(xlstm_cfg)
cfg = from_dict(data_class=xLSTMLMModelConfig, data=OmegaConf.to_container(cfg), config=DaciteConfig(strict=True))
xlstm_stack = xLSTMLMModel(cfg)

x = torch.randint(0, 50304, size=(4, 256)).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)
y.shape[1:] == (256, 50304)

实验

展示 sLSTM 相对于 mLSTM 优势以及反之亦然的最佳合成实验是奇偶校验任务和多查询联想回忆任务。奇偶校验任务只能通过 sLSTM 的内存混合提供的状态跟踪能力来解决。而多查询联想回忆任务则衡量记忆能力，此时 mLSTM 的矩阵内存和状态扩展非常有益。 两者结合在两项任务中都表现出色。

要运行每个实验，请在 experiments 文件夹中运行 main.py，例如：

PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm01.yaml   # xLSTM[0:1], 仅 sLSTM
PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm10.yaml   # xLSTM[1:0], 仅 mLSTM
PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm11.yaml   # xLSTM[1:1], mLSTM 和 sLSTM

请注意，训练循环不包含早停或测试评估。

引用

如果您使用此代码库，或者以其他方式认为我们的工作有价值，请引用 xLSTM 论文：

@inproceedings{beck:24xlstm,
  title = {xLSTM: 扩展的长短期记忆}, 
  author = {Maximilian Beck 和 Korbinian Pöppel、Markus Spanring、Andreas Auer、Oleksandra Prudnikova、Michael Kopp、Günter Klambauer、Johannes Brandstetter 和 Sepp Hochreiter},
  booktitle = {第三十八届神经信息处理系统会议},
  year = {2024},
  url = {https://arxiv.org/abs/2405.04517}, 
}

@article{beck:25xlstm7b,
  title = {{xLSTM 7B}: 用于快速高效推理的循环 LLM},
  author = {Maximilian Beck、Korbinian Pöppel、Phillip Lippe、Richard Kurle、Patrick M. Blies、Günter Klambauer、Sebastian Böck 和 Sepp Hochreiter},
  booktitle = {第四十二届国际机器学习会议},
  year = {2025},
  url = {https://arxiv.org/abs/2503.13427}
}

xLSTM 快速上手指南

xLSTM（Extended Long Short-Term Memory）是一种基于原始 LSTM 理念改进的新型循环神经网络架构。它通过指数门控、归一化稳定技术及新的矩阵记忆机制，克服了传统 LSTM 的局限性，在语言建模任务中展现出媲美 Transformer 和状态空间模型（SSM）的性能。

环境准备

系统要求

• 操作系统: Linux (推荐), macOS, Windows
• Python 版本: 建议 Python 3.8+
• GPU 支持:
  • NVIDIA GPU (推荐): 需安装 CUDA。若使用 sLSTM CUDA 加速内核，要求 Compute Capability >= 8.0 (如 A100, RTX 3090/4090 等)。
  • AMD GPU: 支持 Triton 内核。
  • Apple Silicon: 建议使用原生 PyTorch 实现或社区驱动的 xLSTM-metal (MLX) 版本。

前置依赖

• PyTorch: 版本 >= 1.8 (推荐使用较新版本以获得最佳性能)。
• 可选加速库:
  • mlstm_kernels: 运行 xLSTM Large 7B 模型或追求极致推理速度时必需。
  • flashrnn: 用于独立的更快 sLSTM 内核。

安装步骤

方法一：使用 Conda 环境（推荐）

这是最稳定的安装方式，可确保所有依赖版本匹配。

# 克隆仓库
git clone https://github.com/NX-AI/xlstm.git
cd xlstm

# 创建并激活 conda 环境 (基于 environment_pt240cu124.yaml)
conda env create -n xlstm -f environment_pt240cu124.yaml
conda activate xlstm

# 以开发模式安装 xlstm 包
pip install -e .

提示：国内用户若下载 conda 包较慢，可配置清华或中科大镜像源：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/

方法二：使用 Pip 直接安装

如果你已有合适的 PyTorch 环境，可直接通过 pip 安装。

# 安装基础 xlstm 包
pip install xlstm

# 【可选】若需使用 xLSTM Large 7B 模型，请额外安装加速内核
pip install mlstm_kernels

特殊配置：CUDA 编译问题

若在使用 sLSTM CUDA 内核时遇到编译错误，请尝试设置以下环境变量指定 GPU 架构列表：

export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"

若需指定额外的 CUDA 头文件路径：

export XLSTM_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'

基本使用

场景一：使用 xLSTM Large 7B 架构 (推理优化版)

适用于需要高性能推理的场景。此实现依赖 mlstm_kernels 和 Triton。

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

# 1. 配置模型 (使用 TFLA Triton 内核加速)
xlstm_config = 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", # xl_chunk == TFLA kernels
    sequence_kernel="native_sequence__triton",
    step_kernel="triton",
)

# 2. 实例化模型并移至 GPU
xlstm = xLSTMLarge(xlstm_config)
xlstm = xlstm.to("cuda")

# 3. 创建输入数据 (batch_size=3, seq_len=256)
input_ids = torch.randint(0, 2048, (3, 256)).to("cuda")

# 4. 执行前向传播
output = xlstm(input_ids)

# 输出形状验证: (batch, seq_len, vocab_size) -> (3, 256, 2048)
print(output.shape) 

非 NVIDIA 硬件用户注意：若在 Apple Metal 或其他不支持 Triton 的平台，请将配置中的 kernel 选项改为 native 版本：

chunkwise_kernel="chunkwise--native_autograd",
sequence_kernel="native_sequence__native",
step_kernel="native",

场景二：使用标准 xLSTM 模块 (研究与自定义)

适用于将 xLSTM 作为骨干网络集成到其他项目中，或进行学术研究。支持混合使用 sLSTM 和 mLSTM 块。

import torch
from xlstm import (
    xLSTMBlockStack,
    xLSTMBlockStackConfig,
    mLSTMBlockConfig,
    mLSTMLayerConfig,
    sLSTMBlockConfig,
    sLSTMLayerConfig,
    FeedForwardConfig,
)

# 1. 定义配置
cfg = xLSTMBlockStackConfig(
    mlstm_block=mLSTMBlockConfig(
        mlstm=mLSTMLayerConfig(
            conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4
        )
    ),
    slstm_block=sLSTMBlockConfig(
        slstm=sLSTMLayerConfig(
            backend="cuda", # 使用 CUDA 加速后端
            num_heads=4,
            conv1d_kernel_size=4,
            bias_init="powerlaw_blockdependent",
        ),
        feedforward=FeedForwardConfig(proj_factor=1.3, act_fn="gelu"),
    ),
    context_length=256,
    num_blocks=7,
    embedding_dim=128,
    slstm_at=[1], # 在第 1 层使用 sLSTM，其余默认为 mLSTM
)

# 2. 初始化模型
xlstm_stack = xLSTMBlockStack(cfg)
xlstm_stack = xlstm_stack.to("cuda")

# 3. 准备输入 (batch=4, seq_len=256, dim=128)
x = torch.randn(4, 256, 128).to("cuda")

# 4. 前向传播
y = xlstm_stack(x)

# 输出形状应与输入一致 (除了可能的内部状态处理)
assert y.shape == x.shape

场景三：构建语言模型 (LM)

xLSTMLMModel 是包含 Token 嵌入层和 LM Head 的完整封装，适合直接进行语言建模训练。

from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm import xLSTMLMModel, xLSTMLMModelConfig
import torch

# 使用 YAML 字符串配置模型
xlstm_cfg_str = """ 
vocab_size: 50304
mlstm_block:
  mlstm:
    conv1d_kernel_size: 4
    qkv_proj_blocksize: 4
    num_heads: 4
slstm_block:
  slstm:
    backend: cuda
    num_heads: 4
    conv1d_kernel_size: 4
    bias_init: powerlaw_blockdependent
  feedforward:
    proj_factor: 1.3
    act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""

# 解析配置并创建模型
cfg_dict = OmegaConf.create(xlstm_cfg_str)
cfg = from_dict(
    data_class=xLSTMLMModelConfig, 
    data=OmegaConf.to_container(cfg_dict), 
    config=DaciteConfig(strict=True)
)

model = xLSTMLMModel(cfg)
model = model.to("cuda")

# 模拟输入 token IDs
input_ids = torch.randint(0, 50304, size=(4, 256)).to("cuda")

# 前向传播
logits = model(input_ids)

# 输出形状: (batch, seq_len, vocab_size)
print(logits.shape) # torch.Size([4, 256, 50304])