PaperlayerView exact revision on arXiv

Source-linked AI summary

The Devil Is in the Reconstruction Loss Scale: Rethinking Optimization in LLM Quantization

Chao Li, Shigeng Wang, Anbang Yao

arXiv:2610.00983v1cs.CLcs.AI

TL;DR

The paper investigates how reconstruction loss affects optimization across stages of learning-based PTQ, where MSE’s stage-wise optimization behavior has been little examined. It traces this imbalance to MSE’s loss-dependent gradient scaling and finds that RMSE variants improve quantization performance across methods, models, and tasks.

  • Problem

    How reconstruction loss shapes optimization across sequential PTQ stages, and whether each unit receives sufficient optimization, have not been systematically examined.

    Source: §1 · p. 3

  • Method

    The paper tracks losses, gradients, and parameter updates across sequential PTQ stages and studies RMSE variants that normalize gradients at sample, channel, token, or element levels.

    Source: §1 · p. 3

  • Results

    RMSE variants consistently outperform MSE across language modeling, reasoning, and coding tasks, with absolute average accuracy gains of 5.44%–10.83% under W2A16 and 0.75%–8.05% under W4A4.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Optimization strength across stages should be decoupled from reconstruction loss scale, and RMSE provides a drop-in replacement for MSE without additional trainable parameters or hyperparameters.

    Source: §1 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

This paper examines how MSE shapes optimization in sequential learning-based PTQ, where the behavior of the loss and whether every unit receives sufficient optimization have been little studied. It finds that MSE couples optimization strength to reconstruction loss scale and that RMSE variants improve quantization performance across diverse settings.

Motivation

  • MSE’s optimization behavior across sequential quantization stages remains largely unexplored, and whether every unit receives sufficient optimization has not been systematically examined.Sequential quantization processes one unit at a time while keeping previously quantized units fixed.Source: §1 · pp. 2–3

Approach

  • The study tracks reconstruction losses, gradients, and parameter updates across stages and tests four PTQ methods across model families, scales, architectures, and quantization settings.The evaluation includes OmniQuant, DuQuant, FlatQuant, and SliderQuant, and LLaMA2, LLaMA3.1, and Qwen3.5 models.Source: §1 · p. 3

Finding

  • Across the tested settings, reconstruction losses span a broad cross-stage range, which the paper terms the reconstruction loss scale, and MSE accompanies this with uneven gradients and parameter updates.The largest loss magnitude often exceeds the smallest by more than 10,000×.Source: §1 · p. 3

Analysis and proposal

  • MSE couples optimization strength to loss magnitude, whereas RMSE implicitly normalizes output-space gradient magnitude and enables more balanced optimization across stages.The paper presents RMSE variants defined at sample, channel, token, and element levels.Source: §1 · p. 3

Results

  • RMSE consistently improves quantization across language modeling, reasoning, and coding tasks, with absolute average accuracy gains of 5.44%–10.83% under W2A16 and 0.75%–8.05% under W4A4.All four RMSE variants outperform MSE, and RMSE requires no additional trainable parameters or hyperparameters while adding negligible runtime and memory overhead.Source: §1 · p. 3

2 UNDERSTANDING OPTIMIZATION IMBALANCE IN LEARNING-BASED PTQ

The section formulates learning-based PTQ as stage-wise reconstruction with trainable auxiliary parameters, then examines how loss magnitudes and optimization vary across units. It explains how MSE’s loss-dependent gradient scaling contributes to Optimization Imbalance.

2.1 LEARNING-BASED PTQ PRELIMINARIES

  • Each sequential quantization stage reconstructs a unit’s high-precision output while keeping its pretrained weights frozen and optimizing only auxiliary quantization parameters.A unit may be a layer, block, or group of consecutive blocks.Source: §2.1 · pp. 3–4

2.2 OBSERVATION: OPTIMIZATION IMBALANCE IN SEQUENTIAL QUANTIZATION

  • Across units, reconstruction losses span a broad range, often exceeding 10,000× from the smallest to the largest, as prior quantization errors accumulate and unit output magnitudes vary.The paper calls this cross-stage range the reconstruction loss scale.Source: §2.2 · p. 4
  • Units with larger losses receive larger relative parameter updates, while early units tend to receive the weakest updates despite affecting more subsequent units.Because quantized units remain fixed, errors left by under-optimized units propagate and accumulate through later stages.Source: §2.2 · p. 4

2.3 WHY MSE LEADS TO OPTIMIZATION IMBALANCE

  • MSE’s output-space gradient magnitude is proportional to the square root of loss magnitude, so a loss range above 10,000× yields an output-gradient range exceeding 100×.This scaling couples optimization strength to loss magnitude, though network parameterization also influences the ultimate optimization dynamics.Source: §2.3 · p. 5

