$β$-OPSD: Deriving with Policy Optimization, Training with Self-Distillation

Overview§

$\beta$-OPSD reinterprets vanilla on-policy self-distillation (OPSD) as the $\beta=1$ member of a KL-regularized policy optimization family. Instead of running RL on that objective, the authors derive the closed-form optimal distribution and distill it into the student, making the procedure efficient and stable.

Objective§

For a context $x$ and prefix $y_{<t}$, let $p_{\mathrm{ref}}$ be the reference policy (e.g., the frozen base model) and $p_{\mathrm{tea}}$ the privileged teacher. The generalized objective is

$$ \max_{p_\theta} \mathbb{E}_{y_t \sim p_\theta(\cdot|x,y_{<t})} \left[ \log p_{\mathrm{tea}}(y_t|x,y_{<t}) \right] - \beta \, \mathrm{KL}\left(p_\theta(\cdot|x,y_{<t}) \| p_{\mathrm{ref}}(\cdot|x,y_{<t})\right). $$

Solving this variational problem over unconstrained policies gives

$$ p^\star(y_t|x,y_{<t}) = \frac{1}{Z} p_{\mathrm{ref}}(y_t|x,y_{<t}) \, p_{\mathrm{tea}}(y_t|x,y_{<t})^{1/\beta}, $$

i.e., a geometric interpolation between the reference and teacher. $\beta=1$ recovers vanilla OPSD; larger $\beta$ anchors the solution to the reference, smaller $\beta$ trusts the teacher more.

Rather than optimize this objective by policy gradient (which is high-variance and slow), $\beta$-OPSD uses $p^\star$ as a distillation target. In log space the target is additive:

$$ \log p^\star = \log p_{\mathrm{ref}} + \frac{1}{\beta}\log p_{\mathrm{tea}} + C, $$

so efficient token-level targets are obtained by mixing raw decoder logits:

$$ z_{\mathrm{target}} = z_{\mathrm{ref}} + \frac{1}{\beta} z_{\mathrm{tea}}. $$

Training loss§

Student rollouts are generated on-policy, and target logits are computed from reference/teacher logits at the same positions. Return-to-go $G_t=\sum_{t'\ge t} r_{t'}$ weights each token loss, aligning the distillation objective with the sequence-level return:

$$ \mathcal{L}(\theta)= \mathbb{E}_{x\sim\mathcal{D}, y\sim p_\theta(\cdot|x)} \left[ \sum_t G_t \, \mathrm{KL}\left( p^\star(\cdot|x,y_{<t}) \, \middle\| \, p_\theta(\cdot|x,y_{<t}) \right) \right]. $$

This approximates policy optimization while keeping OPSD's simplicity.

Reference implementation sketch§

beta = 2.5
# student-generated context
input_ids, mask, rtg = sampled_rollout()

with torch.no_grad():
    z_ref = ref_model(input_ids, attention_mask=mask).logits
    z_tea = tea_model(input_ids, attention_mask=mask).logits
    z_target = z_ref + (1.0 / beta) * z_tea
    target_probs = torch.softmax(z_target, dim=-1)

z_stu = student_model(input_ids, attention_mask=mask).logits
log_q = torch.log_softmax(z_stu, dim=-1)
kl = torch.sum(torch.kl_div(log_q, target_probs, reduction='none'), dim=-1)
loss = (kl * rtg.unsqueeze(-1) * mask).sum() / mask.sum()
loss.backward()
optimizer.step()

Results§

Experiments on mathematics reasoning benchmarks show that $\beta$-OPSD improves optimization stability and final accuracy over vanilla OPSD at negligible extra cost, confirming that selecting $\beta$ lets practitioners trade reference retention for teacher guidance.