PaperlayerView exact revision on arXiv

Source-linked AI summary

QUASAR: Lowering the Loss Floor of Quantization-Aware Training with Loss-Aware Reconstruction

Vincent Counathe, Ben Athiwaratkun, Christopher De Sa, Tianyi Zhang

arXiv:2608.13966v1cs.LGcs.CLstat.ML

TL;DR

Low-bit QAT lacks a practical way to keep reconstructed weights aligned with evolving latent weights without costly repeated optimization. QUASAR continuously minimizes loss-aware reconstruction error during training, achieving the lowest held-out KL divergence across evaluated models and bit widths, with especially strong INT2 accuracy gains.

  • Problem

    Applying loss-aware reconstruction to QAT is computationally difficult because evolving weights make repeated PTQ optimization prohibitively expensive.

  • Method

    QUASAR continually minimizes saliency-weighted reconstruction error by searching clipping ranges and fitting dequantizers with online gradient-based saliencies.

  • Results

    QUASAR achieves the lowest training and held-out KL loss across six settings, improving INT2 average task accuracy by 3.5–4.3 points over the strongest QAT baseline.

  • Takeaways & Limitations

    Loss-aware reconstruction is the only reconstruction-dependent term in the QAT convergence bound and, under an additional PL condition, affects the final quantized-model loss.

  • Takeaways & Limitations

    Repeated PTQ-style reconstruction is prohibitively expensive, and the theoretical guarantee additionally assumes a well-behaved loss surface along QUASAR’s candidate directions.

Abstract

from arXiv · show

As large language model inference shifts toward lower precision, post-training quantization (PTQ) becomes increasingly brittle, making quantization-aware training (QAT) essential for preserving model quality. However, QAT computes the loss and surrogate gradients using a lossy reconstruction of latent full-precision weights, while applying updates to the latent weights themselves. This mismatch can lead to suboptimal training trajectories and a higher loss floor. Second-order PTQ methods mitigate a similar gap by minimizing loss-aware reconstruction error, but doing it once for a frozen model can take hours; repeating this process throughout QAT as the weights evolve is impractical. We introduce QUASAR, a QAT method that continuously performs lightweight, loss-aware reconstruction in the training loop to lower the loss floor and improve the resulting low-bit model. At each training step, QUASAR uses the exponential moving average of squared gradients as online saliency estimates, searches over a small set of clipping ranges, and fits affine dequantizers via saliency-weighted least squares. Our analysis shows that the loss-aware reconstruction error is the only reconstruction-dependent term in the QAT convergence bound and controls the loss of the final quantized model, establishing QUASAR's objective as a principled optimization target. QUASAR modifies only the training procedure and supports standard deployment formats, including integer quantization and NVFP4, with no inference-time changes or overhead. Across Qwen3 and Llama-3.1, QUASAR achieves the lowest held-out KL divergence among competitive QAT methods at 2, 3, and 4 bits, reducing KL by at least 10% at 3 and 4 bits and by 29% at 2 bits. At 2 bits, it improves average accuracy across eight tasks by 3.5-4.3 percentage points over strong QAT and PTQ baselines.

1 Introduction

Low-precision inference motivates training models under serving precision because PTQ is increasingly brittle, while QAT suffers from a reconstruction mismatch that raises its loss floor. QUASAR addresses this gap by continuously minimizing loss-aware reconstruction error, with theory and experiments showing improved trajectories and lower loss across precisions and model families.

  • Motivation: Low-precision inference reduces memory footprint, decoding latency, and enables higher concurrency and throughput, but PTQ becomes brittle on newer long-context reasoning and agentic models.Behavioral shifts and biases can also emerge relative to full-precision models.
  • Problem: QAT trains with quantized–dequantized reconstruction weights r while updating latent full-precision weights w, and its STE can produce a suboptimal optimization trajectory.This mismatch creates a loss-floor gap in which QAT reaches higher final training and held-out loss than full-precision training.
  • Method: QUASAR continually minimizes loss-aware reconstruction error during QAT to improve the training trajectory and lower the loss floor.It decomposes reconstruction into quantization with a free clipping range and dequantization back to reconstructed weights.
  • Theory: Only the loss-aware reconstruction error in the QAT convergence bound depends on the reconstruction map, making it QUASAR’s direct optimization target and, under a PL condition, a lever on final quantized loss.The other bound terms are initialization and minibatch noise.
  • Empirical results: At INT4 and INT3, QUASAR reduces evaluation loss by at least 10%, and at INT2 by at least 29%, versus Standard QAT, LSQ, Denoising QAT, and BitDistiller on Qwen3-4B-Thinking.QUASAR achieves lower final evaluation loss than those baselines at every tested bit width, with gains also extending to NVFP4.

