HyperAIHyperAI

Command Palette

Search for a command to run...

強化学習
LLM

Waypoint-OPD: ゲート付きオンポリシー蒸留によるRL誘発方策シフトの転移

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

概要

方策シフト蒸留は、強化学習(RL)の前後における教師の対数確率の差を、生徒が生成した軌跡上での教師信号として再利用する。しかし、この端点方策シフトは、教師の訓練中に局所的な行動選好が一貫して変化したかどうかを明らかにできない。同じ端点シフトは、互いに強め合う、または部分的に打ち消し合う早期と後期のシフトから生じうる。我々は、教師の早期と後期の方策シフト間の一致が選択的転移を導けるかどうかを検討する。我々は、Waypoint-guided On-Policy Distillation(Waypoint-OPD)を導入する。これは、同じRL実行からの中間教師チェックポイントを用いて、生徒が訪問したプレフィックスにおける相対的行動選好の早期と後期のシフトを比較する。我々は、早期と後期のシフト間の方向的一致に基づいて端点シフトをゲートし、それらが反対方向を指す場合にそれを抑制する。端点のみの転移と比較して、本手法は凍結した教師によるスコアリングパスを1回追加するだけで済む。RLで訓練したQwen3-4B教師を用い、Qwen3-8B、Qwen3-14B、Qwen3-30B A3Bを生徒モデルとして数学的推論とコード推論で本手法を評価する。6つのドメイン・生徒設定すべてにおいて、Waypoint-OPDは比較した蒸留手法の中で最高のドメイン平均スコアを達成する。Qwen3-8B生徒を用いたアブレーションでは、ゲートを除去またはシャッフルすると数学的推論性能が低下することが示される。これらの結果は、生徒の訓練中に端点シフトを重み付けするために、教師の早期と後期の方策シフト間の方向的一致を利用することを支持する。コードは 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.


AIでAIを構築

アイデアからローンチまで — 無料のAIコーディング支援、すぐに使える環境、最高のGPU価格でAI開発を加速。

AI コーディング補助
すぐに使える GPU
最適な料金体系

HyperAI Newsletters

最新情報を購読する
北京時間 毎週月曜日の午前9時 に、その週の最新情報をメールでお届けします
メール配信サービスは MailChimp によって提供されています