Weak-to-Strong Generalization via Direct On-Policy Distillation
By Shiyuan Feng, Huan-ang Gao, Haohan Chi, Hanlin Wu, Zhilong Zhang, Zheng Jiang, Bingxiang He, Wei-Ying Ma, Ya-Qin Zhang, Hao Zhou
"Transfers RL-induced policy shifts from weak to strong LLMs via on-policy distillation of implicit reward from teacher checkpoint pairs, achieving rapid gains without target-model RL."
Abstract
Reinforcement learning with verifiable rewards (RLVR) is a powerful recipe for improving language-model reasoning, but it is expensive to repeat on every new strong model because the target model must generate many rollouts during training. As models scale, post-training itself becomes a bottleneck. We study a weak-to-strong alternative: run RL on a smaller model where rollouts are cheaper, then reuse what that RL run learned to improve a stronger target model. Directly distilling the post-RL weak teacher is not enough, because the teacher's final policy mixes useful RL gains with the limitations of the smaller model. We propose Direct On-Policy Distillation (Direct-OPD), which transfers the teacher's RL-induced policy shift instead. Direct-OPD compares the post-RL teacher with its own pre-RL reference and treats their log-ratio as a dense implicit reward for the student. In plain terms, the checkpoint pair tells us which actions RL made the weak model more or less likely to take, and Direct-OPD applies that signal on the stronger student's own on-policy states. This directly reuses the weak model's RL supervision signal without training an explicit reward model or running sparse-reward RL on the target model. Empirically, Direct-OPD consistently leverages weaker teachers to improve stronger target models; notably, it boosts Qwen3-1.7B from 48.3% to 62.4% on AIME 2024 in just 4 hours on 8 A100 GPUs. It outperforms step-matched direct RL and enables the sequential composition of multiple policy shifts. Our results show that RL outcomes can be reused across model scales as implicit reward signals, not merely as final models to imitate.
Technical Analysis & Implementation
Direct On-Policy Distillation (Direct-OPD)§
Core Idea§
RL with verifiable rewards (RLVR) improves reasoning but is expensive for large models. Direct-OPD reuses the RL-induced policy change of a smaller model as an implicit reward signal to train a larger student model, avoiding explicit reward models or sparse-reward RL on the student.
Methodology§
Let $\pi_{\text{weak}}^{\text{post}}$ be the weak teacher after RLVR, and $\pi_{\text{weak}}^{\text{pre}}$ its pre-RL checkpoint. For a state $s$ (e.g., a partial answer prefix), the teacher's policy shift is captured by the log-ratio: $$ r(s, a) = \log \frac{\pi_{\text{weak}}^{\text{post}}(a|s)}{\pi_{\text{weak}}^{\text{pre}}(a|s)} $$ This real-valued $r(s,a)$ serves as a dense implicit reward. Direct-OPD trains the student $\pi_{\text{strong}}$ using on-policy rollouts: the student generates a sequence $\tau = (s_1, a_1, \ldots, s_T, a_T)$ and receives per-token reward $r(s_t, a_t)$ from the teacher's log-ratio. The student is then updated by policy gradient (e.g., REINFORCE) to maximize the total discounted return: $$ \nabla J = \mathbb{E}_{\tau \sim \pi_{\text{strong}}} \left[ \sum_{t=1}^T \nabla \log \pi_{\text{strong}}(a_t|s_t) \cdot R_t \right], \quad R_t = \sum_{k=t}^T \gamma^{k-t} r(s_k, a_k) $$ Critically, no reward model is trained; the teacher's two checkpoints provide the signal. The student learns from its own trajectories (on-policy), not from teacher-generated data, avoiding distribution mismatch.
Implementation Details§
- Teacher: Qwen3-0.6B, pre-RL checkpoint and post-RL checkpoint (trained with RLVR on MATH-like problems).
- Student: Qwen3-1.7B, initialized from its own pre-RL checkpoint.
- Training: Student generates 16 rollouts per prompt; teacher computes log-ratio at each step; advantage estimation with GAE($\lambda=0.95$); PPO clip $\epsilon=0.2$; learning rate 1e-6; 4 hours on 8×A100 GPUs.
- Reward clip: $r$ clipped to $[-5,5]$ for stability.
Code Snippet (PyTorch-like)§
# Direct-OPD training loop
for prompts in dataloader:
# Student on-policy rollouts
student_responses, log_probs_student = student.generate(prompts, return_log_probs=True)
# Teacher log-probabilities for each token (post and pre)
with torch.no_grad():
log_probs_post = teacher_post.log_prob(student_responses)
log_probs_pre = teacher_pre.log_prob(student_responses)
# Implicit reward = log-ratio
rewards = log_probs_post - log_probs_pre # shape: [batch, seq_len]
rewards = torch.clamp(rewards, -5.0, 5.0)
# Discounted returns
returns = compute_gae(rewards, dones, gamma=0.99, lam=0.95)
# Policy gradient update (PPO)
ratio = torch.exp(log_probs_student - log_probs_old)
surrogate = -torch.min(ratio * returns,
torch.clamp(ratio, 1-eps, 1+eps) * returns)
loss = surrogate.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()Results§
- AIME 2024: Qwen3-1.7B improved from 48.3% to 62.4% (surpasses direct RL on the same student by 10+ points).
- Composable shifts: Multiple weak-to-strong transfers can be composed by adding log-ratios from different teachers.
- Efficiency: Training cost ~4 hours on 8 A100s vs. days for RLVR on the strong model.
Key Insight§
The implicit reward from the weak teacher's policy shift captures "which tokens RL made more likely" — a transferable signal that aligns the student's behavior without needing the teacher's absolute quality.
Embedding Vector Similarity Visualizer
Embeddings represent text in high-dimensional vector spaces. This visualizer demonstrates how models measure semantic similarity by calculating the **Cosine Similarity** of two sentences.
Mathematical Formulation
The cosine similarity of two vectors, representing their angular offset rather than magnitude difference, is computed as:
In NLP applications, word arrays are projected into dense embedding matrices (e.g. 1536 dimensions). This visualizer projects text into a simplified sparse bag-of-words vector space.