📊 statistics

Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

This paper introduces Kernelized Linear Attention (KATA), a novel framework that leverages symmetric cones and rank-one PSD features to resolve the capacity-interference tradeoff in linear attention, achieving superior associative recall and significantly higher throughput than FlashAttention-2 while maintaining near-perfect long-range performance with reduced KV-cache overhead.

Original authors: Ayoub Ghriss, Sourav Chakraborty

Published 2026-07-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ayoub Ghriss, Sourav Chakraborty

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a super-smart robot that can read a book and remember every single detail, from the name of a minor character to the exact color of a door mentioned three chapters ago. In the world of artificial intelligence, this is the job of a "Transformer," a type of model that powers many of the chatbots and tools we use today. The secret sauce that makes these robots so good at remembering is something called "attention." Think of attention like a spotlight: when the robot reads a new sentence, the spotlight shines on the most important words it has seen before to help it understand the current one.

However, there is a catch. The traditional spotlight is incredibly powerful but also incredibly heavy. As the story gets longer, the spotlight has to scan through every single previous word to find the right one. This is like trying to find a specific needle in a haystack by checking every single piece of hay one by one; it takes forever and requires a massive amount of storage space (memory) to keep all those needles handy. Scientists have been trying to build a "linear" spotlight that is faster and lighter, one that can remember things without needing to re-scan the whole book every time. But these faster versions often have a terrible memory: they forget the important details or get confused when too many things look similar. They are fast, but they aren't smart enough to handle complex stories.

This is where a new idea called Kernelized Linear Attention (KATA) comes in. The researchers behind this paper, Ayoub Ghriss and Sourav Chakraborty, decided to solve the memory problem by looking at it through the lens of geometry and packing. They realized that the reason fast models forget things is that they are trying to jam too many memories into a tiny, crowded box. To fix this, they invented a new way to organize memories using a mathematical shape called a "symmetric cone."

Think of a memory as a unique key. In the old, fast models, these keys were like flat, 2D shapes that could easily overlap and get mixed up. KATA, however, uses a special 3D shape (specifically, a "positive semi-definite cone") to turn those flat keys into something more robust. It's like taking a flat piece of paper and folding it into a complex origami crane. Even if two pieces of paper look similar when flat, their folded cranes might be totally different and easy to tell apart. By using this "folding" trick, KATA can pack exponentially more unique memories into the same amount of space without them crashing into each other.

The paper shows that this geometric trick works beautifully. They built a new type of attention mechanism that doesn't need to store a massive list of every word it has ever seen (which saves a ton of memory). Instead, it keeps a compact, organized summary. When tested on tasks that require remembering specific details from long texts—like finding a hidden word in a sea of distractions—KATA performed almost as well as the heavy, slow traditional models but with a fraction of the memory. In fact, in some tests, it could remember details from texts 16 times longer than it was trained on, something that other fast models usually fail at.

The researchers didn't just stop at theory; they built the actual computer code to run this on modern graphics cards. They found that their new method is incredibly fast. In some scenarios, it runs up to 11 times faster than the current standard for fast attention, while still keeping the memory accurate. They also discovered that while this new method is great at pure memory, it sometimes needs a little help to understand the flow of a story, suggesting that the best future models might combine this super-efficient memory with other tools to handle both facts and fluency.

In short, KATA is like giving the robot a super-organized filing cabinet where every file has a unique, 3D shape that prevents it from getting lost in the shuffle. It proves that you don't have to choose between a fast robot and a smart one; with the right geometric shape, you can have both.

Technical Summary: Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

1. Problem Statement

Linear attention mechanisms offer constant-time recurrent inference and fixed-size state representations, addressing the $O(T^2)$ training cost and $O(T)$ KV-cache memory of standard softmax attention. However, existing linear attention models suffer from a sharp degradation in associative recall as sequence length increases. This failure is attributed to "memory collision," where many key-value bindings must share a single fixed state, causing retrieval performance to collapse on long-context tasks.

Current gated linear RNNs (e.g., DeltaNet, Mamba2, Gated DeltaNet) attempt to mitigate this via learned forget gates and content-addressed erase operations. Yet, the theoretical limits of their feature geometries remain opaque. Specifically, there is no principled understanding of:

What limits recall capacity at a fixed state size?
Which finite feature maps preserve the isotropy of softmax?
Must feature maps be nonnegative, and if so, what geometries are admissible?

2. Methodology

2.1 Theoretical Framework: Symmetric Cones and KATA

The authors formulate attention recall as a spherical-packing problem. They introduce Kernelized Linear Attention Activations (KATA), a framework derived from first principles by certifying nonnegative attention weights through self-dual homogeneous cones.

Geometric Assumptions: The framework assumes (1) nonnegative attention weights via a self-dual cone ( $K = K^*$ ), (2) isometric invariance of the feature map, and (3) homogeneity of the latent geometry.
Koecher–Vinberg Classification: Under these assumptions, admissible feature geometries are restricted to the irreducible factors of the Koecher–Vinberg theorem: the positive orthant, the Lorentz cone, and the symmetric positive semi-definite (PSD) cone.
Output Gating: The authors prove that the denominator of normalized linear attention naturally emerges as a parameter-free, token-conditioned convex output gate. This gate performs a convex interpolation between the previous readout and the current value, eliminating the need for explicit learned gating mechanisms in the base formulation.

2.2 Capacity Theory

The paper analyzes the "associative capacity" (the number of distinguishable key addresses) for each cone geometry:

