The Key to Going Linear: Analysis-Driven Transformer Linearization

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.

• The paper demonstrates that delta-style updates, particularly via GDN, effectively mimic softmax attention’s key-dependent geometric properties, enabling linear approximations.
• The paper presents a rigorous evaluation in a strict frozen-backbone regime, showing that targeted structural interventions reduce approximation errors in long-context inference.
• The paper bridges theoretical analyses with empirical results by examining key concentration and projection residuals, leading to state-of-the-art performance in linearized transformer models.

Analysis-Driven Transformer Linearization: A Detailed Technical Summary

Introduction

The paper “The Key to Going Linear: Analysis-Driven Transformer Linearization” (2607.07706) addresses the challenge of rationalizing the inference costs of autoregressive transformers over long contexts. While transformers’ causal self-attention delivers state-of-the-art results, its $O(n^2)$ complexity with respect to sequence length presents a scaling bottleneck. Prior literature has explored post hoc linearization pipelines to enable efficient inference; however, the main technical question—what structural aspects of attention-based models need to be preserved to maximize performance post-linearization—remains underexplored in isolation, largely due to confounds from joint interventions (LoRA, distillation, hybrid attention, etc.).

This work adopts a rigorous, isolationist methodology: it systematically analyzes and compares linear attention variants under a “strict frozen-backbone” regime, training only newly introduced parameters and never fine-tuning the pretrained Q, K, V projections. It further presents a first-order theoretical approximation, establishing why delta-style updates such as Gated Delta Networks (GDN) mimic key-dependent geometrical aspects of softmax attention more effectively than pure gating or feature-map approaches. Finally, it introduces targeted structural interventions (sliding window attention, sink tokens, short convolutions) to close the remaining performance gap and provides extensive empirical evidence demonstrating state-of-the-art post hoc linearization across multiple LLM backbones.

Theoretical Analysis: Softmax Attention Linearization

A central contribution is the construction of a formal link between softmax attention and linearized state-update rules. The paper proves that softmax attention implements key-dependent, rank-1 orthogonal projection operations over value vectors, a property that is well-approximated by GDN-style delta updates, but poorly matched by key-independent gating mechanisms such as Gated Linear Attention (GLA). This insight is crucial: linear attention updates that lack these “delta” corrections cannot preserve the spatial geometry induced by exact softmax, even in the absence of normalization.

The theoretical apparatus is presented as a sequence of lemmas. First, a first-order Taylor expansion of the softmax weights reveals a bias plus a term dependent on the deviation of each key from the running mean, motivating centering or projection operations. The analysis then rigorously bounds the approximation error incurred by substituting mean subtraction with a rank-1 projection, showing that the approximation is tight when keys are concentrated in direction. Further, the delta-rule recurrence in GDN models accumulates a product of such projections, mimicking the dynamics of softmax attention on well-behaved key distributions.

Figure 1: Attention replacement blocks in the strict frozen-backbone regime—only new mechanism-specific parameters are trained, all original Q, K, V, and RMSNorm parameters are kept frozen.

Empirical Evaluation and Ablations

Strict Frozen-Backbone Evaluation

To fully disentangle the effects of various design choices, all linearized models are evaluated with completely frozen pretrained backbones—no modification of Q, K, V matrices, no feature adaptation, and no distillation. In this strict setting, GDN consistently achieves the closest approximation to full attention both in terms of normalized mean-squared error and zero/few-shot downstream task accuracy.

Notably, GDN surpasses other linear mechanisms by a substantial margin, particularly in the approximation of attention activations across layers and tokens. The empirical evidence confirms the theoretical prediction: highest approximation error occurs at sequence prefixes (small $t$), where the omitted softmax bias is significant and key concentration is low. This error profile, and the difficulty in matching sharp prefix attention (“sink tokens” effect), is robustly observed across various LLM backbones.

Figure 2: Per layer normalized MSE for Llama3.1-8b linearization, illustrating the superior performance of GDN.

