Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
Jacqueline He, Howard Yen, Shuyue Stella Li, Margaret Li, Hanqing Zeng, Yinglong Xia, Benyu Zhang, Zhuokai Zhao, Qiang Zhang, Pang Wei Koh, Luke Zettlemoyer, Wen-tau Yih
cs.CL
2026-09-02
Meta 在 OLMo-2 上发现:预训练时前向 KL 蒸馏让推理和事实一起涨,中训阶段却拖慢事实。Switch Distillation 只在教师低熵 token 上蒸馏,中训后推理达 NTP 的 1.61–1.71 倍,事实保留 96.7–96.8%。
中训夹在预训练和后训练之间,用更少、更高质量的语料继续做自监督,用来抬推理、编码和指令跟随。token 比预训练少一个数量级,每条信号都更贵,蒸馏看起来很合适:用更强教师的预测分布补 next-token 的硬标签。蒸馏的文献几乎全写在预训练或后训练上,有没有原样搬到中训,没人系统测过。
在 OLMo-2 上测完,答案是不能原样搬。前向 KL 蒸馏在预训练里推理和事实一起涨;到了中训,推理还在涨,事实获取变慢。作者把这叫推理–事实权衡。
学生是 OLMo-2 1B。预训练从随机初始化跑 1000 亿 token,超过 Chinchilla 最优点;中训从已在 4 万亿 token 上预训练过的 Stage-1 检查点再跑 600 亿,数据是 Dolmino Mix 1124。教师是 OLMo-2 的 1B/7B/13B Instruct。评测用 OLMES,切成推理、事实回忆、知识与常识三组宏平均。
机制分析看三件事。教师预测熵在数学、指令这类程序性域更低,在网页知识域更高,且域内熵越低、top-1 和语料标签越一致。按教师熵把事实题分五档,预训练结束时最低熵档已学会 67%,最高熵档只有 5%;中训开始时没学会的事实堆在高熵档,正是教师最虚的地方。前向和反向 KL 对金标签 logit 的梯度,在高熵事实上大约只有 NTP 的 0.5 倍。
Switch Distillation 用教师熵做路由:batch 里熵最低的 q=20% token 走反向 KL,其余走交叉熵。路由信号来自蒸馏本来就要算的教师 logits,不增加前向。反向 KL 的寻模行为和低熵教师分布匹配。
中训后,7B 教师的 Switch Distillation 把推理宏平均从 NTP 的 26.1% 拉到 44.7%,13B 教师到 42.1%,即 1.71 倍和 1.61 倍。事实回忆 29.3% 对 NTP 的 30.3%,保留 96.7–96.8%。知识与常识 49.3% / 46.5%,约 1.13–1.19 倍。同一 α=0.5 的前向、反向蒸馏和 TRKD 都更偏推理、更伤事实,没有点能在中训阶段同时帕累托压过 NTP。7B 教师整体好于 13B,符合容量差距会伤蒸馏的旧结论。
后训练走完 SFT、DPO 和两轮 RLVR 之后,7B 教师的推理宏平均到 50.6%,是 NTP 后训练 38.5% 的 1.32 倍;13B 教师 48.0%,约 1.25 倍。事实缺口被补上,Switch Distillation 最终事实还略高于 NTP,知识与常识仍高 1.13–1.20 倍。消融里换前向 KL 推理掉 2.9 个点;按教师是否答对、按域、按随机掩码路由都明显更差;只用教师 top-1 当标签,事实小涨、推理掉 6.4 个点。熵路由是主因,软分布比硬标签更有用。
SmolLM2 上能看到同一方向的阶段差,不单是 OLMo-2 的配方巧合。
中训正在变成前沿配方里的固定阶段,Llama 3、SmolLM、OLMo-2 都有。蒸馏如果按预训练的用法直接灌进去,会系统性地少学高熵事实。Switch Distillation 是几乎零额外成本的改法:同一套教师 logits,按熵把监督切成「信教师」和「信语料」。对正在做中训的团队,这是可以明天就试的目标函数,不必改数据配比。
它解决的是固定数据上怎么抽监督,和改数据池的工作是互补的。容量差距仍然在:7B 教 1B 比 13B 教 1B 更合适。
受控实验几乎全在 OLMo-2 1B 学生上,因为同时公开中间检查点、中训数据和完整配方的家族极少。SmolLM2 只提供定性复现。教师固定为 Instruct,学生固定为 1B,这些选择没有穷尽网格。路由是按 batch 分位数的硬开关,q=20% 来自小扫描,没有学过的路由器。作者猜想晚期预训练也会出现同类权衡,这篇没有直接测。事实组主要是 TriviaQA、Natural Questions、SimpleQA,覆盖不到长尾知识的全部形态。