Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning
Introduced On-Policy Distillation (OPD) for low-bit reasoning, significantly improving mathematical and code reasoning performance.
Key Findings
Methodology
The study employs Quantization-aware Distillation (QAD) combined with On-Policy Distillation (OPD). QAD restores broad capabilities, while OPD places teacher supervision on trajectories the model generates, using task verifier rewards to enhance reasoning performance.
Key Results
- On the MATH-500 dataset, OPD increased BF16 performance retention from 35% to 70%, and on HumanEval from 66% to 91%.
- OPD significantly exceeds reasoning gains of continued teacher-forced QAD across four models.
- OPD improves mathematical and code reasoning performance at 2.79 and 1.88 effective bits while preserving short-form performance.
Significance
This research provides a comprehensive solution for low-bit quantized models, restoring long-form reasoning capabilities while maintaining broad abilities. It addresses the shortcomings of traditional QAD in long-generation tasks, offering new insights for academia and industry.
Technical Contribution
Technical contributions include applying on-policy distillation to the actual generation trajectories of quantized models, combined with task verifier rewards, significantly enhancing long-form reasoning capabilities. It offers new theoretical guarantees and engineering possibilities compared to existing methods.
Novelty
First to apply on-policy distillation to reasoning recovery in quantized models, providing a more effective solution compared to traditional teacher-forced distillation methods.
Limitations
- OPD requires additional computational resources and time for distillation, which may not be suitable for resource-constrained environments.
- Recovery effects may not be as strong in extremely low-bit scenarios.
Future Work
Future work can explore reducing the computational cost of on-policy distillation and applying the method to a broader range of tasks.
AI Executive Summary
Quantized models often face performance degradation in low-bit scenarios, especially in long-form generation tasks. Existing Quantization-aware Distillation (QAD) methods can recover some performance but still show significant deficiencies in mathematical and code reasoning.
This paper introduces a novel On-Policy Distillation (OPD) method, which places teacher supervision on the trajectories the model actually generates, combined with task verifier rewards, significantly improving the reasoning performance of low-bit models. Experimental results show that OPD significantly enhances performance across multiple datasets while maintaining short-form task capabilities.
This approach offers a comprehensive solution for low-bit quantized models, restoring long-form reasoning capabilities while maintaining broad abilities. Future research can further optimize the computational cost of this method and validate its effectiveness across more tasks.
Deep Analysis
Background
Low-bit quantized models face performance degradation in reasoning tasks, especially in long-form generation. Quantization-aware Distillation (QAD) is a method to recover performance but still shows significant deficiencies in long-form reasoning.
Core Problem
Low-bit quantized models in long-form reasoning tasks tend to enter repetitive loops, exhausting decoding budgets without completing solutions. This is due to quantization-amplified exposure bias, where models are not trained on the trajectories they actually generate.
Innovation
Introduced On-Policy Distillation (OPD), which places teacher supervision on the trajectories the model actually generates, combined with task verifier rewards, significantly improving reasoning performance. OPD provides a more effective solution compared to traditional QAD methods.
Methodology
- �� Use QAD to initialize the model and restore broad capabilities.
- �� Apply On-Policy Distillation (OPD) on the model's actual generation trajectories.
- �� Combine task verifier rewards to provide dense token-level guidance.
Experiments
Experiments conducted on Qwen3-0.6B, Qwen3-1.7B, Qwen3-4B, and Falcon3-1B models, evaluated using MATH-500 and HumanEval datasets. Effectiveness validated by comparing BF16 performance retention.
Results
OPD significantly improves mathematical and code reasoning performance, restoring long-form reasoning capabilities across multiple datasets while maintaining short-form task capabilities.
Applications
This method can be used for reasoning recovery in low-bit quantized models, suitable for deploying efficient models in resource-constrained environments.
Limitations & Outlook
OPD requires additional computational resources and time for distillation, which may not be suitable for resource-constrained environments.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen, and a quantized model is like a chef using low-bit tools to prepare dishes. Traditional QAD methods are like giving the chef a cookbook, but he only gets guidance on fixed steps. On-Policy Distillation (OPD) is like providing real-time guidance to the chef while cooking, helping him make the right choices at every step.
ELI14 Explained like you're 14
Imagine you're playing a game, and your character has a cool skill but performs poorly because of bad gear. A quantized model is like this character, and QAD methods are like giving the character new gear, but it doesn't perform well in long battles. OPD is like giving the character real-time guidance, helping him perform at his best during the fight!
Glossary
Quantization-aware Distillation
A method where a teacher model supervises a quantized model to recover its performance.
Used to restore broad capabilities of low-bit quantized models.
On-Policy Distillation
A method that places teacher supervision on the trajectories the model actually generates.
Used to enhance long-form reasoning capabilities of low-bit models.
Exposure Bias
Occurs when a model is trained on fixed prefixes but conditions on its own generated prefixes during inference.
Leads to performance degradation in long-generation tasks for quantized models.
Task Verifier
A method used to evaluate the correctness of model-generated results.
Combined with OPD to enhance reasoning performance.
BF16
A floating-point format used in deep learning models, providing higher precision.
Used as a benchmark for performance comparison.
Open Questions Unanswered questions from this research
- 1 How to further reduce the computational cost of on-policy distillation for application in resource-constrained environments.
- 2 Validate the effectiveness of on-policy distillation across a broader range of tasks.
Applications
Immediate Applications
Low-bit Reasoning Recovery
Can be used to deploy efficient models in resource-constrained environments, recovering reasoning capabilities.
Long-term Vision
Broad Task Application
Explore methods to apply on-policy distillation across more tasks, further enhancing model performance.
Abstract
Quantization-aware distillation (QAD) restores much of the short-form question-answering performance lost to sub-3-bit quantization, yet leaves mathematical and code reasoning substantially impaired. Long generations often degenerate into repetitive loops, exhausting the decoding budget without completing a solution. We trace this gap to quantization-amplified exposure bias: QAD trains on fixed corpus prefixes, while quantization-induced deviations compound along the model's own autoregressive trajectories. To address this mismatch, we introduce an on-policy distillation (OPD) stage that places teacher supervision where the quantized model actually goes. Starting from a QAD checkpoint, the student generates through the quantized forward path used at deployment and receives feedback from a frozen full-precision teacher on its own prefixes, combining dense token-level guidance with task-verifier rewards. Across four models at 2.79 and 1.88 effective bits, OPD raises average BF16 performance retention from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval while preserving short-form performance, with reasoning gains substantially exceeding those of continued teacher-forced QAD in matched-budget comparisons. By coupling QAD's stable low-bit initialization with OPD's on-policy reasoning recovery, our framework provides a comprehensive sub-3-bit solution that preserves broad capabilities while restoring long-form reasoning.