Paper deep dive
OnlineQAT: On-Policy Distillation for Ultra-Low-Bit Large Language Models
Wenjun Wang, Heng Li, Yanggan Gu, Hongxia Yang
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 10/8/2026, 7:34:10 AM
Summary
The paper introduces OnlineQAT, a two-stage framework for ultra-low-bit quantization-aware training (QAT) of large language models. It addresses the mismatch between offline recovery data and deployed model states by using on-policy distillation (OPD). Stage 1 performs block-wise QAT for initialization, and Stage 2 uses a frozen full-precision teacher to provide reverse-KL signals on student-generated prefixes. Evaluated on Qwen3-1.7B, OnlineQAT achieves state-of-the-art results at W3A16 (57.28 average) and W2A16 (32.52 average), outperforming ReasoningQAT.
Entities (18)
Relation Signals (12)
OnlineQAT → evaluatedon → Qwen3-1.7B
confidence 98% · On Qwen3-1.7B, OnlineQAT obtains the best average among the compared quantized methods
OnlineQAT → achievesscore → 32.52
confidence 95% · 32.52 at W2A16
OnlineQAT → achievesscore → 57.28
confidence 95% · 57.28 at W3A16
OnlineQAT → improves → ReasoningQAT
confidence 95% · improving over ReasoningQAT by 2.90 and 0.44 points, respectively.
OnlineQAT → uses → On-Policy Distillation
confidence 95% · performs on-policy distillation (OPD) on student-generated responses.
OnlineQAT → appliedto → W3A16
confidence 90% · 57.28 at W3A16
OnlineQAT → appliedto → W2A16
confidence 90% · 32.52 at W2A16
OnlineQAT → uses → Reverse-KL
confidence 90% · a frozen full-precision teacher provides a sampled reverse-KL training signal.
OnlineQAT → usesdata → OpenThoughts
confidence 85% · OnlineQAT samples from 32,757 OpenThoughts prompts
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2610.09346v1
- Canonical: https://arxiv.org/abs/2610.09346v1
Trouble viewing inline? Open PDF directly →
Full Text
13,522 characters extracted from source content.
Expand or collapse full text
OnlineQAT: On-Policy Distillation for Ultra-Low-Bit Large Language Models Wenjun Wang Heng Li Affiliation: The Hong Kong Polytechnic University Sun Yat-sen University Yanggan Gu Hongxia Yang Affiliation: PolyU-Daya Bay Technology and Innovation Research Institute *Equal contribution. Abstract Quantization-aware training (QAT) can recover much of the accuracy lost when large language models are compressed below four bits. Existing recovery stages, however, are commonly optimized on fixed completions or teacher-generated answers, whereas the deployed quantized model conditions 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 prefix, 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, improving 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 Weight-only quantization reduces the memory footprint and bandwidth cost of large language models (LLMs), but two- and three-bit weights still incur a substantial accuracy loss. For reasoning models, a small perturbation to an early token can also change the states encountered throughout the remaining generation. Recovering useful behavior at such low precision therefore remains an open problem. BitDistiller, EfficientQAT, and ReasoningQAT recover low-bit models through QAT and distillation (Du et al., 2024; Chen et al., 2025; Okoshi et al., 2026). We use ReasoningQAT as our strongest offline baseline. Its recovery stage evaluates the student and teacher on gold completions and therefore visits reference prefixes. At inference time, however, the quantized model conditions on its own previous tokens. Once quantization changes an early decision, later prefixes can leave the distribution covered by the fixed-completion objective. Figure 1 illustrates this distinction: offline recovery and deployment use different prefix distributions, whereas OPD trains on the prefixes visited by the deployed student. Figure 1: Stage 1 produces a W2/W3 student. In Stage 2, offline recovery queries the teacher on fixed-completion prefixes, while deployment follows student-generated prefixes. OnlineQAT queries the frozen teacher on the student trajectory and applies sampled reverse KL, aligning the training and deployment states. To address this mismatch, we propose OnlineQAT, a two-stage recovery framework. Stage 1 performs block-wise QAT on mixed-domain calibration data, producing a student that can generate meaningful trajectories. Stage 2 samples responses from this student and uses a frozen full-precision teacher to provide a sampled reverse-KL signal at the visited prefixes. On Qwen3-1.7B, OnlineQAT improves the six-task average over ReasoningQAT by 2.90 points at W3A16 and 0.44 points at W2A16, giving the best average among the compared methods in both settings. The W3A16 gain is broad, whereas the smaller W2A16 gain remains task-dependent. 2 Related Work Ultra-low-bit QAT. QAT adapts model or quantization parameters while simulating low-precision weights. BitDistiller combines asymmetric quantization with confidence-aware self-distillation for sub-4-bit LLMs (Du et al., 2024). EfficientQAT introduces block-wise all-parameter training and a lightweight end-to-end quantization-parameter stage (Chen et al., 2025). ReasoningQAT targets reasoning models with mixed-domain block-wise calibration followed by teacher-guided reward rectification and forward-KL distillation (Okoshi et al., 2026). These methods establish effective low-bit initialization and recovery recipes. When recovery is evaluated on fixed calibration sequences or reference completions, however, the objective does not cover prefixes induced by the deployed quantized policy. OnlineQAT retains block-wise QAT but moves end-to-end recovery to those student-visited prefixes. On-policy distillation. On-policy distillation trains a student on its own sampled sequences while querying a teacher at the visited states. GKD formalizes this idea for language-model distillation and shows that student-generated outputs can reduce the mismatch between training and generation (Agarwal et al., 2024). Recent analysis studies when OPD succeeds and attributes its progress to alignment on high-probability tokens at student-visited states (Li et al., 2026); large-scale systems such as Nemotron 3 Ultra further use multi-teacher OPD during post-training (NVIDIA, 2026). Our setting is different: teacher and student share the same architecture, and their gap is created by ultra-low-bit weight quantization. We study whether OPD can recover this precision-induced gap. 3 Method OnlineQAT separates low-bit recovery into local initialization and end-to-end behavioral recovery. Let πT _T be the frozen full-precision model and πθq _θ^q its quantized counterpart. Stage 1 reconstructs one transformer block at a time. Stage 2 samples y y from πθq _θ^q and evaluates both models at student-generated prefixes ht=(x,y^<t)h_t=(x, y_<t), rather than the fixed prefixes (x,y<t)(x,y_<t) used by offline distillation. 3.1 Stage 1: block-wise low-bit initialization Following EfficientQAT and ReasoningQAT (Chen et al., 2025; Okoshi et al., 2026), Stage 1 uses sequential block-wise QAT on calibration data drawn from pre-training and reasoning corpora. Let BlB_l and BlqB_l^q denote the full-precision and fake-quantized versions of block l, and let HlH_l and HlqH_l^q be their respective input streams after the preceding blocks. For a sequence of T tokens and hidden width C, we optimize ℒblock(l)=1TC‖Blq(Hlq)−Bl(Hl)‖F2.L_block^(l)= 1TC \|B_l^q(H_l^q)-B_l(H_l) \|_F^2. (1) The fake quantizer uses asymmetric affine weight quantization with group size 128. We update the block’s master weights, per-group scales and zero-points, and RMSNorm weights through the straight-through estimator, then propagate both streams to the next block. All end-to-end methods start from the resulting checkpoint. 3.2 Stage 2: on-policy reverse-KL recovery For each prompt x∼x , the quantized student samples a response y^∼πθ,τq(⋅∣x) y _θ,τ^q(· x) at temperature τ. At prefix ht=(x,y^<t)h_t=(x, y_<t), let pθ,tτp_θ,t^τ and pT,tτp_T,t^τ denote the temperature-scaled student and teacher distributions. The target objective is the reverse KL averaged over response tokens: ℒonline(θ)=x∼,y^∼πθ,τq[1|y^|∑t=1|y^|DKL(pθ,tτ∥pT,tτ)].L_online(θ)=E_x ,\, y _θ,τ^q [ 1| y| _t=1^| y|D_KL\! (p_θ,t^τ\,\|\,p_T,t^τ ) ]. (2) For the sampled token at=y^ta_t= y_t, a one-sample policy-gradient estimator uses the stop-gradient log-ratio AtRKL=sg[logpθ,tτ(at)−logpT,tτ(at)],A_t^RKL=sg\! [ p_θ,t^τ(a_t)- p_T,t^τ(a_t) ], (3) where sgsg stops gradients through the log-ratio. The training loss is ℒOPD(θ)=1N∑t∈responseAtRKLlogpθ,tτ(at),L_OPD(θ)= 1N _t A_t^RKL p_θ,t^τ(a_t), (4) with N response tokens in the batch. The constant +1+1 term in the reverse-KL policy gradient is omitted because its expectation is zero. Minimizing Eq. 4 raises the probability of a sampled token when the teacher assigns it higher probability than the student, and lowers it in the opposite case. We use τ=0.6τ=0.6 and no cross-entropy term. 4 Experiments 4.1 Setup Benchmarks. We report GSM8K, MATH-500, MMLU-Redux (MMLU-R), IFEval, LiveCodeBench (LCB), and GPQA-Diamond (GPQA-D), covering mathematical reasoning, knowledge, instruction following, and code generation. AVG is their unweighted mean. Model and quantization. We evaluate Qwen3-1.7B in thinking mode under W3A16 and W2A16. Transformer linear weights use asymmetric affine quantization along the input dimension with group size 128 and integer ranges [0,7][0,7] and [0,3][0,3]. Activations, embeddings, and the LM head remain floating point; the teacher and zero-points are frozen during Stage 2. Stage 1 calibration. We use 4,096 sequences of 2,048 tokens from 80% OpenThoughts and 20% FineWeb. Each block is trained for two epochs with batch size 2 and AdamW. Quantization parameters use learning rate 10−410^-4; W3A16 and W2A16 master weights use 10−510^-5 and ×10−52\!×\!10^-5. All recovery methods start from the same Stage 1 checkpoint within each bit width. Baselines and Stage 2. Stg1 has no end-to-end recovery. KD is a fixed-prefix distillation baseline with 0.2ℒCE+ℒJSD0.2L_CE+L_JSD. ReasoningQAT follows Okoshi et al. (Okoshi et al., 2026), using 0.2ℒRSFT+ℒFKL0.2L_RSFT+L_FKL on gold completions, where the forward KL retains the teacher’s top-20 logits. It trains for 512 W3A16 or 1,536 W2A16 steps. OnlineQAT samples from 32,757 OpenThoughts prompts at temperature 0.6 and uses no supervised term. Sequences contain at most 8,192 total tokens and stop at EOS; only response tokens incur loss. All Stage 2 runs use AdamW on eight GPUs with effective batch 64. OnlineQAT trains W3A16 for 40 steps at 5×10−65× 10^-6. W2A16 warms up for 30 steps from 2×10−72× 10^-7 to 2×10−62× 10^-6, followed by 120 steps at 2×10−62× 10^-6. Table 1: Qwen3-1.7B results. ReasoningQAT denotes the method of Okoshi et al. (2026). All entries are percentages; AVG is the unweighted mean. Bold marks the best quantized result within each bit width. Method GSM8K MMLU-R MATH IFEval LCB GPQA-D AVG BF16 (Think) 89.9 73.9 93.4 72.5 33.2 40.1 67.17 W3-Stg1 84.2 65.3 64.8 57.3 22.3 18.2 52.02 W3-KD 85.4 66.0 67.4 58.3 26.7 18.2 53.67 W3-ReasoningQAT 86.0 65.8 71.6 60.3 24.4 18.2 54.38 W3-OnlineQAT 86.3 68.0 73.2 59.5 27.4 29.3 57.28 W2-Stg1 24.40 17.20 21.80 33.20 0.50 11.60 18.12 W2-KD 55.00 38.40 31.40 41.30 4.10 14.10 30.72 W2-ReasoningQAT 56.60 43.80 37.80 33.80 5.30 15.20 32.08 W2-OnlineQAT 59.82 46.33 36.00 32.49 4.83 15.66 32.52 4.2 Main results W3A16 recovery. OnlineQAT averages 57.28, exceeding ReasoningQAT by 2.90 points and KD by 3.62 points. Relative to ReasoningQAT, it improves GSM8K, MMLU-R, MATH-500, LiveCodeBench, and GPQA-Diamond; the largest changes are +3.0+3.0 on LiveCodeBench and +11.1+11.1 on GPQA-Diamond. IFEval decreases by 0.8 points. W2A16 recovery. At W2A16, OnlineQAT improves the ReasoningQAT average from 32.08 to 32.52. It gains on GSM8K, MMLU-R, and GPQA-Diamond but trails on MATH-500, IFEval, and LiveCodeBench; the improvement is therefore modest and task-dependent. Scope. Relative to ReasoningQAT, OnlineQAT changes both the prefix distribution and KL direction, so the comparison does not isolate either factor. The study covers one model without multiple training seeds or an offline reverse-KL control, and OPD incurs rollout generation. Table 1 therefore supports an empirical comparison, not a causal or compute-efficiency claim. 5 Conclusion OnlineQAT retains the mixed-domain block-wise initialization of reasoning-oriented QAT and replaces fixed-sequence recovery with sampled-token OPD. On Qwen3-1.7B, this second stage improves the six-task average over ReasoningQAT by 2.90 points at W3A16 and 0.44 points at W2A16. These initial results motivate a broader evaluation of on-policy recovery for ultra-low-bit language models. References Agarwal et al. (2024) R. Agarwal, N. Vieillard, Y. Zhou, P. Stanczyk, S. Ramos, M. Geist, and O. Bachem On-policy distillation of language models: learning from self-generated mistakes. In International Conference on Learning Representations, Cited by: §2. Chen et al. (2025) M. Chen, W. Shao, P. Xu, J. Wang, P. Gao, K. Zhang, and P. Luo EfficientQAT: efficient quantization-aware training for large language models. In Annual Meeting of the Association for Computational Linguistics, Cited by: §1, §2, §3.1. Du et al. (2024) D. Du, Y. Zhang, S. Cao, J. Guo, T. Cao, X. Chu, and N. Xu BitDistiller: unleashing the potential of sub-4-bit LLMs via self-distillation. In Annual Meeting of the Association for Computational Linguistics, Cited by: §1, §2. Li et al. (2026) Y. Li, Y. Zuo, B. He, J. Zhang, C. Xiao, C. Qian, T. Yu, H. Gao, W. Yang, Z. Liu, and N. Ding Rethinking on-policy distillation of large language models: phenomenology, mechanism, and recipe. External Links: 2604.13016 Cited by: §2. NVIDIA (2026) NVIDIA Nemotron 3 ultra: open, efficient mixture-of-experts hybrid mamba-transformer model for agentic reasoning. External Links: 2606.15007 Cited by: §2. Okoshi et al. (2026) Y. Okoshi, H. Otsuka, D. Fujiki, and M. Motomura Towards quantization-aware training for ultra-low-bit reasoning LLMs. In The Fourteenth International Conference on Learning Representations, Cited by: §1, §2, §3.1, §4.1, Table 1.