PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuanteng Chen, Zhilei Liu, Peisong Wang, Yuantian Shao, Chuangyi Li, Weining Wang, Shuang Qiu, Gang Li, Jing Liu, Jian Cheng

arXiv:2609.26708v1cs.LGcs.AI

TL;DR

Teacher-forced QAD recovers short-form performance after sub-three-bit quantization but leaves long-form reasoning impaired as quantized deviations compound along student-generated trajectories. The paper adds on-policy distillation, supervising rollouts through the deployment quantized path with teacher feedback and task-verifier rewards. Across four models, it improves reasoning retention while preserving short-form performance.

  • Problem

    Teacher-forced QAD supervises fixed corpus prefixes, while quantized students generate from their own predictions, leaving long-form mathematical and code reasoning substantially impaired.

    Source: §1 · p. 1, §3.2 · p. 4

  • Method

    OPD starts from QAD and trains on completions sampled through the deployment quantized path, using frozen-teacher token feedback on student prefixes and task-verifier rewards.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    Across four models, average BF16 retention rises from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval while short-form performance is preserved.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Coupling QAD initialization with OPD recovery preserves broad capabilities while restoring long-form reasoning below three bits.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Sub-three-bit quantization can reduce deployment costs, but teacher-forced QAD recovers short-form performance much better than extended reasoning. The paper attributes this gap to supervision on fixed prefixes and introduces on-policy distillation (OPD), which improves reasoning retention while preserving short-form performance.

1 INTRODUCTION

  • Sub-three-bit quantization offers substantial compression, while low-bit errors can severely degrade or collapse performance without quantization-aware training.The paper frames quantization as a way to reduce memory and bandwidth demands for deployment on constrained hardware.Source: §1 · p. 1
  • Across four models and two bit widths, QAD retains 82% of BF16 performance on short-form question answering but only 35% on MATH-500.Long generations can enter repetitive loops and exhaust the decoding budget without completing a solution.Source: §1 · p. 1
  • OPD follows QAD initialization by sampling through the deployment quantized path and combining frozen-teacher token feedback on student prefixes with task-verifier rewards.Verifiers reward final-answer correctness in mathematics and successful test execution in code.Source: §1 · p. 2
  • Across four models, OPD raises average BF16 retention from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval while preserving short-form performance.Reasoning gains substantially exceed continued teacher-forced QAD under matched corpora and optimizer-step budgets.Source: §1 · p. 2

2 RELATED WORK

Related work covers extreme low-bit quantization and recovery, then situates OPD within research on exposure bias and student-generated trajectories. The paper adapts on-policy distillation to extreme low-bit quantization.

2.1 EXTREME LOW-BIT QUANTIZATION AND REASONING RECOVERY

2.2 EXPOSURE BIAS AND ON-POLICY DISTILLATION

  • Exposure bias is the mismatch between training on reference prefixes and inference conditioned on a model’s own predictions, addressed in sequence modeling by methods including scheduled sampling and sequence-level training (Ross et al., 2011; Bengio et al., 2015; Ranzato et al., 2016).Source: §2.2 · p. 3
  • On-policy distillation supervises student-generated trajectories with teacher targets on visited prefixes, adding token-level feedback to reward-only training (Agarwal et al., 2024; Ko et al., 2024).The paper adapts this approach, established in full-precision post-training, to extreme low-bit quantization (Lu & Thinking Machines Lab, 2025).Source: §2.2 · p. 3

3 WHAT TEACHER-FORCED RECOVERY LEAVES BEHIND

The section shows that teacher-forced QAD restores useful low-bit behavior unevenly: retention falls as tasks require longer generation. Quantization compounds deviations from fixed training prefixes, and long reasoning outputs are especially prone to repetition and budget exhaustion.

