The Key to Going Linear: Analysis-Driven Transformer Linearization
2026-07-08 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how to make transformers work faster with very long input sequences by reducing the costly math involved in their attention mechanism. They focus on one part called the state update and find that a special type of projection used by softmax is important for keeping the model accurate. By adding new structures like sink tokens and short convolutions, they reduce errors caused by simpler approximations. Their method works well on big models like LLaMA and Qwen and performs better than previous techniques on certain tests.
transformerscausal self-attentionsoftmaxstate updatelinearizationLLaMAQwenlong-context inferenceMMLUadaptive caching
Authors
Anna Kuzina, Paul N. Whatmough, Babak Ehteshami Bejnordi
Abstract
The quadratic cost of causal self-attention severely bottlenecks long-context transformer inference. While numerous post hoc linearization pipelines exist, it is difficult to identify which components preserve model quality. This work isolates the effect of state update design in a strict frozen-backbone regime. We show that softmax relies on key-dependent, rank-1 orthogonal projections, elucidating why delta-style networks outperform purely gated accumulation. We identify a potential source of approximation errors and introduce structural interventions, specifically sink tokens, short convolutions, and fixed-budget cache routing, which reduces the remaining gap. We scale this linearization approach across LLaMA and Qwen models up to 32B parameters, outperforming prior post hoc baselines on MMLU and matching the long-context retrieval of complex adaptive-caching frameworks.