3 MITIGATING OPTIMIZATION IMBALANCE VIA RMSE

RMSE is presented as a way to balance optimization across quantization stages: its implicit gradient normalization decouples gradient magnitude from reconstruction-loss scale. The paper extends this principle into four grouped variants.

3.1 RMSE FOR BALANCED OPTIMIZATION

  • RMSE normalizes overall gradient magnitude while preserving gradient direction and relative element magnitudes, so variation in reconstruction-loss scale does not translate into imbalanced gradients across stages.This normalization arises from the loss function itself and requires neither additional hyperparameters nor changes to the PTQ pipeline.Source: §3.1 · p. 5

3.2 A UNIFIED FAMILY OF RMSE LOSS FUNCTIONS

  • Grouping reconstruction errors by sample, channel, token, or element yields four RMSE variants that apply gradient normalization within different feature groups.RMSE-Element treats each feature element as a group and reduces to L1 loss.Source: §3.2 · pp. 5–6
  • All four RMSE variants outperform MSE across the evaluations, although their relative performance differs across model families.Different groupings also show different degrees of within-stage reconstruction-loss variation, making feature grouping a design dimension.Source: §3.2 · p. 6

4 EXPERIMENTS

The experiments test RMSE against MSE across learning-based PTQ methods, models, and quantization settings. RMSE improves quantization results, while its variants produce more balanced optimization across stages, with gains varying by model family and setting.

4.1 EXPERIMENTAL SETUP

  • The evaluation spans five LLMs, four learning-based PTQ methods, five quantization settings, and language-modeling and downstream accuracy benchmarks.The methods use differing learnable designs, and SliderQuant reconstructs sliding windows of consecutive blocks rather than individual transformer blocks.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • Under W2A16 on Qwen3.5-27B, RMSE reduces perplexity from 37.51 to 22.20 and raises downstream accuracy from 60.45% to 70.96%.Across OmniQuant configurations, RMSE reduces perplexity and improves downstream accuracy, with larger gains under more aggressive quantization.Source: §4.2 · p. 6, Table 1 · p. 7
  • Under MSE, broad loss magnitudes accompany uneven output-space gradients, parameter-space gradients, and updates; RMSE keeps output-space gradient magnitudes constant across stages.Figure 2 averages the statistics over optimization steps; parameter-space gradients and updates cover trainable parameters only.Source: §4.2 · p. 7, Figure 2 · p. 7
  • Across learning-based PTQ methods, RMSE improves performance and balances gradients and parameter updates across stages, unlike MSE's uneven pattern.The methods use substantially different quantization designs, while sharing the MSE-based feature reconstruction objective.Source: §4.2 · p. 8, Table 2 · p. 8, Figure 3 · p. 8

4.3 ABLATION STUDIES

  • All four RMSE variants outperform MSE, while LLaMA models generally favor RMSE-Token and Qwen models generally favor RMSE-Channel.The authors use RMSE-Token for LLaMA models and RMSE-Channel for Qwen models in the main experiments.Source: §4.3 · pp. 8–9, Table 3 · p. 8
  • RMSE improves mathematical-reasoning and coding performance overall, but alone cannot fully recover severe degradation under W2A16G128.The evaluation uses OmniQuant models without task-specific training or adaptation.Source: §4.3 · p. 9, Table 4 · p. 9
  • RMSE maintains lower perplexity than MSE as quantization progresses, with the gap generally widening across stages.Figure 4 covers OmniQuant with W4A4 across diverse LLMs and representative PTQ methods on Qwen3.5-9B.Source: §4.3 · p. 9, Figure 4 · p. 9
  • RMSE adds negligible runtime and memory overhead as a drop-in replacement for MSE.The efficiency comparison measures the full sequential quantization process on one NVIDIA L40 GPU with batch size 1.Source: §4.3 · p. 9, Table 5 · p. 9

5 DISCUSSION

The discussion frames optimization behavior as a design consideration alongside quantization design, emphasizing its systematic connection to MSE and sequential quantization.

5 DISCUSSION

  • MSE's loss-dependent gradient scaling couples optimization strength to reconstruction loss magnitude, making Optimization Imbalance systematic rather than incidental.The discussion says its impact becomes more pronounced as reconstruction becomes more difficult.Source: §5 · p. 10
  • Under sequential quantization, insufficient optimization at an early stage may influence subsequent stages because previously quantized units remain fixed.This makes the stage-wise setting a consequential context for optimization imbalance.Source: §5 · p. 10
  • Different feature-grouping strategies behave distinctly, making feature grouping a design dimension for balanced optimization.The discussion points to a broader design space beyond the RMSE variants studied.Source: §5 · p. 10

6 CONCLUSION