3.1 SHORT ANSWERS RECOVER, LONG DERIVATIONS DO NOT

  • Across six settings, QAD retains 86% of BF16 performance on QA9, 56% on GSM8K, 59% on MBPP, 27% on MATH-500, and 16% on AMC23.RTN yields zero accuracy on all five generative benchmarks, while QAD restores useful behavior.Source: §3.1 · p. 3, §3 · p. 3
  • QAD retention declines from short- to medium- to long-form tasks in every model–bit-width setting, alongside longer BF16 generations on the more demanding tasks.QA9 requires no autoregressive generation; BF16 generates roughly 130 tokens on GSM8K and MBPP, compared with about 410 on MATH-500 and 3,200 on AMC23.Source: §3.1 · pp. 3–4, Figure 2 · p. 4

3.2 QUANTIZATION AMPLIFIES EXPOSURE BIAS

  • Teacher-forced QAD supervises fixed corpus prefixes, but quantization changes student trajectories; longer generations and lower precision allow deviations to compound.Even a plausible alternative token can move the student beyond trajectories covered during training.Source: §3.2 · p. 4

3.3 QAD LOSES THE TRAJECTORY, NOT THE ARITHMETIC

  • On MATH-500, 95% of QAD generations exhaust the decoding budget and 70% end in repeated 8-grams, versus 27% and 12% for BF16.On GSM8K, budget exhaustion is 32% for QAD versus 5% for BF16, with repetition rates of 30% and 2%, respectively; the curves separate as generation continues.Source: §3.3 · pp. 4–5, Figure 3 · p. 5

4 ON-POLICY DISTILLATION FOR LOW-BIT REASONING

OPD follows QAD initialization with teacher supervision on quantized student-generated trajectories, combining token-level feedback with task-verifier rewards. On MATH-500, it reduces looping and budget exhaustion while improving accuracy; GSM8K behavioral failures also fall.

4 ON-POLICY DISTILLATION FOR LOW-BIT REASONING

  • QAD establishes a viable low-bit policy, then OPD restores trajectory control by training on deployment-path rollouts with frozen BF16 teacher guidance and verifier rewards.A round-to-nearest model at these widths has nothing worth sampling and earns no reward to bootstrap from.Source: §4 · p. 5

4.1 FROM TEACHER FORCING TO STUDENT FORCING

  • Teacher-forced QAD uses fixed reference-response prefixes, whereas student forcing samples completions and places teacher supervision on the resulting student-generated prefixes.The student policy is evaluated through the quantized forward path used at deployment.Source: §4.1 · p. 5
  • Sampling through the quantized path puts low-bit effects into training, and reverse KL penalizes student-favored continuations that the teacher assigns low probability.This places guidance on continuations shaped by quantization, including departures from reference trajectories.Source: §4.1 · p. 6

4.2 LEARNING FROM QUANTIZED ROLLOUTS

  • For each prompt, the student generates G quantized-path completions; the teacher scores them and verifiers reward mathematical answer correctness or code test execution.Sampled reverse-KL feedback supplies token-level guidance, while group-relative advantage favors completions reaching verified solutions; β=1 in every reported run.Source: §4.2 · p. 6

4.3 OPD RESTORES TRAJECTORY CONTROL

  • 70% to 17% loop rate, 95% to 53% budget exhaustion, and 10.4% to 24.2% accuracy on Qwen3-0.6B W2.79 MATH-500 after OPD.BF16 reference rates were 12% for looping and 27.2% for accuracy.Source: §4.3 · p. 6
  • On GSM8K, OPD cuts loop rate from 30% to 2% and budget exhaustion from 32% to 5%, matching BF16 on both measures.The larger MATH-500 improvement mirrors greater disruption on longer generations.Source: §4.3 · p. 6

5 EXPERIMENTS

The experiments evaluate QAD-initialized OPD across four models, two effective bit widths, and reasoning, code, and short-form QA tasks. OPD restores reasoning while preserving broad capabilities, uses fewer training resources than QAD initialization, and outperforms low-bit baselines on long-form tasks.

