AI Research Portal

This paper addresses the quadratic cost of causal self-attention in transformers, which limits long-context inference. The authors propose an analysis-driven approach to linearize attention mechanisms by studying state updates in a frozen-backbone regime. They demonstrate that delta-style updates better approximate softmax attention than purely gated accumulation. By introducing structural interventions like sink tokens, short convolutions, and projection adaptation, they significantly reduce the performance gap. Their method scales to large models (up to 32B parameters) and achieves competitive results on benchmarks like MMLU and long-context retrieval tasks.

The quadratic computational cost of causal self-attention is a significant bottleneck for long-context transformer inference. Existing post-hoc linearization methods often combine multiple interventions, obscuring the effectiveness of individual components and the fundamental question of which linear state update best approximates softmax attention. This makes it difficult to identify the most effective strategies for achieving linear-time inference while preserving model quality.

The paper studies post-hoc linearization of pretrained language models in a strict frozen-backbone setting, empirically comparing kernelized, gated, and delta-based linear attention mechanisms.
It provides a first-order approximation demonstrating that delta-based mechanisms faithfully capture the key-dependent, rank-1 dynamics of softmax attention, unlike pure gated updates.
The work empirically shows how practical interventions, including sliding window attention (SWA), sink tokens, and limited Q/K/V adaptation (e.g., short convolutions), can compensate for approximation gaps.
The proposed approach outperforms prior post-hoc baselines on MMLU, matches the long-context retrieval of complex adaptive token-caching frameworks, and scales effectively across LLaMA and Qwen models up to 32B parameters.

The core of the method involves analyzing and approximating causal self-attention using linear-time mechanisms within a frozen-backbone setting.

Analysis of Softmax Attention: The paper derives a first-order approximation of softmax attention, showing that it naturally applies key-dependent, rank-1 orthogonal projections. This is mathematically represented by the attention weights Pfull.

Delta-Style Updates: Delta-style updates, such as those in Gated Delta Networks (GDN), are shown to align with this key-dependent, rank-1 projection structure. The paper derives that GDN's update rule approximates the softmax attention weights.

Comparison with Gated Accumulation: Purely gated accumulation methods, like Gated Linear Attention (GLA), are shown to rely on key-independent decay, missing the crucial geometric correction provided by key-dependent projections.

Practical Interventions: To bridge the remaining approximation gap, the paper introduces and evaluates several practical interventions:

Disjoint Sliding Window Path: Incorporating sliding window attention (SWA) to capture local context.
Sink Tokens: Using a small budget of sink tokens to maintain essential long-range retrieval.
Projection Adaptation: Employing short convolutions or LoRA on Q, K, V projections to adapt the linear mechanism.

Key Equations: The paper uses the following key equations to represent attention mechanisms and their approximations:

Softmax Attention: Yt = Σ_{i<t} exp(qᵢkᵢ/√d) / Σ_{j<t} exp(qᵢkⱼ/√d) * Vᵢ
First-order Taylor approximation of softmax: pᵢ(x) ≈ 1/t + 1/t * (xᵢ - x̄) + R(x)
Gated DeltaNet (GDN) state update: St = (I – βtktkᵀ)GtSt−1 + βtkł Ut

Training: The method involves replacing attention blocks in pretrained models with linearized mechanisms and training only the newly introduced parameters, keeping the backbone frozen.

