跳到主要内容
返回时间线
arXiv来源发表:

控制变量蒸馏实验显示:同策略 rollout 并非普遍更优,KL 方向与学习率才是主导因素

核心概要

该研究在强到弱蒸馏中独立改变 rollout 策略、token 级 KL 方向与学习率,跨 Llama 3 与 Qwen2.5 及科学、医学、算术推理任务发现:同策略 rollout 在最终准确率、灾难性遗忘和参数更新稀疏度上没有一致优势,前向 KL 对 rollout 策略稳健而反向 KL 更敏感并偏好学生生成的 rollout,学习率主导遗忘与稀疏度,同策略数据仅在向更难 Countdown 变体泛化时带来提升且该优势在后续 RLVR 后不可靠保留。

AI-generated editorial illustration: On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics

深度剖析

在强到弱蒸馏这一受控测试平台上,独立改变 rollout 策略、token 级 KL 方向和学习率后,同策略与离策略蒸馏在最终任务准确率上几乎无差别,最佳平均准确率为 OnPD 72% 与 OffPD 73%。 以往关于同策略优势的证据多来自 SFT 与 RLVR 的比较,同时改变了目标函数、监督来源与密度、优化流程等多个因素;该工作把 rollout 策略与 KL 方向解耦,打破了前向 KL 配离策略、反向 KL 配同策略的常规配对。 Llama-3.1-8B 教师蒸馏到 Llama-3.2-1B 学生,三个数据集(MedReason、Science、Countdown),每个配置三个随机种子,150 步全参数微调,并在 Qwen2.5-7B 到 Qwen2.5-1.5B 上重复验证。

KL 方向比 rollout 策略更清晰地决定任务表现与输出覆盖:前向 KL 在所有 rollout 策略和学习率下达到 71–73%,反向 KL 则在 35% 到 72% 之间波动且对学习率高度敏感;前向 KL 在重复采样下带来更大的 pass@10 增益。 该工作把 KL 方向作为独立设计维度加以操纵,并给出梯度层面的解释:前向 KL 对学生 logits 的导数坐标有界,其更新差异受 rollout 分布总变差线性约束;反向 KL 的导数被学生与教师概率比加权,可任意放大。 基于 token 级 KL 梯度的推导(附录 B)与 rollout 稳定性定理(附录 C),并在连续的学生–教师 rollout 谱上验证:前向 KL 在整个谱上保持 80% 以上准确率,仅变化 5.2 个百分点,反向 KL 则波动大且在高学习率下偶发崩溃。

灾难性遗忘与参数更新稀疏度主要由学习率决定,而非 rollout 策略:较低学习率下 OOD 平均表现变化至多 1.3 个百分点、稀疏度 85.3–89.6%,较高学习率下 OOD 下降 11.2–14.0 个百分点、稀疏度降至 51.9–60.0%。 这直接检验了同策略学习能减少遗忘、产生更稀疏更新的既有假设,并在匹配比较中发现 OffPD 的更新至少与 OnPD 一样稀疏。 七个 OOD 基准(MMLU-Pro、TruthfulQA、IFEval、HumanEval、EQ-Bench、BBQ、ToxiGen)共 3,935 个样本,加上从低到高的学习率扫描,显示稀疏度随学习率近似线性下降。

同策略数据在向更难任务泛化时确有帮助:在 Countdown-4E 上,两种 KL 方向下更多同策略 rollout 都持续带来更高的 pass@k;但这一初始优势在后续 RLVR 后不能可靠保留,最终最强的持续 RLVR 表现来自离策略检查点。 该工作区分了最终同分布准确率、输出覆盖、更难变体泛化和后续 RLVR 起点质量这几类不同评价,指出同策略的价值依赖于评价设置。 Countdown-3 检查点在 Countdown-4 上做 300 步 RLVR,每个蒸馏配置三个种子共 24 次运行;此外在平均 622 token 的 Numina–MATH 长推理任务上,低学习率下 OnPD 配反向 KL 取得最高 MATH-500 准确率,但该实验每个条件仅一次运行,作者将其视为提示性证据。

启示与展望

该研究面向强到弱蒸馏这一受控设定,适用于学生模型不超过 1.5B 参数、推理轨迹不超过约 2,000 token 的范围,覆盖 Llama 3 与 Qwen2.5 两个模型家族以及科学、医学、算术三类推理任务。它使实践者能够在给定 KL 方向与学习率时判断是否值得为同策略 rollout 支付额外生成成本,并提示把离策略蒸馏纳入标准基线;对希望抑制附带教师风格迁移的场景,同策略配反向 KL 显示出保留学生英语回答风格的倾向。

读者仍需关注:结论建立在不超过 1.5B 参数的学生模型与较短推理轨迹上,更大模型与更长 rollout 是否呈现同样格局尚待检验;长推理 Numina–MATH 实验每个条件仅一次运行,作者本人将其视为提示性证据;教师模型在研究中保持固定,不同学生–教师组合下是否出现 OnPD 与 OffPD 的差异仍是开放问题;此外,本次加载的文本为完整正文与附录,但图表以引用形式出现,具体曲线数值只能依据正文与附录中的文字描述来理解。

来源