Taylor-Calibrate : Initialisation fondée sur des principes pour la distillation d'attention linéaire hybride

Résumé

Les modèles d'attention linéaire hybrides offrent une voie intéressante pour accélérer l'inférence sur de longs contextes : ils réduisent le coût quadratique et la charge du cache KV de l'attention softmax complète tout en conservant une grande partie de la qualité des modèles de type Transformer. Une approche pratique pour obtenir de tels modèles consiste à convertir un Transformer pré-entraîné plutôt que d'en pré-entraîner une nouvelle architecture de zéro, mais cette conversion reste fragile. Le simple fait de copier les projections d'attention du professeur dans un étudiant Gated DeltaNet (GDN) ne spécifie pas les nouvelles dynamiques de décroissance récurrente, d'écriture et de contrôle de la porte de sortie. En conséquence, le modèle converti démarre souvent dans un régime dynamique défavorable et doit consacrer de nombreux tokens de distillation à réparer l'initialisation plutôt qu'à apprendre le comportement restant du professeur. Nous proposons Taylor-Calibrate, une méthode d'initialisation légère pour les étudiants GDN hybrides. Cette méthode utilise les statistiques d'attention du professeur guidées par Taylor pour définir la projection de valeur, l'échelle de temps de la mémoire, les portes d'écriture et la porte de sortie, puis applique une courte étape d'alignement par couche pour faire correspondre chaque couche convertie à la sortie du professeur. Sur quatre configurations de professeur et trois politiques de couches conservées, Taylor-Calibrate produit des étudiants zéro-shot nettement plus performants, avec une amélioration allant jusqu'à 88× dans une ablation représentative, et atteint des objectifs de récupération appariés avec 4,9 à 9,2 fois moins de tokens d'entraînement qu'une conversion naïve.

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.