OPD 极简笔记，方便回顾。

回顾

记号

$s_t=(x,y_{<t}),\qquad a_t=y_t$ $\pi_\theta=\text{student / current policy}$ $\pi_T=\text{teacher}$ $\pi_{\text{old}}=\text{rollout policy / sampling policy}$

SFT

Hard-label SFT:

$L_{\text{SFT}} = -\mathbb{E}_{(x,y^*)\sim\mathcal D} \sum_t \log \pi_\theta(y_t^*\mid x,y_{<t}^*)$

Soft-label / distillation SFT:

$L_{\text{SFT-soft}} = \mathbb{E}_{s\sim d_T} D_{\text{KL}} \left( \pi_T(\cdot\mid s) \|\pi_\theta(\cdot\mid s) \right)$

RL

目标：

$J_{\text{RL}}(\theta) = \mathbb{E}_{\tau\sim\pi_\theta} [R(\tau)]$

Policy gradient:

$\nabla_\theta J = \mathbb{E}_{\tau\sim\pi_\theta} \sum_t A_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)$

GRPO

同一 prompt 采样多个 answers:

$o_1,\dots,o_G\sim\pi_{\text{old}}(\cdot\mid q)$

reward:

$r_1,\dots,r_G$

Group-relative advantage:

$A_i = \frac{ r_i-\operatorname{mean}(r_1,\dots,r_G) }{ \operatorname{std}(r_1,\dots,r_G) }$

Ratio:

$\rho_{i,t}(\theta) = \frac{ \pi_\theta(o_{i,t}\mid q,o_{i,<t}) }{ \pi_{\text{old}}(o_{i,t}\mid q,o_{i,<t}) }$

PPO-style clipped loss:

$L_{\text{GRPO}} = -\frac{1}{G} \sum_i \frac{1}{|o_i|} \sum_t \min \left( \rho_{i,t}A_i, \operatorname{clip}(\rho_{i,t},1-\epsilon,1+\epsilon)A_i \right)$

OPD

目标

OPD 目标：student-prefix 上的 reverse KL。

$L_{\text{OPD}} = \mathbb{E}_{s\sim d_{\pi_\theta}} D_{\text{KL}} \left( \pi_\theta(\cdot\mid s) \|\pi_T(\cdot\mid s) \right)$

固定一个 prefix:

$D_{\text{KL}} \left( \pi_\theta(\cdot\mid s) \|\pi_T(\cdot\mid s) \right) = \mathbb{E}_{a\sim\pi_\theta(\cdot\mid s)} \left[ \log\pi_\theta(a\mid s)-\log\pi_T(a\mid s) \right]$

Full-vocabulary OPD

Full-vocabulary loss:

$L_{\text{full}}(s) = \sum_{a\in V} \pi_\theta(a\mid s) \left[ \log\pi_\theta(a\mid s)-\log\pi_T(a\mid s) \right]$

Gradient:

$\nabla_\theta L_{\text{full}}(s) = \sum_{a\in V} \pi_\theta(a\mid s) \left[ \log\pi_\theta(a\mid s)-\log\pi_T(a\mid s) \right] \nabla_\theta\log\pi_\theta(a\mid s)$

Sampled-token OPD

Student rollout 采一个 token:

$a_t\sim\pi_{\text{old}}(\cdot\mid s_t)$

Sampled-token reverse KL:

$\widehat{\mathrm{RKL}}_t = \log\pi_{\text{old}}(a_t\mid s_t) - \log\pi_T(a_t\mid s_t)$

OPD advantage:

$A_t = -\widehat{\mathrm{RKL}}_t = \log\pi_T(a_t\mid s_t) - \log\pi_{\text{old}}(a_t\mid s_t)$

Surrogate loss

真正目标：

$\min_\theta D_{\text{KL}}(\pi_\theta\|\pi_T)$

RL 实际训练用 surrogate。

最简单 PG surrogate:

$L_{\text{PG}} = -A_t\log\pi_\theta(a_t\mid s_t)$

Importance-sampling ratio:

$\rho_t(\theta) = \frac{ \pi_\theta(a_t\mid s_t) }{ \pi_{\text{old}}(a_t\mid s_t) } = \exp \left( \log\pi_\theta(a_t\mid s_t) - \log\pi_{\text{old}}(a_t\mid s_t) \right)$

IS surrogate loss:

$L_{\text{IS}} = -\sum_t\rho_t(\theta)A_t$

把 OPD advantage 塞进去：

$A_t = \log\pi_T(a_t\mid s_t) - \log\pi_{\text{old}}(a_t\mid s_t)$

得到最终 OPD-IS loss:

$L_{\text{OPD-IS}} = -\sum_t \exp \left( \log\pi_\theta(a_t\mid s_t) - \log\pi_{\text{old}}(a_t\mid s_t) \right) \left[ \log\pi_T(a_t\mid s_t) - \log\pi_{\text{old}}(a_t\mid s_t) \right]$

基于 Tinker 的代码实现

通用 Tinker importance-sampling loss 需要：

new_logprobs      # log π_θ(a_t | s_t)
old_logprobs      # log π_old(a_t | s_t)
advantages        # A_t
mask              # valid token mask

importance_sampling_loss:

def importance_sampling_loss(new_logprobs, old_logprobs, advantages, mask):
    ratio = torch.exp(new_logprobs - old_logprobs)
    loss = -(ratio * advantages * mask).sum() / mask.sum()
    return loss

OPD 只负责制造 advantages:

1 2	reverse_kl = old_logprobs - teacher_logprobs advantages = -reverse_kl

等价于：

1	advantages = teacher_logprobs - old_logprobs

然后传给已有 RL loss:

loss = importance_sampling_loss(
    new_logprobs=new_logprobs,
    old_logprobs=old_logprobs,
    advantages=advantages,
    mask=mask,
)

sampled-token OPD 如下：

sampled_logprobs = trajectories.loss_fn_inputs["logprobs"]
teacher_logprobs = teacher_client.compute_logprobs(trajectories)

reverse_kl = sampled_logprobs - teacher_logprobs
advantages = -reverse_kl

training_client.forward_backward(
    trajectories,
    loss_fn="importance_sampling",
)

直观看一下梯度

令：

$z_t=\log\pi_\theta(a_t\mid s_t)$

单 token loss:

$L_t=-\exp(z_t-\log\pi_{\text{old}}(a_t\mid s_t))A_t$

求导：

$\frac{\partial L_t}{\partial z_t} = -\rho_t A_t$

梯度下降：

$z_t \leftarrow z_t- \eta\frac{\partial L_t}{\partial z_t} = z_t+ \eta\rho_tA_t$

因此：

$A_t>0\Rightarrow \log\pi_\theta(a_t\mid s_t)\uparrow$ $A_t<0\Rightarrow \log\pi_\theta(a_t\mid s_t)\downarrow$

Reference

On-Policy Distillation - Thinking Machines Lab

thinking-machines-lab/tinker-cookbook 的 on_policy_distillation.py