5.1 EXPERIMENTAL SETUP

  • The evaluation covers four models at 2.79 and 1.88 effective bits, with OPD training in mathematics followed by code and evaluation across reasoning, code, and QA tasks.The code phase resumes from the selected mathematics checkpoint; training prompts come from splits that do not overlap evaluation sets.Source: §5.1 · pp. 6–7, Table 1 · p. 6

5.2 RECOVERING REASONING WHILE PRESERVING BROAD CAPABILITIES

  • 35.3% to 69.7% average BF16 retention on MATH-500, while mean QA9 retention rises from 88.1% to 92.1% with OPD.At 2.79 bits, Qwen3-1.7B improves MBPP from 35.3% to 52.0% and HumanEval from 41.5% to 59.1%.Source: Table 2 · p. 7, §5.2 · p. 8

5.3 RECOVERY EFFICIENCY

  • OPD recovery uses approximately 14–23× fewer GPU-hours than QAD initialization in the reported configurations.The OPD stage uses 57 GPU-hours versus 820 for QAD on Qwen3-4B at 1.88 bits, and 11 versus 253 on Falcon3-1B at 2.79 bits.Source: §5.3 · p. 8, Table 3 · p. 8

5.4 COMPARISON WITH QUANTIZATION BASELINES

  • At 2.79 bits, OPD retains 84% on MATH-500 versus 33–58% for other baselines; below two bits, it reaches 64% on GSM8K and 77% on MBPP.Below two bits, post-training and quantization-aware baselines score zero on MBPP and HumanEval, while retaining 67–80% on QA9.Source: Figure 4 · p. 8, §5.4 · p. 9

6 ABLATIONS

Under matched training conditions, OPD outperforms continued teacher-forced QAD across all four benchmarks for two Qwen3 models at both bit widths.

6 ABLATIONS

  • OPD wins all 16 comparisons; for Qwen3-1.7B at 2.79 bits, MATH-500 accuracy reaches 45.6% versus 23.2% with continued QAD.The arms share starting checkpoints, corpora, and optimizer-step budgets; extending QAD to two or three times the matched mathematics budget yields essentially no further gain.Source: §6 · p. 9

7 CONCLUSION

OPD extends teacher-forced QAD supervision to prefixes generated through the deployment quantized path, recovering long-form reasoning while preserving broad capabilities restored by QAD.

7 CONCLUSION

  • Across four models, OPD doubles average BF16 retention on MATH-500 from 35% to 70%, improves code generation, and preserves QAD-restored broad capabilities.The conclusion also reports improved termination and reduced repetition in a few hundred optimizer steps.Source: §7 · p. 10

A THE EDGERAZOR QAD INITIALIZATION

  • Every reported arm starts from a reproduced EdgeRazor mixed-precision QAD checkpoint, so that initialization and its reproduction bound the reported results.The reproduction distills a general instruction mixture from the student's BF16 copy using online logit KL plus a small task loss.Source: A THE EDGERAZOR QAD INITIALIZATION · pp. 12–13
  • The initialization mixes ternary and INT4 output rows for 2.79- or 1.88-bit weights, while leaving the KV cache at 16 bits, looser than the released configuration.The passage notes that the KV-cache deviation favors the study's arms.Source: A THE EDGERAZOR QAD INITIALIZATION · p. 12

B QA9 PER TASK

  • The QA9 expansion reports nine task scores, with MMLU's 57-subject aggregate counted as one task in the equally weighted mean.The QAD reproduction uses an effective batch of 768 across eight GPUs per optimizer step.Source: B QA9 PER TASK · p. 13, Table 5 · p. 13

C BASELINE ALIGNMENT

The baseline comparison aligns quantization settings where possible and describes width and implementation differences that qualify the comparisons.

