Source-linked AI summary
OnlineQAT: On-Policy Distillation for Ultra-Low-Bit Large Language Models
Wenjun Wang, Heng Li, Yanggan Gu, Hongxia Yang
TL;DR
Fixed-completion recovery can miss prefixes visited by a quantized model during generation, complicating recovery at two and three bits. OnlineQAT combines block-wise initialization with on-policy teacher distillation on student-generated prefixes. On Qwen3-1.7B, it improves the six-task average over ReasoningQAT at both tested precisions.
Problem
Two- and three-bit quantization incurs substantial accuracy loss, while fixed-completion recovery does not cover prefixes generated by the deployed quantized policy.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
These initial results motivate broader evaluation of on-policy recovery for ultra-low-bit language models.
Source: §5 · p. 4The one-model study lacks multiple training seeds and an offline reverse-KL control, and its comparison changes both prefix distribution and KL direction.
Source: §4.2 · p. 4
Abstract
from arXiv · showhide
Quantization-aware training (QAT) can recover much of the accuracy lost when large language models are compressed below four bits. Existing re- covery stages, however, are commonly optimized on fixed completions or teacher-generated answers, whereas the deployed quantized model condi- tions on prefixes generated by itself. Quantization errors can therefore move the model into states that are absent from offline recovery data. We introduce OnlineQAT, a two-stage framework that first obtains a usable low-bit initialization through block-wise QAT and then performs on-policy distillation (OPD) on student-generated responses. At each visited pre- fix, a frozen full-precision teacher provides a sampled reverse-KL training signal. On Qwen3-1.7B, OnlineQAT obtains the best average among the compared quantized methods: 57.28 at W3A16 and 32.52 at W2A16, im- proving over ReasoningQAT by 2.90 and 0.44 points, respectively. The results suggest that student-visited states provide a useful recovery signal beyond fixed-completion training, particularly at three bits.
1 INTRODUCTION
Ultra-low-bit quantization leaves substantial accuracy loss, while fixed-completion recovery may miss prefixes the quantized model visits during deployment. OnlineQAT combines block-wise initialization with student-trajectory distillation and improves the six-task average over ReasoningQAT at both tested precisions.
1 INTRODUCTION
- Weight-only quantization reduces LLM memory footprint and bandwidth cost, but two- and three-bit weights still incur substantial accuracy loss.For reasoning models, perturbing an early token can change states throughout the rest of generation.Source: §1 · p. 1
- Fixed-completion recovery visits reference prefixes, whereas deployment follows the quantized model’s own tokens and may enter prefixes outside the offline objective’s distribution.Figure 1 contrasts offline recovery, deployment, and OnlineQAT’s student-trajectory teacher queries.Source: §1 · p. 1, Figure 1 · p. 1
- +2.90 points at W3A16 and +0.44 points at W2A16 on the six-task average over ReasoningQAT, with the best average among compared methods in both settings.The W3A16 gain is broad, while the W2A16 gain remains task-dependent.Source: §1 · p. 2
- OnlineQAT uses block-wise QAT to initialize a student, then applies sampled reverse-KL distillation from a frozen full-precision teacher at student-visited prefixes.The student samples responses for Stage 2 training.Source: §1 · p. 2
2 RELATED WORK
Prior work develops low-bit QAT and distillation recipes, while on-policy distillation trains on student-generated sequences. OnlineQAT applies that latter approach to the precision-induced gap between a quantized model and its full-precision counterpart.
Ultra-low-bit QAT
- BitDistiller combines asymmetric quantization with confidence-aware self-distillation, EfficientQAT uses block-wise all-parameter training, and ReasoningQAT adds reward rectification and forward-KL distillation.These methods establish low-bit initialization and recovery recipes.Source: §2 · p. 2
- Fixed calibration sequences or reference completions do not cover prefixes induced by the deployed quantized policy, so OnlineQAT moves end-to-end recovery to student-visited prefixes.It retains block-wise QAT while changing the recovery data distribution.Source: §2 · p. 2
On-policy distillation
- On-policy distillation trains on student-sampled sequences while querying a teacher at visited states; GKD studies this approach, and later work examines its mechanisms and scale.OnlineQAT studies OPD where teacher and student share an architecture and differ through ultra-low-bit quantization.Source: §2 · p. 2
3 METHOD
OnlineQAT separates low-bit recovery into block-wise initialization and end-to-end behavioral recovery. It then samples from the quantized student and uses teacher guidance at the prefixes that student visits.
3 METHOD
- OnlineQAT separates local low-bit initialization from end-to-end behavioral recovery, evaluating teacher and student on student-generated rather than fixed prefixes.The framework samples responses from the quantized student for Stage 2.Source: §3 · p. 2
3.1 STAGE 1: BLOCK-WISE LOW-BIT INITIALIZATION
- Stage 1 uses sequential block-wise QAT on mixed pre-training and reasoning calibration data, matching fake-quantized and full-precision block outputs.It updates master weights, per-group scales and zero-points, and RMSNorm weights, then propagates both streams to the next block.Source: §3.1 · p. 2
3.2 STAGE 2: ON-POLICY REVERSE-KL RECOVERY
- Stage 2 samples student responses at temperature τ and evaluates temperature-scaled teacher and student distributions at the resulting student-generated prefixes.The target is reverse KL averaged over response tokens.Source: §3.2 · pp. 2–3
- A one-sample policy-gradient estimator uses the sampled token and a stop-gradient log-ratio; the loss raises a token’s probability when the teacher favors it over the student and lowers it otherwise.The method uses τ = 0.6 and no cross-entropy term.Source: §3.2 · p. 3
4 EXPERIMENTS
The experiments evaluate OnlineQAT on Qwen3-1.7B across two low-bit settings and six tasks. Its average exceeds ReasoningQAT in both settings, while the study’s scope limits what the comparison establishes.
4.1 SETUP
- The evaluation uses Qwen3-1.7B in thinking mode at W3A16 and W2A16 across six benchmarks spanning reasoning, knowledge, instruction following, and code generation.AVG is the unweighted mean of GSM8K, MATH-500, MMLU-Redux, IFEval, LiveCodeBench, and GPQA-Diamond.Source: §4.1 · p. 3
- Stage 1 trains on 4,096 sequences of 2,048 tokens, sampled 80% from OpenThoughts and 20% from FineWeb, for two epochs per block.Source: §4.1 · p. 3
4.2 MAIN RESULTS
- At W3A16, ONLINEQAT improved five tasks over ReasoningQAT, with the largest gains on LiveCodeBench (+3.0) and GPQA-Diamond (+11.1), while IFEval fell by 0.8 points.Source: §4.2 · p. 3
- At W2A16, ONLINEQAT gained on GSM8K, MMLU-R, and GPQA-Diamond but trailed on MATH-500, IFEval, and LiveCodeBench.The passage characterizes the improvement as modest and task-dependent.Source: §4.2 · p. 3
- At W3A16, ONLINEQAT averaged 57.28, 2.90 points above ReasoningQAT; at W2A16, it improved ReasoningQAT from 32.08 to 32.52.Table 1 reports percentages and defines AVG as the unweighted mean.Source: §4.2 · p. 3, Table 1 · p. 4
- OnlineQAT’s comparison with ReasoningQAT changes both prefix distribution and KL direction, and the study covers one model without multiple seeds or an offline reverse-KL control.The authors describe Table 1 as an empirical comparison, not a causal or compute-efficiency claim; OPD also incurs rollout generation.Source: §4.2 · p. 4
5 CONCLUSION
OnlineQAT retains mixed-domain block-wise initialization and replaces fixed-sequence recovery with sampled-token OPD. On Qwen3-1.7B, it improves the six-task average over ReasoningQAT at both evaluated bit widths, motivating broader evaluation.
5 CONCLUSION
- OnlineQAT retains mixed-domain block-wise initialization and replaces fixed-sequence recovery with sampled-token OPD.Source: §5 · p. 4
- On Qwen3-1.7B, OnlineQAT improves the six-task average over ReasoningQAT by 2.90 points at W3A16 and 0.44 points at W2A16.Source: §5 · p. 4
- The authors describe these as initial results that motivate broader evaluation of on-policy recovery for ultra-low-bit language models.Source: §5 · p. 4