PaperlayerView exact revision on arXiv

Source-linked AI summary

Taylor-Calibrate: Principled Initialization for Hybrid Linear Attention Distillation

Zhongzhu Zhou, Qingyang Wu, Junxiong Wang, Mayank Mishra, Shuaiwen Leon Song, Ben Athiwaratkun, Chenfeng Xu

arXiv:2606.16429v1cs.LGcs.CL

TL;DR

Converting pretrained Transformers to hybrid Gated DeltaNet students leaves recurrent dynamics under-specified, making initialization brittle. Taylor-Calibrate uses teacher attention statistics and layer-local alignment to initialize those dynamics, yielding stronger zero-shot students and faster recovery than naive conversion.

  • Problem

    Transformer-to-hybrid conversion copies projections but leaves recurrent timescales and gates under-specified, making student quality sensitive to initialization and distillation conditions.

  • Method

    Taylor-Calibrate uses teacher attention statistics to initialize GDN dynamics, then briefly aligns each converted layer’s output with its teacher counterpart.

  • Results

    Across four teacher settings and three layer-selection policies, Taylor-Calibrate achieves higher zero-shot averages and faster downstream recovery than naive initialization.

  • Takeaways & Limitations

    Initialization quality affects both zero-shot quality and subsequent distillation trajectories in converted GDN students.

  • Takeaways & Limitations

    At 700M training tokens, short-context behavior nearly matches teachers, but long-context retrieval remains incompletely restored.

Abstract

from arXiv · show

Hybrid linear attention models offer an appealing path to faster long-context inference: they reduce the quadratic cost and KV-cache burden of full softmax attention while retaining much of the quality of Transformer models. A practical way to obtain such models is to convert a pretrained Transformer instead of pretraining a new architecture from scratch, but this conversion is still brittle. Simply copying the teacher attention projections into a Gated DeltaNet (GDN) student does not specify the new recurrent decay, write, and output-gating dynamics. As a result, the converted model often starts in a poor dynamical regime and must spend many distillation tokens repairing initialization rather than learning the remaining teacher behavior. We propose Taylor-Calibrate, a lightweight initialization method for hybrid GDN students. The method uses Taylor-guided teacher attention statistics to set the value projection, memory timescale, write gates, and output gate, then applies a short per-layer alignment step to match each converted layer to the teacher output. Across four teacher settings and three retained-layer policies, Taylor-Calibrate gives substantially stronger zero-shot students, with up to an 88x improvement in a representative ablation, and reaches matched recovery targets with 4.9x--9.2x fewer training tokens than naive conversion.

1 Introduction

Taylor-Calibrate frames hybrid linear-attention conversion as an initialization problem for recurrent dynamics, not merely teacher-projection copying. Its two-stage GDN method calibrates dynamical parameters from teacher attention statistics and then performs brief layer-local alignment, improving zero-shot quality and downstream recovery across diverse conversion settings.

  • Motivation: Hybrid conversion is brittle because copied Q/K/V/O projections leave recurrent timescales, write gates, and output gates unspecified.These parameters control information flow in GDN, so random initialization can place the student in a poor dynamical regime.
  • Method: Taylor-Calibrate uses teacher attention statistics to initialize GDN value scale, recurrent decay, and write-gate settings.Value scale comes from attention-derived value statistics, average attention distance sets state decay, and attention concentration sets the write gate.
  • Method: A short layer-local calibration step then aligns each converted GDN layer with the teacher output.This forms the method’s second stage and complements the Taylor-guided parameter initialization.
  • Empirical motivation: Zero-shot quality varies sharply with GDN gate initialization even when converted variants share identical teacher Q/K/V/O projections.Better initialization also starts from lower loss and maintains a more favorable optimization trajectory during distillation.
  • Evaluation: Across four teacher settings and three layer-selection policies, Taylor-Calibrate yields higher zero-shot Avg and faster downstream recovery than naive initialization.The evaluation supports initializing dynamical parameters rather than only copying projection weights.

2 Preliminaries

Softmax attention requires a KV cache that grows with sequence length, motivating fixed-size recurrent alternatives. Gated DeltaNet removes this cache but makes initialization of decay and gating dynamics critical for preserving useful context.

  • Softmax attention: The KV cache retains all previous keys and values during autoregressive decoding, creating a central long-context inference bottleneck.Its size grows linearly with sequence length.
  • Calibration lens: Theoretical expansion of the softmax numerator motivates linear-recurrent approximation and calibration through effective scaling terms.A second-order correction may be summarized as γ2 = 1 + µ3/(2µ2).
  • Gated DeltaNet: Gated DeltaNet replaces softmax attention with a recurrent memory state and compresses context into fixed-size state.This provides constant-memory inference with respect to context length.
  • Gated DeltaNet: Poorly initialized GDN decay and gating can erase useful history too aggressively or inject unstructured noise into the residual stream.The recurrent block therefore places substantial burden on decay and gating parameters.

