HyperAIHyperAI

Command Palette

Search for a command to run...

WAYPOINT-OPD: TRANSFER RL-INDUZIERTER POLICY-VERSCHIEBUNGEN DURCH GEGATETE ON-POLICY-DISTILLATION

Bingxi Hou Guochao Jiang Guofeng Quan Weiqing Li Wenfeng Feng Guohua Liu Yuewei Zhang

Zusammenfassung

Policy-Shift-Destillation verwendet die Differenz zwischen den Log-Wahrscheinlichkeiten eines Lehrers vor und nach dem Reinforcement Learning (RL) als Supervision für vom Schüler generierte Trajektorien. Diese Endpunkt-Policy-Verschiebung kann jedoch nicht zeigen, ob sich lokale Aktionspräferenzen während des Lehrertrainings konsistent entwickelt haben. Dieselbe Endpunktverschiebung kann aus frühen und späten Verschiebungen resultieren, die sich gegenseitig verstärken oder teilweise aufheben. Wir untersuchen, ob die Übereinstimmung zwischen frühen und späten Policy-Verschiebungen des Lehrers einen selektiven Transfer leiten kann. Wir führen Waypoint-guided On-Policy Distillation (Waypoint-OPD) ein, die einen zwischenzeitlichen Lehrer-Checkpoint aus demselben RL-Lauf verwendet, um frühe und späte Verschiebungen in den relativen Aktionspräferenzen an vom Schüler besuchten Präfixen zu vergleichen. Wir steuern die Endpunktverschiebung über ein Gate gemäß der richtungsbezogenen Übereinstimmung zwischen frühen und späten Verschiebungen und unterdrücken sie, wenn sie in entgegengesetzte Richtungen zeigen. Im Vergleich zum Transfer nur mit dem Endpunkt benötigt unsere Methode nur einen zusätzlichen Bewertungsdurchlauf mit eingefrorenem Lehrer. Mit RL-trainierten Qwen3-4B-Lehrern evaluieren wir unsere Methode mit Qwen3-8B, Qwen3-14B und Qwen3-30B A3B als Schülermodellen im mathematischen und Code-Schlussfolgern. In allen sechs Domänen-Schüler-Einstellungen erzielt Waypoint-OPD den höchsten domänendurchschnittlichen Wert unter den verglichenen Destillationsmethoden. Ablationen mit dem Qwen3-8B-Schüler zeigen, dass das Entfernen oder zufällige Umsortieren der Gates die mathematische Schlussfolgerungsleistung verringert. Diese Ergebnisse stützen die Verwendung der richtungsbezogenen Übereinstimmung zwischen frühen und späten Policy-Verschiebungen des Lehrers zur Gewichtung der Endpunktverschiebung während des Schülertrainings. Unser Code ist verfügbar unter https://github.com/hhhbxloveai/WaypointOPD.

One-sentence Summary

Alibaba Cloud Computing researchers introduce Waypoint-guided On-Policy Distillation (Waypoint-OPD), which uses an intermediate teacher checkpoint to compare early and late RL policy shifts, gates endpoint supervision when they disagree, requires only one extra frozen-teacher scoring pass, and yields the highest domain-average scores across Qwen3-8B, Qwen3-14B, and Qwen3-30B A3B students on mathematical and code reasoning.

Key Contributions

  • The paper introduces Waypoint-guided On-Policy Distillation (Waypoint-OPD), which uses an intermediate teacher checkpoint to compare early and late shifts in relative action preferences at student-generated prefixes and gates the centered Initial-to-Final policy shift by their directional agreement.
  • Directional agreement between consecutive teacher RL shifts serves as prefix-level evidence for selective endpoint transfer, suppressing updates where early and late shifts oppose each other while requiring only one additional frozen-teacher scoring pass and no verifier calls during distillation.
  • Experiments with RL-trained Qwen3-4B teachers and Qwen3-8B, Qwen3-14B, and Qwen3-30B A3B students in mathematical and code reasoning demonstrate that Waypoint-OPD achieves the highest domain-average score in all six domain-student settings, and ablations show that removing or shuffling the gates reduces mathematical reasoning performance.

Introduction

