Taylor-Calibrate: Principled Initialization for Hybrid Linear Attention Distillation
作者: Zhongzhu Zhou, Qingyang Wu, Junxiong Wang, Mayank Mishra, Shuaiwen Leon Song, Ben Athiwaratkun, Chenfeng Xu
分类: cs.LG, cs.CL
发布日期: 2026-06-15
备注: 24 pages, 9 figures
💡 一句话要点
提出Taylor-Calibrate以解决混合线性注意力蒸馏初始化问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 混合线性注意力 蒸馏训练 模型初始化 泰勒引导 Gated DeltaNet 长上下文推理 自然语言处理 计算机视觉
📋 核心要点
- 现有的混合线性注意力模型在转换过程中存在脆弱性，导致初始化不良，影响学习效果。
- Taylor-Calibrate方法通过泰勒引导的教师注意力统计来优化初始化，确保模型更好地学习教师行为。
- 实验显示，使用Taylor-Calibrate的学生模型在多个设置下性能显著提升，训练代币需求大幅减少。
📝 摘要（中文）
混合线性注意力模型为快速长上下文推理提供了有效路径，减少了全软最大注意力的二次成本和KV缓存负担，同时保留了Transformer模型的质量。将预训练的Transformer转换为混合模型的过程仍然脆弱，简单复制教师注意力投影到Gated DeltaNet（GDN）学生模型中并未明确新的递归衰减、写入和输出门动力学，导致转换后的模型常常处于不良的动态状态。为此，本文提出了Taylor-Calibrate，一种轻量级的初始化方法，利用泰勒引导的教师注意力统计来设置值投影、记忆时间尺度、写入门和输出门，并通过短的逐层对齐步骤将每个转换层与教师输出匹配。实验结果表明，Taylor-Calibrate显著提升了零-shot学生模型的性能，训练代币减少4.9到9.2倍。
🔬 方法详解
问题定义：本文旨在解决混合线性注意力模型在从预训练Transformer转换时的初始化不良问题。现有方法在简单复制教师模型的注意力投影时，未能有效设置新的动力学，导致模型学习效率低下。
核心思路：论文提出的Taylor-Calibrate方法利用泰勒引导的教师注意力统计信息，优化学生模型的初始化过程，确保模型在学习过程中能够更有效地接近教师模型的行为。
技术框架：该方法主要包括四个步骤：首先，利用教师模型的注意力统计设置学生模型的值投影和记忆时间尺度；其次，配置写入门和输出门；然后，通过逐层对齐步骤，确保每个转换层与教师输出一致；最后，进行模型训练。
关键创新：Taylor-Calibrate的创新在于其利用泰勒引导的统计信息进行初始化，这一方法显著改善了模型的初始动态状态，与传统的简单复制方法相比，提供了更为稳健的学习基础。
关键设计：在参数设置上，方法通过教师模型的注意力统计来确定值投影和记忆时间尺度，写入门和输出门的配置则通过短期对齐步骤实现，确保每层的输出与教师模型的输出匹配。
🖼️ 关键图片
📊 实验亮点
实验结果表明，使用Taylor-Calibrate的学生模型在多个教师设置下表现出显著的性能提升，零-shot学习能力增强，训练代币需求减少4.9到9.2倍，达到更高的恢复目标，展示了该方法的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、计算机视觉等需要长上下文推理的任务。通过优化模型初始化，Taylor-Calibrate能够提升模型在实际应用中的学习效率和推理速度，具有重要的实际价值和广泛的应用前景。
📄 摘要（原文）
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.