OnlineQAT: On-Policy Distillation for Ultra-Low-Bit Large Language Models

# OnlineQAT: On-Policy Distillation for Ultra-Low-Bit Large Language Models
Source: [https://arxiv.org/html/2610.09346](https://arxiv.org/html/2610.09346)
Heng LiAffiliation:The Hong Kong Polytechnic University Sun Yat\-sen UniversityYanggan GuHongxia YangAffiliation: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 introduceOnlineQAT, 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,OnlineQATobtains 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\.

## 1Introduction

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](https://arxiv.org/html/2610.09346#bib.bib1);[Chen et al\., 2025](https://arxiv.org/html/2610.09346#bib.bib2);[Okoshi et al\., 2026](https://arxiv.org/html/2610.09346#bib.bib3)\)\. 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](https://arxiv.org/html/2610.09346#S1)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 proposeOnlineQAT, 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,OnlineQATimproves 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\.

## 2Related 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](https://arxiv.org/html/2610.09346#bib.bib1)\)\. EfficientQAT introduces block\-wise all\-parameter training and a lightweight end\-to\-end quantization\-parameter stage\([Chen et al\., 2025](https://arxiv.org/html/2610.09346#bib.bib2)\)\. ReasoningQAT targets reasoning models with mixed\-domain block\-wise calibration followed by teacher\-guided reward rectification and forward\-KL distillation\([Okoshi et al\., 2026](https://arxiv.org/html/2610.09346#bib.bib3)\)\. 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\.OnlineQATretains 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](https://arxiv.org/html/2610.09346#bib.bib4)\)\. Recent analysis studies when OPD succeeds and attributes its progress to alignment on high\-probability tokens at student\-visited states\([Li et al\., 2026](https://arxiv.org/html/2610.09346#bib.bib5)\); large\-scale systems such as Nemotron 3 Ultra further use multi\-teacher OPD during post\-training\([NVIDIA, 2026](https://arxiv.org/html/2610.09346#bib.bib6)\)\. 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\.

## 3Method

OnlineQATseparates low\-bit recovery into local initialization and end\-to\-end behavioral recovery\. LetπT\\pi\_\{T\}be the frozen full\-precision model andπθq\\pi\_\{\\theta\}^\{q\}its quantized counterpart\. Stage 1 reconstructs one transformer block at a time\. Stage 2 samplesy^\\hat\{y\}fromπθq\\pi\_\{\\theta\}^\{q\}and evaluates both models at student\-generated prefixesht=\(x,y^<t\)h\_\{t\}=\(x,\\hat\{y\}\_\{<t\}\), rather than the fixed prefixes\(x,y<t\)\(x,y\_\{<t\}\)used by offline distillation\.

### 3\.1Stage 1: block\-wise low\-bit initialization

Following EfficientQAT and ReasoningQAT\([Chen et al\., 2025](https://arxiv.org/html/2610.09346#bib.bib2);[Okoshi et al\., 2026](https://arxiv.org/html/2610.09346#bib.bib3)\), Stage 1 uses sequential block\-wise QAT on calibration data drawn from pre\-training and reasoning corpora\. LetBlB\_\{l\}andBlqB\_\{l\}^\{q\}denote the full\-precision and fake\-quantized versions of blockll, and letHlH\_\{l\}andHlqH\_\{l\}^\{q\}be their respective input streams after the preceding blocks\. For a sequence ofTTtokens and hidden widthCC, we optimize

ℒblock\(l\)=1T​C​‖Blq​\(Hlq\)−Bl​\(Hl\)‖F2\.\\mathcal\{L\}\_\{\\mathrm\{block\}\}^\{\(l\)\}=\\frac\{1\}\{TC\}\\left\\\|B\_\{l\}^\{q\}\(H\_\{l\}^\{q\}\)\-B\_\{l\}\(H\_\{l\}\)\\right\\\|\_\{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\.2Stage 2: on\-policy reverse\-KL recovery

For each promptx∼𝒳x\\sim\\mathcal\{X\}, the quantized student samples a responsey^∼πθ,τq\(⋅∣x\)\\hat\{y\}\\sim\\pi\_\{\\theta,\\tau\}^\{q\}\(\\cdot\\mid x\)at temperatureτ\\tau\. At prefixht=\(x,y^<t\)h\_\{t\}=\(x,\\hat\{y\}\_\{<t\}\), letpθ,tτp\_\{\\theta,t\}^\{\\tau\}andpT,tτp\_\{T,t\}^\{\\tau\}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τ\)\]\.\\mathcal\{L\}\_\{\\mathrm\{online\}\}\(\\theta\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{X\},\\,\\hat\{y\}\\sim\\pi\_\{\\theta,\\tau\}^\{q\}\}\\left\[\\frac\{1\}\{\|\\hat\{y\}\|\}\\sum\_\{t=1\}^\{\|\\hat\{y\}\|\}D\_\{\\mathrm\{KL\}\}\\\!\\left\(p\_\{\\theta,t\}^\{\\tau\}\\,\\\|\\,p\_\{T,t\}^\{\\tau\}\\right\)\\right\]\.\(2\)For the sampled tokenat=y^ta\_\{t\}=\\hat\{y\}\_\{t\}, a one\-sample policy\-gradient estimator uses the stop\-gradient log\-ratio

AtRKL=sg⁡\[log⁡pθ,tτ​\(at\)−log⁡pT,tτ​\(at\)\],A\_\{t\}^\{\\mathrm\{RKL\}\}=\\operatorname\{sg\}\\\!\\left\[\\log p\_\{\\theta,t\}^\{\\tau\}\(a\_\{t\}\)\-\\log p\_\{T,t\}^\{\\tau\}\(a\_\{t\}\)\\right\],\(3\)wheresg\\operatorname\{sg\}stops gradients through the log\-ratio\. The training loss is

ℒOPD​\(θ\)=1N​∑t∈responseAtRKL​log⁡pθ,tτ​\(at\),\\mathcal\{L\}\_\{\\mathrm\{OPD\}\}\(\\theta\)=\\frac\{1\}\{N\}\\sum\_\{t\\in\\mathrm\{response\}\}A\_\{t\}^\{\\mathrm\{RKL\}\}\\log p\_\{\\theta,t\}^\{\\tau\}\(a\_\{t\}\),\(4\)withNNresponse tokens in the batch\. The constant\+1\+1term in the reverse\-KL policy gradient is omitted because its expectation is zero\. Minimizing Eq\.[4](https://arxiv.org/html/2610.09346#S3.E4)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\\tau=0\.6and no cross\-entropy term\.

## 4Experiments

### 4\.1Setup

#### 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 rate10−410^\{\-4\}; W3A16 and W2A16 master weights use10−510^\{\-5\}and×10−52\\\!\\times\\\!10^\{\-5\}\. All recovery methods start from the same Stage 1 checkpoint within each bit width\.

#### Baselines and Stage 2\.

Stg1has no end\-to\-end recovery\.KDis a fixed\-prefix distillation baseline with0\.2​ℒCE\+ℒJSD0\.2\\mathcal\{L\}\_\{\\mathrm\{CE\}\}\+\\mathcal\{L\}\_\{\\mathrm\{JSD\}\}\.ReasoningQATfollows Okoshi et al\.\([Okoshi et al\., 2026](https://arxiv.org/html/2610.09346#bib.bib3)\), using0\.2​ℒRSFT\+ℒFKL0\.2\\mathcal\{L\}\_\{\\mathrm\{RSFT\}\}\+\\mathcal\{L\}\_\{\\mathrm\{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\.OnlineQATsamples 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 at5×10−65\\times 10^\{\-6\}\. W2A16 warms up for 30 steps from2×10−72\\times 10^\{\-7\}to2×10−62\\times 10^\{\-6\}, followed by 120 steps at2×10−62\\times 10^\{\-6\}\.

Table 1:Qwen3\-1\.7B results\. ReasoningQAT denotes the method of[Okoshi et al\. \(2026\)](https://arxiv.org/html/2610.09346#bib.bib3)\. All entries are percentages; AVG is the unweighted mean\. Bold marks the best quantized result within each bit width\.

### 4\.2Main results

#### W3A16 recovery\.

OnlineQATaverages 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\.0on LiveCodeBench and\+11\.1\+11\.1on GPQA\-Diamond\. IFEval decreases by 0\.8 points\.

#### W2A16 recovery\.

At W2A16,OnlineQATimproves 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](https://arxiv.org/html/2610.09346#S4.T1)therefore supports an empirical comparison, not a causal or compute\-efficiency claim\.

## 5Conclusion

OnlineQATretains 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

- Agarwalet al\.\(2024\)R\. Agarwal, N\. Vieillard, Y\. Zhou, P\. Stanczyk, S\. Ramos, M\. Geist, and O\. BachemOn\-policy distillation of language models: learning from self\-generated mistakes\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px2.p1.1)\.
- Chenet al\.\(2025\)M\. Chen, W\. Shao, P\. Xu, J\. Wang, P\. Gao, K\. Zhang, and P\. LuoEfficientQAT: efficient quantization\-aware training for large language models\.InAnnual Meeting of the Association for Computational Linguistics,Cited by:[§1](https://arxiv.org/html/2610.09346#S1.p2.1),[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2610.09346#S3.SS1.p1.1)\.
- Duet al\.\(2024\)D\. Du, Y\. Zhang, S\. Cao, J\. Guo, T\. Cao, X\. Chu, and N\. XuBitDistiller: unleashing the potential of sub\-4\-bit LLMs via self\-distillation\.InAnnual Meeting of the Association for Computational Linguistics,Cited by:[§1](https://arxiv.org/html/2610.09346#S1.p2.1),[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px1.p1.1)\.
- Liet al\.\(2026\)Y\. Li, Y\. Zuo, B\. He, J\. Zhang, C\. Xiao, C\. Qian, T\. Yu, H\. Gao, W\. Yang, Z\. Liu, and N\. DingRethinking on\-policy distillation of large language models: phenomenology, mechanism, and recipe\.External Links:2604\.13016Cited by:[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px2.p1.1)\.
- NVIDIA \(2026\)NVIDIANemotron 3 ultra: open, efficient mixture\-of\-experts hybrid mamba\-transformer model for agentic reasoning\.External Links:2606\.15007Cited by:[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px2.p1.1)\.
- Okoshiet al\.\(2026\)Y\. Okoshi, H\. Otsuka, D\. Fujiki, and M\. MotomuraTowards quantization\-aware training for ultra\-low\-bit reasoning LLMs\.InThe Fourteenth International Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2610.09346#S1.p2.1),[§2](https://arxiv.org/html/2610.09346#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2610.09346#S3.SS1.p1.1),[§4\.1](https://arxiv.org/html/2610.09346#S4.SS1.SSS0.Px4.p1.1),[Table 1](https://arxiv.org/html/2610.09346#S4.T1)\.