Mismatch Matters: On-Policy Distillation Beyond Token Agreement

TL;DR

TIDE method corrects teacher-student mismatch via bounded Hellinger shaping and top-K injection, boosting Avg@8 from 6.9% to 20.3%.

cs.AI 🔴 Advanced 2026-08-11 92 views
Zichao Yu Chengzhi Yu Shengze Xu Yujin Han Bingqing Jiang Xu Wang Difan Zou
Deep Learning Model Distillation Large Language Models Mismatch Correction Algorithm Innovation

Key Findings

Methodology

This paper introduces TIDE (Token-level Independent Deficit–Excess correction), a novel approach designed to address the 'degenerate agreement' failure mode in on-policy distillation (OPD). The core insight is that teacher-student divergence manifests as two distinct mismatch types: excess tokens, generated by the student but assigned near-zero probability by the teacher, and deficit tokens, favored by the teacher but rarely sampled by the student. TIDE employs a bounded Hellinger transformation to suppress extreme excess tokens, preventing gradient instability. Simultaneously, it utilizes an analytic top-K injection mechanism to restore missing probability mass for deficit tokens without requiring their sampling. The method involves: • Quantile-based selection of severe excess tokens, applying Hellinger shaping for bounded suppression; • Analytic identification of underrepresented teacher-favored tokens via top-K support, with direct probability mass injection. Empirical results on multiple mathematical reasoning benchmarks with Qwen3 as teacher-student pairs demonstrate that TIDE consistently outperforms standard OPD and recent token selection or reward shaping baselines, especially under strong mismatch scenarios, with Avg@8 improving from 6.9% to 20.3%, response length reduced by 3.6 times, and formatting errors significantly decreased.

Key Results

  • In multiple mathematical reasoning benchmarks, TIDE achieved a remarkable increase in Avg@8 from 6.9% to 20.3% on Qwen3-8B models, outperforming baseline OPD and recent token selection methods. The approach also shortened average response length by 3.6 times, effectively mitigating length inflation issues. Furthermore, TIDE reduced formatting failures and irrelevant content, leading to more accurate and reliable outputs. Ablation studies confirmed that both the Hellinger suppression and top-K injection components are essential for optimal performance, with their combination yielding the best results.
  • Analysis of teacher-student mismatch revealed that TIDE accurately identifies excess tokens that cause unstable gradients and deficit tokens that carry critical supervisory signals. The method's adaptive correction mechanism ensures stable training and better reasoning capabilities. Across various tasks, TIDE demonstrated robustness against severe mismatch scenarios, maintaining high performance where traditional OPD struggled. The experimental data showed that TIDE's targeted correction strategies significantly enhance the model's ability to generalize and reason effectively.
  • Across different model sizes and tasks, TIDE maintained consistent superiority, especially in complex reasoning and long-text generation. Its capacity to dynamically address both over- and under-represented tokens ensures comprehensive coverage of the teacher’s knowledge, leading to substantial improvements in downstream performance. These results underscore TIDE's potential as a general framework for improving on-policy distillation, with broad applicability in large-scale language model training and deployment.

Significance

This research addresses a fundamental challenge in large language model training—how to effectively distill knowledge without falling into the trap of degenerate agreement, which masks true model capabilities. By shifting focus from mere token agreement to explicit mismatch correction, TIDE introduces a principled, theoretically grounded framework that enhances the stability and informativeness of teacher-student training. Its ability to identify and correct both excess and deficit tokens ensures that the student model captures the teacher’s reasoning patterns more faithfully, leading to improved reasoning, reduced response length inflation, and fewer formatting errors. These advancements have significant implications for both academia and industry, enabling more efficient, reliable, and interpretable large model training pipelines. Moreover, TIDE’s approach opens new avenues for research into divergence-aware training strategies, pushing the boundaries of what is achievable in model compression and knowledge transfer.

Technical Contribution

The key technical innovations of this work include: • The development of a bounded Hellinger shaping function that replaces unbounded log-ratio corrections, ensuring stable gradients during excess token suppression; • An analytic top-K injection mechanism that directly restores the missing probability mass for teacher-favored tokens, bypassing the need for sampling these rare tokens; • A dual-branch correction framework that independently identifies and rectifies excess and deficit mismatch directions, supported by rigorous theoretical proofs of boundedness and local fidelity. The method leverages divergence geometry principles, ensuring that corrections are both stable and faithful to the underlying distributions. Empirical validation across multiple reasoning tasks demonstrates that these techniques significantly outperform existing baselines, establishing a new standard for mismatch-aware distillation.

