Technical Summary: Asymmetric On-Policy Distillation (AOPD)

1. Problem Statement

On-Policy Distillation (OPD) has emerged as a powerful technique for training student Large Language Models (LLMs) by leveraging token-level feedback from a teacher model on student-generated trajectories. Unlike off-policy distillation, which suffers from exposure bias, OPD aligns the student with the teacher's distribution during inference-time rollouts. However, the authors identify that standard OPD, which relies on an advantage-weighted policy gradient, suffers from three critical structural weaknesses in non-positive advantage regions:

1. High Variance Updates: In regions where the student's probability is significantly lower than the teacher's (negative advantage), the logarithmic nature of the advantage formulation creates unbounded scalar values. This leads to heavy-tailed gradient distributions and unstable optimization.
2. Vanishing Gradients: A vast majority of generated tokens often cluster around a zero advantage. In these "neutral" regions, the standard advantage-weighted update vanishes, discarding the rich conditional distribution information provided by the teacher regarding alternative reasoning paths.
3. Exploration Bottlenecks (Black Holes): When a student makes an error (negative advantage), the standard update suppresses the sampled token but redistributes probability mass strictly according to the student's current prior. If the correct token has a near-zero prior probability in the student model, the corrective update is negligible. This traps the model in an "exploration black hole," preventing it from discovering correct alternatives without external guidance.

2. Methodology: Asymmetric On-Policy Distillation (AOPD)

To address these limitations, the authors propose Asymmetric On-Policy Distillation (AOPD). This framework bridges reinforcement learning (exploitation) and supervised learning (imitation) by dynamically switching learning modes at the token level based on the advantage signal.

Core Mechanism

AOPD decomposes the training objective into two distinct regimes:

• Positive Advantage Regions (Exploitation): When the student token receives a positive advantage ($A_t > 0$), the method retains the standard policy gradient update. This reinforces successful exploration where the student's reasoning aligns with or exceeds the teacher's preference.
• Non-Positive Advantage Regions (Imitation): When the advantage is zero or negative ($A_t \le 0$), AOPD replaces the ineffective policy gradient with localized divergence minimization.

The Intervention Objective

Instead of relying on the scalar advantage, AOPD minimizes the Forward KL Divergence between the teacher and student distributions, but restricted to a Teacher-Defined Support.

• Teacher Support ($S_t$): Defined as the top-$K$ tokens with the highest probability under the teacher distribution $P_T(\cdot | c_t)$.
• Loss Function: On this support, the loss is formulated as:
  $$L_{FKL}^t = \frac{1}{K} \sum_{v \in S_t} P_T(v | c_t) (\log P_T(v | c_t) - \log P_S(v | c_t))$$
• Generalized Objective: The total AOPD loss combines the positive reinforcement term ($L_{Pos}$) with the Forward KL guidance for non-positive regions. A threshold parameter $\tau$ can be used to trigger intervention, with $\tau=0$ being the default setting to target exactly the non-positive regime.

Theoretical Justification

• Bounded Updates: By switching to Forward KL on a bounded support, the update magnitude is constrained by the probability simplex ($P_T - P_S$), eliminating the heavy-tailed variance caused by unbounded log-probability differences.
• Gradient Recovery: In zero-advantage regions, the Forward KL term utilizes the full conditional distribution of the teacher, providing meaningful correction signals even when the scalar advantage is zero.
• Escaping Black Holes: By explicitly matching the teacher's distribution on the support, AOPD injects directional signals toward tokens favored by the teacher but underestimated by the student, effectively breaking the exploration bottleneck.

3. Key Contributions

1. Analysis of OPD Limitations: The paper provides a theoretical and empirical analysis demonstrating that standard advantage-based updates fail in negative and zero-advantage regions due to high variance, vanishing gradients, and exploration bottlenecks.
2. AOPD Framework: The proposal of an asymmetric training framework that adaptively shifts from advantage-weighted policy gradient to localized teacher distribution matching. This design preserves the benefits of RL-based exploitation while applying strong, directed supervision at optimization bottlenecks.
3. Empirical Validation: Comprehensive experiments across multiple model scales and initialization settings show that AOPD consistently outperforms standard OPD and other baselines (SeqKD, GKD, ExOPD).

4. Experimental Results

The authors evaluated AOPD on competition-level mathematical reasoning benchmarks (AIME 2024, AIME 2025, HMMT 2025) using Qwen3 models (4B and 8B) with varying SFT warm-up steps.

• Performance Gains: AOPD consistently achieved the best overall performance.
  • Under strong initialization (3K SFT warm-up), AOPD improved average Pass@1 by 4.09% over standard OPD.
  • Under weak initialization (1K SFT warm-up), the gain was even more significant at 8.34%, demonstrating AOPD's robustness when student models start with weaker foundational capabilities.
• Training Dynamics: AOPD exhibited stable training curves without the initial performance degradation observed in OPD and ExOPD. It maintained higher policy entropy throughout training, indicating a less concentrated policy space.
• Continual Learning: In sequential tool-use adaptation tasks, AOPD demonstrated superior capability retention. While other methods suffered performance drops in mathematical reasoning after learning new tools, AOPD actually improved its math retention (from 67.43 to 68.04 Pass@4) while achieving competitive tool-use success rates.

5. Significance and Claims

The paper claims that AOPD represents a significant step forward in distilling reasoning capabilities by resolving the structural inefficiencies of standard on-policy distillation. The authors argue that:

• Asymmetry is Key: Treating positive and non-positive advantage regions differently is crucial. Positive regions benefit from RL-style exploration, while non-positive regions require the direct, high-fidelity guidance of supervised imitation to avoid optimization stagnation.
• Robustness to Initialization: The method is particularly effective for models with weaker initial capabilities, where standard OPD often fails to recover from early optimization degradation.
• Preservation of Capabilities: By localizing teacher intervention to difficult positions rather than reshaping the entire policy, AOPD better preserves previously acquired reasoning skills during continual learning.

The authors conclude that combining exploitation (RL) and imitation (supervised learning) at the token level offers a principled path to improving both training efficiency and the final capabilities of student models, particularly in complex reasoning tasks.