Summary (Overview)

  • Waypoint-OPD is a novel on-policy distillation method that uses an intermediate teacher checkpoint (Waypoint) from the same RL run to gate the transfer of RL-induced policy shifts to student models.
  • The key insight is that the endpoint (Initial-to-Final) policy shift alone cannot distinguish whether early and late teacher policy changes reinforce or cancel each other; Waypoint-OPD uses directional agreement between consecutive shifts to weight the endpoint signal.
  • The method requires only one additional frozen-teacher scoring pass compared to endpoint-only transfer, with no new teacher training or verifier calls during student transfer.
  • Experiments with Qwen3-4B teachers and Qwen3-8B, Qwen3-14B, and Qwen3-30B-A3B students across mathematical and code reasoning show the highest domain-average scores in all six settings compared to baselines.
  • Ablations confirm that temporal gating, prefix-specific gate assignment, and batch-scale normalization each contribute to the method's performance gains.

Introduction and Theoretical Foundation

Background and Motivation

Reinforcement learning improves language-model reasoning using verifiable rewards, but requires repeated generation, verification, and optimization of on-policy responses. Reusing a teacher's RL improvements offers a way to share this training effort across target models. The central question is: which local changes in the teacher's policy should guide a student on its own generated trajectories?

Theoretical Basis

Standard on-policy distillation (OPD) trains the student to match the teacher's token distributions on student-generated prefixes:

LOPD(θ)=Ex∼D,y∼πθ(⋅∣x)[∑t=1TDKL(πθ(⋅∣st)∥π1(⋅∣st))](1)\mathcal{L}_{\mathrm{OPD}}(\theta) = \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_{\theta}(\cdot | x)} \left[ \sum_{t=1}^{T} D_{\mathrm{KL}}(\pi_{\theta}(\cdot | s_t) \| \pi_1(\cdot | s_t)) \right] \tag{1}

Direct-OPD uses the teacher's RL-induced policy shift as an implicit reward, decomposing the Initial-to-Final shift into token-level log-probability ratios:

Δ(0,1)(y∣x)=log⁡π1(y∣x)π0(y∣x)=∑t=1Tδt(yt)(2)\Delta^{(0,1)}(y \mid x) = \log \frac{\pi_1(y \mid x)}{\pi_0(y \mid x)} = \sum_{t=1}^{T} \delta_t(y_t) \tag{2}

where

δt(a)=log⁡π1(a∣st)−log⁡π0(a∣st)(3)\delta_t(a) = \log \pi_1(a \mid s_t) - \log \pi_0(a \mid s_t) \tag{3}

Key limitation addressed: The Initial-to-Final log-probability ratio captures only the change between pre- and post-RL policies. The same endpoint signal can arise from early and late policy changes that point in similar or opposing directions. Endpoint-only transfer cannot distinguish these cases.

Hypothesis: Agreement between early and late teacher policy shifts provides evidence for deciding how much weight to assign an endpoint shift at a student-generated prefix.

Methodology

3.1 Student-Local Trajectory Consistency

At a student-visited prefix sts_t, the top-k most probable next-token candidates are selected via Kt=TopK(πθ(⋅∣st),k)\mathcal{K}_t = \mathrm{TopK}(\pi_{\theta}(\cdot | s_t), k). The normalized student distribution on this support is:

qt(a)=πθ(a∣st)∑b∈Ktπθ(b∣st),a∈Kt(5)q_t(a) = \frac{\pi_{\theta}(a \mid s_t)}{\sum_{b \in \mathcal{K}_t} \pi_{\theta}(b \mid s_t)}, \qquad a \in \mathcal{K}_t \tag{5}

The two consecutive teacher shifts (early and late) are:

δt(1)(a)=log⁡πm(a∣st)−log⁡π0(a∣st)(6)\delta_t^{(1)}(a) = \log \pi_m(a \mid s_t) - \log \pi_0(a \mid s_t) \tag{6} δt(2)(a)=log⁡π1(a∣st)−log⁡πm(a∣st)(7)\delta_t^{(2)}(a) = \log \pi_1(a \mid s_t) - \log \pi_m(a \mid s_t) \tag{7}

These telescope to δt=δt(1)+δt(2)\delta_t = \delta_t^{(1)} + \delta_t^{(2)}.

To remove action-independent components (uniform shifts across candidates), a qtq_t-centering operator is applied:

Cqt[z](a)=z(a)−Eb∼qt[z(b)](8)C_{q_t}[z](a) = z(a) - \mathbb{E}_{b \sim q_t}[z(b)] \tag{8}

yielding ut=Cqt[δt(1)]u_t = C_{q_t}[\delta_t^{(1)}], vt=Cqt[δt(2)]v_t = C_{q_t}[\delta_t^{(2)}], and dt=Cqt[δt]=ut+vtd_t = C_{q_t}[\delta_t] = u_t + v_t.

3.2 Agreement Gating and Batch-Scale Normalization

