变分线性注意力：用于长上下文 Transformer 的稳定联想记忆

# 变分线性注意力：长上下文 Transformer 的稳定联想记忆
来源：https://arxiv.org/html/2605.11196
Vishal Pandey
独立研究员
伦敦，英国
pandeyvishal\.mlprof@gmail\.com &Gopal Singh
Metriqual
雅典，GR
gopal@metriqual\.com

###### 摘要
线性注意力将 softmax 注意力的二次成本降低至 $O(T)$，但其内存状态在 Frobenius 范数下以 $O(T)$ 增长，导致存储的关联之间产生渐进式干扰。我们引入了**变分线性注意力（Variational Linear Attention, VLA）**，它将内存更新重构为一个在线正则化最小二乘问题，并通过 Sherman-Morrison 秩-1 公式维护自适应惩罚矩阵。我们证明，将写入方向归一化为单位长度可以使递推 Jacobian 的谱范数在所有序列长度和头维度下恰好为 1（命题 2），并且状态范数在输入有界时具有自限性（命题 1）。在经验上，VLA 在 $T=1,000$ 时将 $\|S_t\|_F$ 减少了 $109$ 倍，相对于标准线性注意力，在每头有效内存容量范围内实现了多查询联想回忆（MQAR）的近完美精确匹配准确率。

$\lambda_0 > 0$ 是初始化正则化器（我们使用 $\lambda_0=0.1$，给出 $A_0=10I$）。因为 $u_s$ 是单位归一化的，每次秩-1 更新会在确切的一个方向上以有限的量耗尽 $A_t = M_t^{-1}$。积累大量惩罚质量的子空间接收到的未来更新较小；未触及的子空间保持全幅度。这是 VLA 将新关联写入当前状态尚未利用的方向的机制。

### 3.3 递归更新规则
公式 (7) 的解可以在每步 $O(d_h^2)$ 中精确维护。对公式 (8) 应用 Sherman-Morrison 公式：

$$
z_t = A_{t-1} u_t, \quad \delta_t = 1 + u_t^\top z_t, \quad A_t = A_{t-1} - \frac{z_t z_t^\top}{\delta_t}. \tag{9}
$$

由于 $A_{t-1} \succ 0$，我们总有 $\delta_t \ge 1$；因此除法是无条件安全的。有了 $A_t$，内存更新如下：

$$
e_t = v_t - S_{t-1} \hat{k}_t, \quad \hat{\alpha}_t = \frac{A_t \hat{k}_t}{\|A_t \hat{k}_t\|}, \quad S_t = S_{t-1} + e_t \hat{\alpha}_t^\top. \tag{10}
$$

将 $\hat{k}_t$ 和 $\hat{\alpha}_t$ 都归一化为单位向量至关重要：我们在 §4（命题 2）中证明，当两者都进行单位归一化时，Jacobian $\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}$ 的梯度放大，与我们在早期消融实验中观察到的 NaN 损失一致。

每个位置的输出为：

$$
o_t = \frac{S_t \phi(q_t)}{\max\bigl(\phi(q_t)^\top z_t^{\text{key}}, \;\varepsilon\bigr)}, \quad z_t^{\text{key}} = \textstyle\sum_{s \le t} \phi(k_s), \tag{11}
$$

其中 $\phi = \text{ELU}(\cdot) + 1$ 是标准线性注意力的特征映射，$\varepsilon = 10^{-4}$ 防止除以零。

公式 (9)-(11) 构成了一个完整的注意力头，无需矩阵求逆。

### 3.4 并行公式
$S_t$ 递推 (10) 是一个*线性递推*：

$$
S_t = F_t S_{t-1} + G_t, \quad F_t = I - \hat{\alpha}_t \hat{k}_t^\top, \quad G_t = e_t \hat{\alpha}_t^\top. \tag{12}
$$

对 $(F, G)$ 在组合 $(F_r, G_r) \circ (F_l, G_l) = (F_r F_l, \; F_r G_l + G_r)$ 下是关联的，这使得可以在 $O(\log T)$ 并行步骤中以 $O(T)$ 总工作量执行 Blelloch 并行前缀扫描 [2]。$A_t$ 循环有一个数据相关的分母 $\delta_t$，阻碍了直接的并行化；相反，我们将所有 $T$ 步融合到一个单一的 Triton 内核中，消除了朴素 Python 实现中主导延迟的每 token 内核分发开销。由此产生的 VLA-Triton 内核在 $T=4096$ 时比顺序 Python 实现了 14 倍的速度提升（§6）。

### 3.5 复杂度及与先前工作的关系
VLA 与标准线性注意力和 DeltaNet 具有相同的 $O(T d_h^2)$ 时间复杂度和 $O(d_h^2)$ 内存复杂度；常数因子包括 SM 更新的额外五次每 token 操作。表 1 总结了模型家族。

