Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning
On-policy distillation restores math and code reasoning lost to sub-3-bit model quantization.
Quantization-aware distillation recovers short-form question answering after sub-3-bit quantization but leaves math and code reasoning impaired, often causing repetitive loops. The authors add an on-policy distillation stage in which the quantized student generates its own prefixes and a frozen full-precision teacher supplies token-level feedback plus task-verifier rewards. Across four models at 2.79 and 1.88 effective bits, average BF16 retention rose from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval, exceeding continued teacher-forced QAD at a matched budget.
- QAD restores short-form QA but leaves math and code reasoning impaired.
- Quantization errors compound along the model's own autoregressive trajectories.
- OPD supervises the quantized student on prefixes it actually generates.
- MATH-500 retention rose from 35% to 70%; HumanEval from 66% to 91%.
- Reasoning gains beat continued teacher-forced QAD at matched budget.
Full article188 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.26708