Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning

An on-policy distillation stage recovers long-form reasoning in sub-3-bit quantized LLMs by supervising the student on its own quantized trajectories.

Yuanteng Chen; Zhilei Liu; Peisong Wang; Yuantian Shao; Chuangyi Li; Weining Wang; Shuang Qiu; Gang Li; Jing Liu; Jian Cheng· 2026· DOI 10.48550/arXiv.2609.26708

The core problem

Large language models (LLMs) are increasingly deployed under aggressive quantization to reduce memory and compute, but sub-3-bit quantization severely degrades complex reasoning. Quantization-aware distillation (QAD) recovers much of the lost short-form question-answering performance, yet mathematical and code reasoning remain substantially impaired. The authors identify a key failure mode: long generations often degenerate into repetitive loops, exhausting the decoding budget without completing a solution. The paper traces 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. Because the student is never trained on the states it actually visits during deployment, errors accumulate and reasoning collapses. To address this mismatch, the authors 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 prefixe

Innovation

The authors report substantial improvements in long-form reasoning across multiple benchmarks. The key results are summarized below. | Model | Effective Bits | MATH-500 Retention (BF16) | HumanEval Retention (BF16) | |-------|----------------|---------------------------|----------------------------| | Baseline QAD | 2.79 / 1.88 | 35% | 66% | | + OPD | 2.79 / 1.88 | **70%** | **91%** | - **MATH-500:** Average BF16 performance retention increases from 35% to 70%. - **HumanEval:** Average BF16 performance retention increases from 66% to 91%. - **Short-form performance:** Preserved, indicating that OPD does not sacrifice broad capabilities. - **Comparison with continued QAD:** Reasoning gains substantially exceed those of continued teacher-forced QAD in matched-budget comparisons. The results demonstrate that OPD effectively recovers long-form reasoning while maintaining the benefits of QAD's stable low-bit initialization. The gains are consistent across four models and two effective bit-widths (2.79 and 1.88 bits).

Large language models (LLMs) are increasingly deployed under aggressive quantization to reduce memory and compute, but sub-3-bit quantization severely degrades complex reasoning. Quantization-aware distillation (QAD) recovers much of the lost short-form question-answering performance, yet mathematical and code reasoning remain substantially impaired. The authors identify a key failure mode: long generations often degenerate into repetitive loops, exhausting the decoding budget without completing a solution.

The paper traces 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. Because the student is never trained on the states it actually visits during deployment, errors accumulate and reasoning collapses. To address this mismatch, the authors 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.

Why it matters

The paper's central insight is that **quantization-amplified exposure bias** is the primary cause of reasoning degradation in sub-3-bit models. QAD trains on fixed corpus prefixes, but at inference the model generates its own tokens, and quantization errors compound along these self-generated trajectories. This mismatch leads to repetitive loops and incomplete solutions. OPD addresses this by training on the student's own quantized trajectories. The teacher provides dense token-level supervision on these on-policy prefixes, while task-verifier rewards encourage correct final answers. This combination ensures that the student learns to recover from its own errors, effectively closing the gap between training and deployment conditions. The framework's success across four models at 2.79 and 1.88 effective bits suggests that on-policy distillation is a general and scalable solution for low-bit reasoning. By coupling QAD's stable initialization with OPD's on-policy recovery, the authors provide a comprehensive sub-3-bit solution that preserves broad capabilities while restoring long-form reasoning. **Limitations and future work.** The paper does not specify the exact model architectures or the full set of benchmarks used. Future work could explore the applicability of OPD to even lower bit-widths, other reasoning domains, and different quantization schemes. Additionally, the computational cost of on-policy rollouts and teacher feedback may be a practical consideration for large-scale deployment.

Who should read this

CS practitioners and researchers

Opening member content…