Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

2026-07-19

AI Summary

This work addresses the challenge that linear attention mechanisms face in balancing efficient inference with high-capacity memory in associative recall tasks. The authors formulate attention-based recall as a spherical coding problem and propose the Kernelized Activated Linear Attention (KATA) framework, which derives—based on geometric principles—the first-order positive semidefinite feature map. By incorporating self-dual homogeneous cone constraints, KATA enforces non-negative attention weights and implements a parameter-free convex output gate. The approach reveals that the Welch bound governs the fundamental limit of associative capacity and enables state expansion without increasing model parameters, achieving exponential growth in key storage capacity. Experiments demonstrate that KATA attains 11× higher forward throughput than FlashAttention-2 at 131k tokens, reaches 0.985 accuracy on MQAR tasks, extrapolates sequence length by 16×, and outperforms Gated DeltaNet in long-range association and repeated-key coverage benchmarks.

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.

Problem

• linear attention
• associative recall
• capacity-interference tradeoff
• spherical packing
• attention capacity

Innovation

• Kernelized Linear Attention
• Self-dual Homogeneous Cone
• Positive Semi-definite Features
• Associative Recall
• Spherical Packing