2 Related Work

QUASAR bridges PTQ and QAT by bringing loss-aware reconstruction into the training loop. It addresses the heuristic latent-to-reconstructed weight map while preserving standard RTN-style deployment without inference overhead.

  • PTQ: PTQ is typically one-shot, using calibration data and often second-order loss approximations to choose low-precision weights that minimize damage.Optimal Brain Surgeon, OBQ, and GPTQ use curvature- or loss-aware perturbation measures for quantization.
  • PTQ: QuIP, QuIP#, QTIP, and AQLM strengthen two- and three-bit PTQ with transforms or codebooks, but their inference operations reduce throughput.These methods use random orthogonal transforms and lattice, trellis, or additive codebooks.
  • QAT: Below four bits, overhead-free PTQ degrades sharply, whereas QAT methods improve quality through learned step sizes, full-precision supervision, or low-rank adapters.LSQ, LLM-QAT, BitDistiller, and QLoRA exemplify these approaches, but they do not optimize the latent-to-reconstructed weight map.
  • QUASAR: QUASAR searches clipping ranges and fits dequantizers with saliency-weighted least squares at every step, effectively embedding PTQ reconstruction as weights adapt.The resulting reconstruction error controls reconstruction-induced gradient mismatch and the final quantized model’s loss; deployment requires only a scale and offset.

3 Preliminaries

Standard QAT trains latent full-precision weights but computes the forward pass, loss, and STE gradients on reconstructed quantized weights, so reconstruction error affects every update. QUASAR targets this mismatch through loss-aware reconstruction, whose tractable approximation is tied to QAT convergence and final deployable-model loss.

  • QAT reconstruction: Standard QAT updates latent weights w while computing the forward pass and loss on transient reconstructed weights r derived at each step.The reconstruction maps latent weights to integer codes and dequantizes them back to real-valued weights.
  • Standard reconstruction: Standard QAT selects clipping ranges and dequantization parameters from groupwise extrema without considering the loss.Extremes are placed at the quantization-grid endpoints, and the resulting scale and zero-point are reused for dequantization.
  • QAT reconstruction: The forward pass, loss, and STE gradient are computed at r, but the optimizer applies the update to w, allowing reconstruction errors to affect every optimization step.Rounding has no useful derivative, motivating the straight-through estimator.
  • Loss-aware objective: QUASAR’s tractable approximation bS_t is the explicit reconstruction-dependent term in the QAT convergence bound and controls an upper bound on final deployable-model loss under a PL condition.This motivates minimizing loss-aware reconstruction error rather than relying only on weight extrema.
  • Computational limitation: Repeating full-Hessian PTQ reconstruction during QAT is prohibitively expensive because per-layer linear-algebra solutions become stale after every weight update.Such solvers can already take minutes or hours for billion-parameter frozen models.

4 Methodology

