Variational Linear Attention: Stable Associative Memory for Long-Context Transformers

arXiv:2605.11196v1 Announce Type: new Abstract: Linear attention reduces the quadratic cost of softmax attention to O(T), but its memory state grows as O(T) in Frobenius norm, causing progressive interference between stored associations. We introduce extbf{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 109imes 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_ext{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 14imes speedup over sequential Python and O(T) scaling, crossing below softmax attention latency at approximately 43,000 tokens.

Source: arXiv cs.LG | 2026-05-13