多年来，Transformer 一直是序列建模领域无可争议的王者，为从 GPT 等大型语言模型到科学与多模态 AI 的各种突破提供了动力。然而，即使是王者也有弱点——Transformer 在效率上存在挑战。其计算成本随序列长度呈二次方增长，这意味着处理一本书的成本远高于处理一个句子。随着研究人员不断推动模型去理解整段代码、长期对话，甚至持续的感官数据流，这一局限已成为严重的瓶颈。

此时，状态空间模型 (SSM) 登场。这类架构有望实现与 Transformer 相当的性能，同时具有一个关键优势: 计算量随序列长度线性增长。换句话说，将序列长度加倍只会使计算开销加倍，而不是四倍。最近的研究成果，如 Mamba，表明 SSM 在语言建模、视觉乃至更多领域能够与 Transformer 相媲美甚至超越。

然而，SSM 的设计至今仍带有一定的“艺术性”。开发者往往依靠经验调整参数与递归结构，而缺乏统一的理论来解释为什么某些设计有效。这一过程更多依赖直觉，而非系统的原则。

来自德克萨斯大学奥斯汀分校与 Helixon 的研究团队在最新论文 《LONGHORN: 状态空间模型是摊销式在线学习器》 中，为 SSM 的设计提供了坚实的理论基础。论文提出了一个令人耳目一新的观点: 将 SSM 视作一种在线学习系统，能够在处理数据时逐步更新其内部状态——一次处理一个词元。

这一思想催生了一个全新且优雅的架构——Longhorn，直接源自在线联想回忆的数学形式。结果令人惊艳: Longhorn 不仅超过了当前最先进的模型如 Mamba，还在样本效率与长度泛化方面表现突出。

展示 Longhorn 卓越样本效率和长度外推能力的关键结果。

图 1: (左) Longhorn 在下游任务中的样本效率比 Mamba 提高了 1.8 倍。 (右) 在上下文长度为 2048 的训练下，Longhorn 在推理阶段可泛化到长达 16 倍 (32,768 词元) 的序列，展现出卓越的长度外推能力。

目录

  • 状态空间模型概览
  • 如何将 SSM 框架化为在线学习器
  • Longhorn 架构与其闭式更新
  • 实验结果与主流模型比较
  • 未来方向及更广泛影响

背景: 状态空间模型速览

在深入了解 Longhorn 之前，让我们简要回顾现代序列模型的关键运作原理。

多数大型模型——包括 Transformer 与基于 SSM 的网络——通常由堆叠模块构成，每个模块执行两项核心操作:

  1. 序列混合 (Sequence Mixing): 信息在序列的不同位置之间流动。在 Transformer 中，这一过程由自注意力机制实现，使每个词元能与序列中的所有其他词元交互。
  2. 通道混合 (Channel Mixing): 信息在每个词元表示的内部进行处理，通常由多层感知机 (MLP) 实现。

SSM 被设计为自注意力机制的高效替代方案，专注于实现序列混合的功能。

Mamba 模块的高层视图，展示了序列混合与通道混合的独立路径。

图 2: 基于 SSM 的模型 (如 Mamba) 中的典型模块。SSM (红色路径) 负责序列混合，而 MLP 类组件 (蓝色路径) 执行通道混合。

SSM 的核心是一条简单但极具力量的递归关系:

\[ S_t = A(x_t) \cdot S_{t-1} + B(x_t) \]

其中，\(x_t\) 是第 \(t\) 个时间步的输入词元，\(S_t\) 则汇总此前所有信息。矩阵 \(A_t\) 与 \(B_t\) 决定了状态的演化方式——即保留多少旧信息、引入多少新信息。

虽然该递归形式看似顺序执行，但在训练过程中，SSM 借助前缀扫描算法实现并行计算，从而可以一次性高效地获得所有状态。这一特性使其能够像 Transformer 一样进行训练 (完全并行) ，但在推理阶段则如 RNN 一般高效 (线性扩展) 。

不同的 SSM 在设计 \(A_t\)、\(B_t\) 及其更新机制上各有差异。许多设计依靠启发式方法来平衡性能与计算可行性。Longhorn 论文旨在以一个通用的理论原则取代这种临时性的经验方法。

核心思想: 将状态空间模型视为在线学习器

作者指出，任何 SSM 的递归关系都可以被解释为解决某种在线学习问题的最优解。

在在线学习中，智能体按时间序列进行预测。每一步都会观察新数据、产生损失，并根据结果更新其内部状态，以更好地预测未来。在线学习器需要在*稳定性 (防止遗忘) 与可塑性 *(快速适应新信息) 之间找到平衡。

这种平衡可形式化为在线凸规划 (OCP):

\[ s_t = \underset{s}{\arg\min}\; L_t(s), \quad L_t(s) = D_{\phi}(s, s_{t-1}) + \beta_t \ell_t(s) \]

其中两部分相辅相成:

  • 稳定性项 \(D_{\phi}(s, s_{t-1})\): 使新状态保持与旧状态的接近度，避免灾难性遗忘。
  • 可塑性项 \(\beta_t \ell_t(s)\): 促进模型从新数据中学习，由类似学习率的系数 \(\beta_t\) 控制。

Longhorn 将 SSM 的状态更新视作一次隐式在线学习步骤——即状态 \(S_t\) 针对序列信息优化上述目标函数。

一个概念图，展示如何将序列混合框架化为在线学习问题。

图 3: Longhorn 框架示意图。 (左) 序列模型中的信息混合。 (中) 将该更新过程框架化为在线学习。 (右) Longhorn 的递归结构源于一个在线联想回忆目标。

通过这种视角，SSM 的设计不再是“调参艺术”，而成为理论上可解释、数学上可推导的过程。我们基于一个合理的学习目标来设计更新机制，而更新公式自然从优化过程导出。