QUASAR integrates lightweight, loss-aware weight reconstruction into every QAT step by combining online saliency estimation, clipping-range search, and optimal dequantization. The same two-stage procedure extends beyond integer quantization to production formats such as NVFP4.

  • Core method: QUASAR performs loss-aware weight reconstruction throughout QAT, using scale search and optimal dequantization to improve the final model.The method changes the reconstruction used during the forward pass rather than activation quantization or backpropagation.
  • Online saliency: An exponential moving average of squared gradients provides a diagonal-Fisher curvature proxy, reusable from Adam or AdamW without additional memory overhead.The estimate approximates the Hessian as H_t ≈ diag(h).
  • Online saliency: Saliency-weighted reconstruction allocates lower error to high-saliency weights and permits larger errors for low-saliency weights within each quantization group.Jointly optimizing code assignment and dequantization keeps the reconstructed-weight loss close to the latent-weight loss.
  • Two-stage reconstruction: QUASAR searches candidate clipping ranges for code assignments and fits affine dequantizers by weighted least squares, selecting the reconstruction with the lowest loss-aware error.The clipping range determines which weights share codes, while the fitted scale and zero-point determine represented values.
  • Scale search: 99.6% of Qwen3-4B groups selected narrower-than-full-range clipping during 3-bit integer quantization, reducing bS_t to 69% of the full-range baseline.This result was obtained during quantization-aware distillation.
  • NVFP4 extension: QUASAR applies the same code-assignment search and dequantization fitting to NVFP4 while satisfying its grid and scaling constraints.For NVFP4, it searches candidate group scales and fits the FP8 dequantization scale, while retaining the tensor-level FP32 scale from the absolute maximum.

5 Theoretical Analysis: Loss-Aware Reconstruction Bounds the Loss of the Final Quantized Model

The analysis identifies QUASAR’s loss-aware reconstruction error as the only reconstruction-dependent term in the QAT convergence bound. Under an additional PL condition, it further connects lower reconstruction error during training to a tighter loss guarantee for the final quantized model.

  • Convergence Bound: QUASAR’s loss-aware reconstruction error is the only reconstruction-dependent term in the QAT convergence bound, making it the precise optimization lever.The analysis links this objective to both the training trajectory and the final model’s loss.
  • Convergence Bound: Theorem 1 separates the convergence guarantee into initialization, minibatch-noise, and reconstruction-error terms, with QUASAR minimizing the last term at each step.Lower reconstruction error yields a tighter guarantee that both latent and reconstructed weights approach stationarity.
  • Final Quantized Model: Under Assumptions 1–4 and a PL condition, the final quantized reconstruction’s loss is directly bounded through the loss-aware reconstruction error.This result formalizes why better reconstructions throughout training translate into a tighter guarantee for the deployed model.
  • Empirical Validation: Across model families and tested bit widths, QUASAR has the lowest reconstruction error and held-out KL among Standard QAT, Denoising QAT, and QUASAR.Figure 3 reports that reconstruction error tracks held-out KL to the full-precision counterpart in every setting.

6 Experiments

Experiments show that QUASAR lowers held-out loss and improves downstream capabilities across low-bit INT and NVFP4 settings. Its gains require reconstruction to remain active during training, while adding only 1.4% step-time overhead.

  • Healing: QUASAR reaches the lowest training and held-out KL loss in all six healing settings, with every baseline plateauing above it.At INT2, it reduces KL by about 30%, raises top-1 agreement by about 2 points, and improves average task accuracy by 3.5–4.3 points.
  • Healing: At INT2 on GSM8K, QUASAR scores 68.8 on Qwen and 66.4 on Llama, versus 0.0–0.2 for same-format PTQ.The strongest competing QAT method reaches 49.0 and 53.1, respectively; on Llama, QUASAR is only 3.7 points below BF16.
  • Healing: At 2 bits, an LLM judge prefers QUASAR over every quantized opponent in all 14 comparisons.The judge used 128 WildChat prompts and preferred only the full-precision teacher over QUASAR.
  • Adaptation: QUASAR reaches the lowest training and held-out perplexity at INT4, INT3, and INT2 when adapting Qwen3-4B-Base.At INT2, it reaches 29.6 average accuracy, 10.9 points above the strongest QAT baseline, with 60.7 on MATH-500 and 75.4 on GSM8K.
  • NVFP4: Across Qwen3-8B and Qwen3.5-9B, QUASAR reduces held-out KL by approximately 30% and improves average downstream accuracy by 0.5–2.0 points in NVFP4.The evaluation uses simulated NVFP4 during training and real quantized NVFP4 checkpoints for downstream evaluation.
  • Ablations and overhead: Reconstruction must remain in the training loop: applying it only at initialization gives up nearly all improvement, while Adam’s second moment matches a separate Fisher estimate at lower cost.The median selected range factor falls from 0.95 at INT4 to 0.60 at INT2, and QUASAR adds only 1.4% to matched INT3 optimizer-step time.

