The Devil Is in the Reconstruction Loss Scale: Rethinking Optimization in LLM Quantization
Post-training quantization (PTQ) for LLMs is typically performed sequentially, block by block, by minimizing a reconstruction loss. The authors show that the standard MSE loss creates an Optimization Imbalance: because gradient magnitude scales with the square root of loss, early stages receive much weaker updates than later ones. Replacing MSE with RMSE implicitly normalizes gradients across stages, improving quantization accuracy—especially at aggressive bit widths like W2A16 and W4A4—with no extra parameters or meaningful compute overhead.
The paper identifies a genuine optimization pathology in sequential PTQ: MSE inherently couples gradient magnitude to the reconstruction loss scale, causing early stages to be under-optimized. The RMSE fix is theoretically clean, requires no extra parameters, and is validated across four PTQ methods and multiple models, with especially strong gains under aggressive quantization. However, simpler alternatives such as per-stage learning-rate tuning are not ruled out, and the preference for specific RMSE groupings across model families remains unexplained.
The empirical sweep is unusually thorough: the authors track optimization dynamics across four distinct PTQ methods, multiple model families, dense and MoE architectures, and bit widths from W4A16 down to W2A16 and W4A4. The gradient statistics directly verify the theoretical prediction that MSE output-space gradients scale as $2\sqrt{\mathcal{L}_{\mathrm{MSE}}/N}$ while RMSE fixes them to $1/\sqrt{N}$. Because the change is a strict drop-in replacement with negligible runtime overhead—Table 4 reports under $3\%$ difference—the practical barrier to adoption is minimal.
The paper does not rule out simpler ways to decouple optimization strength from loss scale; for example, per-stage learning-rate tuning could counteract MSE's gradient scaling without changing the loss function, yet this baseline is never discussed. Additionally, while Appendix E shows that high-precision block outputs already grow with depth—meaning the large loss scale is partly intrinsic to the model—the text still frames the 10,000x range as unexpected without fully reconciling these two perspectives. Finally, the family-dependent preference for RMSE-Token on LLaMA and RMSE-Channel on Qwen is noted but not explained.
The evidence strongly supports the central mechanistic claim: Figures 2 and 4 show that RMSE flattens gradient magnitudes across stages, and Table 1 links these balanced dynamics to consistent perplexity and accuracy improvements. Because RMSE is evaluated as a drop-in replacement within identical pipelines, the comparison to MSE is inherently fair. However, the paper does not benchmark against analytical PTQ methods such as GPTQ or AWQ, nor does it compare against other loss-function modifications like the cosine term used in LRQuant, leaving open whether balanced optimization delivers a better frontier than orthogonal quantization advances.
The experiments use publicly available models, standard WikiText2 calibration, and the LM-Evaluation-Harness. The four RMSE variants are defined with exact equations in Appendix A.2, and the authors report following official hyperparameters for each baseline method. The main obstacle to independent reproduction is that the code repository is listed as forthcoming, and exact random seeds or per-method learning-rate schedules are not provided in the main text or appendix. Because the gains are measured relative to reproduced MSE baselines rather than published numbers, reproducing the comparison requires faithfully replicating the entire pipeline.
Post-training quantization (PTQ) methods typically use sequential quantization that partitions a pre-trained LLM into a series of units (e.g., transformer blocks), with one unit quantized at each stage. State-of-the-art PTQ methods are predominantly learning-based, optimizing auxiliary quantization parameters (e.g., scaling factors, rotation matrices, clipping thresholds, and adapters) via gradient descent to minimize a reconstruction loss. A common practice is to use mean squared error (MSE) as the reconstruction loss function, yet its induced optimization behavior remains largely unexplored. In this work, we take a holistic view of sequential quantization and systematically investigate how optimization evolves from the first quantization stage to the last, aiming for a deep understanding of optimization in learning-based PTQ schemes. Through extensive empirical studies spanning representative learning-based PTQ methods, LLM families, model scales, architectures, quantization settings and various tasks, we consistently uncover Optimization Imbalance: reconstruction loss magnitudes vary dramatically across stages, accompanied by highly uneven gradient magnitudes and parameter updates under MSE. We term the cross-stage range of loss magnitudes the reconstruction loss scale, and reveal that MSE translates the unexpectedly large reconstruction loss scale into highly uneven gradient magnitudes, which in turn lead to uneven optimization strength across quantization stages. This finding suggests a general principle for improving learning-based PTQ: optimization strength across stages should be decoupled from the reconstruction loss scale. Theoretically, we show that root mean squared error (RMSE) variants defined at the sample, channel, token, and element levels naturally realize this principle through implicit gradient normalization, outperforming MSE significantly as a drop-in replacement.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.