PaperlayerView exact revision on arXiv

Source-linked AI summary

Asymmetric On-Policy Distillation: Bridging Exploitation and Imitation at the Token Level

Nan Jia, Haojin Yang, Xing Ma, Jiesong Lian, Shuailiang Zhang, Weipeng Zhang, Ke Zeng, Xunliang Cai, Zequn Sun

arXiv:2605.06387v4cs.LGcs.AI

TL;DR

Standard OPD can produce high-variance updates, vanishing signals, and exploration bottlenecks, limiting correction during on-policy distillation. AOPD combines positive policy-gradient reinforcement with localized teacher-distribution matching in non-positive advantage regions, and consistently improves mathematical reasoning while retaining capabilities during tool-use adaptation. Its scope is constrained by exploration failures when essential tokens lie in the student’s low-probability tail, especially after weaker warm-up.

  • Problem

    Standard OPD has high-variance negative updates, nearly vanishing zero-advantage updates, and exploration bottlenecks that limit corrective learning.

  • Method

    AOPD preserves policy-gradient updates where reliable and uses localized teacher-distribution alignment for negative-advantage tokens.

  • Results

    AOPD consistently improves mathematical reasoning across benchmarks, model scales, and initialization settings while better retaining previously acquired capabilities.

  • Takeaways & Limitations

    Combining exploitation and imitation at the token level can preserve policy entropy and improve capability retention during adaptation.

  • Takeaways & Limitations

    OPD’s corrective update can remain negligible for essential tokens with near-zero student prior probability, especially in weaker models with less warm-up.

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student on its own trajectories with token-level teacher feedback and often outperforms off-policy distillation and standard reinforcement learning. However, we find that its standard advantage weighted policy gradient suffers from three structural weaknesses, including high variance updates, vanishing gradients in zero-advantage regions, and exploration bottlenecks when corrective signals are insufficient. We therefore propose Asymmetric On-Policy Distillation (AOPD), which replaces ineffective negative reinforcement with localized divergence minimization in non-positive advantage regions while preserving positive reinforcement learning. Experiments on mathematical reasoning benchmarks show that AOPD consistently outperforms standard OPD, with average gains of 4.09 / 8.34 under strong / weak initialization, respectively. AOPD also maintains higher policy entropy during training and better capability retention during sequential tool-use adaptation.

1 Introduction

On-policy distillation avoids the exposure bias of off-policy trajectories by training on student rollouts with token-level teacher feedback, but standard OPD has unstable, weak, and exploration-limited updates. AOPD addresses these issues by switching from policy-gradient reinforcement to localized teacher-distribution matching when advantage signals are non-positive.

  • Knowledge distillation compresses large-model capabilities into smaller models, but off-policy training exposes students to training–inference mismatch.
  • OPD trains on student trajectories using token-level rewards from teacher–student log-probability differences.This reinforcement-learning formulation improves training efficiency and supports post-training.
  • Standard OPD suffers from high-variance negative updates, nearly vanishing updates near zero advantage, and blind exploration for weak models.These limitations can destabilize optimization and leave insufficient corrective signals when the student lacks foundational capability.
  • AOPD applies localized KL-based teacher-distribution matching to negative-advantage tokens instead of relying on self-exploration, while retaining reinforcement learning elsewhere.The design combines exploitation through policy gradients with stronger correction at optimization bottlenecks.
  • AOPD consistently improves over baselines on mathematical reasoning benchmarks and better preserves prior capabilities during sequential tool-use adaptation.The evaluation covers competition-level mathematical reasoning and continual-learning capability retention.

2 Related Work

Prior work established off-policy and sequence-level distillation, then shifted toward on-policy learning from student rollouts with teacher feedback. Recent extensions vary in divergence objectives, adaptive weighting, teacher count, and bidirectional or coevolutionary training.

  • Classical knowledge distillation transfers softened teacher outputs to compact students, while SeqKD extends distillation to sequence generation.
  • On-policy distillation trains on student rollouts with teacher feedback through generalized f-divergences or dense token-level reinforcement-learning rewards.
  • Recent methods add perplexity-calibrated weighting, multiple teachers, full-vocabulary reverse KL, or bidirectional expert coevolution.These approaches target adaptive supervision, multi-teacher distillation, or stronger internalization of teacher capabilities.

3 Preliminaries