Novelty

This paper is the first to systematically analyze teacher-student mismatch at the token level within an on-policy distillation framework, moving beyond traditional agreement metrics. The introduction of the TIDE mechanism, combining bounded Hellinger suppression with analytic top-K probability injection, represents a novel approach to addressing the asymmetry in teacher-student divergence. Unlike prior work that relies solely on reward shaping or token filtering, TIDE explicitly models and corrects the two primary mismatch directions, offering a theoretically grounded and practically effective solution. Its dual-branch design and divergence-aware correction strategies distinguish it from existing methods, providing a new paradigm for stable and informative knowledge distillation in large language models.

Limitations

  • While TIDE effectively addresses the identified mismatch issues, it may still face challenges in extremely sparse or highly biased teacher models, where the top-K support may not fully capture the teacher’s reasoning complexity. In such cases, the correction mechanisms might be insufficient to recover critical tokens.
  • The additional computational overhead introduced by the dual correction branches and divergence calculations could impact training efficiency, especially in large-scale, multi-task settings, requiring further optimization for industrial deployment.
  • The method’s reliance on teacher top-K support assumes a relatively stable teacher distribution; in scenarios with high teacher model bias or instability, the correction accuracy may diminish, limiting its robustness across diverse tasks.

Future Work

Future research could explore integrating reinforcement learning techniques to adaptively tune the correction thresholds, further improving robustness. Extending TIDE to multi-modal models, such as vision-language systems, is another promising direction. Additionally, optimizing the computational efficiency of divergence calculations and correction steps will be crucial for scaling to larger models and datasets. Investigating the applicability of mismatch correction in unsupervised or semi-supervised settings, as well as in continual learning scenarios, could broaden its impact. Finally, developing theoretical bounds for correction convergence and stability in more complex divergence geometries remains an open challenge.

AI Executive Summary

Large language models (LLMs) have revolutionized natural language processing, enabling applications from chatbots to complex reasoning systems. However, training these models efficiently remains a significant challenge, especially when employing knowledge distillation techniques such as on-policy distillation (OPD). While OPD has gained popularity for its ability to leverage student-generated trajectories and teacher feedback, it suffers from a critical flaw: the phenomenon of degenerate agreement. This occurs when the student model exploits repetitive loops to achieve near-perfect token agreement with the teacher, masking global deficiencies in reasoning and understanding.

This paper investigates the root causes of this issue, revealing that high token agreement does not necessarily correspond to meaningful learning. Instead, the divergence between teacher and student manifests as two distinct mismatch types: excess tokens, generated excessively by the student, and deficit tokens, underrepresented in the student’s output but highly favored by the teacher. Traditional methods tend to overlook these divergences, focusing instead on token-level agreement metrics that can be artificially inflated by repetitive patterns.

To address this, the authors propose TIDE, a novel correction framework that explicitly models and rectifies teacher-student mismatch at the token level. TIDE employs a bounded Hellinger transformation to suppress the influence of extreme excess tokens, preventing gradient instability. Simultaneously, it uses an analytic top-K injection mechanism to restore the probability mass of under-sampled, teacher-favored tokens without requiring their sampling. This dual approach ensures that the student model not only avoids pathological loops but also faithfully captures the reasoning patterns embedded in the teacher.

Extensive experiments on multiple mathematical reasoning benchmarks, using Qwen3 as the teacher, demonstrate that TIDE significantly outperforms existing OPD baselines. The results show an increase in Avg@8 from 6.9% to 20.3%, a reduction in response length by 3.6 times, and a substantial decrease in formatting errors. These improvements are especially pronounced under conditions of strong teacher-student mismatch, highlighting TIDE’s robustness and effectiveness.

The significance of this work lies in its paradigm shift: moving from a sole focus on agreement to a nuanced understanding of divergence directions. By explicitly correcting both excess and deficit tokens, TIDE enhances the stability, interpretability, and reasoning capacity of large models. This approach opens new avenues for research into divergence-aware training and has broad implications for deploying more reliable, efficient, and intelligent AI systems.

In summary, TIDE represents a major step forward in model distillation, addressing longstanding issues of instability and superficial agreement. Its innovative combination of divergence geometry principles and analytic support injection paves the way for next-generation large language model training, with potential impacts spanning academia and industry alike.