Reinforcement learning improves language models on instruction following with human preference rewards and on mathematical and code reasoning with verifiable rewards, but reusing a teacher model’s RL gains across target models remains costly. Prior on-policy distillation methods transfer either the teacher’s absolute distribution, a teacher-to-base change, or the endpoint log-probability ratio between pre-RL and post-RL checkpoints. The endpoint approach is limited because the same Initial-to-Final shift can combine early and late policy changes that point in compatible or opposing directions, and endpoint-only signals cannot distinguish these cases. The authors leverage an intermediate Waypoint checkpoint from the teacher’s RL run to measure agreement between the Initial-to-Waypoint and Waypoint-to-Final shifts, then use that agreement to gate the centered Initial-to-Final signal on student-generated prefixes. Waypoint-OPD retains aligned RL-induced changes, suppresses reversals, and requires only one additional frozen teacher scoring pass with no new teacher training or verifier calls during student transfer.

Method

The authors propose Waypoint-OPD, a method that augments the endpoint signal of direct on-policy distillation by incorporating a Waypoint checkpoint from the same teacher reinforcement learning run. Instead of using the Waypoint model as an imitation target, the framework leverages it to supply consistency evidence. Agreement between the early and late policy shifts is used to weight the endpoint signal on student-visited prefixes.

As shown in the figure below:

The training pipeline consists of four primary stages: on-policy rollout, teacher shift computation, agreement gating with batch scaling, and the student update.

Student-Local Trajectory Consistency At a student-visited prefix sts_tst​, the method selects the top-kkk token candidates with the highest student probabilities and computes a normalized student distribution qtq_tqt​ over this support. These weights focus the comparison on likely student actions. The authors then score these candidates using three frozen teacher checkpoints: the initial π0\pi_0π0​, Waypoint πm\pi_mπm​, and final π1\pi_1π1​ models. Two consecutive teacher shifts are calculated as:

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

These shifts telescope to the total endpoint shift δt=δt(1)+δt(2)\delta_t = \delta_t^{(1)} + \delta_t^{(2)}δt​=δt(1)​+δt(2)​.

To remove action-independent components that might affect the angle between raw shift vectors, a centering operator is applied using the normalized student distribution. The centered shifts for the early segment utu_tut​, late segment vtv_tvt​, and total endpoint dtd_tdt​ are then used to compute local temporal agreement. This agreement is measured as the cosine similarity between the centered early and late shift vectors, weighted by the student probabilities:

ct=⟨ut,vt⟩qt∥ut∥qt∥vt∥qtc_t = \frac{\langle u_t, v_t \rangle_{q_t}}{\| u_t \|_{q_t} \| v_t \|_{q_t}}ct​=∥ut​∥qt​​∥vt​∥qt​​⟨ut​,vt​⟩qt​​​

Agreement Gating and Batch-Scale Normalization To filter the endpoint signal, a positive-cosine gate is applied. If the centered segment norms exceed a small numerical threshold ε\varepsilonε, the gate value is set to:

gt=max⁡(0,ct)g_t = \max(0, c_t)gt​=max(0,ct​)

This mechanism suppresses the endpoint signal when the early and late shifts are misaligned and preserves it when they are aligned. If either norm falls below the threshold, a low-norm fallback sets the gate to one.

To stabilize the average gate scale across different batches, all gated rewards within a batch are multiplied by a shared scaling factor κB\kappa_BκB​:

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

This batch scaling preserves relative gate weights and zero gates while preventing the reward-vector norm from collapsing.

Dense Candidate Objective The final coefficient for each candidate action aaa is derived by combining the batch scaling factor, the normalized student probability, the agreement gate, and the centered total endpoint shift:

At(a)=sg[κBqt(a)gtdt(a)]A_t(a) = \text{sg}[\kappa_B q_t(a) g_t d_t(a)]At​(a)=sg[κB​qt​(a)gt​dt​(a)]

A stop-gradient operator sg\text{sg}sg is applied to detach all these factors during differentiation.

The local on-policy surrogate objective maximizes the weighted log-probabilities of the candidate actions while regularizing the student policy toward its own initialization using a KL divergence penalty:

JWP(θ)=Ex∼D[Ey∼πθ(⋅∣x)[∑t=1T∑a∈KtAt(a)log⁡πθ(a∣st)]−βD^KL(πθ(⋅∣x)∥πS(⋅∣x))]\mathcal{J}_{\text{WP}}(\theta) = \mathbb{E}_{x \sim \mathcal{D}} \left[ \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}_{\text{KL}}(\pi_\theta(\cdot \mid x) \| \pi_S(\cdot \mid x)) \right]JWP​(θ)=Ex∼D​[Ey∼πθ​(⋅∣x)​[t=1∑T​a∈Kt​∑​At​(a)logπθ​(a∣st​)]−βDKL​(πθ​(⋅∣x)∥πS​(⋅∣x))]