The Role of Key Concentration and Projection Residuals

Key metrics—projection residual and key concentration—are analyzed to empirically validate the theoretical assumptions. High residuals correlate tightly with spikes in approximation error, primarily at early sequence indices where both key alignment and mean magnitude are low. These signals are later exploited in hybrid settings for selecting which layers to keep as full attention.

Structural Interventions

The paper systematically introduces and quantifies the effects of progressively relaxed assumptions:

• Sliding Window Attention (SWA): Incorporating a fixed-size local window for exact attention restores fine-grained local context and dramatically decreases the approximation gap, especially in context-based tasks.
• Sink Tokens: Explicit allocation of attention budget for hard-to-approximate prefix tokens further closes the gap, improving performance on both reasoning and context-retrieval benchmarks.
• Short Convolutions and LoRA: Selective Q/K/V adaptation using short convolutions or low-rank adapters increases representational flexibility, nearly saturating performance in most benchmarks but with a measurable risk of catastrophic forgetting or overfitting on knowledge-intensive tasks.

Figure 3: Comparative latency characteristics highlight linear attention’s ability to scale gracefully in context size.

Final Results and Scaling

With the full recipe (delta-style updates, SWA, sink tokens, and optional short convolution adaptation), the GDN-based linearization method achieves performance at parity with full attention models on common reasoning tasks and outperforms all prior post hoc baselines on challenging benchmarks, such as 5-shot MMLU. It matches or exceeds long-context retrieval performance of recent adaptive-caching frameworks, all while requiring fewer additional parameters and dramatically less fine-tuning data.

Figure 4: Per-layer normalized MSE for models with full linearization recipes, showing robust and uniform approximation.

Crucially, the methodology is validated across Qwen and LLaMA backbones up to 32B parameters. Linearization preserves downstream scaling trends, evidencing high transferability of the approach to larger and more diverse foundation models.

Hybrid Model Insights

The analysis extends to hybrid attention models where only a subset of layers retain full attention. Heuristic layer selection, based on key concentration/hardness measures, is shown to improve context-based metrics (e.g., Lambada) compared to naïve uniform allocation, confirming the predictive utility of the proposed theoretical diagnostics.

Implications and Future Directions

Practical Implications

The prescribed recipe enables efficient, linear-time inference for pretrained transformers without the overhead of full retraining or complex knowledge distillation. This provides direct benefit for cost-sensitive deployment, especially in memory-constrained long-context use cases. The explicit separation of architectural and adaptation effects presents a principled template for future post hoc model surgery in industrial settings.

Theoretical Significance

The first-order, rank-1 analysis bridges the gap between heuristic empirical recipes and a rigorous understanding of the essential structures that linearized attention mechanisms must preserve. This theoretical clarity demarcates the attainable performance envelope for linearized architectures and provides precise failure modes (e.g., low key concentration, prefix errors) that future architectures can target.

Open Directions

Future research could explore adaptive caching/retention strategies explicitly tailored to the geometry of the GDN family, as well as dynamic layerwise allocation (e.g., for hybrid models) based on real-time key concentration signals. The extension to instruction-tuned or specialized backbone variants (e.g., for retrieval or code) remains a promising avenue. Finally, a thorough investigation of spatially structured convolutional adaptations versus low-rank approaches in very large LLMs would help optimize performance-cost tradeoffs.

Conclusion

This paper brings a rigorous, analysis-driven approach to the post hoc linearization of large pretrained transformers. Through a sequence of theoretical, empirical, and architectural contributions, it establishes that delta-style updates most faithfully mimic the geometrical corrections implicit in softmax attention, especially under strict adaptation constraints. When combined with minimal structural compensations, the method matches (and in some cases surpasses) the performance of much more complex or compute-intensive linearization frameworks. The implications are both practically and theoretically substantial, offering a new standard for efficiency-oriented LLM model surgery and linearization.