技术摘要：变分线性注意力（VLA）

1. 问题陈述

长上下文序列建模在计算效率与内存稳定性之间面临根本性的权衡。标准 Transformer 注意力机制实现了最优的关联检索，但其时间复杂度为 $O(T^2)$，内存占用为 $O(Td)$，导致其无法处理超过$10^5$ 个 token 的序列。线性注意力通过核特征映射和循环状态更新将复杂度降低至 $O(T)$，从而解决了计算瓶颈。然而，它引入了一种关键的失效模式：由于无条件的累加，内部记忆状态 $S_t$ 的弗罗贝尼乌斯范数（Frobenius norm）会无界增长（$O(T)$）。这种无界增长导致存储的关联之间产生渐进式干扰，随着上下文长度的增加，检索精度逐渐下降。

尽管 DeltaNet 试图通过标量门控机制来缓解这一问题，该机制对状态进行均匀衰减，但它缺乏方向选择性控制。标量门无法优先遗忘最近被覆盖的方向，同时保留旧的、稳定的关联，导致在高内存负载下精度迅速崩溃。作者认为，核心问题并非计算成本，而是内存稳定性，具体而言是内存更新的几何结构。

2. 方法论：变分线性注意力

本文提出了变分线性注意力（VLA），将线性注意力的内存更新重新框架化为一个在线正则化最小二乘问题。

2.1 核心公式

VLA 不再进行无条件累加（$S_t = S_{t-1} + v_t k_t^\top$），而是寻求一个记忆矩阵 $S_t$，以最小化惩罚后的预测误差：
$$S^*_t = \arg \min_S \sum_{s=1}^t \|v_s - S\hat{k}_s\|^2 + \text{tr}(S M_t S^\top)$$
其中 $\hat{k}_s$ 是归一化的键（key），$M_t$ 是一个随时间变化的惩罚矩阵，编码了先前所见键的几何结构。

2.2 自适应惩罚几何

VLA 利用Sherman-Morrison 秩 -1 公式，以每步 $O(d^2)$ 的成本精确维护逆惩罚矩阵 $A_t = M_t^{-1}$。

• 机制：$M_t$ 累积学习到的惩罚方向的外积。因此，$A_t$ 在最近接收到“惩罚质量”（即被写入）的方向上发生收缩。
• 效果：最近被写入的子空间接收更小的更新，而未触及的子空间保持可用。这构成了 DeltaNet 标量门控的矩阵值类比，实现了方向选择性的内存更新。

2.3 递归更新规则

由此公式推导出的精确递归更新规则如下：

1. 更新逆惩罚矩阵：$A_t = A_{t-1} - \frac{z_t z_t^\top}{1 + u_t^\top z_t}$（Sherman-Morrison 公式）。
2. 计算门控向量：$\hat{\alpha}_t = \frac{A_t \hat{k}_t}{\|A_t \hat{k}_t\|}$。
3. 更新状态：$S_t = S_{t-1} + (v_t - S_{t-1}\hat{k}_t) \hat{\alpha}_t^\top$。

关键在于，键 $\hat{k}_t$ 和门控向量 $\hat{\alpha}_t$ 均被归一化为单位长度。输出计算方式类似于标准线性注意力的递归比率，利用 ELU+1 特征映射。

2.4 实现效率

虽然 Sherman-Morrison 更新引入了常数因子的开销（比标准线性注意力多约 $5\times$ 的操作），但作者实现了并行 Blelloch-scan 公式和融合 Triton 内核。这将所有 $T$ 步融合为单个 GPU 内核启动，消除了每 token 的分发开销，相比顺序 Python 实现实现了 $14\times$ 的加速。

3. 主要贡献

3.1 理论保证

本文对两种稳定性属性提供了严格的证明：

1. 有界状态增长（命题 1）：在有界输入下，弗罗贝尼乌斯范数 $\|S_t\|_F$ 具有自限制性。与标准线性注意力中 $\|S_T\|_F = O(T)$ 不同，VLA 的状态随着预测残差的消失而收敛到一个有限的平台期。
2. 单位雅可比谱范数（命题 2）：递归雅可比矩阵 $J_t = \partial S_t / \partial S_{t-1} = I - \hat{\alpha}_t \hat{k}_t^\top$ 的谱范数对于所有序列长度和头维度精确等于 1（前提是向量已归一化）。这保证了在任意深度下梯度的稳定流动（无爆炸或消失）。若无归一化，谱范数将随 $d_h/\lambda_0$ 增长，导致梯度放大（例如，对于 $d_h=96$，25 步后达到 $10^{32}$）。

3.2 关联记忆容量

VLA 保持了每个头 $d_h$ 个关联的理论容量（其中 $d_h$ 是头维度）。然而，与累加不同，VLA 将新写入路由到与最近使用的子空间正交的方向。这使得它能够在不均匀稀释的情况下保留旧关联，在容量边界内保持精确的回忆。

3.3 效率

所提出的融合 Triton 内核实现了 $O(T)$ 的扩展性，并在大约 43,000 个 token 处低于 softmax 注意力的延迟，使其适用于长上下文推理。

4. 实证结果

实验在共享的双层 Transformer 骨干网络上进行，使用了**多查询关联回忆（MQAR）**任务和复制任务。

• 状态范数稳定性：在 $T=1,000$ 时，标准线性注意力的 $\|S_T\|_F \approx 1,600$，而 VLA 保持在 15 以下（减少了 $109\times$）。DeltaNet 也是有界的，但缺乏 VLA 的方向选择性。
• 关联回忆（容量）：在配对数量 $n_{pairs} < d_h$ 的情况下（具体为 $n=24, d_h=32$）：
  • VLA：保持 1.000 的精确匹配准确率。
  • DeltaNet：崩溃至 0.010。
  • 标准线性注意力：下降至 0.074。
  • Softmax：下降至 0.26。
• 长上下文保留：固定 $n=8$ 并将 $T$ 增加至 512，VLA 保持 1.000 的准确率。DeltaNet 由于其标量门在查询到达前衰减了旧关联，在所有长度下均崩溃至 $\approx 0.01$。
• 过载区域：当 $n > d_h$ 时，VLA 比累加线性注意力退化得更温和，因为自适应惩罚矩阵优先覆盖最近写入的方向，从而部分保留了旧关联。

5. 意义与主张

本文提出，长上下文可靠性是一个内存几何问题，而不仅仅是计算问题。

• 几何控制：主要贡献在于证明控制状态允许更新的方向比单纯降低更新计算成本更为关键。通过将无条件累加替换为由自适应惩罚逆控制的残差误差更新，VLA 实现了数值稳定、恒定内存的序列建模。
• 与经典理论的联系：VLA 将线性注意力与**递归最小二乘（RLS）**自适应滤波器联系起来，提供了先前“快速权重”方法中缺乏的原理性推导和稳定性理论。
• 实际范围：作者声称，VLA 提供了 $d_h$ 个最近强化关联的固定大小、数值稳定的摘要，适用于无需 KV 缓存基础设施的流式推理。他们明确指出，VLA 并未克服每个头 $d_h$ 个关联的根本容量限制，但在该界限内提供了更优越的稳定性和检索能力。

该工作得出结论：尽管 VLA 引入了常数因子的开销，但其理论稳定性（单位雅可比、有界状态）与实证性能（容量内的完美回忆）相结合，使其成为长上下文 Transformer 的稳健解决方案，而标准线性注意力则因干扰而失效。