XLSTM：扩展长短期记忆网络的新研究

FreeGuideOnline 最新 2026-06-21

什么是 XLSTM？扩展长短期记忆网络

XLSTM（Extended Long Short-Term Memory）是2024年由 Sepp Hochreiter 团队提出的一套新架构，其目标是对经典 LSTM 进行深度现代化改造，使其在语言建模等大规模任务上能够媲美甚至超越 Transformer。XLSTM 并非单一模型，而是一个方法家族，主要包含两大变体：sLSTM（标量 LSTM）与 mLSTM（矩阵 LSTM）。它们保留了 LSTM 常数级推理复杂度和长程记忆的核心优势，同时修复了原始 LSTM 无法并行训练、存储混合能力受限等问题。

为什么需要重新审视 LSTM？

LSTM 曾是序列建模的基石，但在 Transformer 出现后逐渐被边缘化。这归结于三个关键局限：

  • 无法并行化训练：标准 LSTM 的时间步循环导致无法像 Transformer 那样利用 GPU 一次处理整个序列。
  • 记忆存储方式僵化：LSTM 的记忆更新受限于逐元素相乘的门控，不能动态修正已存储的信息，也难以精细化地清除或重写。
  • 容量限制：隐藏状态和记忆单元是固定维度的向量，当序列变长、信息密度变高时，存储压力增大。

XLSTM 的设计正是围绕这几点展开突破，同时保留了 LSTM 在推理时恒定的状态大小和计算开销，非常适合流式处理和边缘设备。

sLSTM：重振标量记忆

sLSTM（scalar LSTM）表面看类似经典 LSTM，但引入了多项重要改进，让“老样本”焕发新生。

指数门控与记忆混合化

经典 LSTM 的记忆细胞更新是加法运算：

c_t = f_t ⊙ c_{t-1} + i_t ⊙ z_t

其中遗忘门 f_t 和输入门 i_t 都使用 sigmoid 激活，值域 (0,1)。这种方式缺乏对旧记忆进行重新加权或彻底覆写的能力。

sLSTM 改为指数门控思路。遗忘门和输入门不再强制绑定为 (0,1)，而是允许其指数化表达，使得记忆更新变为：

c_t = f_t · c_{t-1} + i_t · z_t

但这里的门值是直接经过指数函数处理的，从而允许对历史状态选择性衰减或增强。同时 sLSTM 引入了记忆混合（memory mixing），即在更新记忆之前，允许隐状态与记忆单元之间进行更深层的交互，这让 LSTM 能动态决定要记住什么、改写什么。

多头与标量控制

sLSTM 采用多头结构，每个头拥有独立的记忆处理流水线，最后拼接输出。该设计让不同头专注于不同时间尺度或语义维度的信息。尽管存在多个头，每个头的记忆单元依旧是标量维，即每个时间步每个单元存储单一数值，这与 Transformer 中键值对的高维向量截然不同，确保了状态空间的经济性。

mLSTM：将记忆升级为矩阵

mLSTM（matrix LSTM）是 XLSTM 家族中的另一极，它彻底改变了记忆的存储形式：不再是向量，而是矩阵。

从向量到矩阵：扩容存储

在 mLSTM 中，单元状态是一个 d×d 的矩阵 C_t。这个矩阵通过外积更新的方式写入和读取信息，灵感来源于快速权重编程（Fast Weight Programmers）和线性注意力。

更新规则可简写成：

C_t = f_t · C_{t-1} + i_t · (v_t ⊗ k_t)

这里 v_t 是值向量，k_t 是键向量，⊗ 代表外积操作。遗忘门 f_t 和输入门 i_t 控制着矩阵记忆的衰减与新信息的注入。最终，从矩阵记忆中读取信息时，执行矩阵与查询向量 q_t 的乘法：

o_t = C_t · q_t

这一机制使得 mLSTM 能够同时并行训练（因为矩阵更新可以构造成类似线性注意力的形式），又能在推理时保持固定大小的状态，实现真正的递归与并行统一。

与线性注意力的关系