The local temporal agreement (cosine similarity weighted by student probability) is:

ct=⟨ut,vt⟩qt∥ut∥qt∥vt∥qt(9)c_t = \frac{\langle u_t, v_t \rangle_{q_t}}{\|u_t\|_{q_t} \|v_t\|_{q_t}} \tag{9}

When both centered segment norms exceed ε=10−8\varepsilon = 10^{-8}, the positive-cosine gate is:

gt=max⁡(0,ct)(10)g_t = \max(0, c_t) \tag{10}

If either norm is at most ε\varepsilon, a low-norm fallback sets gt=1g_t = 1. The gate gt∈[0,1]g_t \in [0,1] is shared by all candidates at prefix sts_t, removing the endpoint signal when ct≤0c_t \leq 0 and scaling its magnitude by ctc_t otherwise.

Illustrative example: For a nonzero centered endpoint dd: the pair u=v=d/2u = v = d/2 gives c=1c = 1 and g=1g = 1, whereas u=2d,v=−du = 2d, v = -d gives c=−1c = -1 and g=0g = 0. Both yield u+v=du + v = d, so endpoint-only transfer cannot distinguish them.

Batch-scale normalization stabilizes the average gate scale:

κB=clip⁡(τmax⁡(∣Bpref∣−1∑t∈Bprefgt,εnum),κmin⁡,κmax⁡)(11)\kappa_B = \operatorname{clip}\left(\frac{\tau}{\max\left(|B_{\text{pref}}|^{-1} \sum_{t \in B_{\text{pref}}} g_t, \varepsilon_{\text{num}}\right)}, \kappa_{\min}, \kappa_{\max}\right) \tag{11}

3.3 Dense Candidate Objective

The final coefficient for candidate action a∈Kta \in \mathcal{K}_t is:

At(a)=sg⁡[κBqt(a)gtdt(a)](12)A_t(a) = \operatorname{sg}[\kappa_B q_t(a) g_t d_t(a)] \tag{12}

The local on-policy surrogate to maximize:

JWP(θ)=Ex∼D[Ey∼πθ(⋅∣x)[∑t=1T∑a∈KtAt(a)log⁡πθ(a∣st)]−βD^KL(πθ(⋅∣x)∥πS(⋅∣x))](13)\mathcal{J}_{\mathrm{WP}}(\theta) = \mathbb{E}_{x \sim \mathcal{D}} \Bigg[ \mathbb{E}_{y \sim \pi_{\theta}(\cdot | x)} \left[ \sum_{t=1}^{T} \sum_{a \in \mathcal{K}_t} A_t(a) \log \pi_{\theta}(a \mid s_t) \right] - \beta \widehat{D}_{\mathrm{KL}}(\pi_{\theta}(\cdot \mid x) \| \pi_S(\cdot \mid x)) \Bigg] \tag{13}

The implementation uses dual-clipped PPO, which has the same gradient as this surrogate at the rollout policy.

Empirical Validation / Results

Experimental Setup

  • Teachers: Three checkpoints from the same Qwen3-4B RL run: Initial (before RL), Waypoint (step 600), Final (step 1200)
  • Students: Qwen3-8B, Qwen3-14B, Qwen3-30B-A3B
  • Domains: Mathematical reasoning (7 benchmarks) and code reasoning (5 evaluations)
  • Baselines: Base, OPD, OPD², Direct-OPD

Main Results — Mathematics

Table 1: Mathematical reasoning across three student sizes (mean@32 accuracy, %)

MethodAIME24AIME25AIME26HMMT25FebHMMT25NovHMMT26FebAMC23Avg.
Teacher: Qwen3-4B
Initial21.6721.6721.6712.5010.0015.1571.2524.84
Waypoint58.9654.7956.7731.3535.7331.7293.4451.82
Final65.8360.0060.8335.8340.8335.1394.4556.13
Student: Qwen3-8B
Base29.1720.8318.3311.6713.3318.9470.0026.04
OPD63.6549.6954.9029.0638.6532.5888.7551.04
OPD²64.4848.3353.9628.4440.5231.5392.9751.46
Direct-OPD67.2955.3157.4031.7740.8333.7191.4153.96
Waypoint-OPD68.8557.9260.2134.3842.6036.3692.6656.14
Student: Qwen3-14B
Base30.0027.5020.0010.0011.6718.9473.1327.32
OPD69.1755.5258.4436.7745.2131.8296.2556.17
OPD²69.4846.9854.3833.1341.9831.9191.9552.83
Direct-OPD70.1056.4661.9837.9247.6035.1394.0657.61
Waypoint-OPD70.2160.5263.7540.2147.0836.5596.0959.20
Student: Qwen3-30B-A3B
Base31.6722.5022.5010.0013.3321.2175.0028.03
OPD63.5449.0657.5030.8337.7128.0391.8851.22
OPD²64.9051.7760.4231.5638.5430.2191.3352.68
Direct-OPD69.1754.6961.9838.9641.5633.0592.3455.96
Waypoint-OPD70.9458.4463.6539.9044.6932.7794.5357.85