C BASELINE ALIGNMENT

  • All baselines are brought onto the study's setting so score differences reflect quantization method rather than differences in what was quantized.The QA9 table expands results by task; its Ret. column is the mean as a percentage of the BF16 row.Source: C BASELINE ALIGNMENT · p. 13, Table 7 · p. 15
  • EfficientQAT, FlatQuant, and AutoRound use fallback widths of 3 bits against 2.79 and 2 bits against 1.88, rounding up in the baselines' favor.Q-Palette is run natively at the two target widths.Source: C BASELINE ALIGNMENT · p. 13
  • For width-matched QA9, baselines without INT1.58 outputs mix their 2-bit and wider solutions; both candidate mixtures are built and the better is reported.AutoRound and EfficientQAT mix rows, whereas FlatQuant mixes layers because its folded scales change the basis.Source: C BASELINE ALIGNMENT · pp. 13–14
  • On generative benchmarks, baselines retain BF16 embeddings and output heads while the study's arms use INT4, an asymmetry that favors the baselines.On QA9, embeddings and lm heads are also INT4 across arms; the reported isolated GSM8K cost of that override is −0.23 points for BF16 weights and −0.45 for FlatQuant at 4 bits.Source: C BASELINE ALIGNMENT · pp. 13–14
  • Q-Palette's QA9 and MBPP results come from a pure-PyTorch reimplementation, so the comparison reflects the study's reading of the method rather than the authors' code.Calibrating methods share a 10k-document Pile subset.Source: C BASELINE ALIGNMENT · p. 14

D TRAINING EFFICIENCY

Against continued teacher-forced QAD from the same starting checkpoint, OPD yields more GSM8K points per optimizer step across four models.

D TRAINING EFFICIENCY

  • With Qwen3-4B at 2.79 bits, 20 OPD steps gain 4.32 GSM8K points, exceeding the 4.09-point gain from 800 QAD steps.Figure 5 compares points per thousand optimizer steps from the same starting checkpoint; the vertical axis is logarithmic.Source: D TRAINING EFFICIENCY · p. 15, Figure 5 · p. 15

E OPD TRAINING CURVES

The Qwen3-4B 1.88-bit phase-1 curves show rapid early recovery followed by a reward plateau, with distillation dominating the loss scale.

E OPD TRAINING CURVES

  • In the first 30 steps, entropy falls from about 10 to 0.35 nats per token, the distillation term drops from 2.07 to 0.25, and verifier reward rises from −0.39 to roughly +0.15.Across the remaining 90 steps, reward stays between 0.1 and 0.2.Source: E OPD TRAINING CURVES · pp. 15–16
  • At β = 1, the distillation term accounts for almost all total loss while the policy-gradient term stays near 0.05; each step scores 32 sequences, making the smoothed reward curve more informative than raw values.Figure 6 plots per-step values faintly and nine-step exponential moving averages as solid lines.Source: E OPD TRAINING CURVES · p. 16, Figure 6 · p. 16

F TEACHER CHOICE

From Qwen3-1.7B upward, larger teachers do not improve OPD recovery over the student’s own BF16 teacher; Qwen3-0.6B is an exception that needs a stronger teacher. Extending teacher-forced QAD to three times the matched budget still leaves it behind OPD.

F TEACHER CHOICE

  • Across Qwen3-1.7B and Qwen3-4B at both widths, larger teachers do not consistently improve recovery over the student’s own BF16 teacher.At 1.88 bits, Qwen3-1.7B’s self-teacher leads on both tasks; for Qwen3-4B, results are close or mixed across tasks and widths.Source: F TEACHER CHOICE · p. 16, Table 8 · p. 17
  • Qwen3-0.6B needs a stronger teacher: its own BF16 copy is barely ahead of the quantized student, while teachers of 1.7B or above yield larger gains.Scaling beyond a 1.7B teacher adds little or can reduce performance.Source: F TEACHER CHOICE · p. 16, Table 8 · p. 17

G EXTENDED QAD BUDGET

  • At 3× the matched budget, teacher-forced QAD still trails OPD by 8.9–18.1 GSM8K points across configurations.The control extends the mathematics-only QAD phase from the same starting checkpoints and corpus; the Qwen3-0.6B 2.79-bit run also falls below its 1× baseline.Source: G EXTENDED QAD BUDGET · p. 18, Table 9 · p. 18
Loading 2609.26708v1…