Standard OPD uses token-level advantage-weighted reinforcement learning on student trajectories, but its negative and near-zero advantage regions create unstable, weak, or ineffective corrective updates.

  • OPD formulation: OPD optimizes reverse KL divergence through reinforcement learning using token-level teacher–student log-probability differences.The advantage is computed from the teacher and student conditional log-probabilities for each sampled token.
  • OPD formulation: Positive reinforcement strengthens sampled tokens favored by the teacher, while negative reinforcement suppresses disfavored sampled tokens and promotes exploration.The loss separates token positions into positive- and negative-advantage sets.
  • Limitations of OPD: Negative advantages have heavy tails, so logarithmic advantage scaling produces large update variance and fragile optimization.The unbounded scalar can amplify as the teacher or student probability approaches zero.
  • Limitations of OPD: Most sampled tokens lie near zero advantage, where OPD receives an almost zero signal despite the teacher retaining preferences over alternative reasoning paths.This produces an almost vanishing policy-gradient update in the neutral regime.
  • Limitations of OPD: Negative reinforcement can create an exploration black hole by suppressing mistakes without effectively increasing the probability of correct low-prior alternatives.The released probability mass is redistributed according to the student’s existing prior, with weaker warm-up making this pathology more likely.
  • Motivation for asymmetric training: These non-positive regions require a different optimization rule from positive reinforcement because their learning signal is insufficient.The analysis motivates an asymmetric training rule for such tokens.

4 Asymmetric On-Policy Distillation

AOPD dynamically switches between reinforcement learning for exploitation and teacher-guided distribution matching for imitation at the token level. It applies truncated forward-KL guidance in non-positive advantage regions to address OPD’s optimization bottlenecks while preserving reinforcement learning elsewhere.

  • AOPD switches between reinforcement learning for exploitation and soft-label supervised learning for imitation according to the student’s local capability.
  • Truncated Forward-KL Guidance: AOPD replaces the update in non-positive advantage regions with forward KL minimization on teacher-selected support.The support contains the teacher’s K highest-probability tokens, and forward KL is chosen because the teacher defines both the support and reference distribution.
  • Intervention Boundary: The default intervention boundary is zero advantage, while alternative thresholds interpolate between standard OPD and global supervised distribution matching.τ = −1 recovers OPD, τ = 1 recovers GKD, and the method uses τ = 0.
  • High-Variance Negative Advantages: AOPD maintains a substantially smaller early-training gradient norm than OPD by replacing high-variance negative-advantage policy gradients with bounded divergence guidance.The correction is governed by a bounded probability gap constrained by the probability simplex.
  • Zero Advantages and Exploration Bottlenecks: KL intervention in zero-advantage regions provides meaningful teacher corrections, and targeted intervention can improve performance while affecting fewer than 10% of generated tokens.The zero-advantage variant yields clear gains over OPD, while the τ = −0.2 variant achieves gains with sparse intervention.
  • Main Results: Table 1 compares AOPD with SeqKD, OPD, GKD, and ExOPD across multiple mathematical reasoning benchmarks and configuration blocks.Bold values indicate the best result within each configuration block, while underlined values indicate the second best.

5 Experiments

Across reasoning and continual-learning experiments, AOPD consistently provides stronger, more stable performance than baselines across model scales and initialization settings. Its localized guidance also improves capability retention while intervention weighting and location materially affect outcomes.

  • Main reasoning results: AOPD achieves the best overall reasoning performance across model scales and warm-up settings, including the challenging 1k-step initialization.For Qwen3-8B-Base with 1k-step warm-up, standard OPD falls below the warm-up baseline while AOPD remains strongest.
  • Training dynamics: AOPD improves steadily from the start and reaches the highest performance ceiling, whereas OPD and ExOPD show early performance drops.The degradation and recovery difficulty are especially severe under weaker initialization.
  • Continual learning: AOPD achieves a competitive 75.0% tool-call success rate while preserving the strongest mathematical reasoning performance after continual learning.Other methods suffer mathematical-performance degradation, with OPD showing the largest drop.
  • Continual learning: AOPD maintains higher policy entropy during training while applying divergence guidance to a decreasing fraction of tokens, from about 40% to 30%.Most tokens remain optimized by policy gradient, while correction is concentrated at difficult positions.
  • Divergence weighting: AOPD attains its highest Pass@1 at β = 1 and highest Pass@4 at β = 0.9, with performance increasing as divergence becomes more teacher-biased.The result supports stronger teacher guidance during interventions.
  • Support size: Increasing teacher support size K from 8 to 32 improves final Pass@1, while smaller K often produces larger early-training gains.Larger support preserves more of the teacher distribution; smaller support provides denser supervision early.
  • Intervention location: Standard AOPD with τ = 0 performs best, while shrinking the intervention range or restricting guidance to selected positions reduces performance.AOPD-Zero still outperforms τ = −0.2, and accuracy filtering trails standard AOPD.

6 Conclusion

The paper identifies three limitations of OPD and proposes AOPD, which preserves reliable policy-gradient updates while applying direct teacher alignment where stronger correction is needed. Across benchmarks and training settings, AOPD improves reasoning stability and capability retention while preserving policy entropy.

  • Conclusion: AOPD preserves policy-gradient updates where they are reliable and applies direct teacher distribution alignment where stronger correction is needed.This design targets the limitations identified in zero and negative advantage regions.
  • Conclusion: Across model scales and initialization steps, AOPD consistently improves mathematical reasoning and remains stable where standard OPD often struggles.The conclusion emphasizes robustness across the evaluated settings.
  • Conclusion: Combining exploitation and imitation at the token level can preserve policy entropy and better retain previously acquired capabilities.This conclusion is supported within the paper’s mathematical-reasoning and tool-use evaluations.

