‹ 返回 2026-06-21

Taylor-Calibrate:混合线性注意力蒸馏的基于原则的初始化方法

Taylor-Calibrate: Principled Initialization for Hybrid Linear Attention Distillation

▲ 3 💬 5 2026-06-21

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

摘要

混合线性注意力模型为实现更快的长上下文推理提供了有效的途径:它们能够降低全softmax注意力机制所带来的二次计算成本以及KV缓存负担,同时又能保留Transformer模型的大部分特性。实现这种模型的实用方法是将预训练的Transformer模型直接用于模型训练,而非从头开始构建新架构,不过这种转换方式仍然存在缺陷。简单地将教师模型的注意力投影直接应用到Gated DeltaNet模型中,并不能确保新的递归衰减、写入和输出门控机制能够正常运作。因此,转换后的模型往往处于不良的动态状态,需要大量额外的初始化步骤来修复初始状态,而不是学习教师的剩余行为特征。我们提出了Taylor-Calibrate这一轻量级初始化方法,它利用泰勒引导的教师模型注意力统计信息来设定值投影、记忆时间尺度、写入门和输出门的值,然后对每个层进行短暂的对齐处理,以使其与教师模型的输出保持一致。在四种不同的教师模型和三种不同的保留层策略下,Taylor-Calibrate能够产生性能更优的模型,在典型的消融实验中,其性能提升了88倍;而且与简单的转换方式相比,所需的训练次数则减少了4.9倍到9.2倍。

English Abstract

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.