Taylor-Calibrate: ハイブリッド線形注意蒸留のための原理に基づく初期化

要旨

ハイブリッド線形注意モデルは、より高速な長文脈推論への魅力的な経路を提供する。これらは、完全なソフトマックス注意の二次的なコストとKVキャッシュの負担を軽減しつつ、Transformerモデルの品質の多くを保持する。このようなモデルを得る実用的な方法の一つは、新しいアーキテクチャをゼロから事前学習する代わりに、事前学習済みTransformerを変換することであるが、この変換は依然として脆弱である。単に教師の注意射影をGated DeltaNet（GDN）生徒にコピーするだけでは、新しい再帰的減衰、書き込み、出力ゲーティングのダイナミクスを指定できない。その結果、変換されたモデルはしばしば貧弱な動的状態から開始し、残りの教師行動を学習するよりも、多くの蒸留トークンを初期化の修復に費やさざるを得なくなる。本稿では、ハイブリッドGDN生徒のための軽量な初期化手法であるTaylor-Calibrateを提案する。本手法は、テイラー誘導による教師注意統計量を用いて値射影、記憶タイムスケール、書き込みゲート、出力ゲートを設定し、その後、各変換層を教師出力に一致させるための短い層ごとのアライメントステップを適用する。4つの教師設定と3つの保持層ポリシーにわたって、Taylor-Calibrateははるかに強力なゼロショット生徒を与え、代表的なアブレーションでは最大88倍の改善を示し、ナイーブな変換と比較して4.9倍から9.2倍少ない訓練トークンで一致した回復目標に達する。

English

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.