3 Method: Taylor-Calibrate

Taylor-Calibrate initializes converted GDN layers by mapping teacher attention statistics to value, decay, write, and output-gate parameters, then performs brief layer-local alignment before full distillation. The method targets layer-local mismatch rather than solving end-to-end distillation, addressing recurrent dynamics left unspecified by projection transfer.

  • Overview: Taylor-Calibrate uses teacher attention statistics to initialize GDN value scale, memory decay, write gates, and output gating before per-layer output alignment.The pipeline has two stages: Taylor-guided parameter initialization followed by alignment on the same calibration inputs.
  • Problem formulation: The initialization objective is to reduce each converted layer’s mismatch with the teacher on a small calibration corpus before expensive full-model training.It does not attempt to solve end-to-end distillation at this stage.
  • Motivation: The method fills missing recurrent-dynamics information—memory timescale, write strength, and output gating—that teacher projection copying leaves unspecified.Students inheriting the same WQ, WK, WV, and WO can still begin with sharply different PPL and task accuracy depending on GDN gate settings.
  • Phase 1: Taylor-guided initialization: Teacher output amplitude calibrates the value path, average attention distance calibrates decay, entropy calibrates write strength, and a small output gate prevents recurrent-state domination.Longer attention look-back implies slower decay, while lower entropy yields a larger concentration score and write-gate target.
  • Phase 2: Layer-local alignment: The alignment stage updates transferred projections and recurrent parameters to capture nonlinear residual mismatch and adapt the layer to recurrent rather than softmax dynamics.It complements the analytical stage’s low-order scale and gate choices rather than replacing them.

4 Experiments

Experiments across four teachers and three layer-retention policies show that Taylor-Calibrate improves zero-shot initialization and accelerates early short-context recovery. Long-context retrieval improves more slowly, requiring additional training beyond the short recovery schedule.

  • Experimental setup: Four teacher settings and three retained-layer policies evaluate Taylor-Calibrate against projection-copy and output-gate initialization strategies under matched distillation.Teachers include Qwen2.5-1.5B-Instruct, Qwen2.5-3B-Instruct, Llama-3.2-3B-Instruct, and Qwen3-8B; policies are Uniform, AR (LM-PPL), and GA-S2.
  • Zero-shot evaluation: Taylor-Calibrate most consistently improves zero-shot Avg at 50% and 75% retained attention, while RULER gains remain model- and budget-dependent.At aggressive 25% or 50% retention, all methods have very low RULER; output-gate stabilization is especially useful when more softmax layers remain.
  • Recovery training: At 100M tokens, Qwen2.5-1.5B improves from 37.2 to 62.6 Avg under Uniform selection, while Qwen3-8B improves from 39.8–49.4 to 66.5–67.4 across policies.The largest gains occur early in recovery, before Stage 2 can repair poor initialization; by 700M tokens, Taylor-Calibrate still gives the best Avg in every row, but the gap narrows.
  • Long-context recovery: At 700M tokens, Llama-3.2-3B reaches 64.8 vs. 65.6 Avg but trails 63.4 vs. 89.6 on RULER, while Qwen3-8B reaches 76.7 RULER versus 94.0 for its teacher.Short-context behavior largely recovers, whereas long-context retrieval remains below the teacher on the short schedule.
  • Ablations: Phase 2 alignment lowers PPL from 37337.3 for Baseline to 424.1 for full Taylor-Calibrate, which also achieves the best Avg, 32.1.Zero-Gate and Taylor-Only remain around 22470, while Alignment-Only reaches 2015.9, showing that Taylor calibration is most useful as initialization for local alignment.
  • Training efficiency: 4.9×–9.2× fewer training tokens are required for Taylor-Calibrate to reach matched target quality across representative runs.Extended SFTMix v0.2 training beyond 20k optimization steps can bring RULER into the teacher range in the 1.5B setting.

5 Conclusion

Taylor-Calibrate is presented as a lightweight initialization method for converting pretrained Transformers into GDN hybrid students, addressing poorly calibrated recurrent timescales and gates before global distillation. Across representative settings, it reduces the tokens needed to reach the same target quality, while limitations and future work include scaling and finer-grained calibration.

  • Distillation efficiency: 4.9× to 9.2× speedups: Taylor-Calibrate reaches the same target quality with fewer training tokens than Baseline across representative runs.The comparison concerns Stage-2 distillation after initialization.
  • Limitations and future directions: Limitations and future directions include larger-model support, architecture-specific calibration, and finergrained head-wise conversion.These directions are discussed in Appendix C.
  • Contribution: Taylor-Calibrate combines Taylor-derived calibration with short layer-local alignment before global distillation.The method targets poorly calibrated memory timescales and gates left by simply copying teacher projections.