7 Conclusion · Appendix

QUASAR addresses the mismatch between reconstructed weights used in QAT forward passes and latent full-precision weights updated during training. It minimizes loss-aware reconstruction error through saliency-weighted dequantization fitting and analyzes QAT convergence via a three-term bound decomposition.

  • 7 Conclusion: QUASAR is a QAT method designed to mitigate the mismatch between reconstructed forward-pass weights and updated latent full-precision weights.The method targets a discrepancy inherent in the training procedure.
  • 7 Conclusion: QUASAR minimizes loss-aware reconstruction error by searching over quantization scales.Scale search is part of the method’s reconstruction procedure.
  • 7 Conclusion: QUASAR fits dequantization parameters using saliency-weighted least squares.The saliency statistics are derived from the optimizer.
  • 7 Conclusion: The method derives saliency statistics from optimizer information to guide its reconstruction procedure.These optimizer-derived statistics support the saliency weighting used during fitting.
  • 7 Conclusion: QUASAR’s theoretical analysis decomposes the QAT convergence bound into three terms.The supplied passage identifies initialization and minibatch-related terms, while the final term is truncated.
  • Appendix: The supplied Appendix material contains no passage-level evidence for an additional appendix-specific finding.No appendix claim can be reported without introducing unsupported content.

A Proofs

This appendix provides proofs for the results stated in Section 5.

  • The appendix contains proofs of the results stated in Section 5.

A.1 Proofs of the main results · A.2 Validation under the SGD dynamics · A.3 Choice of saliency weighting

The appendix proves QUASAR’s main convergence results using smoothness, conditional expectations, telescoping loss terms, and terminal reconstruction bounds. It also validates the predicted method ordering under plain SGD and shows that linear saliency weighting outperforms squared weighting across tested settings.

  • A.1 Proofs of the main results: The sufficient condition for Assumption 3 follows when L is L-smooth and h_t,i ≥ h_min > 0, yielding C = L^2/h_min.The appendix explicitly states that this establishes equation (9).
  • A.1 Proofs of the main results: Theorem 1’s proof sets d_t = r_t − w_t and applies the fundamental theorem of calculus to establish the first assertion.The argument then uses the definition of r_t to conclude the assertion.
  • A.1 Proofs of the main results: The convergence proof combines smoothness, the update rule, conditional expectations, and the bound E_t ∥c_t + ξ_t∥^2 ≤ (1 + M)∥c_t∥^2 + σ^2.The proof takes expectations, sums over 0 ≤ t < T, and uses telescoping loss terms before dividing by ηT/2.
  • A.1 Proofs of the main results: Corollary 1 uses q = 1 − ηµ, the step-size conditions, the PL inequality, and a terminal comparison between w_T and r_T to derive the stated result.The terminal comparison applies smoothness at w_T, the theorem’s reconstruction argument, and the PL inequality at r_T.
  • A.2 Validation under the SGD dynamics: The SGD validation repeats the INT2 quantization-aware distillation comparison under plain SGD to test whether the predicted ordering depends on the optimizer.It uses no momentum, a constant learning rate, no warmup, no gradient clipping, 1,024 steps, and global batch 32; QUASAR maintains an EMA of squared gradients.
  • A.3 Choice of saliency weighting: The saliency score b_S weights squared perturbations by h_t, which tracks diagonal curvature and targets the linearized gradient mismatch ∥H_t∆_t∥^2.The comparison evaluates the two weighting choices with all other factors fixed under AdamW and plain SGD settings.
  • A.3 Choice of saliency weighting: The h_t weighting reaches lower held-out loss in all six comparisons across INT2, INT3, and INT4 AdamW rates and three stable plain-SGD rates.The squared weighting is unstable, collapsing at INT4 and at the larger SGD learning rates.

B Method Details

