在过去十年的大部分时间里，Transformer 一直定义着序列建模的前沿。其并行处理长上下文的能力开启了大语言模型 (LLM) 时代。但这一进步也令最初的序列引擎——循环神经网络，尤其是 Sepp Hochreiter 共同发明的 LSTM 架构——逐渐淡出人们的视野。

《xLSTM: 扩展长短期记忆网络》这篇论文回顾了这一技术谱系，并提出了一个看似简单的问题: 如果我们借助现代工程手段扩展 LSTM，并消除其已知的弱点，它能走多远？简短的答案是: 非常远。论文介绍了一个 LSTM 扩展家族，这些改进恢复了果断的记忆更新，大幅提升了存储容量，并且——至关重要的是——使得部分架构可并行化，从而能够与现代替代方案竞争。在多项语言建模基准和合成测试中，xLSTM 模型的表现达到或超过了当前最先进的 Transformer 和状态空间模型。

本文将梳理 xLSTM 背后的核心理念、动机和实验。我会将复杂内容拆分为直观的构建模块，展示架构全貌，并总结支撑 xLSTM 的实证证据。

读完本文你将获得:

  • 对 LSTM 核心要点的简明回顾，以方便理解其扩展。
  • 对指数门控的直观理解，以及它为何能解决“优柔寡断”的记忆更新问题。
  • 对标量记忆 (sLSTM) 和矩阵记忆 (mLSTM) 如何应对不同局限的理解。
  • 将这些单元构建为可扩展架构的残差块配方。
  • 关键实验结论: 合成任务、记忆测试、大规模语言建模及推理特性。

首先是一张可视化路线图: 一图看懂 xLSTM 家族。

从 LSTM 到 sLSTM 和 mLSTM 的演进，组装成 xLSTM 块并堆叠成 xLSTM。

图 1: 扩展长短期记忆网络 (xLSTM) 家族。从左到右: 1) 经典 LSTM 记忆单元 (常数误差环路 + sigmoid 门控) ，2) 两个新的记忆单元——sLSTM 和 mLSTM，均配备了指数门控 (sLSTM 保留标量记忆并加入新混合机制；mLSTM 以矩阵记忆和协方差更新取代标量记忆) ，3) 这些单元被包装成残差块 (xLSTM 块) ，4) 堆叠残差块形成 xLSTM 架构。

1 — LSTM 基础快速回顾

标准 LSTM 的核心由两个相互交织的思想构成:

  • 常数误差环路 (constant error carousel) : 对单元状态的加性更新，让梯度能够跨越多个时间步流动。
  • 门控 (sigmoid 函数) : 控制信息的遗忘、写入和输出。

具体来说，在时间步 t，经典标量 LSTM 的更新为:

\[ c_t = f_t \, c_{t-1} + i_t \, z_t, \qquad h_t = o_t \, \psi(c_t), \]

其中 f (遗忘) 、i (输入) 、o (输出) 是 sigmoid 门，z 是候选单元输入 (通常经 tanh 变换) ，ψ 为作用于单元状态的压缩/非线性函数。

这些机制曾推动 LSTM 在多个领域表现卓越: 如语言生成、序列到序列翻译、强化学习以及长时间序列预测。然而，在大规模语言建模中，最终有三大局限制约了它的发展:

  1. 不能果断修正存储决策。Sigmoid 门取值在 (0, 1)，很难稳定地用新记忆完全替换已有记忆。
  2. 存储容量有限。每个记忆单元是标量 (或标量向量) ，压缩了大量信息，对罕见且特定的词元保真度不高。
  3. 缺乏可并行性。循环的隐藏到隐藏连接需要顺序计算，GPU 难以像处理全并行注意力那样高效利用。

xLSTM 同时解决了上述三大问题。

2 — 两大核心创新

xLSTM 引入了两个正交但互补的创新:

  • 指数门控: 用指数激活替换 (或增强) sigmoid 门，并引入稳定机制。
  • 新型记忆结构: 保留具备表达性循环特征的标量型变体 (sLSTM) ，新增具备键值关联存储能力且易并行化的矩阵型变体 (mLSTM) 。

接下来逐一介绍这些变体及其直觉。

2.1 指数门控——让记忆更新更果断

