cs.LGJul 19, 2026

Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

Authors: Ayoub GhrissSourav Chakraborty

Organizations: Department of Computer Science, University of Colorado Boulder

Abstract

Linear attention promises constant-time recurrent inference but degrades sharply on associative recall. We formulate attention recall as a spherical-packing problem and introduce Kernelized Linear Attention Activations (KATA), a framework whose feature maps are derived from first principles by certifying nonnegative attention weights through a self-dual homogeneous cone. Building on this observation, we show that rank-one positive semi-definite (PSD) features offer a favorable capacity--interference tradeoff. KATA recovers a parameter-free convex output gate and characterizes associative capacity through the Welch interference floor. For tolerances above this floor, KATA enlarges the state without adding parameters and admits spherical codes with exponentially many keys in the projection dimension. We implement KATA as fused Triton kernels at two operating points: a flash-attention-style forward up to 1.6×{\sim}1.6\times FlashAttention-2 throughput, and an exact O(T)O(T) chunked-state form that reaches 11×{\sim}11\times FlashAttention-2 forward throughput at 131131k tokens. An associative scan of the first-order feature lowers the inter-chunk recurrence depth to O(log(T/C))O(\log(T/C)) for chunk size CC and averages 2.4×{\sim}2.4\times the throughput of a matched sequential linear-attention baseline. On long-range MQAR and repeated-key overwrite, several KATA variants outperform Gated DeltaNet, with parameter counts and state sizes reported alongside accuracy. Induction preserves near-perfect recall, while kernel benchmarks show that the maps can be implemented efficiently. KATA retains 0.9850.985 MQAR at a 16×16\times out-of-distribution length, approaching the softmax with roughly one quarter of the KV-cache entries. Experiments on 340M-parameter LLMs reveal a feature-dependent fluency trade-off and clarify how positional embeddings, delta rules, and decay gates interact with feature geometry.

Explore similar work

May 11, 2026cs.LG

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

Linear attention reduces the quadratic cost of softmax attention to O(T)\mathcal{O}(T), but its memory state grows as O(T)\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 11 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 StF\|S_t\|_F by 109×109\times relative to standard linear attention at T=1,000T{=}1{,}000, achieves near-perfect exact-match accuracy on multi-query associative recall within the effective per-head memory regime (npairs<dhn_\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×14\times speedup over sequential Python and O(T)\mathcal{O}(T) scaling, crossing below softmax attention latency at approximately 43,000 tokens.
Vishal Pandey, Gopal Singh
Jul 2, 2026cs.AI

A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets

Linear-attention and state-space language models compress the prefix into a fixed-size recurrent state, yielding O(1) memory at the cost of a lossy exact memory: when many key--value associations compete, earlier facts are overwritten and needle recall degrades. Inspired by Complementary Learning Systems, we give linear attention a hippocampal complement. HOLA (Hippocampal Linear Attention) keeps the usual delta-rule state as a compressive memory and adds a bounded exact KV cache, forming a semiparametric test-time memory: the state models linearly compressible structure, while the cache stores associations that should not be forced through that state. The cache writes without a learned eviction module, keeping tokens with large beta * ||e||, the prediction residual actually committed to the state; a decoupled RMSNorm-gamma cache read then turns these exact KV pairs into sharp retrieval rather than soft averaging. At 340M parameters trained on 15B SlimPajama tokens, HOLA lowers Wikitext perplexity from 27.32 to 22.92 (-16.1%), below a full-attention Transformer++ (26.88), and improves LAMBADA perplexity from 30.95 to 30.26. It also achieves the best linear in-context retrieval and remains much more robust than GDN or a matched HOLA+recency cache on RULER needle-in-a-haystack recall out to 32k tokens (16x its training length).
Wanyun Cui
May 9, 2026cs.LG

Kaczmarz Linear Attention

Long-context language modeling remains central to modern sequence modeling, but the quadratic cost of Transformer attention makes scaling computationally prohibitive. Linear recurrent models address this bottleneck by compressing the context into a fixed-size state, making the rule that forgets, writes, and edits information a central design problem. To address state maintenance, Gated DeltaNet (GDN) combines gated state decay with delta-rule residual writes, using a learnable coefficient to balance forgetting and update magnitude. However, this coefficient is learned empirically rather than derived from the underlying objective, which can lead to suboptimal update magnitudes. We revisit the online-regression objective underlying GDN and, inspired by the Kaczmarz projection method, derive the key-norm-normalized dynamic step size βt=ηt/(kt22+ε)β_t = η_t / (\|k_t\|_2^2 + ε) for residual updates. We propose Kaczmarz Linear Attention (KLA), a one-scalar modification of GDN that preserves the state shape, gates, linear recurrence, and chunkwise parallel algorithm. At the 0.4B scale with a 1B-token budget, KLA achieves the lowest validation perplexity among evaluated linear-time baselines, 8.09 versus 8.50 for GDN, and remains stable up to 65K tokens. On controlled tasks, KLA reaches 100% on single-needle-in-a-haystack retrieval, improves 8x multi-query associative recall by 7.03 points over GDN, and delivers 2.1x higher decode throughput at 32K context. These results suggest that the key-norm-normalized Kaczmarz coefficient is a first-order design axis for delta-rule sequence models: it improves accuracy, extrapolation, and decoding efficiency without changing the recurrent state or hardware kernel.
Jiaxuan Zou, Ruifeng Ren, Yong Liu