QUASAR: Lowering the Loss Floor of Quantization-Aware Training with Loss-Aware Reconstruction
Vincent Counathe, Ben Athiwaratkun, Christopher De Sa, Tianyi Zhang
cs.LG, cs.CL, stat.ML
2026-08-14
QUASAR re-fits the quantizer every QAT step with Adam saliency, cutting held-out KL by about 29% at 2 bits and lifting 8-task accuracy 3.5-4.3 points, at 1.4% extra step time.
Inference is moving to INT4 and NVFP4, and post-training quantization gets brittle below four bits. Same-format GPTQ and AWQ drive GSM8K on Qwen3-4B-Thinking to roughly zero. Quantization-aware training is supposed to adapt the model to the serving grid. The optimizer is looking at the wrong weights.
The forward pass and the loss use reconstructed weights r after quantize-dequantize. The optimizer updates latent full-precision weights w. Rounding has no useful gradient, so training copies the gradient at r onto w via a straight-through estimator. When r sits far from w, that borrowed gradient is a poor descent direction for w, and training plateaus above the full-precision loss floor.
Second-order PTQ methods close a similar gap by minimizing Hessian-weighted reconstruction error, once, on a frozen model, often in minutes to hours. In QAT, w moves every step, so rerunning GPTQ on every forward is not viable. Reconstruction error still governs the trajectory. Nobody was cheaply re-optimizing it inside the loop.
QUASAR changes only the map from w to r. Backprop still uses the STE. Deployment is still ordinary grouped quantization. The extra work lives in the training-time reconstruction.
Saliency comes from the second-moment buffer Adam or AdamW already keeps, treated as a diagonal Fisher, one hi per weight. A group of g weights shares one scale and zero-point, so error has to be budgeted inside the group: protect high-hi weights, allow more error on the rest.
Reconstruction has two stages. First, search clipping ranges: scale the group's center and half-width by factors f from 0.30 to 1.00 in steps of 0.05, each f assigning a different integer code vector. With codes fixed, scale s and zero-point z have a closed-form weighted least-squares solution that minimizes the saliency-weighted squared error. Keep the candidate with the smallest error. During INT3 distillation of Qwen3-4B, 99.6% of groups picked a range narrower than min-max, cutting that error to 69% of the full-range baseline.
NVFP4 uses the same split: search E2M1 codes inside a group of 16, then fit the FP8 dequant scale. The tensor-level FP32 scale still comes from the absmax and is not optimized. What gets stored for inference is those scales, so the serving path does not change.
The STE-SGD bound splits into initialization, minibatch noise, and loss-aware reconstruction error. Only the third term depends on the reconstruction map, and that is the quantity minimized every step. Under a Polyak–Lojasiewicz condition, the same term bounds the loss of the final quantized model. Main runs use AdamW; an INT2 rerun under plain SGD kept the same method ranking.
Quantization-aware distillation uses the matching full-precision model as teacher, Open-PerfectBlend data, 4096 steps. Weights are INT2/3/4 with group size 128; activations and embeddings stay BF16.
Held-out KL on Qwen3-4B-Thinking-2507:
| Bits | QUASAR | Best QAT baseline | Best same-format PTQ |
| INT2 | 0.126 | Denoising QAT 0.179 | GPTQ 1.140 |
| INT3 | 0.054 | Denoising QAT 0.060 | GPTQ 0.130 |
| INT4 | 0.016 | Denoising/BitDistiller 0.018 | GPTQ 0.027 |
At INT2, eight-task average accuracy is 53.2 versus BitDistiller 48.9 and the teacher 66.1. GSM8K is 68.8 versus BitDistiller 49.0; PTQ sits at 0.0–0.2. Llama-3.1-8B-Instruct under the same recipe: INT2 KL 0.105 versus BitDistiller 0.151, average accuracy 59.5 versus 56.0, GSM8K 66.4 versus 53.1, 3.7 points behind the teacher at 70.1.
INT3 and INT4 still have the lowest KL, with smaller task gaps. On Qwen INT3, LSQ averages 63.4 against QUASAR 62.7: KL wins, the average score does not. At INT4 both hug the teacher (Qwen 65.6 vs 66.1, Llama 69.4 tying the teacher).
The harder setting trains a new skill directly at low bits. Qwen3-4B-Base fine-tuned on OpenMathReasoning hits 29.6 average across five math benchmarks at INT2, 10.9 points above the best QAT baseline at 18.7; full-precision SFT then GPTQ scores 0.6. MATH-500 is 60.7 versus 39.2 for that QAT baseline and 83.8 for full-precision SFT. AIME'24 is still 4.0 against 22.9, so long traces are partly kept, not fully kept.
On NVFP4, Qwen3-8B and Qwen3.5-9B drop held-out KL from 0.016 to 0.011 and 0.009 to 0.006, about 30%, with 0.5–0.8 point gains on a nine-task average, evaluated through vLLM's real NVFP4 path.
Ablations show scale search doing the heavy lifting: search plus uniform fit cuts KL from 0.175 to 0.114; Adam weighting without search only reaches 0.164. On 8×H100, an INT3 step is 2.972 s versus 2.930 s for standard QAT, +1.4%.
The serving format does not change. The finished checkpoint is ordinary grouped integers or NVFP4, no Hadamard, no codebook. When same-format 2-bit PTQ is unusable, reconstruction quality decides whether the model can still do math. For teams already running QAT, this is a second-order PTQ idea folded into the forward pass at almost no step-time cost, not a new inference stack.
At INT4 the gain is incremental. The split appears at 2 bits, and when teaching a new skill at the target precision instead of training full precision then crushing it.
The Hessian proxy is a diagonal Fisher from Adam's second moment, not GPTQ's intra-layer correlation matrix. Clipping factors are a 0.05 grid, not a continuous optimal code. Main integer runs quantize weights only, with activations in BF16, which is a different constraint from W4A4 NVFP4. Results stop at 4B–9B; there is no 70B number. Qwen INT3 is the awkward case where the lowest KL does not mean the highest average score. INT2 math jumps relative to QAT baselines and still lags full-precision SFT, especially on AIME and HMMT. The paper has no standalone limitations section; these bounds come from the experimental setup.