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

The evolution of large language models (LLMs) such as GPT and BERT has driven significant advances in NLP, enabling tasks like translation, summarization, and reasoning. Knowledge distillation emerged as an effective method to compress these models, transferring knowledge from a large teacher to a smaller student. Early approaches like sequence-level distillation (Kim and Rush, 2016) focused on matching output distributions, but suffered from exposure bias. To mitigate this, on-policy distillation (Lu and Lab, 2025) introduced token-level feedback, leveraging reinforcement learning to optimize the student model directly on its own trajectories. Despite improvements, these methods face challenges: high variance in policy gradients, gradient vanishing in neutral regions, and exploration bottlenecks, especially in complex reasoning tasks like mathematics or code understanding. Recent works have attempted to address these issues through various divergence measures and curriculum strategies, but a comprehensive solution remains elusive.

Standard OPD relies on advantage-weighted policy gradients, which suffer from high variance when advantage signals are large and unstable in zero or negative advantage regions. This leads to inefficient learning, particularly in tasks requiring precise reasoning, where the model often gets trapped in local optima or exploration dead-ends. The core issues include: 1) heavy-tailed advantage distributions causing unstable updates; 2) vanishing gradients in neutral advantage zones, limiting correction signals; 3) exploration black holes where the model cannot escape from suboptimal policies due to insufficient corrective feedback. These bottlenecks hinder the training of large models on complex tasks, reducing their ability to generalize and transfer knowledge effectively.

AOPD introduces a token-level asymmetric training mechanism that dynamically switches from advantage-weighted policy gradients to direct divergence minimization based on the advantage signal. Specifically, it constructs a teacher support set via top-K probabilities and applies forward KL guidance to non-positive advantage tokens, effectively replacing the ineffective gradient updates. The method employs a threshold parameter τ to control intervention scope, enabling flexible adaptation to different task complexities. This approach reduces gradient variance, enhances learning in neutral zones, and facilitates escape from exploration black holes. The integration of local divergence matching with advantage-based reinforcement constitutes a novel contribution, offering a more stable and efficient training paradigm for large language models.

• �� Compute token advantages (At) using the difference in log probabilities between teacher and student. • Identify support set St via top-K highest teacher probabilities. • In positive advantage zones, optimize with advantage-weighted policy gradient. • In non-positive zones, halt policy gradient and minimize forward KL divergence between teacher and student distributions over support set. • Introduce threshold τ to regulate intervention scope, with Gt indicating whether to trigger divergence guidance based on probability differences. • During training, dynamically switch between policy gradient and divergence matching based on token advantage and Gt. • Maintain policy entropy by limiting divergence scope, ensuring exploration. • Use hyperparameters like support size K and τ to balance correction strength and stability.

The evaluation employed datasets such as OpenThoughts for reasoning warm-up, DeepMath for mathematical tasks, and ToolAlpaca for continual tool-use learning. Baselines included SeqKD, GKD, OPD, and ExOPD, with models trained under different initializations (3K and 1K steps). Metrics focused on Pass@1/4 scores and tool success rates. Experiments involved hyperparameter tuning for support size K, threshold τ, and divergence weights. Ablation studies assessed the impact of localized KL guidance, threshold sensitivity, and intervention frequency. The training monitored gradient norms, policy entropy, and task performance across epochs, demonstrating the robustness and efficiency of AOPD in diverse settings.

AOPD achieved the highest average Pass@1 scores: 53.40 for Qwen3-4B and 67.43 for Qwen3-8B, surpassing all baselines. Under weak initialization, it improved reasoning accuracy by 8.34 points over OPD, and by 4.09 under strong initialization. Training curves showed more stable convergence, with lower gradient norms and higher policy entropy, indicating better exploration. Ablation results confirmed that divergence guidance in non-positive advantage regions significantly reduces gradient variance and prevents exploration deadlocks. In continual learning, AOPD maintained or improved prior capabilities, unlike other methods suffering from catastrophic forgetting. These findings demonstrate the method’s effectiveness in complex reasoning and adaptive scenarios.

AOPD can be applied to large-scale language model training, especially in tasks requiring complex reasoning, continual learning, and capability transfer. It is suitable for deploying smaller, efficient models in real-world applications like AI assistants, automated reasoning systems, and code generation tools. Its ability to maintain diverse policies and adapt to new tasks with minimal forgetting makes it valuable for industry applications demanding robustness and flexibility. Additionally, the framework can be integrated into existing reinforcement learning pipelines to improve stability and sample efficiency in high-dimensional action spaces.

The approach depends on the choice of threshold τ, which requires careful tuning for different tasks. Computational overhead increases due to support set construction and divergence calculations, potentially limiting scalability. The method assumes high-quality teacher distributions; noisy or weak teachers may reduce guidance effectiveness. Furthermore, the current implementation may struggle with extremely large vocabularies or real-time constraints, necessitating further optimization. Future work should focus on adaptive parameter tuning, efficiency improvements, and extending robustness to noisy supervision.