A Additional Preliminaries and Derivations · A.1 Hybrid Attention Architectures · A.2 Decay Calibration from Half-Life Matching

The paper studies hybrid models that retain some full softmax-attention layers while replacing others with GDN, and calibrates recurrent decay using teacher attention distance and half-life matching. It assumes replacement layers are already selected and focuses on initializing their decay, write, and output-gate scales.

  • A.1 Hybrid Attention Architectures: Hybrid models partition the L layers into disjoint full-attention and GDN subsets.The partition satisfies S_soft ∪ S_gdn = {1, …, L} and S_soft ∩ S_gdn = ∅.
  • A.1 Hybrid Attention Architectures: Layer selection is treated as separate from the initialization problem.The method assumes the layers for replacement have already been chosen.
  • A.1 Hybrid Attention Architectures: The initialization question is how to calibrate newly introduced GDN decay, write, and output-gate scales before distillation.This framing targets the starting dynamical regime of converted layers rather than layer-selection policy.
  • A.2 Decay Calibration from Half-Life Matching: The decay map uses teacher softmax-attention statistics to set GDN’s recurrent decay parameterization.For head h, the per-step multiplicative decay is defined as ρ_h = exp(g_h).
  • A.2 Decay Calibration from Half-Life Matching: The target memory behavior is a half-life of roughly d_h tokens, where d_h is the teacher’s average attention distance for head h.This connects teacher attention distance to recurrent memory persistence.
  • A.2 Decay Calibration from Half-Life Matching: At initialization, setting A_log = 0 and assuming a small input-dependent term a_proj(x_t) simplifies the GDN decay.These assumptions define the approximation used for decay calibration.
  • A.2 Decay Calibration from Half-Life Matching: The resulting decay-logit approximation is g_h ≈ −softplus(d_tbias,h).This closed-form relation follows immediately from the half-life calibration setup.

A.3 Deriving the Entropy-to-Write-Gate Mapping … B Related Work

The appendices justify Taylor-Calibrate’s write-gate, value-side, and output-gate initialization as stable local calibration procedures. Related work situates the method among efficient hybrid architectures, Transformer conversion, layer selection, and teacher-informed local transfer.

  • A.3 Deriving the Entropy-to-Write-Gate Mapping: Taylor-Calibrate maps attention concentration to a bounded logit and uses its first-order affine approximation to initialize write gates near a neutral regime.The construction avoids under-writing and sigmoid saturation while keeping initialization in a high-gradient region.
  • A.3 Deriving the Entropy-to-Write-Gate Mapping: The write gate modulates update sharpness around neutrality because the decay timescale is handled separately by gh.Median-concentration heads are intended to initialize near a balanced write gate.
  • A.4 Matching the Write-Gate Logit by Row Rescaling: Phase 1 matches each teacher-informed write-gate logit by a scalar row rescaling that preserves the module-default row direction.The scale is chosen using a cheap proxy for typical preactivation magnitude under unit-RMS inputs.
  • A.5 Closed-Form Value-Side OLS Rescaling: The value-side calibration uses a closed-form one-dimensional OLS rescaling to align the student context with the teacher in squared error.Clipping is applied afterward as an implementation safeguard.
  • A.6 RMS-Matched Output-Gate Initialization: After value calibration, output-gate initialization restricts the projection to a one-parameter family and matches its scale to a small fraction of the teacher pre-WO output.The derivation is a calibration proxy for the actual SiLU parameterization.
  • A.6 RMS-Matched Output-Gate Initialization: The output gate uses damping factor λ = 0.01 to remain near-linear and contribute minimally at initialization.This is the regime in which the RMS-matching proxy is most appropriate.
  • B Related Work: Related efficient-sequence work includes production-scale hybrid and linear architectures such as Jamba, Jamba-1.5, MiniMax-01, Falcon-H1, Kimi Linear, and Qwen3.5.Another efficiency direction compresses softmax attention itself.
  • B Related Work: Transformer-conversion methods improve transfer through attention transfer, hidden-state alignment, low-rank finetuning, or progressive distillation, with RADLADS closest to this training protocol.Layer-selection studies show that retained softmax layers affect long-context behavior and recovery, whereas Taylor-Calibrate focuses on initializing converted recurrent dynamics; related local-transfer methods reuse teacher projections or perform layer-local matching.

C Limitations and Future Work