mLSTM 的结构与线性 Transformer 非常相似，但引入了指数门控来弥补纯线性注意力的表达力不足。遗忘门使用指数激活，能产生极其尖锐的衰减曲线，有效防止记忆矩阵被无关信息淹没，这比标准线性注意力中用固定衰减率要强大得多。

XLSTM 架构总览：双轨并行

一个完整的 XLSTM 块通常将 sLSTM 与 mLSTM 以交错残差的方式堆叠。典型配置如下：

  1. 前处理：输入先经过层归一化，再投射到多个头。
  2. mLSTM 层：利用矩阵记忆对序列全局交互模式建模，支持完全并行训练。适合捕捉长期依赖。
  3. sLSTM 层：对 mLSTM 的输出进行深度精炼，利用标量记忆执行逐 token 的精细状态修正。
  4. 门控 MLP：XLSTM 也引入了类似的门控前馈网络，进一步增强非线性。

这种混合架构让 XLSTM 兼具了训练时并行计算与推理时恒定状态的优势，彻底打破了“RNN 式模型无法高效处理长序列”的刻板印象。

关键技术创新点解读

指数门控（Exponential Gating）

这是 XLSTM 最重要的设计。传统门控使用 sigmoid 或 tanh，指数门控则直接输出无界的正值或负值，再经指数映射，可获得近乎阶跃的取舍能力。这意味着模型可以决绝地遗忘，也可以突兀地写入，控制粒度远胜之前。

记忆混合（Memory Mixing）

在更新候选项 z_t 之前，允许输入与前一时刻记忆单元进行多轮交互投影，让 LSTM 具备“三思而后行”的递归深度，相当于在每个时间步内部展开一小段思考链。

归一化技术与初始化策略

为保持数值稳定性，XLSTM 采用精心设计的层归一化与门控初始化。遗忘门初始化为极大的偏置（如 3~6），使得模型在训练初期就具备“保留记忆”的默认倾向，这延续了经典 LSTM 遗忘门偏置技巧，但在指数门控框架下更为关键。

实战：XLSTM 能做什么？

  • 语言建模：在 SlimPajama 等数据集上，XLSTM 在相同参数量下性能已接近甚至超过最先进的 Transformer，并且能处理超长上下文（数万 token），推理内存恒定。
  • 强化学习与时间序列：笛卡尔状态递归特性让 XLSTM 天然适合需要逐步决策或流式预测的场景，如金融异常检测、机器人控制。
  • 并行训练与推理加速：mLSTM 可利用 chunkwise 并行算法，大幅提升训练吞吐量，而推理时无需 KV 缓存，显存占用极低。
  • 多模态与科学应用：矩阵记忆提供了更丰富的关联存储，有望在蛋白质序列、气象预测等结构化时间数据上展现优势。

代码片段：构建一个 XLSTM 块的伪代码

class XLSTMBlock(nn.Module):
    def __init__(self, dim, heads, mlp_ratio=4.0):
        super().__init__()
        self.mlstm = mLSTM(dim, heads)
        self.slstm = sLSTM(dim, heads)
        self.mlp = GatedMLP(dim, int(dim * mlp_ratio))

    def forward(self, x, state_mlstm, state_slstm):
        # 并行训练：state 可由下三角矩阵注意力实现
        x = x + self.mlstm(x, state_mlstm)
        x = x + self.slstm(x, state_slstm)
        x = x + self.mlp(x)
        return x, state_mlstm, state_slstm

在推理时，每个模块只保存并更新自身的记忆状态，无需重新计算历史序列，真正实现了 O(1) 计算和内存的 token 生成。

总结：LSTM 的文艺复兴

XLSTM 证明，递归网络远未过时。通过将标量记忆细化、矩阵记忆扩容，并辅以指数门控和并行化设计，LSTM 重获了与 Transformer 正面竞争的实力。对于追求推理效率、流式处理以及低资源部署的开发者，XLSTM 提供了非常值得关注的现代替代方案。

下一步建议阅读：官方论文《xLSTM: Extended Long Short-Term Memory》，以及xLSTM的实现库（如mlx-lstm、xlstm-pytorch）来动手实验。