Sigmoid 门限制在 [0, 1] 之间，导致更新偏“柔和”。若需完全覆盖旧记忆，必须同时让输入门接近 1、遗忘门接近 0，这在实践中很难稳定实现。

指数门控将输入门 (及可选的遗忘门) 激活函数改为指数函数:

\[ i_t = \exp(\tilde i_t), \quad f_t = \sigma(\tilde f_t)\ \text{ 或 }\ \exp(\tilde f_t), \]

这样，新内容的有效贡献范围可远大于 1。因为 exp() 的取值范围是 (0, ∞)，模型在需要时能使新信息压倒已有内容，实现彻底的记忆修正。

注意: exp() 可能导致数值溢出。论文提出了轻量级稳定技巧: 维护一个名为 m_t 的额外稳定器状态，持续对门控贡献进行对数尺度归一化，并用其重新归一化 i_t 和 f_t，从而保证前向传播数值安全。重要的是，该稳定器不影响网络其余部分的梯度，仅用于缩放组合项避免溢出。

2.2 sLSTM——具备更优混合的标量记忆

sLSTM 是“循环专家”: 保留标量 (或向量) 单元状态，同时结合指数门控以及跟踪门控强度的归一化状态。核心的前向规则为:

\[ c_t = f_t c_{t-1} + i_t z_t, \qquad n_t = f_t n_{t-1} + i_t, \qquad \tilde h_t = \frac{c_t}{n_t}, \qquad h_t = o_t \, \tilde h_t. \]

这里 n_t 是累计门控强度的归一化因子，输出除以它可稳定隐藏状态。指数门控让 sLSTM 在需要时能迅速、大幅地修正存储的标量状态。sLSTM 保留了隐藏到隐藏的循环混合，并可配置多头 (块对角循环矩阵) 来高效地追踪状态不同方面。

适用场景: 需要精细状态跟踪与逐步推理的任务——如形式语言、代码执行、需逐步更新和查询显式状态的问题。

2.3 mLSTM——为容量与并行化而生的矩阵记忆

sLSTM 提升了果断性，但信息仍压缩在标量中。mLSTM 将单元状态升级为矩阵 \(C_t \in \mathbb{R}^{d\times d}\)。该记忆充当相关/关联存储: 对于键向量 k_t 和值向量 v_t，协方差式更新规则为:

\[ C_t = f_t \, C_{t-1} + i_t \, v_t k_t^\top, \]

即将外积 v k^T 存入矩阵。检索时，给定查询向量 q_t:

\[ \tilde h_t = \frac{C_t q_t}{\max\{n_t^\top q_t, 1\}}, \qquad h_t = o_t \odot \tilde h_t, \]

其中 n_t 是归一化器 (键的加权和) ，分母加下界以防不稳定。这与经典关联记忆及早期“快速权重”/外积记忆工作密切相关。关键在于将其嵌入类 LSTM 框架: 遗忘门控制衰减，输入门控制外积更新的学习率。

为何选矩阵？矩阵记忆能并行存储大量不同键值关联，不必压缩至单一标量槽，容量更高，尤其适合存储罕见词元或大规模键值对。

为何可并行？协方差更新可转化为全序列的矩阵运算，且可用与 sLSTM 相同的对数归一化技巧稳定，从而让训练在时间维上并行 (类似注意力实现) 。生成阶段仍可用循环形式高效自回归解码。

3 — 构建 xLSTM 模型 (残差块 + 堆叠)

单个单元只是深层模型的一部分。作者将 sLSTM 和 mLSTM 集成进残差块，借鉴了当代 LLM 工程 (LayerNorm、MLP、升/降维投影) 。两种模式:

  • 后上投影块 (sLSTM) : 输入 → sLSTM → 门控 MLP → 残差连接。类似 Transformer 块，非线性在升维前的原始嵌入空间。
  • 前上投影块 (mLSTM) : 输入先升维，再在高维空间中应用 mLSTM，最后降维。矩阵记忆受益于高维操作。

完整 xLSTM 为残差块堆叠。记法 xLSTM[a:b] 表示 mLSTM 块与 sLSTM 块比例，如 xLSTM[7:1] 表示每 8 块中 7 块 mLSTM、1 块 sLSTM。

