← 返回 2026-05-15

面向 LLM 推理的自蒸馏中自适应教师暴露 Adaptive Teacher Exposure for Self-Distillation in LLM Reasoning

Zihao Han, Tiangang Zhang, Huaibin Wang, Yilun Sun 📅 2026-05-12 👍 7 2026-07-13 08:36
OPSD 强化学习控制器 推理后训练 教师暴露 自蒸馏

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}$ 锚定教师对参考解的可预测性,确保不会被'全屏蔽'式地利用低暴露刷分。

Overview of ATESD. The OPSD backbone samples student continuations from the problem-only prompt. Given an exposure action αt, ATESD truncates only the reasoning prefix of the privileged solution, preserves the boxed answer, and builds the teacher context from the problem, exposed reference solution, and transition prompt.
Figure 3: Overview of ATESD. The OPSD backbone samples student continuations from the problem-only prompt. Given an exposure action αt, ATESD truncates only the reasoning prefix of the privileged solution, preserves the boxed answer, and builds the teacher context from the problem, exposed reference solution, and transition prompt.

实验结果

主结果在 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。

Main results on competition-level mathematical reasoning benchmarks. We follow the OPSD reporting protocol and report Average@12 accuracy (%) under the Qwen3 sampling configuration. Baseline numbers are from Zhao et al. [35]; ATESD is evaluated with the same within-100-step checkpoint selection convention. Best results are in bold; second-best results are underlined.
Table 1: Main results on competition-level mathematical reasoning benchmarks. We follow the OPSD reporting protocol and report Average@12 accuracy (%) under the Qwen3 sampling configuration. Baseline numbers are from Zhao et al. [35]; ATESD is evaluated with the same within-100-step checkpoint selection convention. Best results are in bold; second-best results are underlined.
Controller ablations on AIME 2024 Average@12 (%). (A) Delayed credit for learning exposure; (B) learned exposure versus fixed or uncontrolled alternatives. Both subtables use the same evaluation setting, isolating controller design choices rather than data or checkpoint-budget changes.
Table 2: Controller ablations on AIME 2024 Average@12 (%). (A) Delayed credit for learning exposure; (B) learned exposure versus fixed or uncontrolled alternatives. Both subtables use the same evaluation setting, isolating controller design choices rather than data or checkpoint-budget changes.
Mechanism ablations for exposure control. (A) Keeping the problem, student rollout, and scoring positions fixed, reducing teacher exposure lowers the token-level teacher–student KL spikes on a positive trajectory. (B) The learned Beta policy evolves from broad exploration toward a structured exposure distribution rather than collapsing to either no reference or full exposure.
Figure 4: Mechanism ablations for exposure control. (A) Keeping the problem, student rollout, and scoring positions fixed, reducing teacher exposure lowers the token-level teacher–student KL spikes on a positive trajectory. (B) The learned Beta policy evolves from broad exploration toward a structured exposure distribution rather than collapsing to either no reference or full exposure.
查看结构化数据
任务指标本文基线提升
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 级别资源。