The rollout samples and candidate supports are held fixed when differentiating this surrogate. The implementation utilizes dual-clipped PPO, which shares the same gradient as the surrogate at the rollout policy. Relative to direct on-policy distillation, this approach adds only one frozen teacher checkpoint and a scoring pass, requiring no additional trainable parameters or student rollouts.

Experiment

This work evaluates Waypoint-OPD, which transfers policy shifts from frozen Qwen3-4B RL checkpoints at initial, waypoint, and final stages to Qwen3-8B, Qwen3-14B, and Qwen3-30B-A3B students on math and code benchmarks, using on-policy distillation and endpoint transfer baselines. Waypoint-OPD achieves the highest domain average in all six settings and improves over Direct-OPD across both domains and all student scales, with the larger students surpassing the final teacher on math and code averages. Learning trajectory results show sustained gains over Direct-OPD during training, and ablations on Qwen3-8B math indicate that temporal gating, correct gate-prefix assignment, and mean normalization each contribute to performance.

Waypoint-OPD achieves the highest domain-average score among evaluated student methods across all three student sizes in both math and code. On math, it outperforms Direct-OPD, with the 14B student surpassing the final teacher on all seven benchmarks and the 30B-A3B student also exceeding the teacher average, while the 8B student essentially matches it. Gains are broad and sustained, improving 19 of 21 math benchmark comparisons and holding an advantage from step 60 onward in the 8B learning trajectory. Waypoint-OPD records the highest domain-average score among evaluated student methods in all six settings spanning math and code. For math, Waypoint-OPD improves over Direct-OPD by about 1.6 to 2.2 percentage points across the three student sizes. The 14B and 30B-A3B students exceed the final teacher math average, while the 8B student essentially matches it. Waypoint-OPD improves 19 of 21 math benchmark comparisons over Direct-OPD and sustains gains from step 60 onward in the 8B learning trajectory.

Waypoint-OPD achieves the highest average code reasoning performance among the evaluated student methods across all three student scales. It improves over Direct-OPD in every code comparison, with the largest average gain on the smallest student, while the 14B and 30B-A3B students exceed the final teacher checkpoint average. Waypoint-OPD is the strongest student method for code reasoning across all three student sizes. It improves over Direct-OPD in all 15 code comparisons, and the 14B and 30B-A3B students surpass the final teacher average code score.

Ablations on Qwen3-8B mathematical reasoning show that Waypoint-OPD benefits most from consistency-dependent gating and from matching each gate to the prefix that produced it. Removing the gate lowers accuracy on all seven benchmarks, while shuffled gate assignment is even more harmful. Normalization provides a smaller additional average benefit and improves most benchmarks, though the unnormalized variant still outperforms Direct-OPD on all seven benchmarks. Removing the gate reduces average accuracy and produces lower scores on all seven benchmarks. Shuffled gate assignment is more detrimental than removing the gate, with the same ordering across every benchmark. Mean normalization gives a smaller average benefit and improves six of seven benchmarks, with HMMT February 2025 as the exception. Even without normalization, the method still outperforms Direct-OPD on average and on all seven benchmarks.

Waypoint-OPD consistently outperforms Direct-OPD across math and code reasoning tasks for all evaluated student sizes, with the 14B and 30B-A3B students often surpassing the final teacher average and the 8B student roughly matching it. The gains are broad across benchmarks and sustained during training. Ablations on mathematical reasoning show that consistency-dependent gating and matching each gate to its originating prefix are the most important components, while mean normalization provides a smaller additional benefit and the method remains effective without it.


KI mit KI entwickeln

Von der Idee bis zum Launch – beschleunigen Sie Ihre KI-Entwicklung mit kostenlosem KI-Co-Coding, sofort einsatzbereiter Umgebung und bestem GPU-Preis.

KI-gestütztes kollaboratives Programmieren
Sofort einsatzbereite GPUs
Die besten Preise

HyperAI Newsletters

Abonnieren Sie unsere neuesten Updates
Wir werden die neuesten Updates der Woche in Ihren Posteingang liefern um neun Uhr jeden Montagmorgen
Unterstützt von MailChimp