嵌入在残差结构中的 sLSTM 块 (左) 和 mLSTM 块 (右) 。

图 3: xLSTM 块。左: 带后上投影的残差 sLSTM 块 (类似 Transformer) ——可选小型因果卷积，接门控 MLP。右: 带前上投影的残差 mLSTM 块 (类似部分状态空间模型) ——mLSTM 包裹在 MLP、归一化与逐元素门控中。

4 — 这些设计的意义

  • 指数门控赋予模型在新信息出现时果断覆盖记忆的能力。适用于序列最近邻搜索等任务: 找到更优匹配时，应替换原结果而非柔和混合。
  • 矩阵记忆消除了标量压缩瓶颈: 罕见词元和大规模关联存储成为可能，而不必增加巨大维度。
  • 混合 sLSTM 与 mLSTM 得到实用组合: mLSTM 作大容量并行主干，关键处用 sLSTM 处理复杂顺序推理。
  • 将 mLSTM 循环部分改写为并行矩阵运算 (配合稳定化) ，让训练对 GPU 更友好。

5 — 实验: 能力与规模验证

论文通过广泛实验评估 xLSTM: 合成任务、关联回忆 (MQAR) 、Long Range Arena、消融实验及大规模语言建模 (SlimPajama 数据集上 150 亿与 3000 亿词元规模训练) 。要点如下:

5.1 合成任务——状态跟踪与形式语言
形式语言任务考查模型维护与更新结构化状态的能力 (如奇偶校验、上下文无关语言、类栈行为) 。无循环连接的模型 (Transformer、SSM) 在此类任务常表现欠佳。

带 sLSTM (记忆混合+指数门控) 的 xLSTM 变体能稳定解决多数任务，而普通 Transformer 与并行 SSM 往往失败。这验证了记忆混合对于某些算法泛化的必要性。

5.2 关联回忆 (MQAR) ——记忆容量
在多查询关联回忆 (MQAR) 基准中，Transformer 因注意力的高容量特性可达标杆水平。非 Transformer 模型中，带 mLSTM 的 xLSTM 表现最佳，尤其是 xLSTM[1:1] 与 xLSTM[1:0]。即便键值对增至 256 对、序列长达 2048，mLSTM 架构仍保持较高回忆率。

MQAR 实验中，不同键值对设置下准确率与模型维度关系。Transformer 很快接近满分；xLSTM 变体与最优非 Transformer 模型相当，且随维度良好扩展。

图 5: 多查询关联回忆实验，不同模型在各模型维度、键值数量下的准确率。xLSTM[1:1] 与 xLSTM[1:0] 在非 Transformer 模型中表现最佳。

5.3 Long Range Arena——多样长上下文任务
在长距离竞技场的检索、列表操作、像素级图像任务等基准中，xLSTM 一致表现强劲，达到或超过其他线性时间或长上下文架构。

5.4 组件消融实验
消融结果表明:

  • 为普通 LSTM 添加现代残差骨干 (LayerNorm+残差连接) 显著提升训练。
  • 引入指数门控使困惑度下降明显。
  • 用 mLSTM 替换部分 sLSTM 块继续提升。
  • 门控可学习且依赖输入，可带来额外增益。完整门控设置最佳。

结论: 指数门控与矩阵记忆均是 xLSTM 强性能的必要因素。

5.5 大规模语言建模 (150 亿词元)
在相同 SlimPajama 子集训练下，xLSTM 各变体在验证困惑度方面超过 Transformer、SSM、RWKV 和线性注意力方法。

模型参数量 (百万)困惑度
Llama40714.25
Mamba42313.70
RWKV-545614.25
xLSTM[1:0]40913.43
xLSTM[7:1]40813.48

验证困惑度与模型大小: xLSTM 曲线在所有尺寸上均低于竞争 RNN 与 SSM。

图 6: 在 150 亿 SlimPajama 词元上的扩展性，xLSTM 始终领先。

5.6 完整 LLM 训练 (3000 亿词元) ——外推与下游任务
在 3000 亿词元训练下，测试多种大小模型 (1.25亿–13亿参数) :

  • 长度外推: 训练上下文 2048，测试至 16k。xLSTM 困惑度稳定，Transformer 超范围后困惑度急剧恶化。
  • 下游任务: xLSTM 在各种基准上有竞争力或领先；PALOMA 评估中大多数领域困惑度更低。

