发表机构
Tencent(腾讯)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
本文针对在线自蒸馏的次优问题,提出统一在线自蒸馏(USD)框架,该框架通过耦合学习难度预算与师生 divergence,在多模型规模及推理基准上均优于基线方法。
AI 中文摘要
在线自蒸馏(OPSD)通过将特权上下文内化到模型参数中,提升大语言模型(LLM)的推理能力。近期有两类研究分别从选择学习的token、控制教师接收的特权信息两方面改进基础OPSD,但本文指出两类研究均固定其中一个变量优化另一个,导致次优解。本文提出两类变量通过学生学习能力耦合:特权信息设定教师规定的单token divergence,token加权选择学生需吸收的信息。本文将两类研究形式化为统一优化框架,在学生可吸收的总学习难度预算约束下最大化师生总 divergence。据此提出轻量在线算法统一在线自蒸馏(USD)以求解拉格朗日函数,USD表明单个对偶变量同时决定token选择阈值与特权信息调整方向,使监督匹配学生的动态能力。经大量实验,USD在不同模型规模、不同推理基准上,均优于OPSD及token侧、特权信息侧基线。代码可在指定URL获取。
英文摘要
On-policy self-distillation (OPSD) improves the reasoning abilities of LLMs by internalizing privileged context into model parameters through self-distillation. Two recent research lines promote vanilla OPSD by choosing which tokens to learn from and by controlling how much privileged information the teacher receives, respectively. However, we show that each line optimizes one variable while holding the other fixed, which leads to a suboptimal solution. We argue that the two variables are coupled through the student's learning capacity: the privileged information sets the per-token divergence the teacher prescribes, while token weighting selects which of these the student must absorb. We formalize the two lines of work into a unified optimization framework, which maximizes the aggregate teacher--student divergence, subject to a budget on the aggregate learning difficulty the student can absorb. Under this modelling, we propose Unified On-Policy Self-Distillation (USD), a lightweight online algorithm to solve the Lagrangian. USD reveals that a single dual variable governs both decisions: at one price for learning difficulty, it simultaneously sets the token-selection threshold and the direction of privileged-information adjustment, keeping supervision matched to the student's evolving capacity. Through extensive experiments, USD consistently demonstrates superior performance over OPSD and token- and PI-side baselines across various model scales on various reasoning benchmarks. Code is available at https://github.com/lauvlalala/USD.