Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning
TL;DR - This paper introduces on-policy distillation (OPD) to recover long-form reasoning degraded by sub-3-bit quantization. By supervising quantized models on their own generated trajectories, OPD substantially improves math and code performance over standard teacher-forced quantization-aware distillation.
- OPD targets quantization-amplified exposure bias, where small deviations compound during autoregressive generation and can produce repetitive loops.
- A frozen full-precision teacher provides dense token-level guidance and task-verifier rewards on prefixes generated through the deployment-time quantized path.
- Across four models at 2.79 and 1.88 effective bits, average BF16 performance retention rose from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval.
- The method preserved short-form performance and outperformed continued teacher-forced distillation under matched training budgets.