**表 1：每注意力头的复杂度比较。** $T=$ 序列长度，$d_h=$ 头维度 ($d_h = d/H$)。**Gate** 是每步遗忘机制。VLA 和线性注意力共享 $O(T d_h^2)$ 渐近时间；VLA 的常数因子由于 Sherman-Morrison 更新而大约大 5 倍（通过 Triton 内核缓解，§3）。† Softmax 注意力 KV 缓存随 $T$ 增长；在 $T=100\text{K}$, $d=4096$, $L=32$ 层时，这约为 52GB。所有其他模型使用 $O(d_h^2)$ 固定状态，独立于 $T$。

VLA 严格泛化了上述两个先前模型。设置 $u_t=0$ 并固定 $A_t = \lambda_0^{-1} I$ 将 (9)-(10) 坍缩为标准线性注意力（累加，无几何结构）。用标量门替换 $\hat{\alpha}_t$ 并丢弃 $A_t$ 恢复 DeltaNet。唯一的架构差异是 $d_h \times d_h$ 矩阵 $A_t$，它提供了每个方向的选择性：它可以同时保护编码旧关联的子空间，并为新关联打开一个新方向。标量门必须在两者之间进行权衡。这种区别构成了 VLA 在高内存负载下的优势，我们在 §6 中通过经验证明了这一点。

## 4 理论性质

我们建立了 VLA 的两个核心性质：(1) 内存状态范数具有自限性，(2) 递推 Jacobian 具有单位谱范数，保证稳定的梯度流。然后我们推导容量和长上下文行为作为推论。

### 4.1 有界状态动力学

###### 命题 1（有界状态增长）
设对所有 $t$ 有 $\|v_t\| \le C_v$。在 VLA v3 更新 $S_t = S_{t-1} + e_t \hat{\alpha}_t^\top$ 且 $\|\hat{\alpha}_t\|=1$ 下，Frobenius 范数满足

$$
\|S_t\|_F \le \|S_0\|_F + \sum_{s=1}^t \|e_s\|, \tag{13}
$$

并且当输入有界时，$\|S_t\|_F$ 收敛到一个有限平台。

###### 证明
由于 $\hat{\alpha}_t$ 是单位归一化的，更新的 Frobenius 范数为 $\|e_t \hat{\alpha}_t^\top\|_F = \|e_t\| \|\hat{\alpha}_t\| = \|e_t\|$。由三角不等式，$\|S_t\|_F \le \|S_{t-1}\|_F + \|e_t\|$，这望远镜求和得到 (13)。随着 $S_t$ 越来越适应存储的关联，预测残差 $e_t = v_t - S_{t-1} \hat{k}_t$ 根据 LMS 滤波器的 Widrow-Hoff 收敛性质 [4] 趋向于零：一旦对于观察到的键有 $S_{t-1} \hat{k}_s \approx v_s$，后续更新接近于零，且 $\|S_t\|_F$ 达到平台。$\blacksquare$

**与标准线性注意力的对比。** 在累加更新 $S_t = S_{t-1} + v_t k_t^\top$ 下，范数满足 $\|S_t\|_F \le \|S_0\|_F + \sum_{s=1}^t \|v_s\| \|k_s\|$，这对于有界输入以 $O(T)$ 增长，没有消失的残差来阻止累积。我们在图 2(b) 中通过经验验证了这一点：在 $T=1,000$ 时，标准线性注意力达到 $\|S\|_F \approx 1,600$，而 VLA v3 保持在 15 以下。

### 4.2 单位 Jacobian 谱范数

###### 命题 2（单位 Jacobian）
设 $\hat{k}_t, \hat{\alpha}_t \in \mathbb{R}^{d_h}$ 为单位向量。$S_t$ 递推关于 $S_{t-1}$ 的 Jacobian 为

$$
J_t = \frac{\partial S_t}{\partial S_{t-1}} = I - \hat{\alpha}_t \hat{k}_t^\top. \tag{14}
$$

该矩阵对所有 $\hat{k}_t, \hat{\alpha}_t$ 满足 $\|J_t\|_2 = 1$。

###### 证明
$J_t = I - \hat{\alpha}_t \hat{k}_t^\top$ 是单位矩阵的秩-1 扰动。其特征值为 $1$（重数 $d_h-1$，对应于 $\hat{k}_t$ 的 $d_h-1$ 维补空间）以及 $1 - \hat{k}_t^\top \hat{\alpha}_t$（对应于方向 $\hat{k}_t$）。由于 $\|\hat{k}_t\| = \|\hat{\alpha}_t\| = 1$，柯西-施瓦茨不等式给出 $|\hat{k}_t^\top \hat{\alpha}_t| \le 1$，因此后一个特征值的模在 $[0, 2]$ 中。然而，对于秩-1 投影矩阵 $P = \hat{\alpha}_t \hat{k}_t^\top$ 且 $\|P\|_2 = \|\hat{\alpha}_t\| \|\hat{k}_t\| = 1$，我们一般有 $\|I-P\|_2 \le 1+\|P\|_2 = 2$，但当两个向量都进行单位归一化时（更新是非扩张投影步），$I-P$ 的奇异值满足 $\sigma_{\max}(I-P)=1$。形式上：$\|J_t x\|^2 = \|x\|^2 - 2(\hat{k}_t^\top x)(\hat{\alpha}_t^\top x) + (\hat{k}_t^\top x)^2$。在 $\|x\|=1$ 上最大化并应用柯西-施瓦茨得到 $\sigma_{\max}^2 = 1$。$\blacksquare$

