Variational Linear Attention: Stable Associative Memory for Long-Context Transformers
- URL: http://arxiv.org/abs/2605.11196v1
- Date: Mon, 11 May 2026 20:03:12 GMT
- Title: Variational Linear Attention: Stable Associative Memory for Long-Context Transformers
- Abstract summary: Linear attention reduces the quadratic cost of softmax attention to $mathcalO(T)$, but its memory state grows as $mathcalO(T)$ in Frobenius norm.<n>We introduce textbfVariational Linear Attention (VLA), which reframes the memory update as an online regularised least-squares problem.
- Score: 0.6117371161379209
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: 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.
Related papers
- RunningTensor: Generalizing Linear Attention to Higher-Order Recurrent States [2.2153783542347805]
RunningTensor generalizes recurrent memory to an order-$o$ tensor.<n>It retains both recurrent and parallel forms while remaining linear in sequence length $T$.<n>RunningTensor outperforms linear-attention and SSM baselines on associative recall.
arXiv Detail & Related papers (2026-09-11T13:07:29Z) - Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones [4.847095054351615]
Linear attention promises constant-time recurrent inference but degrades sharply on associative recall.<n>We introduce Kernelized Linear Attention Activations (KATA), a framework whose feature maps are derived from first principles.<n>We show that rank-one positive semi-definite features offer a favorable capacity--interference tradeoff.
arXiv Detail & Related papers (2026-07-19T21:46:09Z) - Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels [0.0]
We present a reformulation of scaled dot-product attention and its numerically stable softmax.<n>The DNF achieves $O(n_dk + n_dv)$ data movement versus $O(n2 + n_dk + n_dv)$ for the standard implementation.<n>Unlike hardware-specific accelerators or empirical tiling schemes such as FlashAttention, MoA simultaneously provides array fusion, shape-transformation correctness, and predictive cost models.
arXiv Detail & Related papers (2026-06-05T14:44:49Z) - Don't Read Everything: A Curvature-Conditioned Query for Linear Attention [42.02790959450039]
Linear attention reduces the quadratic cost of softmax attention by maintaining a recurrent fast-weight state.<n>Existing remedies act on the write side of memory through gating, delta updates, or kernel feature maps.<n>We borrow one specific piece of softmax's geometry to construct a cheap read-time contraction of the query.<n>We call this mechanism Curvature-Conditioned Query (CCQ)<n>It improves perplexity, zero-shot downstream accuracy, S-NIAH retrieval at and beyond the training context, length-extrapolation perplexity from 4K to 20K, and Long
arXiv Detail & Related papers (2026-05-31T15:25:42Z) - Sharp Capacity Thresholds in Linear Associative Memory: From Winner-Take-All to Listwise Retrieval [59.859592671274704]
How many key-value associations can a $dtimes d$ linear memory store?<n>We show that the answer depends not only on the $d2$ degrees of freedom in the memory matrix, but also on the retrieval criterion.
arXiv Detail & Related papers (2026-05-06T17:53:20Z) - Mild Over-Parameterization Benefits Asymmetric Tensor PCA [12.923414933046574]
Asymmetric PCA (ATPCA) is prototypical model for studying the trade-offs between sample complexity, computation, and memory.<n>We focus on the setting where $overlinek geq 4$ is even and consider descent-based algorithms under a limited memory budget.
arXiv Detail & Related papers (2026-04-11T13:34:48Z) - Optimal Decay Spectra for Linear Recurrences [6.508284558109273]
PoST integrates into any diagonal linear recurrence without overhead.<n>We instantiate it across Mamba RWKV-7, Gated DeltaNet, Gated Linear Attention, and RetNet.
arXiv Detail & Related papers (2026-04-08T23:54:57Z) - Scaling Attention via Feature Sparsity [50.64995497733461]
Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n2 d)$ cost of self-attention.<n>We propose Sparse Feature Attention, where queries and keys are represented as $k$-sparse codes that preserve high-dimensional expressivity.<n>Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $2.5times$ and reducing FLOPs and KV-cache by nearly 50%.
arXiv Detail & Related papers (2026-03-17T08:41:50Z) - Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training [0.0]
Attention scores in transformers are bilinear forms $S_ij = x_itop M x_j / sqrtd_h$ whose maximum magnitude governs overflow risk in low-precision training.<n>We derive a emphrank-aware concentration inequality: when the interaction matrix $M = WQ WKtop$ has rank $r ll d$, tail probabilities for $max_i,j|S_ij|$ decay as $exp(-d22/(
arXiv Detail & Related papers (2026-02-21T14:29:22Z) - Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression [53.48692193399171]
Gated KalmaNet (GKA) is a layer that reduces the gap by accounting for the full past when predicting the next token.<n>We solve an online ridge regression problem at test time, with constant memory and linear compute cost in the sequence length.<n>On long-context, GKA excels at real-world RAG and LongQA tasks up to 128k tokens, achieving more than $10$% relative improvement over other fading memory baselines.
arXiv Detail & Related papers (2025-11-26T03:26:37Z) - Scaling Up Liquid-Resistance Liquid-Capacitance Networks for Efficient Sequence Modeling [50.994194925685434]
LrcSSM is a $textitnon-linear$ recurrent model that processes long sequences as fast as today's linear state-space layers.<n>By forcing its Jacobian matrix to be diagonal, the full sequence can be solved in parallel.<n>LrcSSM offers a formal gradient-stability guarantee that other input-varying systems such as Liquid-S4 do not provide.
arXiv Detail & Related papers (2025-05-27T20:02:59Z) - Gated Linear Attention Transformers with Hardware-Efficient Training [60.670102007737476]
This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability.
We then generalize this algorithm to a more expressive variant of linear attention with data-dependent gates.
When used as a replacement for the standard attention layer in Transformers, the resulting gated linear attention Transformer is found to perform competitively.
arXiv Detail & Related papers (2023-12-11T18:51:59Z) - TCNCA: Temporal Convolution Network with Chunked Attention for Scalable
Sequence Processing [52.64837396100988]
MEGA is a recent transformer-based architecture, which utilizes a linear recurrent operator whose parallel computation, based on the FFT, scales as $O(LlogL)$, with $L$ being the sequence length.
We build upon their approach by replacing the linear recurrence with a special temporal convolutional network which permits larger receptive field size with shallower networks, and reduces the computational complexity to $O(L)$.
We evaluate TCNCA on EnWik8 language modeling, long-range-arena (LRA) sequence classification, as well as a synthetic reasoning benchmark associative recall.
arXiv Detail & Related papers (2023-12-09T16:12:25Z) - HyperAttention: Long-context Attention in Near-Linear Time [78.33061530066185]
We present an approximate attention mechanism named HyperAttention to address the computational challenges posed by the growing complexity of long contexts.
Empirically, employing Locality Sensitive Hashing (LSH) to identify large entries, HyperAttention outperforms existing methods.
We validate the empirical performance of HyperAttention on a variety of different long-context length datasets.
arXiv Detail & Related papers (2023-10-09T17:05:25Z) - Sketchy: Memory-efficient Adaptive Regularization with Frequent
Directions [22.09320263962004]
We find the spectra of the Kronecker-factored gradient covariance matrix in deep learning (DL) training tasks are concentrated on a small leading eigenspace.
We describe a generic method for reducing memory and compute requirements of maintaining a matrix preconditioner.
We show extensions of our work to Shampoo, resulting in a method competitive in quality with Shampoo and Adam, yet requiring only sub-linear memory for tracking second moments.
arXiv Detail & Related papers (2023-02-07T21:50:06Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.