Waypoint-OPD improves over Direct-OPD by 2.18, 1.59, and 1.89 percentage points for Qwen3-8B, Qwen3-14B, and Qwen3-30B-A3B, respectively, improving 19 of 21 benchmark comparisons.

Main Results — Code

Table 2: Average code reasoning performance (average mean@8 scores, %)

MethodQwen3-8BQwen3-14BQwen3-30B-A3B
Base41.0347.7250.33
OPD43.9151.6651.91
OPD²51.2155.6055.96
Direct-OPD50.5956.3056.35
Waypoint-OPD53.2557.8257.54

Waypoint-OPD improves over Direct-OPD by 2.66, 1.52, and 1.19 percentage points on average, with gains in all 15 code comparisons.

Learning Dynamics

Waypoint-OPD delivers consistent gains over Direct-OPD across training checkpoints (steps 0–100), leading on all three benchmarks from step 60 onward. On HMMT November 2025, Waypoint-OPD exceeds the Final teacher's 40.83 at steps 60, 80, and 100, reaching 42.60 at step 60, demonstrating ability to surpass the Final teacher's score.

Ablation Results

Table 3: Gate and normalization ablations on Qwen3-8B mathematical reasoning

MethodAIME24AIME25AIME26HMMT25FebHMMT25NovHMMT26FebAMC23Avg.
No gate66.2554.2756.6732.6039.7934.6690.2353.50
Shuffled gate63.3352.9254.5830.0037.0831.0689.4551.20
No normalization68.1356.4658.3334.7941.5635.2392.0355.22
Direct-OPD67.2955.3157.4031.7740.8333.7191.4153.96
Waypoint-OPD68.8557.9260.2134.3842.6036.3692.6656.14

Key ablation findings:

  • Removing the gate: Lowers Avg. from 56.14 to 53.50 (−2.64 points)
  • Shuffling gates: Lowers Avg. to 51.20 (−4.94 points vs. full method), more detrimental than removing gates entirely
  • Removing normalization: Lowers Avg. by 0.92 points to 55.22, though still exceeding Direct-OPD by 1.26 points

Theoretical and Practical Implications

Theoretical Contributions

  1. Temporal agreement as a selection criterion: The paper introduces the use of an intermediate Waypoint checkpoint to assess temporal agreement at student-generated prefixes, revealing directional information unavailable from the endpoint pair alone. This addresses a fundamental limitation of endpoint-only distillation methods.

  2. Decomposition of policy shifts: The method formalizes that the same endpoint shift δt=δt(1)+δt(2)\delta_t = \delta_t^{(1)} + \delta_t^{(2)} can arise from early and late shifts that reinforce or cancel each other, and provides a principled way to distinguish these cases using cosine similarity of centered segment shifts.

  3. Gating as weighting, not filtering: The positive-cosine gate gt=max⁡(0,ct)g_t = \max(0, c_t) both removes conflicting signals (ct≤0c_t \leq 0) and scales aligned signals by their agreement strength, providing a continuous weighting mechanism rather than a binary filter.

Practical Implications

  • Efficiency: Waypoint-OPD requires only one additional frozen-teacher scoring pass compared to Direct-OPD, with no new teacher training or verifier calls during student transfer — making it practical for large-scale deployment.
  • Scalability: The method works across student scales (8B to 30B-A3B) and architectures (dense and MoE), demonstrating robustness.
  • Teacher surpassing: The 14B and 30B-A3B students exceed the Final teacher's average by 3.07 and 1.72 points (math) and 2.45 and 2.17 points (code), showing that students can surpass their teacher's capabilities.
  • Practical guidance: Intermediate checkpoints from the same RL run provide valuable evidence for weighting local policy changes during student learning, suggesting that practitioners should retain and reuse intermediate checkpoints rather than only the final one.

Conclusion

Waypoint-OPD presents a novel approach to policy-shift distillation that leverages directional agreement between consecutive teacher RL shifts to gate endpoint transfer signals. The method addresses a fundamental limitation of endpoint-only transfer: the inability to distinguish between early and late policy changes that reinforce versus cancel each other.

Key takeaways:

  • Intermediate teacher checkpoints provide valuable consistency evidence for weighting local policy changes during student learning
  • The method achieves state-of-the-art transfer performance across mathematics and code reasoning at multiple student scales
  • Prefix-specific gate assignment and batch-scale normalization are both important components

Future directions suggested by the work include:

  • Investigating the role of Waypoint checkpoint position (the paper explores steps 500, 600, and 700 in Appendix F)
  • Exploring whether directional agreement alone guarantees beneficial supervision (the paper notes it does not)
  • Extending the approach to other domains and model families
  • Studying the interaction between gate design and student architecture

Related papers