Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

AI summary

Overview

• Research area: Efficient transformer architectures and linear attention — specifically the theory of associative-recall capacity in fixed-size recurrent states, combined with GPU kernel design and language-model pretraining.
• Technical level: Advanced. The paper leans on convex geometry (symmetric cones, the Koecher–Vinberg classification), spherical-packing bounds, and signal-to-noise analysis, alongside Triton kernel engineering and large-scale pretraining.
• Scope in one sentence: The authors derive admissible feature geometries for linear attention from a self-dual homogeneous cone, analyze how many key addresses each geometry can separate, and validate the resulting KATA (Kernelized Linear Attention Activations) framework with fused Triton kernels and 340M-parameter language-model experiments.

What This Paper Is About

Linear attention replaces softmax's exponential kernel with a finite feature map, giving constant-time recurrent inference but degrading sharply on associative recall because many key–value bindings must share one fixed-size state. The authors reframe this "memory collision" as a spherical-packing problem and ask what feature maps are admissible, how much recall capacity a given feature dimension supports, and how to implement the resulting maps efficiently on modern accelerators. The goal is a principled account of why linear attention loses recall and which cone-based feature geometries recover it.

Key Contributions

1. A cone classification for nonnegative linear attention. Under assumptions of self-dual nonnegative geometry, isometric invariance, and homogeneity, the admissible feature geometries are organized by the Koecher–Vinberg classification into ordinary real factors: the positive orthant, the Lorentz cone, and the symmetric positive semidefinite (PSD) cone. The denominator of normalized linear attention emerges as a parameter-free, token-conditioned convex output gate.

2. A capacity theory for associative recall. Recall is cast as discrete spherical packing, yielding the Lorentz Rankin wall (Theorem 13), dimension-dependent packing bounds for PSD rank-one rays (Theorem 4), explicit finite dictionaries from mutually unbiased bases (Theorem 5) and DeVore frames (Theorem 6), and an optimal-packing pSNR ladder whose argument also covers temperature-scaled softmax attention.

3. Hardware-aligned kernels. Fused Triton kernels are released for the gate-free, delta-rule-free KATA-Mg recurrence at two operating points: a quadratic O(T²) forward reaching up to ~1.6× FlashAttention-2 throughput, and an exact O(T) chunked-state form reaching ~11× FlashAttention-2 forward throughput at 131k tokens. An associative scan of the first-order feature lowers inter-chunk recurrence depth to O(log(T/C)) for chunk size C.

4. Experimental evidence across scales. The Zoology playground is used to evaluate multi-query associative recall (MQAR), repeated-key overwrite, and induction under matched training configurations, with parameter counts and state sizes reported alongside accuracy. Matched 340M-parameter language models are pretrained for 15B tokens and evaluated on standard benchmarks and in-context recall.

Main Findings

• Memory collision is a spherical-packing problem. The paper formalizes interference as the minimum achievable maximum pairwise inner product among T unit-norm feature vectors in a cone, and defines ε-capacity as the largest T whose interference stays at or below ε.

• The orthant inherits the Welch floor. For any T > d unit vectors in R^d, μ² ≥ (T−d)/(d(T−1)); consequently μ ≤ ε < 1/√d forces T ≤ d(1−ε²)/(1−dε²). Real equiangular tight frames attain equality only for special (d, T) pairs and satisfy T ≤ d(d+1)/2.

• The Lorentz cone hits a Rankin wall. For 0 ≤ ε < 1/2, the ε-capacity of the Lorentz cone is bounded by min{d, 1 + 1/(1−2ε)}, so increasing feature width alone cannot buy unbounded capacity at a fixed tolerance.

• Rank-one PSD features provide exponential capacity above the Welch floor. With p ≥ 2, n = p(p+1)/2, and 0 ≤ ε < 1, the greedy spherical-cap construction gives C_ε(S_+^p) ≥ (1/2)(1−ε)^{−(p−1)/2}. Below the Welch floor (ε < 1/p) capacity stays polynomial, T ≤ p(1−ε)/(1−pε); for any fixed ε in (0,1), the lower bound is exponential in p once p > 1/ε.

• Squaring the inner product changes the geometry. PSD interference at most ε is equivalent to |⟨u_i, u_j⟩| ≤ √ε, replacing the mutually obtuse constraint of the Lorentz lift with an unsigned near-orthogonal spherical code. Higher even-order features u^{⊗2r} are noted, with the first nontrivial order r = 1 used because the quadratic expanded state already dominates the hardware budget.

