The Key to Going Linear: Analysis-Driven Transformer Linearization
For practitioners deploying long-context transformers, this work provides a principled, analysis-driven linearization method that matches or exceeds prior post hoc baselines without retraining, enabling efficient inference.
This paper identifies key factors enabling high-quality linearization of causal self-attention for long-context transformer inference, showing that delta-style state updates outperform gated accumulation due to rank-1 orthogonal projections. The proposed structural interventions (sink tokens, short convolutions, fixed-budget cache routing) reduce approximation errors, achieving SOTA results on MMLU and matching complex adaptive-caching methods in long-context retrieval across LLaMA and Qwen models up to 32B parameters.
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.