写在前面：一场关于记忆的复兴之旅
你好，我是 Sepp Hochreiter，或者你可以叫我 Sepp。如果你对深度学习有所了解，可能听过我曾经的“孩子”——长短期记忆网络，也就是我们常说的 LSTM。它在90年代诞生，核心思想简单却极具革命性：通过一个“恒定误差传送带”（Constant Error Carousel）和一个巧妙的“门控”机制，它能记住遥远过去的信息，同时又能有选择地遗忘。这在当时解决了困扰循环神经网络（RNN）的梯度消失问题，让它成为处理序列数据的王者，也成为了第一代大语言模型的基石。
但时代的车轮滚滚向前，2017年，Transformer 横空出世。它凭借并行化的自注意力机制，在大规模数据上展现出无与伦比的性能，像一辆超跑将我那朴实耐用的“小货车”远远甩在身后。整个学术界和工业界的热情都转向了 Transformer，并用它打造了 GPT、Llama、PaLM 等一系列巨型模型。
然而，我的内心始终有一个简单而执着的疑问：如果我们用今天的技术武装 LSTM，去弥补它已知的局限性，它还能走多远？这就是我与我的团队着手进行这项研究的初衷。我们没有选择重新发明轮子，而是站在巨人的肩膀上，将现代大语言模型的成功经验——比如大规模并行、矩阵运算和高级门控——融入到 LSTM 的血脉之中。这场旅程的终点，就是我们今天正式发布的——XLSTM，一种被赋予新生的、面向大模型的扩展长短期记忆网络。
摘要
20世纪90年代，恒定误差传送带（constant error carousel）和门控（gating）作为长短期记忆网络（LSTM）的核心思想被提出。自那时起，LSTM经受住了时间的考验，并为众多深度学习的成功案例做出了贡献，尤其是在它们构成了首批大型语言模型（LLMs）方面。然而，以可并行化自注意力机制为核心的Transformer技术的出现，标志着一个新时代的黎明，并在规模上超越了LSTM。我们现在提出一个简单的问题：当我们将LSTM扩展到数十亿参数的规模，利用现代LLMs的最新技术，但同时缓解LSTM已知的局限性时，我们能在语言建模上取得多大进展？
首先，我们引入了指数门控（exponential gating）以及相应的归一化和稳定化技术。其次，我们修改了LSTM的记忆结构，得到了：(i) sLSTM，它具有标量记忆、标量更新和新的记忆混合（memory mixing）机制；(ii) mLSTM，它具有矩阵记忆和协方差更新规则，并且可以完全并行化。我们将这些LSTM扩展集成到残差块（residual block）骨干网络中，形成了xLSTM块，然后将这些xLSTM块以残差方式堆叠成xLSTM架构。实验证明，指数门控和修改后的记忆结构显著增强了XLSTM的能力，使其在性能和扩展性方面与最先进的Transformers和状态空间模型（State Space Models）相比具有优势。
LSTM 的局限性：记忆的“三宗罪”
在开始介绍 XLSTM 的创新之前，我们先来复盘一下 LSTM 的“三宗罪”：
- **无法改变已存储的决定（Inability to revise storage decisions）**：想象一下，你用一个传统的 LSTM 模型来解决“近邻搜索”（Nearest Neighbor Search）问题。模型需要顺序扫描一串向量，并在序列末尾返回与给定参考向量最相似的那个。但问题是，如果它在序列中间找到一个“A”，并决定把它存起来，但后面又遇到一个更相似的“B”，它很难“忘记”A而“记住”B。它就像一个固执的档案管理员，一旦将一份文件归档，就不太愿意轻易地去替换它。
- **记忆容量有限（Limited storage capacities）**：传统的 LSTM 使用一个单一的标量（scalar）作为记忆单元。这就像你只有一张小小的便利贴来记录所有重要的事。当信息量变大，尤其是像“罕见词汇预测”这类任务时，有限的记忆容量使得模型难以压缩和存储那些不常出现但至关重要的信息。你很难用一张小纸条记录下一整个图书馆的知识。
- **难以并行化（Lack of parallelizability）**：由于 LSTM 的设计是循环式的，即当前时间步的计算依赖于前一个时间步的隐藏状态。这使得训练过程必须像“串珠子”一样顺序进行，无法像 Transformer 那样同时处理整个序列。这种固有的“串行”特性，成为了其在大规模并行计算时代的主要瓶颈。
为了解决这些问题，我们对 LSTM 进行了“深度改造”。
XLSTM 的两大核心创新
XLSTM 的核心是两大武器：指数门控和新型记忆结构。它们相互配合，分别解决了上述的“固执”和“容量”问题，也为并行化铺平了道路。
1. 指数门控：让记忆“敢于取舍”
传统的 LSTM 使用 Sigmoid 门控，其函数形式为 $\sigma(x) = \frac{1}{1 + e^{-x}}$。这个函数的值域是 \([0, 1]\)，可以很好地模拟“忘记”和“输入”的比例。但它存在一个问题：梯度在 Sigmoid 接近 0 或 1 时会变得非常小，导致梯度流受阻，模型很难做出彻底的“忘记”或“记住”决策。
我们引入了指数门控（Exponential Gating），其核心是使用指数函数 \( \exp(x) \) 作为激活函数。这带来了几个优势：
- **更强的决策能力**：指数函数的值域是 \((0, \infty)\)，它能产生比 Sigmoid 门控大得多的值。这使得输入门和遗忘门可以产生非常大的权重，让模型有能力对记忆单元进行“彻底的覆盖”或“完全的保留”。这就像是给那个固执的档案管理员一个权限，让他可以果断地把旧文件扔掉，换上新的。
- **解决“记忆覆盖”问题**：在近邻搜索这类任务中，当找到一个更优的向量时，我们希望模型能完全忘记旧的，这正是指数门控所擅长的。它解决了传统 LSTM “无法修改已存储决定”的痛点。
- **稳定化与归一化**：为了防止指数函数产生过大的值导致数值溢出，我们引入了相应的归一化和稳定化技术，确保了训练过程的稳定性。这就像是在给档案管理员大权力时，也给他配了一套严格的管理系统，以防他乱来。
这套机制可以抽象为以下两个核心公式。传统的 LSTM 单元状态更新公式是：
\[ c_{t} = f_{t} \odot c_{t-1} + i_{t} \odot z_{t} \]
其中，\(f_t\) 是遗忘门，\(i_t\) 是输入门，它们都由 Sigmoid 函数控制。而在我们的 sLSTM 变体中，我们用指数门控 \(exp(\tilde{i_t})\) 替换了输入门，并引入了一个额外的归一化状态 \(n_t\) 来保持数值稳定：
\[ \overline{h}_{t} = c_t / n_t \]
这个 \( \overline{h}_{t} \) 才是真正的记忆内容。这就像用一个动态的“总和”来校准每一笔“收入”，防止单个大额交易让整个账本失衡。
动画标题：门控机制大比拼
生活化类比：想象两种记忆机制。Sigmoid 门控像一个温柔的过滤器，只能决定信息“通过一点”或“通过多一点”，但很难完全阻止或完全放行。而指数门控就像一个强力的阀门，能够完全打开或关闭，做出更果断的决定。
2. 新型记忆结构：从“便利贴”到“档案柜”
为了解决 LSTM 的第二宗罪——记忆容量有限，我们提出了两种新的记忆结构，由此诞生了 XLSTM 的两个新成员：
sLSTM：保留串行，升级记忆混合
sLSTM (scalar LSTM) 仍然沿用传统的标量记忆。但我们将其升级为多头（multiple heads）架构，并引入了一种新的记忆混合（new memory mixing）方式。这意味着每个头都有自己的记忆，它们之间可以进行信息交流。这就像你把一张便利贴升级成了几张不同的笔记，每张笔记记录不同的内容，但你可以在它们之间进行交叉引用和整理。
sLSTM 保留了 LSTM 的核心优势：记忆混合（memory mixing），即前一个时间步的隐藏状态会对当前时间步的计算产生影响。这使得 sLSTM 在处理需要“状态追踪”（state tracking）的问题时表现出色，例如评估一段代码或追踪长篇叙事中的实体。在形式语言任务上，sLSTM 甚至比 Transformer 和 Mamba 等模型更具表达能力。
静态示意图：sLSTM 记忆混合
这个示意图展示了 sLSTM 如何将输入（x）与上一时间步的隐藏状态（h）结合，通过多个门控头（i, f, z, o）来更新其内部的多个记忆细胞（c）。这类似于一个由多个子系统组成的复杂档案室，每个子系统独立工作但又相互影响。
mLSTM：矩阵记忆，拥抱并行化
mLSTM (matrix LSTM) 则是一个更激进的创新。我们直接将 LSTM 的记忆细胞从一个简单的标量 \($c \in \mathbb{R}^d$\) 升级为一个矩阵 \($C \in \mathbb{R}^{d \times d}$\)！这就像把你的便利贴笔记本升级成了一个巨大的、有行有列的表格，可以存储键值对（key-value pairs）。
为了更新这个矩阵记忆，我们引入了协方差更新规则（Covariance Update Rule），这个规则在关联记忆（Associative Memories）领域被证明是有效的。它的核心思想是：当你有一个新的键值对 \((k_t, v_t)\) 时，用它们的外积 \($v_t k_t^{\top}$\) 来更新矩阵记忆 \($C_t$\)。这个过程可以完全并行化，因为每个时间步的更新不再依赖于前一个时间步的隐藏状态。这就像一个多人的大型团队同时往一个共享的数据库里录入数据，而不需要排队等候。
由于 mLSTM 放弃了传统的记忆混合机制，它的训练可以像 Transformer 一样实现全并行，这彻底解决了 LSTM 的“第三宗罪”。
动画标题：矩阵记忆的并行狂想曲
生活化类比：想象一个巨大的电子表格，每一列是"键"，每一行是"值"。当一个新的信息到来时，整个团队可以同时找到对应的行和列，并瞬间更新整个表格。这与需要排队、按部就班更新的旧系统形成鲜明对比。
实验证明：XLSTM 的卓越表现
我们的实验涵盖了多个方面，旨在全面评估 XLSTM 的能力。结果令人振奋：
- **长序列外推**：我们训练了 1.3B 参数的模型，并测试它们在处理比训练时更长的序列时的表现。结果显示，Llama 和 Mamba 等模型在序列长度超过一定阈值后，性能会急剧下降，而 XLSTM 模型的困惑度（Perplexity）始终保持在低位。这证明了 XLSTM 强大的长程记忆和外推能力。
- **合成任务表现**：在多查询关联记忆（Multi-Query Associative Recall）任务中，mLSTM 展现出卓越的记忆容量，能够轻松记住多达 256 个键值对。这直接证明了矩阵记忆结构的有效性。
- **大模型对比**：在 SlimPajama 300B Token 的大规模语言建模实验中，XLSTM 在多个下游任务和基准测试（如 PALOMA）中，其性能均优于或与 Llama、Mamba 和 RWKV 等主流模型持平，甚至在一些任务上取得了最优结果。
- **推理速度**：尽管 mLSTM 的矩阵运算在理论上计算量较大，但由于其可并行化特性，我们的 CUDA 实现可以在 GPU 上实现极高的吞吐量。与 Llama 的二次方复杂度相比，XLSTM 的推理时间与序列长度呈线性关系，使其在长文本生成方面具有巨大优势。
动画标题：长文本，谁主沉浮？
生活化类比：想象在漆黑的隧道中前行。Llama 和 Mamba 像手电筒，光照范围有限，走远了就什么也看不见。而 XLSTM 像一个探照灯，光线能够穿透很远的距离，即使在很长的隧道里也能保持清晰视野。
未来展望：一场记忆革命的序章
我们的研究部分回答了那个简单的问题：“当 LSTM 扩展到数十亿参数时，我们能走多远？”答案是：“至少和目前主流的 Transformer 和状态空间模型一样远，甚至更远。” XLSTM 的诞生，证明了 LSTM 这种经典架构在经过现代化改造后，依然具有强大的生命力，甚至有可能在大语言模型领域掀起一场新的记忆革命。
当然，这仅仅是一个开始。XLSTM 的潜力远不止于此。它在强化学习、时间序列预测和物理系统建模等领域，都可能带来巨大的影响。我们相信，这个“新瓶装旧酒”的 XLSTM，将成为未来大模型研究中一个值得关注的竞争者。
技术细节附录
数学符号与公式（部分）
LSTM 的核心更新规则可以向量化表示为：
\[ c_t = f_t \odot c_{t-1} + i_t \odot z_t \qquad \text{(单元状态)} \]
其中，\(c_t\) 是当前时间步的单元状态，\(c_{t-1}\) 是前一时间步的单元状态。\(f_t\) 和 \(i_t\) 分别是遗忘门和输入门，它们由输入 \($x_t$\) 和前一隐藏状态 \($h_{t-1}$\) 决定，并通过 Sigmoid 函数 \(\sigma(\cdot)\) 归一化。\(z_t\) 则是新的单元输入。
\[ f_t = \sigma(W_f x_t + R_f h_{t-1} + b_f) \qquad \text{(遗忘门)} \]
\[ i_t = \sigma(W_i x_t + R_i h_{t-1} + b_i) \qquad \text{(输入门)} \]
XLSTM 中的一个关键创新是**指数门控**。在 sLSTM 中，我们用 \(\exp(\tilde{i}_t)\) 替代了传统的输入门，并引入了归一化状态 \(n_t\) 来防止数值溢出。
\[ n_t = f_t \odot n_{t-1} + i_t \qquad \text{(归一化状态)} \]
在 mLSTM 中，我们将记忆单元扩展为一个矩阵 \(C \in \mathbb{R}^{d \times d}\)，其更新规则基于**协方差更新**：
\[ C_t = f_t \odot C_{t-1} + i_t v_t k_t^{\top} \qquad \text{(矩阵单元状态)} \]
这里的 \($k_t$\) 和 \($v_t$\) 分别是键（key）和值（value）向量，其外积 \($v_t k_t^{\top}$\) 形成一个秩为 1 的矩阵，用于更新主记忆矩阵 \($C_t$\)。这种更新方式使得整个计算可以高度并行化，是 mLSTM 能够在大规模训练中与 Transformer 竞争的关键。