Get Started
Home
Topics
Search
Library
LLM Training · Reasoning · Jul 6, 2026

Trust Region Policy Distillation

Source: research paper via Hugging Face Daily Papers
0:00 / 6:35
On-policy distillation blows up when the student picks a token the teacher hates: the log-ratio reward diverges to negative infinity. TOP-D replaces it with log(α·ρ+1−α), blending teacher toward student in probability space — bounded reward, zero extra compute, +25.84pp on AIME24 over vanilla OPD.
TL;DR
TOP-D fixes the training instability of on-policy distillation by mixing the teacher’s next-token distribution with the student’s own, so the token-level reward is bounded from below instead of exploding to negative infinity when the teacher assigns near-zero probability, lifting AIME 2024 accuracy by +25.84 pp over standard On-Policy Distillation on an 8B student.
Why It Matters
Say you’re post-training a small reasoning model on your own math or coding traces, using a bigger model as the teacher. The obvious move is On-Policy Distillation: let the student sample its own responses, then reward each token by how much more likely the teacher would have made it. This gives dense per-token feedback, unlike verifier-only RL which only scores whole answers. The problem shows up quickly in practice. When the student produces a token the teacher considers implausible, the log-ratio reward blows up toward negative infinity, one bad token dominates the gradient, and training destabilizes. Prior fixes are mostly empirical patches (mixed sampling, reward clipping, top-p filtering) without guarantees. TOP-D asks: can we bound this reward analytically without adding compute?
How It Works
The trick is to never let the teacher’s probability go to zero in the denominator of the reward. Instead of comparing the student to the raw teacher, compare it to a blend: mostly the student itself, with a small dose of the teacher mixed in. Formally, the proximal teacher is a convex combination in probability space, controlled by a coefficient the authors call Interpolation coefficient α (set to 0.1 or 0.2 in experiments). The per-token reward becomes log(α·ρ + 1−α), where ρ is the teacher/student probability ratio. As ρ → 0 (student picks a token the teacher hates), the reward flattens at log(1−α) instead of diverging. As α → 1, you recover vanilla OPD.
Crucially, no second forward pass is needed. The proximal teacher exists only algebraically. You compute the same teacher and student logprobs you already have, then apply the log(α·ρ + 1−α) transform. Zero extra compute.
On top of this stabilized reward, TOP-D adds a PPO-style inner loop: sample a batch of rollouts, then reuse them for multiple gradient epochs with a clipped importance-sampling ratio, borrowing the sample-efficiency trick from trust-region RL. Advantages are normalized at the token level within each prompt’s group of rollouts, in the Group Relative Policy Optimization (GRPO) style.
while not converged: pi_old = copy(pi_theta) for x in batch_of_prompts: rollouts = sample(pi_old, x, G=8) for y_i in rollouts: r_tilde = log(alpha * (pi_star(y)/pi_old(y)) + (1 - alpha)) R_tilde = r_tilde + mean(future r_tilde) # length-normalized A_hat = (R_tilde - mu) / sigma # token-level, per-prompt group for epoch in range(E): update pi_theta with clipped-ratio objective on (A_hat, rollouts)
Core Insight
The prevailing fix for unstable distillation is engineering: clip the reward, filter the vocabulary, mix in teacher samples. This paper shows the opposite. Reshape the reward analytically by interpolating teacher and student in probability space, and the unbounded log-ratio becomes a smooth, lower-bounded signal for free. The evidence that this is the load-bearing move is the ablation at α=1.0, which strips out the interpolation and recovers OPD’s instability, not the headline benchmark lift.
What They Found
The ablation is the finding that proves the mechanism. Setting α=1.0 removes the proximal teacher and reverts to raw OPD; the training curves become erratic and final accuracy collapses. Disabling the inner trust-region epochs (forcing strictly on-policy updates) doesn’t destabilize training but drastically slows convergence. Both pieces matter, and they matter for the reasons the theory predicts.
Secondary evidence, headline benchmarks on a Qwen3-8B-Base student taught by Qwen3-30B-A3B-Instruct-2507:
•
AIME 2024 avg@32: OPD 24.58% → TOP-D 50.42% (+25.84 pp)
•
AIME 2025: +10.73 pp over OPD; AIME 2026: +18.64 pp
•
Beats the strongest RLVR baseline (DAPO) by +17.5 pp on AIME24
The capacity-gap story sharpens on the smaller Qwen3-1.7B-Base student: raw OPD falls behind verifier-only RL baselines (the unbounded log-ratio bites hardest when student and teacher disagree most), while TOP-D still leads by roughly 10 pp on AIME24/25. Sensitivity to α across {0.1, 0.2, 0.3} is minor.
What’s Useful
Reach for this when you’re distilling a smaller reasoning model from a larger one and OPD-style training either diverges or plateaus below your verifier-only RL run. The change is a one-line reward transform: replace r = log(π*/π) with r̃ = log(α·π*/π + 1−α), with α around 0.1–0.2. Everything else in your PPO/GRPO loop stays the same, including the teacher and student forward passes you were already doing. No extra GPUs, no extra latency.
The paper doesn’t link a public code release. Training uses the DAPO-Math-17k dataset and evaluates on standard AIME and AMC competitions; students are the Qwen3 base models, teachers are Qwen3-14B or Qwen3-30B-A3B-Instruct-2507. Reported hardware is a 4-node H200 cluster, but the authors state the full pipeline reproduces on a single 8-GPU node.
Takeaway
When a log-ratio reward wants to explode, don’t clip it. Blend the reference distribution toward your current policy so the ratio can’t blow up in the first place. The bound comes from math, not heuristics, which is why it survives large teacher-student capacity gaps where clipping-based fixes get shakier.
Caveats
•
Validated only up to 8B students and roughly 200–400 update steps; scaling behavior at 30B+ or over long training horizons is untested by the authors.
•
All experiments are mathematical reasoning with a strong same-family teacher (Qwen3). Whether the proximal-teacher trick helps as much for code, general chat, or cross-family teacher/student pairs is not shown.
•
The interpolation happens in probability space, which requires access to teacher token probabilities over the same vocabulary as the student. Distilling across tokenizers or from an API-only teacher isn’t addressed.
Topics
Don't miss new content
Log in to follow topics and personalize your feed.
By content type
Research Paper178 episodes
AI178 episodes