• The readout gate is free. Under the nonnegativity assumption, the normalized readout is a convex interpolation between the previous readout and the current value, with weights α_t = D_{t−1}/D_t and β_t = c_t/D_t satisfying α_t + β_t = 1 — no additional learned parameters, and the state update remains commutative and additive, unlike Gated DeltaNet's rank-one noncommutative transition.

• pSNR separates linear from PSD retrieval. At the idealized Welch-scale margin ε* = (T−p)/(p(T−1)), linear attention gives pSNR = p/(T−p) and is ideally retrievable only for T < 2p (capacity Θ(p)), while PSD gives pSNR = p²(T−1)/(T−p)² and extends the threshold to Θ(p²).

• Softmax sharpens the same margin, at a temperature cost. With μ* = √ε*, pSNR_Soft = exp(2(1−μ*)/τ)/(T−1), so softmax is ideally retrievable when T−1 < exp(2(1−μ*)/τ). For T ≫ p this gives T_Soft(p, τ) ≈ 1 + exp(2(1−1/√p)/τ), and an inverse temperature β_T = 1/τ = Θ(log T) is needed to keep a fixed positive margin — recovering the logarithmic scale identified for long-context softmax dynamics by Chen et al. (2025).

• Finite dictionaries are constructible. MUB realification gives T = s(s+1) vectors in R^d with d = 2s and μ ≤ 1/√s = √(2/d) for prime power s. DeVore's construction gives T = s^{r+1} vectors in R^d with d = s² and μ ≤ r/s for prime power s and integer 0 ≤ r < s. Table 1 shows for T = 10⁵ that Welch requires d ≥ 399 at μ ≤ 0.05 versus d ≥ 25 at μ ≤ 0.20, while random dictionaries need roughly d ≈ 18,421 and d ≈ 1,151 respectively, MUB needs d = 802 (s = 401) and d = 634 (s = 317), and DeVore needs d = 2,209 (s = 47) and d = 361 (s = 19).

• State size, not just feature width, dominates memory. The full rank-one PSD map has d² outer-product entries or n = d(d+1)/2 unique packed coordinates, and the feature–value state is n·d_v. At the standard head dimension d = 64 the direct and packed widths are already 4096 and 2080, placing the full PSD map at the edge of hardware-friendly state sizes. Two reduced families are introduced: KATA-Σg with state (d/g)(d/g+1)/2 · d_v, and KATA-Mg with state d(d/g+1)/2 · d_v; both recover the full map at g = 1.

• Chunk recurrence avoids the triangular solve. For Gated DeltaNet, the inter-chunk cost is O(C d_k² + d_k² d_v) because each chunk forms and inverts a C×C unitriangular matrix; for KATA the corresponding cost is O(C n_ψ d_v), and SRAM rather than solve conditioning sets the chunk-size limit.

• Kernel throughput is competitive to strongly favorable. The quadratic KATA-Mg forward reaches up to ~1.6× FlashAttention-2 throughput and remains competitive on the training step. In the linear-state form, KATA-M2 reaches parity with its quadratic counterpart near 16k tokens, is faster by 32k, and reaches ~11× FlashAttention-2's forward throughput at 131k; KATA-M1 is faster by 128k. The associative scan averages ~2.4× the throughput of a matched sequential linear-attention baseline.

• Recall benchmarks favor PSD variants. Several KATA variants outperform Gated DeltaNet on long-range MQAR and repeated-key overwrite under matched training configurations, with parameter counts and state sizes reported alongside accuracy; induction preserves near-perfect recall. KATA-M1 retains 0.985 MQAR accuracy at 16× the training context, approaching softmax with roughly one quarter of the KV-cache entries at that length.

• Pretrained language models show a feature-dependent trade-off. On matched 340M-parameter models trained for 15B tokens, the PSD variants remain broadly comparable on zero-shot accuracy while showing a meaningful perplexity spread. They retain much more high-entropy needle signal than Gated DeltaNet, whose UUID recall falls to 0.004.

Methodology in Plain English

The authors start by writing down what linear attention actually does: it swaps softmax's exponential kernel for an inner product of finite feature vectors, which turns token-by-token mixing into a recurrence on a fixed-size state. They then add three mathematical constraints to the feature map — that feature vectors live in a cone that is its own dual (so attention weights can never be negative), that rotating the input rotates the feature the same way, and that the cone has no privileged interior region. Under those constraints, classical convex-geometry theory says the possible geometries must be built from a short list: the positive orthant, the Lorentz cone, and the PSD cone. That gives a principled menu instead of an ad hoc set of activation choices.