QUASAR uses strongly varying within-group saliency to prioritize loss-aware reconstruction and typically selects clipping ranges narrower than min–max. The method also instantiates in NVIDIA Rubin INT3 LUT format while preserving an ordinary deployable checkpoint.

  • Clipping-range search: 99.6% of Qwen3-4B INT3 weight groups select a clipping range narrower than min–max.The searched range clips extreme weights and follows the bulk of each group.
  • Saliency-guided reconstruction: QUASAR’s saliency h varies strongly within each 128-weight group, directing reconstruction accuracy toward the most important weights.The groups run along the input dimension, with 128 weights per group.
  • NVIDIA Rubin INT3 instantiation: Rubin INT3 LUT-QUASAR replaces affine dequantization with an NVIDIA hardware lookup table using 3-bit indices and eight E4M3 codebook entries.For fixed indices, each used entry is set to its loss-aware mean rounded to E4M3; candidate clipping ranges are searched and the best reconstruction retained.
  • NVIDIA Rubin INT3 instantiation: 3.125 bits per weight are stored using 512 packed 3-bit indices and eight 8-bit E4M3 entries per 512-weight block.The exported representation uses packed indices and a codebook.
  • Deployment: After training, QUASAR discards curvature estimation, candidate search, and codebook fitting, exporting an ordinary Rubin INT3 LUT checkpoint.The deployment checkpoint therefore contains the packed hardware-format representation rather than the reconstruction procedure.

C Additional Healing Results

The appendix extends the healing analysis with additional checkpoint metrics, training dynamics, loss-aware reconstruction results, and qualitative INT2–INT4 response examples. It also examines response statistics and preference outcomes for QUASAR against competing methods and full precision.

  • Additional checkpoint metrics: Table 7 reports reverse KL, 99th-percentile per-token KL, and top-1 agreement on uncertain full-precision tokens for INT4, INT3, and INT2 checkpoints.The evaluation covers Qwen3-4B-Thinking-2507 and Llama-3.1-8B-Instruct using the same final checkpoints as Tables 1 and 2.
  • Additional checkpoint metrics: The appendix compares average benchmark accuracy and held-out KL across bit widths, model sizes, and QAT or PTQ methods.Figures 16–19 and 22 include both healing models and mark full-precision references where applicable.
  • Training dynamics: The appendix evaluates top-1 agreement during QAD and reports net pairwise preference for QUASAR with 95% bootstrap confidence intervals over 128 prompts.These comparisons include competing QAT methods and the full-precision model.
  • Training dynamics: QUASAR’s loss-aware reconstruction-error reduction is shown by projection type at initialization and after training across INT4, INT3, and INT2.Figure 21 reports median reductions over modules, measured against full-precision weights, relative to Standard QAT.
  • Response quality: Qualitative INT2–INT4 examples compare RTN, GPTQ, AWQ, and Standard responses, while response statistics measure repetition and termination across 128 held-out prompts.The examples include a stranded-on-an-island prompt and cover both Qwen3-4B-Thinking-2507 and Llama-3.1-8B-Instruct.

D Additional Adaptation Results

This appendix extends the INT2 adaptation analysis by examining answer outcomes, consistency across repeated samples, and divergence from full-precision SFT along reasoning traces.

  • Answer outcomes: The appendix analyzes how INT2 PTQ and QAT checkpoints divide generated samples among correct, wrong, and missing final answers.Results cover Qwen3-4B-Base on MATH-500 and GSM8K, with the un-finetuned base model and full-precision SFT included for reference.
  • Sampling consistency: It measures repeated-attempt consistency for INT2 checkpoints using the pass@8-to-avg@1 ratio on MATH-500.Eight samples are generated per problem at temperature 0.6; a ratio of 1 indicates identical solved problems across attempts, while larger values indicate more variable successes.
  • Reasoning-trace divergence: It tracks each quantized checkpoint’s per-position forward KL divergence from full-precision SFT as reasoning traces unfold.The evaluation feeds identical tokens from 1,000 full-precision SFT traces on MATH-500 to every model and averages divergence within each tenth of the trace.
Loading 2608.13966v1…