Taylor-Calibrate is presented as a lightweight initialization step rather than a complete Transformer-to-recurrent conversion solution. Future work must address scale, architecture and hybridization choices, heuristic calibration, evaluation scope, continued distillation dependence, and broader dual-use implications.

  • Scope: Taylor-Calibrate remains an initialization method, not a complete Transformer-to-recurrent conversion recipe.Its purpose is to provide converted GDN layers with a better starting point before full distillation.
  • Larger Model Support: Larger teachers require memory-efficient forward passes, distributed evaluation, and matched recovery schedules across retained-attention policies.The paper does not quantify these settings because the full end-to-end protocol was limited to smaller teacher configurations.
  • Architecture specificity: The calibration formulas are specific to Gated DeltaNet and should not be treated as architecture-independent recipes for GLA, Mamba-like layers, KDA, or other efficient sequence models.These alternatives use different parameterizations and memory-update or input-dependent-gating mechanisms.
  • Layer-wise rather than head-wise hybridization: Layer-wise hybridization ignores head-level functional differences, including heads specialized for long-range retrieval or copying.Selecting softmax layers keeps the setup controlled and makes retained-attention budgets easier to compare, but overlooks structure within layers.
  • Heuristic calibration rules: Several calibration rules remain heuristic, including the half-life rule, entropy-to-write-gate mapping, clipping ranges, and output-gate scale.They are designed to be simple, cheap, and stable near initialization, but are not globally optimal for nonlinear GDN recurrence interactions.
  • Calibration and evaluation scope: The small calibration corpus and short layer-local alignment stage leave sensitivity to calibration size, alignment steps, learning rate, and domain mismatch incompletely tested.These limits are intentional because calibration is designed to cost little relative to full recovery training.
  • Downstream distillation dependence: Taylor-Calibrate does not eliminate downstream distillation, which remains necessary to repair cross-layer interactions and adapt the residual stream to recurrent blocks.Final quality also depends on layer selection, retained-softmax budget, distillation loss, optimizer settings, and recovery duration.
  • Broader impact: The reduced compute and token budget may broaden access to efficient long-context models while also lowering the cost of harmful dual-use applications.The paper notes that it does not release new model artifacts.

D Extended Experimental Results

The appendix provides supplementary empirical results referenced in Sections 4 and 4.5, including additional teacher settings and gate-statistics diagnostics that interpret initialization behavior.

  • D Extended Experimental Results: The appendix reports additional teacher settings and gate-statistics diagnostics for interpreting initialization behavior.These materials supplement the empirical results referenced in Sections 4 and 4.5.

D.1 Detailed Experimental Results

This section provides detailed zero-shot and recovery results for four teacher settings, covering short-context benchmarks and RULER as a long-context probe. It compares initialization methods across retained-layer policies and softmax budgets, then reports checkpoint performance after staged distillation.

  • Experimental scope: Detailed results cover Qwen2.5, Llama-3.2, and Qwen3 teacher settings across short-context benchmarks and the RULER long-context probe.The short-context suite includes ARC-C/E, HellaSwag, PIQA, MMLU, OpenbookQA, ReArc, WinoGrande, BoolQ, LAMBADA, COPA, and SciQ.
  • Zero-shot details: Table 5 evaluates five initialization methods under three retained-layer selection policies and 25%, 50%, and 75% softmax budgets before downstream training.Avg is computed over available short-context metrics, while RULER is reported separately.
  • Recovery details: Table 6 reports checkpoint benchmarks at 100M tokens after Stage 1 and 700M tokens after Stage 2 for all four teacher settings.Avg appears only when the full short-context set is available, with RULER reported separately.

D.2 Gate Value Statistics

In the Qwen2.5-3B-Instruct uniform-25% conversion, Taylor calibration regularizes recurrent memory timescales and removes extreme random-decay outliers. Layer-local alignment preserves that calibrated timescale while substantially increasing projection-gate magnitudes.

  • Experimental setup: The analysis evaluates four initialization settings across 27 recurrent GDN layers and 432 recurrent heads in the uniform-25% conversion.The retained softmax layers are {0, 4, 8, 12, 16, 20, 24, 28, 32}.
  • Memory timescale: Taylor-calibrated variants concentrate around a median half-life of approximately 71 steps and cap the maximum near 128 steps, unlike baseline outliers up to roughly 3.06×10^4 steps.Half-life measures the memory timescale induced by the initialized recurrence before input-dependent updates.
  • Alignment effects: 0.692 steps: mean half-life changes only from 69.337 to 70.030 steps after alignment, indicating that Phase 2 largely preserves the calibrated timescale.The alignment step instead primarily changes projection gates.
  • Alignment effects: 89.8×: mean absolute gproj magnitude increases from 1.32×10^-4 after Taylor-Only to 0.012 after full Taylor-Calibrate.bproj increases by 1.73×, while write and output projection gates receive most alignment-induced change.
Loading 2606.16429v1…