Deep Dive

Abstract

On-policy distillation (OPD) has emerged as a core component of modern LLM post-training pipelines, yet we reveal a failure mode: degenerate agreement, where students exploit repetitive loops to achieve near-perfect token agreement with the teacher despite globally flawed responses. We therefore shift our focus from agreement to teacher-student mismatch, and find that mismatch tokens can be mainly categorized into two types: student-excess tokens and student-deficit tokens. Student-excess tokens are generated by the student but assigned near-zero probability by the teacher; their log-ratio corrections grow unbounded and destabilize the update. Student-deficit tokens, in contrast, are preferred by the teacher but rarely sampled by the student; their absence blocks the transfer of the teacher's reasoning patterns. To tackle these mismatch directions, we propose TIDE (Token-level Independent Deficit-Excess correction), which applies bounded Hellinger shaping to suppress the most severe sampled excesses and an analytic teacher top-$K$ injection to restore deficient probability mass without requiring deficit tokens to be sampled. Across mathematical reasoning benchmarks with multiple Qwen3 teacher-student pairs, TIDE consistently outperforms standard OPD and recent token-selection and reward-shaping baselines. Moreover, the gains of TIDE are more pronounced under strong teacher-student mismatch, where it improves Avg@8 from 6.9% to 20.3%, reduces average response length by a factor of 3.6, and substantially reduces formatting failures. Code is available at https://github.com/yzc-666/TIDE

cs.AI cs.CL

References (20)

On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou et al.

2023 662 citations ⭐ Influential View Analysis →

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

Nan Jia, Haojin Yang, Xingchen Ma et al.

2026 11 citations ⭐ Influential View Analysis →

PowerOPD: Stabilizing On-Policy Distillation with Bounded Power Transformation

Anhao Zhao, Junlong Tong, Yingqi Fan et al.

2026 3 citations ⭐ Influential View Analysis →

DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models

Zhihong Shao, Peiyi Wang, Qihao Zhu et al.

2024 8137 citations ⭐ Influential View Analysis →

DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

Adam Suma, Sam Dauncey

2025 4954 citations ⭐ Influential

Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation

Yuying Li, Leqi Zheng, Yongzi Yu et al.

2026 8 citations ⭐ Influential View Analysis →

DAPO: An Open-Source LLM Reinforcement Learning System at Scale

Qiying Yu, Zheng Zhang, Ruofei Zhu et al.

2025 2361 citations View Analysis →

Divergence measures and message passing

T. Minka

2005 667 citations

TIP: Token Importance in On-Policy Distillation

Yuanda Xu, Hejian Sang, Zhengze Zhou et al.

2026 46 citations View Analysis →

Rethinking On-Policy Self-Distillation for Thinking Models

Simran Kaur, Narutatsu Ri, Yinghui He et al.

2026 7 citations View Analysis →

Teacher-Guided Policy Optimization for On-Policy Reasoning Distillation under Large Policy Divergence

Xinyu Liu, Kechen Jiao, Chunyang Xiao et al.

2026 3 citations View Analysis →

Fast and Effective On-policy Distillation from Reasoning Prefixes

Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani et al.

2026 26 citations View Analysis →

Qwen3 Technical Report

An Yang, Anfeng Li, Baosong Yang et al.

2025 7149 citations View Analysis →

Proximal Policy Optimization Algorithms

John Schulman, Filip Wolski, Prafulla Dhariwal et al.

2017 30084 citations View Analysis →

PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

Yanli Zhao, A. Gu, R. Varma et al.

2023 821 citations View Analysis →

Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks

Samy Bengio, O. Vinyals, N. Jaitly et al.

2015 2452 citations View Analysis →

GLM-5: from Vibe Coding to Agentic Engineering

GLM-4.5 Team Aohan Zeng, Xin Lv, Zhenyu Hou et al.

2026 315 citations View Analysis →

Trust Region On-Policy Distillation

Xingrun Xing, Haoqing Wang, Boyan Gao et al.

2026 4 citations View Analysis →

Demystifying OPD: Length Inflation and Stabilization Strategies for Large Language Models

Feng Luo, Yu-Neng Chuang, Guanchu Wang et al.

2026 15 citations View Analysis →

HybridFlow: A Flexible and Efficient RLHF Framework

Guangming Sheng, Chi Zhang, Zilingfeng Ye et al.

2024 2116 citations View Analysis →