> ## Documentation Index
> Fetch the complete documentation index at: https://agent-rl.tacetai.com/llms.txt
> Use this file to discover all available pages before exploring further.

# 16. 面向智能体的蒸馏

第 15 章用了示范、筛选后的成功轨迹和偏好对，但这些信号只告诉学生某条路径该怎样走，或者哪次尝试更好。如果教师能返回完整的下一 token 分布，我们还能知道未选中的候选有多大概率。**蒸馏**（distillation）利用这类分布信息训练学生。本章先区分教师生成与学生生成的前缀，再推导相应的学习目标。

蒸馏常用于把较大模型的能力迁移到较小模型。用于智能体时，还要先回答一个关键问题：*我们在谁的前缀上训练？* 智能体的前缀包含它自己之前的动作所产生的观测。只使用教师轨迹训练，通常覆盖不到学生自己犯错后会遇到的前缀。本章在 BugFix gym 上构建两种形式的蒸馏，推导逐 token 的反向 KL（reverse KL）信号及其梯度，直接测量暴露偏差（exposure bias），并在相同的学生步数下比较蒸馏与结果 RL（outcome RL）。

```python theme={null}
#@save
import torch
import torch.nn.functional as F
import agentrl as arl
```

```python theme={null}
import copy
import random
from collections import Counter
tasks = arl.make_tasks('all')
tok = arl.build_tokenizer(tasks)
train, heldout = arl.make_tasks('train'), arl.make_tasks('heldout')
oracle_demos = [arl.scripted_rollout(arl.scripted_policy('oracle'), arl.BugFixEnv(), t, tok) for t in train]
weak_play = [arl.scripted_rollout(arl.scripted_policy(kind, seed=s), arl.BugFixEnv(), t, tok)
             for kind in ['blind', 'random'] for s in range(4) for t in train]
print('oracle demos:', len(oracle_demos), '| weak-play episodes:', len(weak_play))
```

Output:

```text theme={null}
oracle demos: 20 | weak-play episodes: 160
```