The conclusion identifies MSE-linked Optimization Imbalance in sequential learning-based PTQ and proposes RMSE losses as drop-in alternatives. Broad experiments report improvements across tasks and settings, especially under more aggressive quantization.

6 CONCLUSION

  • The paper identifies Optimization Imbalance as stage-wise optimization strength being coupled to reconstruction loss scale under MSE.It proposes RMSE loss functions to decouple those quantities and enable more balanced optimization.Source: §6 · p. 10
  • RMSE losses are drop-in replacements for MSE and naturally decouple optimization strength from reconstruction loss scale.Source: §6 · p. 10
  • Experiments report consistent quantization-performance improvements across tasks and settings, with gains increasing under more aggressive quantization.The reported scope spans PTQ methods, model families, scales, architectures, and quantization settings.Source: §6 · p. 10

AI USE STATEMENT

The authors report that generative AI tools were used only to proofread the manuscript.

AI USE STATEMENT

  • Authors say they reviewed and verified all AI-assisted content and take full responsibility for the work.Source: AI USE STATEMENT · p. 10

A.1 EXPERIMENTAL DETAILS

The experiments follow each PTQ method’s official implementation, and the section defines four RMSE variants by the feature groups over which reconstruction error is normalized.

A.1 EXPERIMENTAL DETAILS

  • All methods use 128 WikiText2 calibration samples of length 2048, with identical samples for each MSE/RMSE comparison.RMSE replaces MSE while other method components remain unchanged; larger-model inference uses two L40 GPUs, while runtime and peak memory are measured during reconstruction on one L40 GPU with batch size 1.Source: §A.1 · p. 14
  • Reasoning tasks use the same quantized models without task-specific training or adaptation.GSM8K and MATH-500 are evaluated with LM-Evaluation-Harness; HumanEval and MBPP use EvalPlus.Source: §A.1 · p. 14

A.2 DEFINITIONS OF RMSE VARIANTS

  • The four RMSE variants normalize reconstruction error over sample, channel, token, or element groups, with each variant averaging its group-wise RMSEs.Grouping is defined within each sample; RMSE-Element is equivalent to mean absolute reconstruction error, and the variants differ only in feature grouping. Their shared gradient normalization depends on group size rather than local reconstruction loss magnitude.Source: §A.2 · pp. 14–15

B FULL DOWNSTREAM EVALUATION RESULTS

This section presents full per-task downstream results for OmniQuant and multiple learning-based PTQ methods, complementing the main paper’s average-accuracy reporting.

B FULL DOWNSTREAM EVALUATION RESULTS

  • Table 6 reports full downstream results for OmniQuant across LLMs and quantization settings, with Avg covering six commonsense tasks.The tasks are PIQA, ARC-Easy, ARC-Challenge, Winogrande, HellaSwag, and BoolQ.Source: Table 6 · p. 15
  • Table 7 reports full downstream results across representative learning-based PTQ methods under W4A4 quantization, with Avg covering the same six tasks.The tasks are PIQA, ARC-Easy, ARC-Challenge, Winogrande, HellaSwag, and BoolQ.Source: Table 7 · p. 16
  • The per-task tables give a more detailed view of MSE and RMSE behavior across models, quantization settings, and PTQ methods.The main paper reports average downstream accuracy for compactness.Source: B FULL DOWNSTREAM EVALUATION RESULTS · p. 16

C ADDITIONAL RESULTS OF RMSE VARIANTS

Across the evaluated configurations, all four RMSE variants outperform MSE, while the favored variant differs by model family and is not necessarily optimal for each configuration.

C ADDITIONAL RESULTS OF RMSE VARIANTS

D ADDITIONAL ANALYSIS OF OPTIMIZATION DYNAMICS

Additional analyses compare optimization statistics across stages, examine high-precision block-output magnitudes, and measure reconstruction-loss variation within stages.

D ADDITIONAL ANALYSIS OF OPTIMIZATION DYNAMICS

E OUTPUT MAGNITUDES ACROSS BLOCKS

  • High-precision block-output magnitudes vary substantially and generally increase with depth across the evaluated models, providing a natural source of broad reconstruction-loss scale.The statistics are collected from unquantized models using 128 WikiText2 calibration samples.Source: E OUTPUT MAGNITUDES ACROSS BLOCKS · p. 20, Figure 8 · p. 20

F RECONSTRUCTION LOSS MAGNITUDE VARIATION WITHIN STAGES

  • Within OmniQuant units trained with MSE, group-wise reconstruction losses vary substantially across samples and become more dispersed under channel, token, and element groupings.The analysis computes group-wise MSE between high-precision and quantized unit outputs using the four feature groupings.Source: F RECONSTRUCTION LOSS MAGNITUDE… · p. 20, Figure 9 · p. 20
Loading 2610.00983v1…