Limitations

The paper’s evaluation and method analysis leave several scope and design boundaries. Results are limited to mathematical reasoning and tool use, while the current intervention design is fixed and other training combinations remain unexplored.

  • Scope: The evaluation is limited to mathematical reasoning and tool-use settings, leaving performance on other task types unexplored.The authors identify broader task coverage as future work.
  • Method boundary: The current AOPD instantiation uses a fixed intervention design, motivating future investigation of more adaptive intervention strategies.This is an explicit limitation of the present method configuration.
  • Future direction: The paper suggests combining AOPD with reinforcement learning driven by verifiable rewards for tasks where outcome-based reward is available.This is proposed as a possible direction rather than an evaluated result.

B.2 Training Dynamics

Training dynamics vary substantially with divergence weighting: student-biased settings produce unstable gradients and entropy-collapse behavior, while teacher-biased settings are smoother. The intervention fraction falls during this collapse, coinciding with reasoning degradation.

  • Training dynamics: Student-biased divergence produces more unstable gradient behavior, with the most severe gradient-norm spikes at β = 0 and 0.1.Teacher-leaning settings β = 0.9 and 1 yield the smoothest training dynamics.
  • Training dynamics: The proportion of tokens receiving divergence guidance plummets during the entropy-collapse phase, when most generated positions register positive advantages.The accompanying ablation results associate this apparent convergence with severe reasoning-capability degradation.
  • Training dynamics: The authors attribute the collapse pattern to reward hacking, in which the student suppresses reverse KL while failing to acquire the intended capability.This explanation is presented for the anomalous reverse-KL training behavior.

C Implementation Details

The experiments use a distributed on-policy distillation setup, with supplementary Qwen3-4B-Base results under 1K-step SFT warm-up showing AOPD achieves the best overall performance.

  • Environment: Training uses VeRL for on-policy distillation and vLLM for teacher deployment across distributed A100-80G GPUs.PyTorch FSDP and Ray coordinate training on 16 GPUs, with 8 additional GPUs serving the teacher.
  • Parameters: Table 3 summarizes the main training and evaluation parameters, with the complete parameter set available in the code repository.
  • Supplementary results: AOPD attains the best overall performance for Qwen3-4B-Base with a 1K-step SFT warm-up initialization.This result is consistent with the trend reported in the main text.

D.2 Detailed Training Curves

Across both weak and strong initialization, AOPD consistently achieves the highest Pass@1/4/8 scores across all evaluated benchmarks. Its advantage is especially pronounced under weak initialization, while it retains the highest performance ceiling under strong initialization.

  • Training curves: AOPD consistently achieves the highest Pass@1/4/8 scores across all benchmarks under both initialization settings.
  • Weak initialization: Under weak initialization, AOPD’s performance advantage over the baselines is particularly pronounced.The larger gap indicates effectiveness for students with weaker foundational capabilities.
  • Strong initialization: Under strong initialization, the margin between AOPD and baselines narrows, but AOPD still reaches the highest performance ceiling.

D.3 Extended Ablation Studies on τ

The τ ablation shows that reducing intervention weakens AOPD toward standard OPD: performance declines and plateaus once τ reaches −0.2, when fewer than 10% of tokens receive supervision. Comparisons with stabilization and adaptive OPD methods further position AOPD as applying forward-KL guidance throughout the non-positive advantage region.

  • Extended ablation studies on τ: Mathematical reasoning performance gradually declines as τ decreases and begins to plateau at τ = −0.2.Lower τ reduces the proportion of positions receiving forward-KL supervision.
  • Intervention ratio: When τ ≤ −0.2, fewer than 10% of generated tokens trigger the supervision loss.The resulting sparse teacher guidance explains the smaller performance gap beyond this threshold.
  • Interpretation: AOPD’s full benefit is realized when teacher guidance covers the zero-advantage and negative-advantage regions.Reducing forward-KL supervision in the negative-advantage region correspondingly weakens AOPD’s effect.
  • OPD stabilization methods: OPD-Norm and OPD-Clip reduce high-variance advantage effects without changing the underlying policy-gradient update.These strategies do not address insufficient corrective directions in non-positive advantage regions.
  • Comparative methods: AOPD applies forward-KL guidance throughout the non-positive advantage region, unlike KDRL’s outcome-level signals and EOPD’s teacher-entropy selection.The comparison includes methods trained for 60 steps under the same experimental configuration.
Loading 2605.06387v4…