To compare the menu items, they pose recall as a packing question: how many unit-length feature vectors can you place in a cone so that no two overlap too much? Standard bounds (Welch, Rankin) cap the orthant and Lorentz cases. The key observation is that if a PSD feature is the outer product of a vector with itself, the PSD inner product is the square of the underlying inner product — so a tolerance ε in PSD space only requires √ε separation among the raw vectors, which a greedy spherical-cap construction turns into exponentially many distinguishable keys once the projection dimension exceeds 1/ε.

They then ask whether a given score margin is actually enough to retrieve a value, under an isotropic-value noise model, and derive a power signal-to-noise ratio (pSNR) that separates linear from PSD readouts and predicts how sharply softmax would need to be temperature-scaled to match. Finally, they instantiate the PSD family in fused Triton kernels at several reduced state sizes (KATA-Σg and KATA-Mg), use an associative scan to parallelize the inter-chunk recurrence, and test the models on synthetic recall tasks and 340M-parameter language-model pretraining.

Why This Matters

Impact on research. The paper offers a first-principles explanation for a widely observed failure mode of linear attention, and it replaces an ad hoc zoo of feature-map and gating variants with a classified set of admissible geometries. It also derives a capacity law for temperature-scaled softmax from the same packing argument, which connects two previously separate lines of work on associative recall and long-context attention dynamics. The prediction that capacity above the Welch floor requires rank-one PSD-style features — rather than simply a wider orthant map — is a concrete, testable structural claim.

Potential applications (implications of the work, not claims tested in the paper):

• Long-context language modeling, where replacing a growing KV cache with a fixed-size PSD state could reduce memory requirements at long sequence lengths.
• Retrieval-style and associative-recall workloads, where the paper directly measures MQAR, repeated-key overwrite, and induction performance.
• Efficient inference and on-device deployment, where constant-time recurrent decoding and the reported throughput gains over FlashAttention-2 matter.
• Memory-constrained and high-throughput training pipelines, where the additive (rather than triangular-solve) inter-chunk recurrence simplifies kernel scheduling.

Industry relevance. The throughput numbers are measured against FlashAttention-2 (Dao, 2023), the de facto production attention kernel, and the reported two operating points — a fast quadratic forward and an exact chunked linear-state form — map onto the two regimes practitioners actually choose between. The paper situates its normalized-address assumptions with respect to RMSNorm and the query–key normalization used by Qwen3, Gemma 3, MiniMax-class models, DeltaNet, and Gated DeltaNet, so the capacity results are stated in terms those architectures already use.

Future Directions

• Complex and exceptional cones. The Koecher–Vinberg classification also contains complex Hermitian, quaternionic Hermitian, and exceptional factors. The paper restricts itself to the ordinary real families and explicitly leaves complex extensions to future work, noting that over C^d the ceiling is d² even though the representation uses 2d real scalars.

• Better finite spherical packings for attention heads. The authors state that their constructions illustrate available finite regimes and that better real spherical packings may exist, leaving the question of which packings are best suited to attention heads open.

• Higher-order features. The framework admits even-order features u^{⊗2r} with symmetric ambient dimension binom(p+2r−1, 2r), but only r = 1 is used because the quadratic expanded state already dominates the hardware budget. Whether higher orders become practical at different state budgets is untested.

• Closing the gap between geometric capacity and empirical recall. The paper notes that ε-capacity measures address-packing quality at a fixed interference tolerance, and that the update rule, noise model, and optimization determine how closely empirical recall realizes that geometric capacity. The feature-dependent perplexity spread and fluency trade-off observed in the 340M-parameter models are exactly the kind of behavior this gap would explain.

Target Audience

This paper is most useful to researchers working on efficient attention, linear-time sequence models, and long-context language modeling who are comfortable with convex geometry and kernel methods; to systems and kernel engineers designing fused GPU implementations of recurrent attention; and to readers who want a theoretical account of why fixed-size states lose associative recall and what class of feature maps restores it. Readers primarily interested in applied results without the geometric derivations will still find the kernel throughput and recall benchmark sections, but the capacity theory requires the mathematical background.

Authors’ 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 ${\sim}1.6\times$ FlashAttention-2 throughput, and an exact $O(T)$ chunked-state form that reaches ${\sim}11\times$ FlashAttention-2 forward throughput at $131$k tokens. An associative scan of the first-order feature lowers the inter-chunk recurrence depth to $O(\log(T/C))$ for chunk size $C$ and averages ${\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.985$ MQAR at a $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.