上面的代码块重建了 [15.2 节](/tutorial/zh/chapters/15-sft-and-rejection-sampling#15-2-热启动：先格式，后技能)的数据：每个训练任务一条 oracle 示范，以及 160 个格式正确但没有技能的行为回合。

## 16.1 off-policy 蒸馏与 on-policy 蒸馏

设 $q$ 为教师，$p_\theta$ 为学生，二者都是同一 token 词表（vocabulary）上的分布。要在智能体轨迹上让 $p_\theta$ 接近 $q$，有两种自然的做法。

**off-policy 蒸馏**（off-policy distillation）从教师采样回合，并用 [15.1 节](/tutorial/zh/chapters/15-sft-and-rejection-sampling#15-1-轨迹-sft-与动作掩码)的带掩码 SFT 损失在这些回合上微调学生。由于环境项相互抵消，它的期望损失为

$$
\mathbb{E}_{\tau\sim q}\big[-\log p_\theta(\tau)\big] = \mathrm{KL}\big(q(\tau)\,\|\,p_\theta(\tau)\big) + H\big(q(\tau)\big),
$$

即轨迹分布之间的**前向** KL（forward KL）加上一个常数。前向 KL 倾向于 *mode-covering*：教师赋予较高概率的 token 若被学生忽略，学生就会受到惩罚。问题出在下标上。期望是对教师的回合取的，所以学生只在教师访问过的前缀上接受训练。

**on-policy 蒸馏**（on-policy distillation）从学生采样回合，在每个学生前缀 $c_i$ 上向教师询问其分布，并在那里最小化逐 token 的散度。本书的默认选择是逐 token 的**反向** KL，

$$
\mathcal{L}_{\text{OPD}}(\theta) = \mathbb{E}_{\tau\sim p_\theta}\Big[\sum_{i\in\text{actions}(\tau)} \mathrm{KL}\big(p_\theta(\cdot\mid c_i)\,\|\,q(\cdot\mid c_i)\big)\Big],
$$

它具有 *mode-seeking* 的倾向：学生在教师概率很小的地方放置概率质量会受到惩罚。现在前缀 $c_i$ 来自学生，包括它自身错误所产生的观测，所以教师的纠正恰好出现在学生部署时会到达的位置。on-policy 蒸馏是 DAgger 的软标签版本；DAgger 是一种模仿学习算法，它在学习者自己访问的状态上向专家查询（[Ross, Gordon, and Bagnell, 2010](https://arxiv.org/abs/1011.0686)）。该论文的动机同样适用：在序列预测中，未来的输入取决于过去的输出，这破坏了监督学习背后的独立同分布假设。

只有当教师知道在学生的前缀上（包括糟糕的前缀）该怎么做时，on-policy 蒸馏才有用。因此我们构建一个见过错误的教师。`with_mistake` 包装 oracle，使它在某一轮犯一个给定的错误，然后就像那一轮没有发生过一样继续。在得到的示范中，出错的那一轮被从损失中掩码掉，所以教师学到的是*从*错误*中恢复*，而不是犯这个错误。

```python theme={null}
def with_mistake(mistake, at):
    oracle = arl.scripted_policy('oracle')
    def policy(observations):
        k = len(observations) - 1
        if k == at:
            return mistake
        return oracle(observations if k < at else observations[:at + 1] + observations[at + 2:])
    return policy

def noop_op(task):
    """A wrong repair op that leaves this task's source unchanged."""
    return next(op for op in arl.FIX_OPS if arl.EDIT_OPS[op](task.source, task.public_tests) == task.source)

def recovery_demo(task, mistake, at):
    traj = arl.scripted_rollout(with_mistake(mistake, at), arl.BugFixEnv(), task, tok)
    traj.action_mask = traj.action_mask * (traj.turn[1:] != at)       # no loss on the mistake itself
    return traj
```

三种错误，各在第 0、1、2 轮注入，共得到 180 条恢复示范：一个没有终止的调用、一个参数无效的调用，以及一个什么也没改变的错误编辑。下面是其中一条。

```python theme={null}
mistakes = lambda t: ['<call> read', '<call> submit read </call>', f'<call> edit {noop_op(t)} </call>']
recovery_demos = [recovery_demo(t, m, at) for t in train for at in range(3) for m in mistakes(t)]
example = recovery_demos[5]
print('demos:', len(recovery_demos), '| all reward 1:', all(d.reward == 1.0 for d in recovery_demos))
print(tok.decode(example.ids[7:].tolist()).replace('\n', ' '))
print('masked-out tokens:', tok.decode(example.ids[1:][(example.action_mask == 0) & (example.owner[1:] == 2)].tolist()))
```

Output:

```text theme={null}
demos: 180 | all reward 1: True
<call> read </call> <obs> def clamp ( value , lower , upper ) : return min ( lower , max ( value , upper ) ) </obs> <call> edit fix_range_end </call> <obs> no change </obs> <call> edit swap_minmax </call> <obs> ok </obs> <call> test </call> <obs> pass ( -6 , -3 , -3 ) pass ( 5 , -2 , 3 ) pass ( 1 , -3 , 2 ) </obs> <call> submit </call>
masked-out tokens: <call> edit fix_range_end </call>
```

在这条 `clamp` 示范中，错误出现在第 1 轮：读取之后，专家应用了 `fix_range_end`，它没有改变源代码（`no change`），随后专家继续执行正确的 `swap_minmax`、一次测试和一次提交。只有出错调用的四个 token 被从损失中掩码掉。全部 180 条示范的奖励都是 1。

教师是一个标准配置的 TinyAgent（157k 参数），先在弱行为上训练以掌握格式，再在 oracle 示范和恢复示范上训练。另一个**干净教师**（clean teacher）只看过 20 条 oracle 示范；16.4 节用它说明恢复数据为什么重要。**学生**是一个更小的 TinyAgent，宽度 32、只有一层，只在弱行为上训练过，对应一个懂格式但没有技能的小型预训练模型。

```python theme={null}
torch.manual_seed(0)
teacher = arl.TinyAgent(tok.vocab_size)
arl.sft_train(teacher, weak_play, tok, steps=150)
clean_teacher = copy.deepcopy(teacher)
arl.sft_train(clean_teacher, oracle_demos, tok, steps=150)
arl.sft_train(teacher, oracle_demos + recovery_demos, tok, steps=300)
torch.manual_seed(1)
base = arl.TinyAgent(tok.vocab_size, d_model=32, n_layers=1, n_heads=2)
arl.sft_train(base, weak_play, tok, steps=150)
evaluate = lambda m: (arl.pass_rate(m, train, tok, 4, seed=11)[0], arl.pass_rate(m, heldout, tok, 4, seed=12)[0])
for name, m in [('teacher', teacher), ('clean teacher', clean_teacher), ('student base', base)]:
    n_params = sum(p.numel() for p in m.parameters())
    print(f'{name:<14} params {n_params:>7,} | train/held-out pass {evaluate(m)[0]:.2f}/{evaluate(m)[1]:.2f}')
```

Output:

```text theme={null}
teacher        params 157,244 | train/held-out pass 0.99/1.00
clean teacher  params 157,244 | train/held-out pass 0.99/0.90
student base   params  41,500 | train/held-out pass 0.14/0.05
```

两个教师在训练任务上的通过率都是 0.99，带恢复数据的教师通过了全部留出任务（干净教师为 0.90）。学生基座在训练集上起点为 0.14，在留出集上为 0.05，等于或低于猜测水平，参数量约为教师的四分之一（41,500 对 157,244）。

## 16.2 作为稠密信号的逐 token 反向 KL

结果 RL 每个回合给出一个数。on-policy 蒸馏为每个动作 token 给出一个数 $\mathrm{KL}(p_\theta(\cdot\mid c_i)\,\|\,q(\cdot\mid c_i))$，并且无论回合是否成功都能得到。要使用它，我们需要它的梯度。

固定一个前缀 $c$，记 $p_v = p_\theta(v\mid c)$，$q_v = q(v\mid c)$。于是 $\mathrm{KL}(p\|q) = \sum_v p_v(\log p_v - \log q_v)$，利用 $\nabla p_v = p_v \nabla\log p_v$，

$$
\nabla_\theta\,\mathrm{KL}(p\,\|\,q) = \sum_v p_v \nabla\log p_v\,(\log p_v - \log q_v) + \sum_v p_v \nabla \log p_v
= \mathbb{E}_{y\sim p}\big[(\log p_y - \log q_y)\,\nabla_\theta\log p_y\big],
$$

因为第二个求和是 $\nabla\sum_v p_v = 0$。这是单个 token 上的策略梯度（policy gradient），其奖励为 $r_y = \log q_y - \log p_y$，即负的对数比，只是要最小化而不是最大化。由此有两种实现。**采样 token** 形式只使用学生实际生成的 token $y$ 以及教师对它的对数概率，把 $r_i = \log q(y_i\mid c_i) - \log p_\theta(y_i\mid c_i)$ 当作普通策略梯度损失中的逐 token 优势（advantage）。它每个 token 只需要教师的一个数。**全词表**形式精确计算对 $v$ 的求和，并让自动微分对其求导。它在单个 token 上没有采样噪声，但在每个位置都需要教师的完整分布。

这个推导在*固定*前缀下成立。完整目标 $\mathcal{L}_{\text{OPD}}$ 还通过学生访问哪些前缀而依赖于 $\theta$，这会额外引入一个得分函数（score function）项，其中每个 token 的学习信号还包括后续 token 的 KL（逐 token KL 的剩余回报，return-to-go）。实际做法会丢掉它。Thinking Machines 的实现使用“折扣因子为零”（“a discount factor of zero”），所以每个 token 只使用当前位置的 KL 信号（16.6 节）。GKD 式训练在学生生成的序列上计算逐 token 散度，并且不对采样过程求导。我们遵循同样的约定，并明确说明它是一个近似：它优化的是学生已访问前缀上的局部分布差异，没有显式优化早期动作对未来 KL 的影响。

在 logit 层面，两个 KL 方向也不同。对学生的 logits $z$ 求导，

$$
\frac{\partial\,\mathrm{KL}(p\|q)}{\partial z_v} = p_v\big(\log p_v - \log q_v - \mathrm{KL}(p\|q)\big), \qquad \frac{\partial\,\mathrm{KL}(q\|p)}{\partial z_v} = p_v - q_v .
$$

反向 KL 在一个 token 上的梯度与学生自己的概率 $p_v$ 成正比。如果学生几乎不给教师偏好的 token 分配概率质量，反向 KL 几乎拉不动它；它主要是压低学生当前偏好的东西。前向 KL 则无论如何都以 $q_v - p_v$ 的力度把教师的 token 拉上来。下一个代码块通过枚举验证采样 token 恒等式，并针对一个错过了教师众数的学生打印两种 logit 梯度。

```python theme={null}
student_z = torch.tensor([3.0, 0.0, -1.0, -3.0], dtype=torch.float64, requires_grad=True)   # [V]
teacher_z = torch.tensor([-2.0, 0.0, 0.0, 3.0], dtype=torch.float64)                       # mode = token 3
logp, logq = student_z.log_softmax(-1), teacher_z.log_softmax(-1)
reverse = (logp.exp() * (logp - logq)).sum()
g_exact, = torch.autograd.grad(reverse, student_z, retain_graph=True)
surrogate = (logp.exp().detach() * (logp - logq).detach() * logp).sum()   # E_y~p[(log p - log q) grad log p]
g_sampled, = torch.autograd.grad(surrogate, student_z, retain_graph=True)
forward = (logq.exp() * (logq - logp)).sum()
g_forward, = torch.autograd.grad(forward, student_z)
print('p:', [round(v, 3) for v in logp.exp().tolist()], '| q:', [round(v, 3) for v in logq.exp().tolist()])
print('reverse KL grad:', [round(v, 4) for v in g_exact.tolist()], '| matches sampled form:', torch.allclose(g_exact, g_sampled))
print('forward KL grad:', [round(v, 4) for v in g_forward.tolist()])
```

Output:

```text theme={null}
p: [0.934, 0.047, 0.017, 0.002] | q: [0.006, 0.045, 0.045, 0.904]
reverse KL grad: [0.3369, -0.2158, -0.0965, -0.0246] | matches sampled form: True
forward KL grad: [0.928, 0.0015, -0.0279, -0.9016]
```

学生把 0.93 的概率质量放在 token 0 上，教师把 0.90 放在 token 3 上。梯度下降使每个 logit 沿其梯度的反方向移动。在反向 KL 下，最大的分量 +0.34 把 token 0 往下压，而教师的众数 token 3 只得到 $-0.025$，因为学生只给它 0.002 的概率。在前向 KL 下，token 3 得到 $-0.90$，几乎就是完整的差距 $q_3 - p_3$。两种梯度最终都会把学生移到教师那里，但反向 KL 先清除学生自身的偏好，然后才找到教师的众数。枚举得到的得分函数形式与精确的反向 KL 梯度一致，与推导相符。

`distill_loss` 保存全词表版本。它接收同一位置上的学生与教师 logits（与全书一样，对齐到被预测的位置），把教师从计算图中分离，在任一方向上计算逐位置散度，并在动作掩码上取平均。

```python theme={null}
#@save
def distill_loss(student_logits, teacher_logits, mask, kind='reverse_kl'):
    """Per-token KL between student and (detached) teacher: logits [B, T, V], mask [B, T] -> masked mean."""
    logp = torch.log_softmax(student_logits.float(), dim=-1)            # [B, T, V]
    logq = torch.log_softmax(teacher_logits.detach().float(), dim=-1)    # [B, T, V]
    if kind == 'reverse_kl':        # KL(p || q): mode-seeking, evaluated on student prefixes
        per_token = (logp.exp() * (logp - logq)).sum(-1)
    elif kind == 'forward_kl':      # KL(q || p): mode-covering, classic logit distillation
        per_token = (logq.exp() * (logq - logp)).sum(-1)
    else:
        raise ValueError(f'unknown kind {kind!r}')
    return (per_token * mask).sum() / mask.sum()                        # [B, T] -> scalar
```

在真实数据上做一个快速检查：在教师采样的回合上，教师与自身的逐 token KL 在两个方向上都是 0，而教师与学生基座之间的 KL 很大且不对称。

```python theme={null}
_, teacher_eps = arl.pass_rate(teacher, train, tok, 2, seed=3)
batch = arl.collate(teacher_eps, tok.pad_id)
with torch.no_grad():
    t_logits, s_logits = teacher(batch['ids'])[:, :-1], base(batch['ids'])[:, :-1]     # [B, T-1, V]
for kind in ['reverse_kl', 'forward_kl']:
    print(f"{kind}: teacher vs itself {distill_loss(t_logits, t_logits, batch['action_mask'], kind):.4f}"
          f" | student base vs teacher {distill_loss(s_logits, t_logits, batch['action_mask'], kind):.3f}")
```

Output:

```text theme={null}
reverse_kl: teacher vs itself 0.0000 | student base vs teacher 3.044
forward_kl: teacher vs itself 0.0000 | student base vs teacher 1.334
```

在相同的位置上，两个方向相差超过两倍（每个动作 token 3.04 对 1.33 nats）。只要学生对教师认为不太可能的东西很有把握，反向 KL 就会很大，而一个没有技能但格式良好的学生大体正是如此。

## 16.3 实验：蒸馏与结果 RL 的比较

现在我们训练学生基座的三个副本，各训练 40 步，批次结构相同：每步 4 个训练任务，每个任务 8 个回合（每步 32 个回合）。

* **off-policy 蒸馏**：教师采样 32 个回合，学生在其上走一步 `sft_loss`。
* **on-policy 蒸馏**：学生采样 32 个回合，教师在一次前向传播中为每个前缀打分，学生走一步 `distill_loss`（反向 KL）。
* **结果 RL**：学生采样，使用带逐任务组基线的 REINFORCE 和隐藏测试奖励，与 [15.5 节](/tutorial/zh/chapters/15-sft-and-rejection-sampling#15-5-实验：从同一个弱基座出发的四种做法)相同。[第 9 章](/tutorial/zh/chapters/09-grpo-and-variants)给出完整的讨论。

两种蒸馏方法的学习率都是 $10^{-2}$，RL 为 $3\times 10^{-3}$。撰写本章时，我们在同一随机种子下为每种方法尝试了 $3\times10^{-3}$、$10^{-2}$ 和 $3\times10^{-2}$。这些是最佳取值，on-policy 蒸馏在 $3\times10^{-2}$ 时坍缩到通过率 0。这次扫描不会由检查器重新运行。

```python theme={null}
def student_step(student, opt, method, rng, gen, group=8):
    """One update of `method` in {'off', 'on', 'rl'} on 4 tasks x group episodes -> teacher tokens used."""
    batch_tasks = [t for t in rng.sample(train, 4) for _ in range(group)]
    sampler = teacher if method == 'off' else student
    trajs = arl.rollout_batch(sampler, [arl.BugFixEnv() for _ in batch_tasks], batch_tasks, tok, generator=gen)
    batch = arl.collate(trajs, tok.pad_id)
    mask = batch['action_mask']                                                    # [B, T-1]
    if method == 'off':
        loss = arl.sft_loss(student, batch)
    elif method == 'on':
        with torch.no_grad():
            t_logits = teacher(batch['ids'])[:, :-1]                               # [B, T-1, V]
        loss = distill_loss(student(batch['ids'])[:, :-1], t_logits, mask)
    else:
        rewards = batch['rewards'].view(-1, group)                                 # [4, G]
        adv = (rewards - rewards.mean(1, keepdim=True)).view(-1, 1)                # [B, 1]
        loss = -(adv * arl.sequence_log_probs(student, batch['ids']) * mask).sum() / mask.sum()
    opt.zero_grad()
    loss.backward()
    opt.step()
    return 0 if method == 'rl' else int(batch['attention_mask'].sum())
```

返回值统计教师在这一步中需要处理的 token 数：off-policy 蒸馏是生成的回合，on-policy 蒸馏是被打分的回合，RL 则为零。下面的循环负责训练、每 5 步评估一次，并保留训练好的学生。

```python theme={null}
def train_student(method, lr, steps=40, seed=0):
    student = copy.deepcopy(base)
    opt = torch.optim.AdamW(student.parameters(), lr=lr)
    rng, gen = random.Random(seed), torch.Generator().manual_seed(seed)
    history, teacher_tokens = [(0, *evaluate(student))], 0
    for step in range(1, steps + 1):
        teacher_tokens += student_step(student, opt, method, rng, gen)
        if step % 5 == 0:
            history.append((step, *evaluate(student)))
    return student, history, teacher_tokens

runs = {name: train_student(method, lr) for name, method, lr in
        [('off-policy', 'off', 1e-2), ('on-policy', 'on', 1e-2), ('outcome RL', 'rl', 3e-3)]}
for name, (_, history, teacher_tokens) in runs.items():
    print(f'{name:<11} teacher tokens {teacher_tokens:>6} |', ' '.join(f'{s}:{a:.2f}/{b:.2f}' for s, a, b in history[::2]))
```

Output:

```text theme={null}
off-policy  teacher tokens  87355 | 0:0.14/0.05 10:0.23/0.20 20:0.51/0.50 30:0.72/0.75 40:0.82/0.78
on-policy   teacher tokens  89345 | 0:0.14/0.05 10:0.24/0.17 20:0.26/0.25 30:0.56/0.65 40:0.64/0.82
outcome RL  teacher tokens      0 | 0:0.14/0.05 10:0.17/0.17 20:0.15/0.17 30:0.20/0.17 40:0.31/0.17
```

![图 16.1：off-policy 蒸馏、on-policy 蒸馏（反向 KL）和结果 RL 的学生学习曲线，三者都从同一个小型学生基座出发。左：训练任务。右：留出变体。每一步使用 32 个回合。](https://mintlify.s3.us-west-1.amazonaws.com/agent-rl/tutorial/figures/ch16-curves.svg)

图 16.1 展示了三条曲线。在同样的 40 步、1,280 个回合内，两种形式的蒸馏都远远领先于 RL：off-policy 蒸馏在训练集上达到 0.82、在留出集上达到 0.78，on-policy 蒸馏为 0.64 和 0.82，结果 RL 为 0.31 和 0.17。原因在于每个回合携带的信息量。在结果混合的组中，一个 RL 回合贡献一个优势值，并分摊到每个动作 token 上；而回合全部失败的组什么也不贡献。一个蒸馏回合在其 10 到 20 个动作 token 的每一个上都贡献一个完整的分布目标，无论它是否成功。

on-policy 曲线起步更慢。它在第 5 步时在两个划分上都跌到 0，到第 10 步恢复到 0.24，到第 20 步在训练集上仍是 0.26，而此时 off-policy 曲线已经达到 0.51。16.2 节预言了这一延迟。学生基座的前缀来自它自己的无技能行为，反向 KL 梯度先压低学生当前偏好的东西，而这时教师偏好的 token 还没获得足够的概率被拉上来。off-policy SFT 从第一步起就对教师的 token 施加前向 KL 的力 $q_v - p_v$。到第 40 步，两条蒸馏曲线在留出集上的差距处于噪声范围内，在训练集上相差 0.18。下一节展示它们真正的差别所在，而那不是平均通过率。

> **证据。** 本书测得：单个随机种子，评估分别使用 80 和 40 个回合，所以低于约 0.1 的差异属于噪声。“40 步时蒸馏 ≫ 结果 RL”这一排序远超噪声。教师使用了恢复示范；16.4 节测量这一点有多重要。

## 16.4 暴露偏差与误差累积

平均通过率掩盖了策略失败的方式。假设当学生的历史看起来与教师的相同时，学生在第 $t$ 轮以概率 $\epsilon$ 犯错，并且在第一次犯错之后，它就处在前缀空间中一个从未训练过的区域，在那里失败的概率接近 1。那么在 $H$ 轮中，它的成功率约为 $(1-\epsilon)^H$，而处于分布之外的期望轮数按 $\epsilon H^2/2$ 而不是 $\epsilon H$ 增长：一次早期错误会毁掉之后的每一轮。这是 [Ross, Gordon, and Bagnell (2010)](https://arxiv.org/abs/1011.0686) 对行为克隆（behavior cloning）的经典分析，而区分 off-policy 训练与 on-policy 训练的，正是第二个假设：*学生在自己犯错之后会做什么？*

```python theme={null}
H = torch.arange(1, 9, dtype=torch.float64)
for eps in [0.02, 0.05, 0.1]:
    no_recovery = (1 - eps) ** H                                  # any mistake is fatal
    one_retry = no_recovery + H * eps * (1 - eps) ** (H - 1)      # survives exactly one mistake
    print(f'eps {eps:.2f}: success at H=4 {no_recovery[3]:.3f} vs {one_retry[3]:.3f}'
          f' | H=8 {no_recovery[7]:.3f} vs {one_retry[7]:.3f}')
```

Output:

```text theme={null}
eps 0.02: success at H=4 0.922 vs 0.998 | H=8 0.851 vs 0.990
eps 0.05: success at H=4 0.815 vs 0.986 | H=8 0.663 vs 0.943
eps 0.10: success at H=4 0.656 vs 0.948 | H=8 0.430 vs 0.813
```

在每轮 5% 的出错率下，一个无法恢复的策略在 4 轮内的成功率为 81%，在 8 轮内为 66%。一个能从一次错误中恢复的策略，成功率分别为 99% 和 94%。差距随时间跨度（horizon）增大，这就是暴露偏差对智能体而言比对单轮生成更严重的原因。

我们直接测量恢复能力。`forced_rollout` 用给定的动作列表开启一个回合，这些动作经由与真实 rollout 相同的 token 路径注入，随后让模型继续。我们注入三种错误类型，并测量后续部分的成功率。

```python theme={null}
def forced_rollout(model, task, forced, gen):
    episode = arl.Episode(arl.BugFixEnv(), task, tok)
    for text in forced:
        episode.act(tok.encode(text))
    while not episode.done:
        episode.act(*arl.generate_call(model, episode.ids, tok, generator=gen))
    return episode.reward

injections = {'none': lambda t: [], 'unterminated call': lambda t: ['<call> read'],
              'read + no-op edit': lambda t: ['<call> read </call>', f'<call> edit {noop_op(t)} </call>']}
models = {'clean teacher': clean_teacher, 'teacher': teacher, 'off-policy': runs['off-policy'][0],
          'on-policy': runs['on-policy'][0]}
recovery = {}
for inj, prefix in injections.items():
    gen = torch.Generator().manual_seed(0)
    recovery[inj] = {name: sum(forced_rollout(m, t, prefix(t), gen) for t in train for _ in range(2)) / 40
                     for name, m in models.items()}
    print(f'{inj:<18}', ' | '.join(f'{name} {v:.2f}' for name, v in recovery[inj].items()))
```

Output:

```text theme={null}
none               clean teacher 1.00 | teacher 0.97 | off-policy 0.80 | on-policy 0.72
unterminated call  clean teacher 0.00 | teacher 1.00 | off-policy 0.03 | on-policy 0.00
read + no-op edit  clean teacher 0.10 | teacher 1.00 | off-policy 0.00 | on-policy 0.47
```

![图 16.2：两个教师和两个蒸馏学生在注入错误之后的成功率。只在干净示范上训练的教师无法恢复；off-policy 学生继承了干净的行为；on-policy 学生学会了从它实际会犯的错误中恢复。](https://mintlify.s3.us-west-1.amazonaws.com/agent-rl/tutorial/figures/ch16-recovery.svg)

上表和图 16.2 区分出三种效应。

1. **只在干净示范上训练的教师自己也无法恢复。** 干净教师在每个干净开局上都成功（1.00），而在注入错误之后只在 0.00 和 0.10 的任务上成功。教师本身也是在某个分布上训练出来的策略，它在从未见过的前缀上的行为并不是专业能力。on-policy 蒸馏的效果上限取决于教师在*学生的*前缀上表现如何。
2. **off-policy 学生继承的是教师的干净行为，而不是它的恢复能力。** 完整教师能从两种注入的错误中恢复（1.00），但它自己的样本几乎从不包含错误，所以 off-policy 蒸馏从未向学生展示过恢复。off-policy 学生在干净开局上的成功率为 0.80，犯错之后为 0.03 和 0.00，与干净教师的模式相同。这就是被测量出来的暴露偏差。
3. **on-policy 学生能从它自己会犯的错误中恢复，而且只限于这些错误。** 在一次无效编辑之后，它的成功率为 0.47，而 off-policy 学生为 0.00，因为学生基座会猜测操作，在训练中做过这类编辑，于是教师在那里纠正了它们。它无法从没有终止的调用中恢复（0.00），因为在格式正确的行为上训练的学生基座几乎从不产生这种调用，所以该前缀从未出现在它的训练批次中。on-policy 蒸馏覆盖的是学生自己的错误分布，而不是所有错误。

干净开局那一行也显示了代价：on-policy 学生在干净开局上略逊于 off-policy 学生（0.72 对 0.80，各 40 个回合，处于噪声范围内）。在这个预算下，它用对教师主路径的一部分模仿换来了对自身偏离路径的覆盖。

## 16.5 全词表与 top-k 教师分布

`distill_loss` 需要教师在每个位置上的完整分布：每个批次 $B \times T \times V$ 个数，对于 Qwen3 的 151,936 词表，fp32 下每个 token 约 0.6 MB（[3.6 节](/tutorial/zh/chapters/03-language-models-as-policies#3-6-真实词表：qwen3-0-6b)）。当教师运行在单独的推理服务上时，这些数据必须经过网络传输。因此，系统只为每个位置发送教师的前 $k$ 个条目，或者只发送被采样 token 的对数概率（16.2 节的采样 token 形式）。截断会改变目标，而改变的方式取决于如何处理缺失的概率质量：

* **重归一化 top-$k$** 把两个分布都限制在教师的前 $k$ 个 token 上并重新归一化。它忽略学生放在该集合之外的任何概率质量，所以一个把概率质量分散到教师排除的 token 上的学生不会因此受到惩罚。当 $k=1$ 时，损失恒为 0。
* **top-$k$ 加尾部桶**（tail bucket）对学生和教师都保留教师的 $k$ 个 token，并把其余所有 token 合并成一个“其他”结果。粗化后分布之间的散度是完整 KL 的下界（数据处理不等式），而学生仍会因为放在教师前 $k$ 个之外的总概率质量受到惩罚。

我们在真实的学生前缀上测量这两种做法：学生基座自己的回合，由教师打分。对每个 $k$，我们报告损失值，以及它对学生 logits 的梯度与全词表梯度之间的余弦相似度。

```python theme={null}
_, student_eps = arl.pass_rate(base, train, tok, 2, seed=3)
batch = arl.collate(student_eps, tok.pad_id)
mask = batch['action_mask']
with torch.no_grad():
    t_logits = teacher(batch['ids'])[:, :-1]                                         # [B, T-1, V]
s_logits = base(batch['ids'])[:, :-1].detach().requires_grad_(True)                 # leaf, [B, T-1, V]
grad_of = lambda loss: torch.autograd.grad(loss, s_logits)[0].flatten()
full = distill_loss(s_logits, t_logits, mask)
g_full = grad_of(full)

def tail_bucket(logits, top):
    """log-probs of the top-k entries plus one merged 'other' entry: [B, T, V] -> [B, T, k+1]."""
    logp = torch.log_softmax(logits.float(), -1).gather(-1, top)
    other = torch.log1p(-logp.exp().sum(-1, keepdim=True).clamp(max=1 - 1e-6))
    return torch.cat([logp, other], -1)
```

`tail_bucket` 返回 $k+1$ 个结果上归一化的对数概率，所以 `distill_loss` 可以直接把它们当作 logits 使用。下面的循环在若干个 $k$ 上比较两种截断方式。

```python theme={null}
topk_rows = []
for k in [1, 2, 4, 8, 32, tok.vocab_size]:
    top = t_logits.topk(k, -1).indices                                              # [B, T-1, k]
    renorm = distill_loss(s_logits.gather(-1, top), t_logits.gather(-1, top), mask)
    bucket = distill_loss(tail_bucket(s_logits, top), tail_bucket(t_logits, top), mask)
    cos = [F.cosine_similarity(grad_of(l), g_full, 0).item() if l.requires_grad and l.item() > 0 else 0.0
           for l in (renorm, bucket)]
    topk_rows.append((k, renorm.item(), bucket.item(), *cos))
    print(f'k={k:>3}: full {full.item():.3f} | renormalized {renorm.item():.3f} (cos {cos[0]:.3f})'
          f' | tail bucket {bucket.item():.3f} (cos {cos[1]:.3f}) | bytes/token {8 * min(k, tok.vocab_size)}')
```

Output:

```text theme={null}
k=  1: full 2.070 | renormalized 0.000 (cos 0.000) | tail bucket 1.759 (cos 0.906) | bytes/token 8
k=  2: full 2.070 | renormalized 1.724 (cos 0.739) | tail bucket 1.945 (cos 0.979) | bytes/token 16
k=  4: full 2.070 | renormalized 1.976 (cos 0.953) | tail bucket 2.035 (cos 0.995) | bytes/token 32
k=  8: full 2.070 | renormalized 2.027 (cos 0.989) | tail bucket 2.060 (cos 0.999) | bytes/token 64
k= 32: full 2.070 | renormalized 2.055 (cos 0.999) | tail bucket 2.068 (cos 1.000) | bytes/token 256
k=316: full 2.070 | renormalized 2.070 (cos 1.000) | tail bucket 2.070 (cos 1.000) | bytes/token 2528
```

![图 16.3：在学生前缀上对教师分布做 top-k 截断。左：逐 token 损失值相对全词表损失的比例。右：截断梯度与完整梯度之间的余弦相似度。](https://mintlify.s3.us-west-1.amazonaws.com/agent-rl/tutorial/figures/ch16-topk.svg)

图 16.3 画出了上表。重归一化 top-$k$ 在 $k$ 较小时有严重偏差。$k=1$ 时它给出恰好为零的损失和梯度，$k=2$ 时它的梯度方向与完整梯度相差很远。尾部桶每个位置只多花一个数，并且在每个 $k$ 下都更接近：它的值低估了完整 KL（$k=1$ 时为 1.76 对 2.07，$k=2$ 时为 1.95），它的梯度与完整梯度的余弦在 $k=1$ 时为 0.91，在 $k=2$ 时为 0.98，而重归一化给出的是 0 和 0.74。差别就在于学生放在教师没有概率质量之处的那部分质量。重归一化恰好丢掉了反向 KL 本来就是要惩罚的那一项。到 $k=32$ 时，两者都接近精确，因为教师在这些前缀上的分布是集中的：它的熵很低，少数几个 token 承载了几乎全部概率质量。熵更高的教师（推理文本，温度 1 采样）需要更大的 $k$ 才能达到同样的精度。

> **陷阱。** 先在教师上计算 top-$k$，再把教师其余的 logits 设为 $-\infty$ 去调用全词表损失，那么只要学生在前 $k$ 个之外有概率质量，反向 KL 就是无穷大，在浮点运算中会变成 `nan` 梯度。要么显式地把学生限制在相同的支撑集上，要么保留一个尾部桶。

## 16.6 近期工作

以下条目是来自公开来源的文献报告结果，查阅日期为 2026-10-13。本书没有复现其中任何一项。

* **GKD**（[Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes”](https://arxiv.org/abs/2306.13649)，v1 2023-06-23；查阅了摘要）。指出了标准蒸馏中训练与推理之间的不一致，并用教师反馈在学生自己生成的输出上训练学生。它允许使用前向 KL 以外的散度，以适应无法表示教师的学生，并且可以与 RL 微调结合。
* **MiniLLM**（[Gu et al.](https://arxiv.org/abs/2306.08543)，v1 2023-06-14；查阅了摘要）。用反向 KL 取代前向 KL，使学生不会高估教师的低概率区域，并以 on-policy 方式优化它。它报告在 120M 到 13B 参数的指令遵循任务上，暴露偏差更低、校准更好。
* **Qwen3 Technical Report**（[Qwen Team](https://arxiv.org/abs/2505.09388)，v1 2025-05-14；查阅了第 4.5 节和表 21）。描述了为较小模型进行的由强到弱蒸馏，分两个阶段：先在教师输出上做 off-policy 蒸馏，再做 on-policy 蒸馏，其中“学生模型生成 on-policy 序列用于微调”（“the student model generates on-policy sequences for fine-tuning”），并以最小化与教师 logits 的 KL 散度为目标进行训练。表 21 从同一个经过 off-policy 蒸馏的 Qwen3-8B 检查点（AIME'24 55.0）出发，比较了耗费 17,920 GPU 小时的 RL（67.6）与耗费 1,800 GPU 小时的 on-policy 蒸馏（74.4）。
* **On-Policy Distillation**（[Kevin Lu and Thinking Machines Lab](https://thinkingmachines.ai/blog/on-policy-distillation/)，2025-10-27；查阅了全文）。使用 16.2 节的采样 token 形式：逐 token 优势是负的反向 KL，由教师对学生所采样 token 的对数概率计算得到，并采用“折扣因子为零”（“a discount factor of zero”）。文章说明它不使用 top-$k$ logit 蒸馏。从在教师推理轨迹上做过 off-policy SFT 的 Qwen3-8B-Base 出发（AIME'24 上为 60%），on-policy 蒸馏报告在约 150 步内达到 70%，并估计相对于外推的 SFT 可节省 9 到 30 倍的计算量，具体取决于是否计入教师采样。

这些来源中反复出现两种模式，并与 16.3 和 16.4 节在小规模实验上测得的结果相符。第一，off-policy 蒸馏在前，on-policy 蒸馏在后：off-policy 数据每个 token 的成本很低，能把学生带到教师附近，然后 on-policy 蒸馏再修正学生自己的错误。第二，与 RL 的比较是在相同起点上进行的，而所报告的蒸馏优势体现在计算量上，这正是 16.3 节的稠密信号论证。这些来源都没有研究多轮工具使用，所以 16.4 节的恢复能力测量是它们留下的、智能体特有的开放问题。[第 20 章](/tutorial/zh/chapters/20-frontier#20-7-蒸馏与评估)跟踪智能体方面的后续工作。

## 16.7 实践中的设计问题

> **问题：教师与学生的分词器不一致。**
>
> **现象：** 蒸馏损失很大且不下降，或者损失下降而学生变差。
>
> **原因：** `distill_loss` 按索引逐一比较分布。如果教师与学生的词表顺序不同，索引 $v$ 对应的是不同的字符串；如果它们切分文本的方式不同，位置 $i$ 对应的是不同的前缀。
>
> **可选方案：** (a) 使用来自同一模型家族、共享分词器（tokenizer）的教师和学生，这是常见选择。(b) 对于字符串相同但 id 不同的词表，把教师条目映射到学生 id，对一一映射而言这是精确的。(c) 对于切分方式不同的情况，在序列层面蒸馏（在教师文本上做 off-policy SFT，由学生重新分词），或者对 token 片段做显式对齐，这两种做法都会失去逐 token 信号。
>
> **建议做法：** (a)。下一个代码块展示前两种失败模式。

```python theme={null}
perm = torch.randperm(tok.vocab_size, generator=torch.Generator().manual_seed(0))   # teacher id -> different word
same = t_logits[..., perm]                                  # identical distributions, teacher-side id order
inverse = torch.argsort(perm)
print(f'reverse KL, same model, same ids:   {distill_loss(t_logits, t_logits, mask):.4f}')
print(f'reverse KL, same model, permuted:   {distill_loss(t_logits, same, mask):.2f}')
print(f'reverse KL, permuted then mapped:   {distill_loss(t_logits, same[..., inverse], mask):.4f}')
from transformers import AutoTokenizer
qtok = AutoTokenizer.from_pretrained('Qwen/Qwen3-0.6B')
call = '<call> edit swap_minmax </call>'
print(f'{call!r}: TinyAgent {len(tok.encode(call))} tokens | Qwen3', qtok.tokenize(call))
```

Output:

```text theme={null}
reverse KL, same model, same ids:   0.0000
reverse KL, same model, permuted:   16.83
reverse KL, permuted then mapped:   0.0000
'<call> edit swap_minmax </call>': TinyAgent 4 tokens | Qwen3 ['<', 'call', '>', 'Ġedit', 'Ġswap', '_min', 'max', 'Ġ</', 'call', '>']
```

置换后的教师对各个 token 的预测概率完全没变，而按索引计算的 KL 却报告了很大的散度。把 id 映射回去就恢复为 0。切分方式不同的情况更糟：同一个调用在 TinyAgent 中是 4 个 token，在 Qwen3 中是一串更长的子词 token，这时同一位置不再对应同一个预测事件，不能直接按 token 位置计算 KL。

> **问题：教师成本。**
>
> **现象：** on-policy 蒸馏的每一步比 RL 的每一步更慢，或者教师服务成为集群的瓶颈。
>
> **原因：** on-policy 蒸馏在每一步都要对每个学生 token 运行一次教师前向传播，而教师通常比学生大得多。off-policy 蒸馏每个回合只为教师付一次费用，其数据可以跨轮次、跨学生复用。RL 完全不需要教师，但需要环境和验证器。
>
> **可选方案：** (a) 只对动作 token 打分（工具 token 不需要目标），在观测占据上下文主体时节省最多。(b) 返回带尾部桶的 top-$k$，而不是完整分布（16.5 节）。(c) 使用采样 token 形式，每个 token 一个数。(d) 先做 off-policy 蒸馏，只在最后阶段做 on-policy 蒸馏，Qwen3 报告和 Thinking Machines 的文章都是这样做的。
>
> **建议做法：** (d) 加 (a) 和 (c)。下一个代码块根据 16.3 节的 token 计数估计每一步的前向传播成本。

```python theme={null}
student_params = sum(p.numel() for p in base.parameters())
teacher_params = sum(p.numel() for p in teacher.parameters())
for name, (_, _, teacher_tokens) in runs.items():
    per_step = teacher_tokens / 40
    teacher_flops = 2 * teacher_params * per_step                  # forward only
    student_flops = 6 * student_params * per_step if per_step else 0
    print(f'{name:<11} teacher tokens/step {per_step:>6.0f} | teacher FLOPs/step {teacher_flops:.2e}'
          + (f' = {teacher_flops / student_flops:.2f} x student train FLOPs' if per_step else ''))
```

Output:

```text theme={null}
off-policy  teacher tokens/step   2184 | teacher FLOPs/step 6.87e+08 = 1.26 x student train FLOPs
on-policy   teacher tokens/step   2234 | teacher FLOPs/step 7.02e+08 = 1.26 x student train FLOPs
outcome RL  teacher tokens/step      0 | teacher FLOPs/step 0.00e+00
```

这里教师的参数量是学生的 3.8 倍，所以教师在批次上只做前向传播的成本，是学生前向加反向传播合计成本的 1.26 倍。两种蒸馏方法每步都处理了约 2,200 个教师 token。在生产中这个比例更大（Qwen3 的报告从大得多的旗舰模型蒸馏到小模型），而且 off-policy 数据还有一个这个估计忽略了的优势：它只生成一次就可以复用。对于 off-policy 那一行，该数字统计的是生成回合中的教师 token，真实的推理引擎会借助 KV cache 但以顺序方式处理它们，利用率远低于 on-policy 的 prefill。

> **问题：教师在某些任务类型上是错的。**
>
> **现象：** 学生的平均分接近教师，而评估中的某个切片对两者都是零。
>
> **原因：** 蒸馏损失衡量的是与教师的一致性，而不是正确性。只要教师自信地出错，学生就会被训练成以同样的方式自信地出错，而且用的正是让蒸馏高效的那种稠密信号。
>
> **可选方案：** (a) 蒸馏前用验证器逐切片验证教师，并在教师失败的切片上丢弃教师信号。(b) 混入结果奖励，$\mathcal{L} = \mathcal{L}_{\text{distill}} + \lambda\,\mathcal{L}_{\text{RL}}$。(c) 在该切片上修正或替换教师。
>
> **建议做法：** (a)，并把被丢弃的切片当作单独的 RL 问题或数据问题来处理。我们注入这一失败：一个教师副本在 `zero` 任务类型上使用 `hardcode_tests` 的示范上做了微调，这是一个在某类任务上用可被投机利用的奖励训练出来的教师的简化示例。

```python theme={null}
def hack_zero(observations):
    action = arl.scripted_policy('oracle')(observations)
    return '<call> edit hardcode_tests </call>' if 'add_zero_guard' in action else action

bad_demos = [arl.scripted_rollout(hack_zero, arl.BugFixEnv(), t, tok) for t in train]
bad_teacher = copy.deepcopy(teacher)
arl.sft_train(bad_teacher, bad_demos, tok, steps=60)

def by_family(model, n=8):
    _, trajs = arl.pass_rate(model, train, tok, n, seed=21)
    return {f: round(sum(t.reward for t in trajs if t.info['family'] == f) / (4 * n), 2) for f in arl.FAMILIES}
print('bad teacher:', by_family(bad_teacher))
```

Output:

```text theme={null}
bad teacher: {'minmax': 1.0, 'range': 1.0, 'compare': 1.0, 'zero': 0.0, 'init': 1.0}
```

这个坏教师能解决四个任务类型，在 `zero` 上则完全失败。我们用三种方式把它 on-policy 地蒸馏到学生基座中：普通的反向 KL；反向 KL 加上在同一批次上权重为 $\lambda = 1$ 的 REINFORCE 损失；以及用上面的逐任务类型验证来**门控**（gated）同样的混合，使失败任务类型的回合只得到 RL 信号而没有教师信号。

```python theme={null}
def distill_with(teacher_model, rl_weight=0.0, gated=(), steps=40, lr=1e-2, seed=0, group=8):
    student = copy.deepcopy(base)
    opt = torch.optim.AdamW(student.parameters(), lr=lr)
    rng, gen = random.Random(seed), torch.Generator().manual_seed(seed)
    for _ in range(steps):
        batch_tasks = [t for t in rng.sample(train, 4) for _ in range(group)]
        batch = arl.collate(arl.rollout_batch(student, [arl.BugFixEnv() for _ in batch_tasks], batch_tasks, tok,
                                              generator=gen), tok.pad_id)
        mask = batch['action_mask']
        trusted = torch.tensor([[float(t.family not in gated)] for t in batch_tasks])   # [B, 1]
        logits = student(batch['ids'])[:, :-1]
        with torch.no_grad():
            t_logits = teacher_model(batch['ids'])[:, :-1]
        loss = distill_loss(logits, t_logits, mask * trusted) if trusted.sum() else 0.0
        if rl_weight:
            rewards = batch['rewards'].view(-1, group)
            adv = (rewards - rewards.mean(1, keepdim=True)).view(-1, 1)
            logp = arl.selected_log_probs(logits, batch['ids'][:, 1:])
            loss = loss - rl_weight * (adv * logp * mask).sum() / mask.sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return student
```

```python theme={null}
wrong = {'KL only': distill_with(bad_teacher), 'KL + RL': distill_with(bad_teacher, rl_weight=1.0),
         'gated + RL': distill_with(bad_teacher, rl_weight=1.0, gated={'zero'})}
zero_tasks = [t for t in train if t.family == 'zero']
for name, student in wrong.items():
    _, trajs = arl.pass_rate(student, zero_tasks, tok, 8, seed=4)
    top_ops = Counter(' '.join(t.info['ops']) for t in trajs).most_common(1)[0]
    print(f'{name:<10}', by_family(student), f'| zero-family top ops: {top_ops[0]!r} ({top_ops[1]}/32)')
```

Output:

```text theme={null}
KL only    {'minmax': 0.94, 'range': 0.88, 'compare': 0.22, 'zero': 0.0, 'init': 1.0} | zero-family top ops: 'hardcode_tests' (27/32)
KL + RL    {'minmax': 1.0, 'range': 0.97, 'compare': 0.59, 'zero': 0.0, 'init': 0.91} | zero-family top ops: 'hardcode_tests' (29/32)
gated + RL {'minmax': 0.44, 'range': 0.84, 'compare': 1.0, 'zero': 0.0, 'init': 1.0} | zero-family top ops: '' (11/32)
```

错误被原样传递了下去。使用普通反向 KL 时，学生在 `zero` 上复制了这个漏洞利用：它在那里的 32 个样本中有 27 个使用 `hardcode_tests`，通过率为 0。（它在 `compare` 上 0.22 的分数是一个独立的、单随机种子下由短预算造成的效应；KL + RL 运行在那里达到 0.59。）

以相同权重混入结果 RL 项，并不能在 40 步内修复该任务类型，投机行为反而变得更普遍（32 个中有 29 个）。原因在于组基线：在 `zero` 任务上，学生的八个样本全是教师的投机行为，全部失败，优势全为 0，所以 RL 项没有可以推动的方向，而 KL 项一直把它拉向投机行为。

在该任务类型上关闭教师信号消除了投机行为，学生在 `zero` 上最常见的行为变成了根本不做编辑（32 个中有 11 个），但通过率仍是 0，而且这一次运行在 `minmax` 上也丢失了准确率（0.44）：没有教师信号，组内也没有任何成功样本，就没有什么能教会它 `add_zero_guard`。图 16.4 汇总了这四个策略。

如果教师在某类任务上系统性地出错，单靠模仿它不能纠正这些错误。它必须先由独立检查发现，然后用一个能够探索的信号来解决：在这个切片上，这意味着付出 RL 的成本并保证足够的探索以找到一次成功（练习 6），或者提供新的示范。

![图 16.4：一个在 zero 任务类型上出错的教师，以及用普通反向 KL、加入结果 RL 项、以及在失败任务类型上门控掉教师这三种方式从它蒸馏得到的学生，在各任务类型上的隐藏测试通过率。](https://mintlify.s3.us-west-1.amazonaws.com/agent-rl/tutorial/figures/ch16-wrong-teacher.svg)

## 小结

* off-policy 蒸馏是在教师回合上做 SFT，在教师前缀上最小化轨迹分布之间的前向 KL。on-policy 蒸馏在学生生成的前缀上最小化逐 token 散度，默认使用反向 KL。它就是带软教师的 DAgger。
* 在固定前缀下，$\nabla\,\mathrm{KL}(p\|q) = \mathbb{E}_{y\sim p}[(\log p_y - \log q_y)\nabla\log p_y]$：这是一个奖励为 $\log q - \log p$ 的单 token 策略梯度。实际做法忽略未来前缀对当前 token 的依赖（折扣为 0）。反向 KL 的 logit 梯度随学生自己的概率缩放，所以当学生错过教师的众数时，起步会变慢。
* 从同一个小型学生出发，在每步 32 个回合、共 40 步的条件下，两种蒸馏方法都远远胜过结果 RL（本书测得）。两者之间的差别体现在恢复能力上，而不是平均通过率：off-policy 学生继承了教师的干净行为，无法从错误中恢复；on-policy 学生能从它自己犯的错误中恢复，但无法从它从未犯过的错误中恢复。
* 教师自身必须在学生的前缀上胜任。只在干净示范上训练的教师无法从注入的错误中恢复，因此也不可能教会恢复。
* top-$k$ 教师分布会使目标产生偏差。重归一化 top-$k$ 丢弃了学生在教师支撑集之外的概率质量，并在 $k=1$ 时给出零损失。尾部桶只多花一个数，在每个 $k$ 下都与完整梯度保持接近。
* 蒸馏信号衡量的是一致性，而不是正确性。教师在一个任务类型上的错误被原样传递给了学生，而额外的 RL 项和门控都没能在预算内修复它，因为该任务类型上的每个组奖励方差都为零。

## 练习

1. 实现 on-policy 蒸馏的采样 token 形式（逐 token 优势 $\log q(y_i) - \log p(y_i)$，从计算图中分离，用于 REINFORCE 损失），并在相同步数下比较它与 `distill_loss` 的学习曲线和梯度方差。采样形式在每个前缀上需要多少个样本，才能使其与真实梯度的余弦达到 16.5 节中全词表梯度的水平？
2. 改变一个假设：让学生与教师一样大。off-policy 蒸馏和 on-policy 蒸馏是否仍然殊途同归？然后把学生变得更小（`d_model=16`），比较前向 KL 与反向 KL，并用mass-covering 和 mode-seeking来解释其中的差异。
3. 加回实际做法丢掉的剩余回报项：用某个 token 之后各 token 的逐 token KL 之和（折扣 $\gamma$）为该 token 的得分函数加权，并在 16.4 节的恢复测量上比较 $\gamma = 0$、$0.5$ 和 $1$。
4. 在 on-policy 蒸馏过程中，以每轮 0.1 的概率向*学生的* rollout 注入没有终止的调用这一错误（并把注入的 token 从损失中掩码掉），然后重新测量恢复能力。把这一做法与 DART 式的噪声注入以及 16.4 节发现的局限联系起来。
5. 在教师温度为 0.5、1 和 2 时（在 softmax 之前把教师 logits 除以温度）运行 top-$k$ 测量。要使梯度余弦达到 0.99 所需的 $k$ 如何随教师的熵变化？
6. 修复 16.7 节的 `zero` 任务类型：门控掉教师，并只在被门控的任务上提高学生的 rollout 温度。什么样的温度能让组奖励方差足够频繁地非零，从而让 RL 找到 `add_zero_guard`？这对其他任务类型有什么影响？
7. 设计评审：你必须用一个同家族的 32B 教师、一个沙箱化的基于测试的验证器以及固定的 GPU 预算，对一个 1.5B 参数的编程智能体做后训练。列出各个阶段（数据、off-policy 蒸馏、on-policy 蒸馏、RL），说明你会测量什么来决定何时从一个阶段进入下一个阶段，以及如何在学生复制之前发现一个在某个代码仓库或某种语言上出错的教师。


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.