Longhorn 架构: 学习回忆

在上述原则指导下，研究人员选择一个简单而强大的目标:** 在线联想回忆**。

这一概念与 Transformer 的“归纳头 (induction head)”能力直接相关——该模式负责上下文学习。模型处理 (键，值) 对后，学会在给定键时预测正确的值。Longhorn 将这一行为明确内嵌于其递归公式之中。

在每个时间步，它观测键 \(k_t\) 与值 \(x_t\)，并根据如下公式更新状态 \(S_t \in \mathbb{R}^{d \times m}\):

\[ S_t = \underset{S \in \mathbb{R}^{d \times m}}{\arg\min} \left\{ \|S - S_{t-1}\|_F^2 + \|S k_t - x_t\|_{\mathrm{diag}(\beta_t)}^2 \right\} \]

其中，\(\| \cdot \|_F\) 表示弗罗贝尼乌斯范数，而 \(\beta_t\) 控制新信息影响更新的幅度。

该优化问题存在一个闭式解:

\[ S_{t,i} = (I - \varepsilon_{t,i} k_t k_t^\top) S_{t-1,i} + \varepsilon_{t,i} k_t x_{t,i}, \quad \varepsilon_{t,i} = \frac{\beta_{t,i}}{1 + \beta_{t,i} k_t^\top k_t} \]

为提高计算效率，作者将 \(k_t k_t^\top\) 替换为其对角近似 \(k_t^{\odot 2}\)，使更新过程可与标准 SSM 并行化方式一致。最终公式遵循常见的 SSM 模板:

\[ S_t = A_t \odot S_{t-1} + B_t, \quad A_t = (1_{d \times m} - \varepsilon_t \otimes k_t^{\odot 2}), \quad B_t = (\varepsilon_t \odot x_t) \otimes k_t \]

一个极具美感的结果是: Longhorn 的遗忘门由数学推导自然产生，无需手动参数化。遗忘与记忆的平衡随键的动态变化自然实现。

实验与结果

多查询联想回忆 (MQAR)

为了验证理论基础，研究者首先在 MQAR 基准上测试了 Longhorn，该任务衡量模型检索已存储 (键，值) 对的能力。

在多查询联想回忆基准测试上的准确率。

图 4: 在 MQAR 测试中，Longhorn (青色) 实现了近乎完美的准确率，在长序列与小维度场景下优于 Mamba 与其他 SSM。

Longhorn 即使在序列长度为 512、隐藏维度较小的情况下，也能保持几乎完美的召回率，验证了其更新规则中有效的联想记忆机制。

语言建模扩展规律

随后，研究人员在 OpenWebText 数据集上对 Longhorn 进行语言建模评估，训练了参数量介于 1.2 亿至 3.5 亿的模型，使用上下文长度分别为 1024 与 4096。

在 OpenWebText 数据集上验证损失与模型大小的关系。

图 5: Longhorn 的验证损失始终低于其他 SSM，表现可与强大的 LLaMA Transformer 基线媲美。

不同模型和上下文长度的验证损失表。

表 1: OpenWebText 的详细结果显示，在 350M 规模下，Longhorn 在所有模型中达成最佳验证损失。

在所有配置中，Longhorn 的性能均优于 Mamba、RWKV 与 GLA——部分情况下甚至超越了基于 Transformer 的 LLaMA。

在 SlimPajama 上的大规模训练

进一步扩展规模后，团队在 SlimPajama 数据集上使用 1000 亿词元训练了一个 13 亿参数的 Longhorn 模型，并在八项下游任务上与其他模型进行了比较。

13 亿参数模型在下游基准测试上的结果。

表 2: Longhorn 在八项下游任务中取得最高平均分，以更少参数超越 Mamba。

在这些大规模实验中，Longhorn 不仅保持了强劲的整体性能，还展示了比 Mamba 高 1.8 倍的样本效率——使用不到一半训练数据即可达到竞争性的困惑度。

长度外推

众所周知，Transformer 在超出训练上下文长度时泛化能力较弱，而 Longhorn 的在线学习框架则赋予其卓越的外推性能。当在 2048 词元长度下训练时，Longhorn 在最长达 32K 词元的序列上仍能保持稳定困惑度——**是训练上下文长度的 16 倍 **(见图 1 右侧) 。

视觉任务

为验证跨领域能力，作者将 Longhorn 改编用于图像分类任务 (“Vision Longhorn”，简称 ViL) ，并在 ImageNet 上与 Vision Mamba (ViM) 对比。

在 ImageNet 分类任务上的结果。

表 3: Vision Longhorn (ViL) 在 ImageNet 上的 Top-1 准确率略高于 Vision Mamba (ViM)。

即使未进行额外调优，Vision Longhorn 也能匹敌甚至略胜 ViM，展现其跨模态的稳健性。

结论与未来方向

Longhorn 论文不仅提出了新架构，更建立了一个基于在线学习理论设计状态空间模型的系统化框架。

主要启示:

  1. 统一设计原则: SSM 可被理解为优化稳定性与可塑性平衡的在线学习器。
  2. 简洁与高效: Longhorn 的递归源自闭式解，无需手工门控设计，参数量更少。
  3. 卓越性能: Longhorn 在语言、视觉及合成任务上均达当前最优水平，具备非凡的样本效率与上下文外推能力。

展望未来，这一理论视角启示了新的探索方向: 围绕推理、工具使用或持续学习的其他在线优化目标。Longhorn 不只是 Transformer 的又一对手——它揭示了高效、长上下文模型的下一步进化。

如今，领域终于拥有了一个指导原则。不再凭直觉设计，我们可以学会学习——正如 Longhorn 所展现的那样。