###### 推论 1（稳定梯度流）
标量损失 $\mathcal{L}$ 通过完整 $T$ 步递推的梯度满足 $\|\partial \mathcal{L} / \partial S_0\|_F \le \|\partial \mathcal{L} / \partial S_T\|_F$。梯度既不爆炸也不消失。

**经验证实。** 如果没有归一化（$\|\hat{k}_t\|, \|\hat{\alpha}_t\| \ne 1$），Jacobian 谱范数随着 $d_h / \lambda_0$ 增长：在 $d_h=96$ 时，$\|J_t\|_2 \approx 20.5$，在 25 步后产生约 $10^{32}$ 的梯度放大。这解释了在 $d_h \ge 96$ 时运行未归一化 VLA 时观察到的 NaN 损失。表 2 报告了两种公式的经验测量谱范数。

**表 2：** 未归一化（VLA v2）和归一化（VLA v3）公式下的 Jacobian 谱范数 $\|J_t\|_2$。经过 $T=25$ 步后的梯度放大为 $\|J_t\|_2^{25}$。

### 4.3 联想记忆容量
在有界输入和精确算术下，当键线性无关时，RLS 目标 (7) 可以精确恢复存储的关联：

###### 命题 3（精确恢复）
设 $\{\hat{k}_i\}_{i=1}^n$ 线性无关且 $n \le d_h$。在处理所有 $n$ 对之后，对于所有 $i \le n$，$S_n \hat{k}_i = v_i$。因此每头容量为 $d_h$ 个关联。

VLA 相对于线性注意力或 DeltaNet 并未增加此界限，三者都有 $d_h \times d_h$ 状态矩阵。VLA 的优势在于它在顺序覆盖下*保持*这种容量：因为 $A_t$ 将新写入路由到与最近使用的子空间正交的方向，旧关联不会像在累加下那样被均匀稀释。这解释了图 2(c) 中的经验结果：当 $n_{\text{pairs}} > d_h$ 时（见表 4，§6.6），在 $n=24$ 时，VLA 保留 1.000，而 DeltaNet 降至 0.010，softmax/线性注意力 plateau 在 0.07–0.08 附近。图 5 显示了评估准确率作为 $n_{\text{pairs}}$ 的函数。在 $n=4$ 时，所有模型都很接近：VLA 1.000，DeltaNet 0.97，softmax 0.26，linear 0.27。随着 $n$ 增加，DeltaNet 急剧下降，在 $n=8$ 时为 0.73，在 $n=12$ 时为 0.10，在 $n=24$ 时为 0.010，而 VLA 保持在 1.000 不变。线性注意力和 softmax 都 plateau 在 0.08–0.15 附近，低于随机基线 $1/128 \approx 0.008$ 仅因为交叉熵头利用了嵌入几何。

有两个限定条件。首先，所有测试值满足 $n_{\text{pairs}} \le 24$ 且 $d_h > d_h$ 的情况在 §6.6 中报告。其次，DeltaNet 崩溃是实质性发现：尽管有残差误差更新，当最近的写入主导相同方向时，标量门无法保护旧关联。VLA 的 $d_h \times d_h$ 矩阵门将新写入正交路由，直到容量边界保留所有先前的关联。

### 6.5 长上下文 MQAR

**图 6 标题：** 带有 $n=8$ 对和递增 $T \in \{64, 128, 256, 512\}$ 的 MQAR 评估准确率。VLA 在所有序列长度下保持 1.000。线性注意力和 softmax 注意力 plateau 在 $\approx 0.14$–0.15；DeltaNet 在所有测试长度下崩溃至 $\approx 0.01$。

图 6 保持 $n=8$。$n > d_h$ 的容量曲线超出理论边界 $d_h=32$ 至 $n \in \{32, 48, 64, 96\}$。每个工作点的准确率报告在表 4 的“Overload”列中。所有模型在 $n=d_h$ 之后都会退化；问题是*如何*退化。VLA 比线性注意力退化得更缓慢，因为 $A_t$ 优先覆盖最近写入的方向，通过过载机制部分保留旧关联。

#### 组件消融：
我们移除单独的 VLA 组件并报告 $n=16$ 时的 MQAR 准确率。移除键归一化（$\hat{k}_t \leftarrow k_t$）在 $d \ge 96$ 时产生 NaN 损失（Jacobian 爆炸，参见表 2）。固定 $A_t = 10I$（无 SM 更新）将准确率降低到匹配标准线性注意力，证实改进的来源是惩罚几何结构，而不仅仅是残差误差。完整结果见表 4。