Orthant & Lorentz: These geometries are subject to the Welch bound and the Rankin wall, respectively. They offer linear capacity ( $O(d)$ ) or are strictly limited by interference thresholds (e.g., Lorentz interference cannot drop below $1/2$ without forcing mutually obtuse keys).
PSD Cone (Rank-One): The authors propose using rank-one positive semi-definite (PSD) features, $\psi(u) = uu^\top$ $ψ (u) = u u^{⊤}$ . This map squares the raw inner product ( $\langle u_i, u_j \rangle^2$ $⟨ u_{i}, u_{j} ⟩^{2}$ ), converting the interference geometry.
- Theoretical Result: While the raw vectors are subject to the Welch floor, the squared interference allows for an exponential capacity in the projection dimension $p$ for any fixed tolerance $\epsilon > 0$ . Specifically, the capacity scales as $T \geq \frac{1}{2}(1-\epsilon)^{-(p-1)/2}$ .
- Signal-to-Noise: The authors derive a power signal-to-noise ratio (pSNR) analysis showing that the PSD map extends the idealized retrieval threshold from $\Theta(p)$ (linear) to $\Theta(p^2)$ , effectively enlarging the state without adding parameters.

2.3 Hardware-Aligned Implementation

To realize these theoretical gains, the authors implement KATA using fused Triton kernels on NVIDIA GPUs.

Chunked Parallelism: The sequence is split into chunks. Intra-chunk attention uses tensor cores, while inter-chunk state propagation is handled via an associative scan (parallel prefix sum).
Complexity Reduction: This reduces the inter-chunk recurrence depth from $O(N_C)$ (sequential) to $O(\log N_C)$ , where $N_C$ is the number of chunks.
Two Operating Points:
1. Quadratic ( $O(T^2)$ ) Forward: A FlashAttention-style kernel that avoids materializing the full state, achieving up to 1.6 $\times$ the throughput of FlashAttention-2.
2. Linear ( $O(T)$ ) State: An exact chunked-state form that reaches $\sim$ 11 $\times$ FlashAttention-2 throughput at 131k tokens.
Variants: The framework includes KATA-Mg (multiple PSD factors) and KATA- $\Sigma$ g (summed PSD factors) to balance state size and packing efficiency.

3. Key Contributions

Cone Classification for Linear Attention: The paper characterizes the admissible nonnegative feature geometries (Orthant, Lorentz, PSD) under the constraints of nonnegative weights and isometric invariance, grounding linear attention in the Koecher–Vinberg classification.
Associative Recall Capacity Theory: By casting recall as discrete spherical packing, the authors prove the Lorentz Rankin wall and dimension-dependent packing bounds for PSD rank-one rays. They demonstrate that rank-one PSD features offer an exponential capacity advantage over linear features below the Welch floor.
Hardware-Optimized Kernels: The authors release fused Triton kernels for the KATA recurrence. These kernels utilize an associative scan to lower inter-chunk depth, achieving significant throughput gains over sequential linear attention baselines and FlashAttention-2 at long contexts.
Empirical Validation: Experiments on synthetic tasks (MQAR, repeated-key overwrite, induction) and 340M-parameter language models show that KATA variants (particularly those on the PSD cone) outperform Gated DeltaNet and other linear baselines in associative recall, retaining near-perfect recall at 16 $\times$ out-of-distribution length extrapolation.

4. Results

Synthetic Benchmarks:
- MQAR (Multi-Query Associative Recall): KATA-M1 retains 0.985 accuracy at 16 $\times$ the training context length, approaching softmax performance while using roughly one-quarter of the KV-cache entries.
- Repeated-Key Overwrite: Variants with the delta rule (DeltaKATA) achieve near-perfect overwrite accuracy, correcting stale bindings at repeated addresses.
- Needle-in-a-Haystack (NIAH): On high-entropy tasks (UUIDs), KATA-M2 and DeltaKATA-M1 significantly outperform Gated DeltaNet, which degrades to near-zero recall (0.004) on UUIDs.
Language Model Pretraining:
- Pretrained 340M-parameter models show that PSD variants maintain competitive zero-shot accuracy on standard benchmarks (ARC, PIQA, etc.) while exhibiting a "fluency trade-off" (higher perplexity than softmax but better recall than other linear models).
- Induction: KATA preserves near-perfect induction capabilities, a task where linear attention often fails.
Performance:
- The linear-state KATA-M2 kernel achieves $\sim$ 11 $\times$ the forward throughput of FlashAttention-2 at 131k tokens.
- The quadratic KATA-M1 forward reaches $\sim$ 1.6 $\times$ FlashAttention-2 throughput.
- Associative scan implementations average $\sim$ 2.4 $\times$ the throughput of matched sequential linear-attention chunks.

5. Significance and Claims

The paper claims to resolve the "capacity wall" of linear attention by identifying the correct geometric substrate (symmetric cones, specifically the PSD cone) for nonnegative attention.

Theoretical Insight: It establishes that the "price" of constant state in linear attention is not inherent to linearity itself, but to the choice of feature geometry. By moving to rank-one PSD features, the capacity can be exponentially increased without increasing the number of parameters, only the state dimension.
Practical Impact: The work demonstrates that high-capacity associative recall is achievable with linear-time inference and fixed-size states, provided the feature map is aligned with hardware (via fused kernels) and the geometry is theoretically sound (via symmetric cones).
Limitations: The authors note that while the geometric capacity is high, learnability is a factor; models may struggle to learn the full high-dimensional spherical packing (e.g., KATA-M1 has a larger theoretical capacity but sometimes underperforms KATA-M2 in contextual tasks due to training dynamics). Additionally, the current implementation focuses on real-valued cones, leaving complex extensions for future work.

In summary, KATA provides a principled bridge between the theoretical limits of associative memory and the practical constraints of hardware-efficient inference, offering a path to linear-time transformers that do not sacrifice long-range recall.