Architecture: The study evaluates several linear attention mechanisms as replacements for causal self-attention, including Hedgehog, Gated Linear Attention (GLA), and Gated Delta Networks (GDN). The pretrained backbone parameters (Wq, Wk, Wv, RMSNorm) are kept frozen.
Key Equations:
- Softmax Attention: Yt = Σ_{i<t} exp(qᵢkᵢ/√d) / Σ_{j<t} exp(qᵢkⱼ/√d) * Vᵢ
- Kernelized Linear Attention: Yt = φ(qt)St / φ(qt)Zt, where St = Σ_{i<t} φ(ki) and Zt = Σ_{i<t} φ(ki).
- Hedgehog feature map: φ(x) = σ(xW) ⊕ σ(−xW).
- Gated Linear Attention (GLA) state update: St = GtSt−1 + kt ut, where Gt = diag(at).
- Gated Delta Net (GDN) state update: St = (I – βtktkᵀ)GtSt−1 + βtkł Ut.
- First-order Taylor approximation of softmax: pᵢ(x) ≈ 1/t + 1/t * (xᵢ - x̄) + R(x).
Loss/Objective: The models are trained using the standard next-token prediction loss (cross-entropy).
Optimization: AdamW optimizer is used with a learning rate of 2e-3 and cosine learning rate scheduler. LoRA training uses a learning rate of 5e-4. Gradient clipping is set to 1.
Augmentation: Not stated in the paper.
Data Handling: Sequences are packed to 4096 tokens. The DCLM-Edu dataset is used for training.
Implementation Details: FLA2 library is used for GLA and GDN kernels. Hedgehog is implemented using FLA for feature maps and a chunk-wise-parallel form for training and recurrent form for inference. PEFT library is used for LoRA.

Datasets: DCLM-Edu for training. Downstream evaluation includes common reasoning benchmarks (PIQA, ARC-e, ARC-c, HellaSwag, WinoGrande), 5-shot MMLU, Lambada, S-NIAH, and extended RULER benchmarks.
Metrics: Accuracy (%), Normalized MSE, Latency (sec.), GPU Memory Usage (GB).
Baselines: Full Attention, Hedgehog, Gated Linear Attention (GLA), Gated Delta Networks (GDN), LoLCATs, Liger-GLA, Lizard, LoLA, STILL.
Numeric Results:
- In the strict frozen-backbone setting, GDN consistently achieves lower MSE than other linear kernels.
- With SWA and sink tokens, the gap between linearized models and full attention is significantly reduced. For Llama3.1-8b, the gap with full attention is 0.17% on CR, 6.37% on MMLU, and 5.97% on Lambada.
- With short convolutions, the gap between models with and without the delta rule almost disappears.
- The proposed method (GDN with SWA and sink tokens) matches or exceeds prior linearization baselines on common reasoning benchmarks and outperforms them on 5-shot MMLU.
- Scalability tests on Qwen3 models (0.6B to 32B) show that linearized performance scales consistently with the original base model.
- In a hybrid setting, keeping a subset of layers with full attention and replacing others with GDN shows competitive results.
Runtime/Training Cost: Training GLA and GDN variants takes 1 hour on a single H100 GPU. Training Hedgehog takes approximately 8 hours.

The first-order theory relies on approximate assumptions that may not hold uniformly across layers or very short prefixes.
The paper leaves the extension of adaptive caching strategies to GDN linearization for future work.

Method evidence: "We derive a first-order approximation connecting softmax attention to linear state updates. Our analysis proves that softmax naturally applies key-dependent, rank-1 orthogonal projections. Delta-style updates (like Gated Delta Networks) natively implement these corrections through rank-1 operators constructed from keys, suppressing components aligned with previous keys. Conversely, pure gated accumulation (like Gated Linear Attention) relies on key-independent decay, missing this vital geometric correction." (Section: 1 Introduction)
Result evidence: "We outperform prior post hoc baselines on MMLU, match the long-context retrieval of complex adaptive-caching frameworks, and remain competitive across all downstream tasks, all while demonstrating effectiveness in hybrid settings and consistent scaling across LLaMA and Qwen models up to 32B parameters." (Section: Abstract)
Limitation evidence: "Limitations: Our first-order theory relies on approximate assumptions that may not hold uniformly across layers or very short prefixes. We leave the extension of adaptive caching strategies to GDN linearization to future works." (Section: 6 Conclusion and Limitations)

The Key to Going Linear: Analysis-Driven Transformer Linearization

Abstract

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Additional Figures