不同 token 位置的困惑度: xLSTM 在长上下文保持低困惑度，Transformer 出现上升。

图 7: 长度外推测试 (训练上下文 2048，评估至 16k) ，13 亿参数模型，3000 亿词元训练。

5.7 推理速度与吞吐量
xLSTM 推理优势:

  • 循环解码生成时间与长度线性相关；Transformer KV 缓存可能导致更高复杂度。
  • 内存占用与序列长度无关 (每个头的矩阵记忆固定) ，推理时可用更大批量，吞吐更高。

左: 生成时间 vs 长度——xLSTM 线性增长；右: 吞吐量 vs 批量——xLSTM 支持更大批量且吞吐更高。

图 9: 13 亿参数模型推理速度与吞吐量。左: 循环模型生成时间线性扩展；Transformer 增速更快。右: 不同批量下 tokens/s，xLSTM 恒定内存占用支持更大批量。

6 — 局限性与实践考量

  • sLSTM 因循环混合无法并行，尽管 CUDA 内核优化，但仍慢于全并行。
  • mLSTM 的 \(d \times d\) 矩阵运算计算量大，可通过高效 GPU 内核或类似 FlashAttention 的优化缓解。
  • 遗忘门偏置的初始化对训练稳定性重要。
  • 更大规模 xLSTM 模型可从更充分的超参调优与内核工程中获益。

简言之: 架构潜力巨大，但工程优化仍有空间，以缩小与高优化 Transformer 内核的性能差距。

7 — 总结: xLSTM 的定位

  • 当任务需要显式且可修正的状态 (如算法类任务、部分推理、长上下文状态跟踪) ，sLSTM 风格循环独具威力。
  • 当容量是关键 (大量键值、罕见词元记忆) ，矩阵记忆 (mLSTM) 是有吸引力的选择。
  • 混合 xLSTM 在表达力和并行性上取得平衡: 大部分用 mLSTM 批处理，少量 sLSTM 应对难处理的顺序现象。
  • 大规模下，xLSTM 在困惑度和泛化上匹敌或超越竞争架构，并具备实用推理优势 (恒定内存、线性生成成本) 。

8 — 最后的思考: RNN 回归了吗？

xLSTM 并非对旧架构的怀旧翻版，而是将 LSTM 核心思想 (常数误差环路与门控) 与现代创新 (指数门控、稳定归一化、矩阵记忆、残差骨干) 结合，打造出适合 LLM 时代的竞争性序列模型家族。

从“设计原语”的角度看，xLSTM 带来两个值得记住的要点:

  • 当任务需要修正时，让记忆更新果断 (指数门控) 。
  • 为模型提供合适的记忆基底: 状态推理用标量循环，大容量回忆用矩阵关联记忆——并在残差堆叠中结合。

Transformer 革命体现了特定的并行与关联能力平衡。xLSTM 展示了另一种平衡——更智能的循环 + 关联矩阵记忆——同样可达顶尖水平，甚至在部分场景超越。这为架构设计开辟了新方向，尤其适合长上下文、高容量记忆、恒定内存推理场景。

想进一步探索，可参考论文附带的代码，内含详细伪代码、CUDA 实现及消融实验。

参考文献与注释

  • 内容与实验结果来自 “xLSTM: Extended Long Short-Term Memory” (Maximilian Beck 等) 。论文提供了完整推导、伪代码，以及关于数值稳定性和并行公式的附录。
  • 关于 LSTM 核心思想的历史背景，请参阅 Hochreiter & Schmidhuber (1997) 及后续注意力、状态空间模型、快速权重的相关文献。
  • 数据集: SlimPajama、PALOMA、MQAR、Long Range Arena。

拓展阅读

  • 原始 LSTM 论文 (Hochreiter & Schmidhuber) 。
  • 快速权重编程、Hopfield 网络、外积关联记忆文献。
  • 最新长上下文与 SSM 论文 (S4、Mamba、Retention) 。
  • RWKV 等现代 RNN LLM 工作。