论文
变分线性注意力:长上下文的稳定联想记忆 Transformer
Variational Linear Attention: Stable Associative Memory for Long-Context Transformers
摘要
线性注意力将softmax注意力的二次成本降低到$\mathcal{O}(T)$,但其记忆状态随着Frobenius范数的$\mathcal{O}(T)$增长,导致存储的关联之间逐渐干扰。我们引入了 \textbf{变分线性注意}(VLA),它将内存更新重新构建为在线正则化最小二乘问题,并通过 Sherman-Morrison 秩 1 公式维护自适应惩罚矩阵。我们证明,将写入方向归一化为单位长度,对于所有序列长度和头尺寸,递归雅可比谱范数恰好为 $1$(命题 2),并且状态范数在有界输入下是自我限制的(命题 1)。根据经验,VLA 相对于 $T{=}1{,}000$ 处的标准线性注意力减少了 $\|S_t\|_F$ $109\times$,在有效的每头内存机制 ($n_\text{pairs} < d_h$) 内实现了多查询关联召回的近乎完美的精确匹配精度,在内存负载增加的情况下保持比 DeltaNet 和标准线性注意力更高的检索性能,并保持每人容量边界的准确度为 62%。 Triton 融合内核比顺序 Python 和 $\mathcal{O}(T)$ 缩放实现了 $14\times$ 加速,在大约 43\,000 个 token 时低于 softmax 注意力延迟。