Summary (Overview)
- Core Problem: Converting pretrained Transformers to hybrid linear attention models (specifically Gated DeltaNet, GDN) by simply copying teacher Q/K/V/O projections leaves the new recurrent decay, write, and output-gating dynamics uninitialized, resulting in poor starting quality and slow recovery during distillation.
- Proposed Solution: Taylor-Calibrate, a lightweight two-stage initialization method that (1) extracts Taylor-guided statistics from teacher attention maps to calibrate value scale, decay timescale, write gate, and output gate, and (2) performs a brief per-layer alignment step to match each converted GDN layer to the teacher output.
- Key Results: Taylor-Calibrate achieves – fewer training tokens than naive conversion to reach teacher-like quality, with consistent improvements in zero-shot perplexity (Avg) and RULER scores across multiple teacher models (Qwen3-8B, Llama-3.2-3B, Qwen2.5-7B).
- Insight: Initialization quality mainly affects recovery speed rather than the final converged quality; long-context retrieval (RULER) requires longer training budgets and is governed by a slower timescale than short-context perplexity.
- Broader Message: Linear-attention transfer needs to initialize dynamical parameters, not just copy projection weights.
Introduction and Theoretical Foundation
Background
Transformers remain the default for LLMs, but long-context decoding suffers from quadratic attention cost and linearly growing KV cache. Hybrid models (e.g., Jamba, Samba, Hymba, MiniMax-01) replace some softmax layers with cheaper recurrent/linear modules to improve throughput and memory trade-offs. However, pure recurrent models struggle on retrieval-heavy tasks, motivating hybrid designs that retain some softmax layers.
Motivation
Recent pipelines convert pretrained Transformers into hybrid students by copying teacher projections. The central issue: hybrid transfer is not only a projection-copying problem. Copying weights preserves useful projections but does not specify:
- Recurrent timescales (decay)
- Write gates
- Output gates
For GDN, the recurrent output is additionally modulated by a learned output gate before the final projection. If these new parameters are left random, the converted student can have very poor zero-shot quality.
Theoretical Foundation: Taylor View of Softmax Attention
The method is motivated by expanding the softmax numerator around a neutral logit regime, following the view that softmax attention admits linear-attention and recurrent approximations:
Truncating this expansion motivates a fixed-size linear-recurrent approximation of softmax attention, making teacher attention statistics (value amplitude, average look-back distance, attention concentration) operational proxies for GDN's value scale, decay timescale, and write strength.
Gated DeltaNet (GDN)
GDN replaces softmax attention with a recurrent state update:
where is the decay, is the write gate, and the output is further modulated by an output gate . GDN provides constant-memory inference but places substantial burden on decay and gating parameters.
Methodology
Stage 1: Taylor-Guided Analytical Calibration
For each converted layer , the method reads statistics from teacher attention maps:
| Parameter/statistic | Standard GDN transfer | Taylor-Calibrate initialization |
|---|---|---|
| Not used | Teacher attention statistics | |
| Copied | Copied |
Value scale — Teacher output amplitude calibrates : the value-side OLS factor is computed per head by matching the student value-path output to the teacher output in least squares (Appendix A.5).
Decay timescale — The distance statistic is the attention-weighted average look-back length. If a head should preserve half its information after roughly tokens, the target decay magnitude is:
Write gate — Attention entropy is converted into a write-gate target via a concentration score ; higher entropy (more diffuse attention) leads to weaker writes.
Output gate — Initialized small enough (RMS-matched) to avoid dominating the layer output, keeping the recurrent branch in a low-contribution, high-gradient regime.
Stage 2: Layer-Local Alignment
A short per-layer distillation step matches each converted GDN layer output to the teacher layer output:
This is performed on calibration data before global distillation. It allows transferred projections to adapt to recurrent dynamics and updates gate directions rather than just scales. Notably, successful alignment does not dramatically increase the magnitude of ; instead, it rotates parameters into directions producing teacher-like outputs.
Empirical Validation / Results
Setup
- Teachers: Qwen3-8B, Llama-3.2-3B, Qwen2.5-7B
- Conversion: Replace a fraction (25%, 50%, 75%) of softmax layers with GDN, following RADLADS-style projection transfer
- Baselines: Baseline (random recurrent params), Zero-Gate, Small-Gate, Taylor-Only (Stage 1 only), Taylor-Calibrate (Stages 1 + 2)
- Metrics: Avg PPL (short-context), RULER (long-context retrieval)
Zero-Shot Results (Table 2)
Key trends:
- Aggressive conversion is hard: At 25–50% retained attention, all methods have very low RULER; Baseline Avg is near chance-level.
- Output-gate stabilization helps most when enough softmax remains: At 75% retention, Zero-Gate and Small-Gate recover much of the short-context Avg.
- Full Taylor-Calibrate most consistently improves Avg, especially at 50% and 75% retention. RULER gains are more model- and budget-dependent.
Recovery Dynamics (Table 3)
| Setting | Baseline | Taylor-Only | Alignment-Only | Taylor-Calibrate |
|---|---|---|---|---|
| PPL (initial) | 37337.3 | ~22470 | 2015.9 | 424.1 |
| Avg (best) | — | — | — | 32.1 |
Taylor calibration is most useful as a better starting point for local alignment, not a standalone replacement. The largest benefit appears early in recovery, before Stage 2 has enough tokens to repair poor initialization.
Training Token Savings
- – fewer tokens to reach teacher-like quality compared to naive conversion
- For Llama-3.2-3B, the baseline recovers reasonably after Stage 1, but Taylor-Calibrate still gives consistent short-context gains
- By 700M tokens, the Avg gap narrows as distillation restores short-context behavior for both initializations
Long-Context Recovery
RULER recovery is governed by a slower timescale than PPL or short-context Avg. Restoring retrieval behavior requires SFTMix v0.2 and more than 20k optimization steps, beyond the short recovery schedule. Taylor-Calibrate improves the starting point and early recovery, but full long-context restoration remains a slower process.
Theoretical and Practical Implications
Theoretical Implications
- Initialization as architecture transfer: The results reframe hybrid conversion as an initialization problem for recurrent dynamics, not just projection copying. The Taylor view of softmax attention provides a principled bridge between teacher attention statistics and GDN's decay/gating structure.
- Layer-local alignment matters: Analytical calibration alone (Taylor-Only) is insufficient; the short alignment step is critical for rotating gate directions into teacher-like configurations.
- Recovery speed vs. final quality: Initialization primarily accelerates recovery rather than changing the converged optimum, suggesting the loss landscape is well-behaved but the initialization determines the path.
Practical Implications
- Token efficiency: Taylor-Calibrate reduces distillation token budgets by up to , directly lowering training compute costs for model conversion.
- Hybrid model deployment: Enables faster conversion of pretrained Transformers into efficient hybrid architectures for long-context inference, reducing KV-cache and quadratic attention costs.
- Design guidance: Output-gate stabilization is a cheap but effective safeguard; full calibration is most valuable when aggressive layer replacement is desired.
Conclusion
Main takeaways:
- Taylor-Calibrate is a lightweight, two-stage initialization method for GDN-based hybrid distillation that combines Taylor-derived gate calibration with brief layer-local alignment.
- It significantly improves zero-shot quality and reduces recovery tokens (–) across multiple teacher models.
- Initialization quality mainly affects recovery speed; long-context retrieval requires additional training beyond short-context recovery.
Future directions:
- Extending calibration to other linear-attention families (e.g., Mamba, RWKV)
- Larger-model support and architecture-specific calibration
- Finer-grained head-wise conversion strategies
- Understanding the slower timescale of long-context retrieval recovery
Related papers
- Data Scarcity and Model Sparsity: Mixtures-of-Experts Overfit More to Repeated Data
Mixture-of-Experts models overfit repeated data faster than dense Transformers, losing their advantage by 32x repetition, but dropout and output masking can restore it even at 64x.
- Evidence for Shared Routing Geometry and Dynamics in Sparse Mixture-of-Experts
After orthogonal gauge alignment, router-control states across MoE layers share a low-dimensional dynamical core, where a single linear transition captures most of layer-specific predictive power.
- Good Pretraining, Bad SFT: Checkpoint Quality Across the Training Stack
Pretraining loss and benchmark scores fail to predict post-SFT performance, as a checkpoint with worse pretraining metrics yields better final models after supervised fine-tuning.