提出新型线性注意力机制,解决长序列记忆干扰问题。
Variational Linear Attention: Stable Associative Memory for Long-Context Transformers

- 将记忆更新建模为自适应正则化最小二乘问题,稳定存储状态
- 在1000长度下记忆范数降低109倍,准确率达近完美
- 适合高容量长序列任务,性能超越现有线性注意力方法
线性注意力将软注意力的二次复杂度降至 $/mathcal{O}(T)$,但其记忆状态的Frobenius范数随 $T$ 线性增长,导致存储关联的渐进干扰。本文提出变分线性注意力(VLA),将记忆更新重构为带有自适应惩罚矩阵的在线正则化最小二乘问题,并通过Sherman-Morrison公式维护该矩阵。理论证明:归一化写入方向至单位长度可使递推雅可比谱范数恒为1(命题2),且在有界输入下状态范数自我受限(命题1)。实验表明,当 $T=1{,}000$ 时,VLA相较标准线性注意力将 $\|S_t\|_F$ 降低 $109\times$;在有效每头记忆范围($n_ ext{pairs} < d_h$)内实现近乎完美的精确匹配召回准确率;在内存负载增加时,仍显著优于DeltaNet与标准线性注意力,在每头容量边界处保持62%准确率。基于Triton的融合内核实现相较于串行Python代码14倍加速,$/mathcal{O}(T)$ 扩展性使其在约43,000个词元时低于软注意力延迟。
原文摘要 · Abstract (English)
Linear attention reduces the quadratic cost of softmax attention to $\mathcal{O}(T)$, but its memory state grows as $\mathcal{O}(T)$ in Frobenius norm, causing progressive interference between stored associations. We introduce \textbf{Variational Linear Attention} (VLA), which reframes the memory update as an online regularised least-squares problem with an adaptive penalty matrix maintained via the Sherman-Morrison rank-1 formula. We prove that normalising the write direction to unit length gives the recurrence Jacobian spectral norm exactly $1$ for all sequence lengths and head dimensions (Proposition 2), and that the state norm is self-limiting under bounded inputs (Proposition 1). Empirically, VLA reduces $\|S_t\|_F$ by $109\times$ relative to standard linear attention at $T{=}1{,}000$, achieves near-perfect exact-match accuracy on multi-query associative recall within the effective per-head memory regime ($n_\text{pairs} < d_h$), maintaining substantially higher retrieval performance than DeltaNet and standard linear attention under increasing memory load, and maintains 62\% accuracy at the per-head capacity boundary. A Triton-fused kernel achieves $14\times$ speedup over sequential Python and $\mathcal{O}(T)$ scaling, crossing below softmax attention latency at approximately 43\,000 tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。