面向 LLM 推理的自蒸馏中自适应教师暴露 Adaptive Teacher Exposure for Self-Distillation in LLM Reasoning
ATESD:用 Beta 策略控制器把教师暴露比例变成可学习变量,动态调节特权 CoT 的可见量。
前置知识
On-Policy Self-Distillation (OPSD)
单模型自蒸馏:同一模型既是教师也是学生。学生从问题 $x$ 采样轨迹 $\hat y$,教师以参考解 $y^\star$ 为条件提供 per-token 监督,损失是 forward KL。
本文直接建立在 OPSD 之上,并质疑其'教师总能看完整参考解'的默认假设,ATESD 正是要打破这个默认。
Beta 分布策略 (Beta-policy controller)
定义在 $[0,1]$ 上的连续分布,参数 $(a,b)$ 控制形状。论文用 2 层 MLP 把训练状态映射到 $a_t,b_t>1$,采样 $\alpha_t\sim\mathrm{Beta}(a_t,b_t)$。
把'教师看到多少参考解'从离散选择变为连续可微手段,比固定超参或 ε-greedy 更贴合语义。
REINFORCE 与延迟信用分配
REINFORCE 用采样动作的总回报更新策略。在本文中 $\alpha$ 对学生的影响要 $L$ 步后才显现,故用折扣 look-ahead 回报 $G_{lp}=\sum_i\gamma^{i-1}\max(0,\ell_{t+i-1}-\ell_{t+i})$。
不解决延迟信用,控制器只会学 myopic 规则。Table 2 (A) 显示从即时奖励 52.22 提升到全延迟 59.17 的关键。
Token-level forward KL 失配
度量 $p_T^\alpha$ 与 $p_S$ 在同条轨迹上每位的差距 $M(\alpha)=\mathbb{E}\frac{1}{|\hat y|}\sum_n\mathrm{KL}(p_T^\alpha\|p_S)$。$\alpha$ 越大教师越尖锐,KL 越大。
Figure 2B 显示 KL 随 $\alpha$ 单调增长,是'教师侧暴露失配'的可量化证据。
研究动机
OPSD 及其后续变体虽然消除了学生侧的分布失配,但都默认教师始终看到完整的参考解 $y^\star$——这一'全暴露'假设从未被严肃质疑。论文通过在 AIME 2024 上用 Qwen3-1.7B 跑 3 颗种子的对照实验发现:固定 $\alpha=0.5$ 反而比 $\alpha=1.0$ 表现更好,'全暴露'并非最优。更糟的是,把 $\alpha$ 从 0 扫到 1,per-token forward KL 尾值与 top-1 预测分歧都单调增长——即教师参考解越长,监督目标越尖锐,越超出学生当前可吸收范围。Figure 2C 的难度分桶分析更揭示:简单题偏好 $\alpha=1.0$,中等题偏好 0.5,难题偏好 0.0,'一刀切'的全暴露在所有学习阶段都不合适。
本文的目标是把教师暴露比例 $\alpha\in[0,1]$ 从固定超参升级为训练时可学习、随训练状态动态调节的控制变量,使监督强度实时匹配学生当前能力,从而在 OPSD 框架下取得更稳定的推理后训练效果。具体目标包含三点:在 AIME 24/25、HMMT 25 三个数学竞赛基准上稳定超过 OPSD 原始报告,在 Qwen3-1.7B/4B/8B 三种规模上均取得正收益,并通过消融证明延迟信用分配与学到的暴露策略都是必要组件。换言之,ATESD 想证明'教师特权信息量'本身就是一个值得被端到端学习的新维度,而不是被工程师手工拍脑袋调的常数。
与已有工作不同的是,已有工作要么调学生侧 (scheduled sampling、DAgger、importance reweighting) 来缓解分布失配,要么用静态课程/动态温度去改变蒸馏目标,但都把教师特权上下文当成不可调的黑箱。ATESD 是第一个把'教师能看到多少参考解'本身当作决策变量的方法,通过一个 Beta 策略控制器配合延迟信用分配来端到端学习这个控制信号,而不是靠手工网格搜索或事后选最优固定值。
核心方法
ATESD 保留 OPSD 的学生采样管线不变,只替换教师的特权上下文:把参考解的推理前缀按比例 $\alpha_t$ 截断,但始终保留最终 boxed 答案。然后用一个轻量 Beta 策略控制器 $\pi_\phi(\alpha|s_t)$,每隔 $H$ 步根据六个训练状态统计量 (训练进度、损失 EMA、暴露 EMA、错配 EMA、probe-NLL EMA、批次学生置信度) 采样一个 $\alpha_t$ 作用于整段 hold 窗口。控制器的训练采用双时间尺度:学生每步更新蒸馏梯度,控制器则在 $L=20$ 步 look-ahead 后才用 REINFORCE 配合折扣学习进步奖励 $R=G_{lp}+\lambda_g G_{gt}$ 一次性回填信用。直觉上:先告诉系统'暂时只给教师看一半的解',再观察学生后续 $L$ 步是真进步还是停滞,最后用结果反推这个暴露决策好不好。
把 on-policy 蒸馏中的教师特权信息量从二元 (有/无) 提升为连续可学习控制变量,并用延迟信用分配 (discounted lookahead reward + REINFORCE) 训练一个训练状态条件化的 Beta 策略控制器。和 OPSD 的本质区别不在损失函数,而在教师上下文生成方式:$\alpha=1$ 时退化为 OPSD,$\alpha=0$ 时退化为'只给答案',中间值由策略在线决定。
方法步骤详情
七步流程:(1) 截断算子按比例切推理前缀,保留 boxed 答案;(2) 教师 prompt $q_T^{\alpha_t}=[x;\tilde y^\star_{\alpha_t};\tau]$;(3) 学生从 $x$ 采样 $\hat y\sim p_S(\cdot|x)$,teacher-force 同条轨迹通过学生与教师上下文,按 Eq.(6) per-token KL 更新 $\theta$;(4) 2 层 MLP $f_\phi$ 把 6 维状态 $s_t$ 映射到 $a_t,b_t>1$,采 $\alpha_t\sim\mathrm{Beta}(a_t,b_t)$ 后 clip;(5) hold $\alpha_t$ 固定 $H$ 步;(6) $L=20$ 步后用 $\max(0,\ell_{t+i-1}-\ell_{t+i})$ 折扣求和得 $G_{lp}$,加 $G_{gt}$ 组合 $R$;(7) 中心化优势 $A_j$,按 Eq.(9) 一次 REINFORCE 更新 $\phi$。
技术新颖性
新颖性主要在三方面。(1) 概念上首次把'教师特权信息量'形式化为连续可学习控制变量,并把'监督端失配'和'学生端失配'对称地放上议程。(2) 训练范式上引入 hold/lookahead 双时间尺度:学生快速更新、控制器慢周期 REINFORCE,避开了把 on-policy 蒸馏当作 i.i.d. 监督学习时常见的 myopic 错误。(3) 奖励设计上用 $\max(0, \ell_{t+i-1}-\ell_{t+i})$ 的非负学习进步代替原始 loss 差,避免了'损失下降但泛化退步'的虚假信号;再加 $G_{gt}$ 锚定教师对参考解的可预测性,确保不会被'全屏蔽'式地利用低暴露刷分。
实验结果
主结果在 AIME 24/25、HMMT 25、Qwen3-{1.7B,4B,8B} 上用 Average@12 报告。ATESD 在三基准均值上取得 44.35 (+0.95)、65.65 (+2.05)、67.13 (+2.33),均超 OPSD,4B/8B 提升更显著。Figure 4A 把 $\alpha$ 从 1.0 降到 ATESD 学到的 0.3,mean KL 从 0.0136 降到 0.0061、max KL 从 0.2432 降到 0.0645、第 26 位 token 尖峰从 0.2432 跌到 0.0098,证明暴露控制确实削减过强教师监督。Figure 4B 显示 Beta 策略从较宽分布逐步集中到中段,没有塌缩到 0 或 1,证明学到真策略。Table 2 (A) 信用消融:即时 52.22 → 短延迟 56.11 → 折扣 look-ahead 58.06 → 全延迟 59.17;Table 2 (B) 策略消融:OPSD 满暴露 57.20、最优固定 57.44、随机 54.94、学到的策略 59.17。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME 2024 (Qwen3-1.7B) | Average@12 准确率 (%) | 59.17±0.8 | OPSD 57.2 | +1.97 |
| AIME 2024 (Qwen3-4B) | Average@12 准确率 (%) | 78.06±0.43 | OPSD 76.4 | +1.66 |
| AIME 2024 (Qwen3-8B) | Average@12 准确率 (%) | 80.56±0.70 | OPSD 77.8 | +2.76 |
| AIME 2025 (Qwen3-1.7B) | Average@12 准确率 (%) | 44.72±0.28 | OPSD 43.9 | +0.82 |
| AIME 2025 (Qwen3-4B) | Average@12 准确率 (%) | 71.39±0.48 | OPSD 68.3 | +3.09 |
| AIME 2025 (Qwen3-8B) | Average@12 准确率 (%) | 72.50±0.32 | OPSD 70.8 | +1.70 |
| HMMT 2025 (Qwen3-1.7B) | Average@12 准确率 (%) | 29.17±1.37 | OPSD 29.2 | -0.03 (持平) |
| HMMT 2025 (Qwen3-4B) | Average@12 准确率 (%) | 47.50±0.89 | OPSD 46.1 | +1.40 |
| HMMT 2025 (Qwen3-8B) | Average@12 准确率 (%) | 48.33±0.28 | OPSD 45.8 | +2.53 |
| 三基准平均 (Qwen3-1.7B) | Average@12 准确率 (%) | 44.35±0.23 | OPSD 43.4 | +0.95 |
| 三基准平均 (Qwen3-4B) | Average@12 准确率 (%) | 65.65±1.04 | OPSD 63.6 | +2.05 |
| 三基准平均 (Qwen3-8B) | Average@12 准确率 (%) | 67.13±0.32 | OPSD 64.8 | +2.33 |
局限与改进
作者明确承认三点限制:(1) 控制器是全局粒度的——一个 hold 窗口内对所有样本共用同一个 $\alpha$,缺乏逐题/逐难度的差异化暴露;(2) 延迟信用依赖固定的 look-ahead 窗口 $L$,没有用反事实或基于模型的奖励估计;(3) 实验只覆盖数学竞赛基准,未在代码生成、科学推理、对话等更广推理域验证。我自己的观察还包括:(a) 1.7B 上 HMMT 25 几乎与 OPSD 持平 (29.17 vs 29.2),说明在容量有限时暴露控制的红利被压扁;(b) 控制器只用 6 维状态统计量与一个 2 层 MLP,状态表征能力有限,可能错过细粒度的训练信号;(c) 论文没有分析不同 hold 长度 $H$、look-ahead $L$、折扣 $\gamma$ 对最终性能的影响,只说'all main runs use L=20'。
独立分析的弱点
几个独立判断的弱点。第一,'全局共享 $\alpha$'与 Figure 2C 揭示的'不同难度需要不同暴露'之间的张力没解决,控制器的 $\alpha$ 是 batch-level scalar,无法对难题自动调到 $\alpha\to 0$;改进方向是按学生置信度或 RFT 难度分桶做 per-sample $\alpha$。第二,look-ahead 窗口 $L=20$ 是硬编码超参,训练早期/晚期最优 $L$ 可能不同;改进方向是 schedule $L$ 或用 TD($\lambda$)。第三,奖励中 $\max(0, \ell_{t+i-1}-\ell_{t+i})$ 鼓励'损失上升被避免',但 on-policy KL 自身波动大,可能奖励噪声;可考虑用验证集准确率等更稳定的信号做 credit。第四,控制器训练用 6 维人工特征,未来可换成从模型内部表征派生的特征,去除手工设计。第五,方法只在 OpenThoughts 数学数据上、100 步预算内训练,缺少与 GRPO/RLHF 阶段级联时的兼容性证据。
未来方向
作者提了三类未来方向:(1) 把暴露控制从全局扩展到 per-sample 或难度感知,并明确指出 Figure 2C 的分桶分析是这条路的强动机。(2) 用反事实估计 (counterfactual) 或 model-based reward 替代固定 $L$ 窗口的 look-ahead,以更好地做延迟信用分配。(3) 在更大模型、代码生成、科学推理上验证泛化性。基于结果可进一步延伸:(a) 把 $\alpha$ 进一步分解为 reasoning-prefix 截断比例 + answer 注入时机两个维度;(b) 与 DPO/GRPO 等偏好/RL 训练阶段级联,研究是否需要不同的 $\alpha$ 调度;(c) 探索 Beta 之外的分布族 (e.g. Truncated Normal) 在更高方差训练阶段是否更稳;(d) 引入对训练过程不可逆的 safety bound,避免控制器把 $\alpha$ 推到 0 学到'全屏蔽教师'的 trivial solution。
复现评估
复现难度中等。论文在方法部分给出完整算法描述 (Eq. 5-9)、控制器结构 (2 层 MLP、6 维状态、$L=20$)、训练细节 (100 步 budget、LoRA、on-policy 采样、OPSD 同等 student rollout),并把 ATESD 表格与 OPSD 原始报告对齐。但有潜在阻力:(1) Baseline 行直接抄自 OPSD 论文 (Zhao et al.) 而非自跑,跨论文对比可能有采样温度/超参差异;(2) 训练状态 6 维特征的具体计算公式没有全部展开;(3) 没有公开代码或权重链接,控制器 checkpoint、reward 曲线、状态 EMA 衰减系数等关键实现细节缺失;(4) 训练需同时维护 student、teacher、controller 三个计算图,显存与 LoRA rank 选择需要自己摸索。算力上,Qwen3-8B + LoRA + 100 步 on-policy 蒸馏大约需要 8 卡 A100 级别资源。
论文图表
左半 (A) 展示两个场景:简单题 'Solve 2+3' 时教师特权 CoT 是 '从 2 数三个',与学生能力匹配,蒸馏成功;难题 'Solve $2x^2+5x-3=0$' 时教师特权 CoT 包含完整的判别式与求根公式,远超学生当前水平,产生不可吸收的监督目标。右半 (B) 展示 ATESD 引入一个 Beta 策略控制器 $\pi_\phi$,按学习进步奖励选择 $\alpha\in[0,1]$,把教师特权 CoT 限制在学生可消化的范围内。
这是论文的核心概念图,把抽象的'教师侧暴露失配'翻译成两个具体例子,让读者立刻理解为什么需要自适应暴露。
三个子图全部基于 3 种子、Qwen3-1.7B 在 AIME 2024 上的对照实验。(A) 展示固定 $\alpha\in\{0,0.25,0.5,0.75,1.0\}$ 下的平均准确率:$\alpha=0.5$ 反而是峰值,$\alpha=1.0$ 反而次优。(B) 画 KD loss tail 与 top-1 失配随 $\alpha$ 单调上升。(C) 按 easy/medium/hard 难度分桶画最佳网格 $\alpha$:简单题 1.0、中等题 0.5、难题 0.0。
这是论文诊断阶段的实验铁证,三幅子图分别支撑了 '全暴露非最优'、'失配单调增长'、'暴露应随学习阶段变化' 三个核心论点。