PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Outliers: A Study of Optimizers Under Quantization

Georgios Vlassis, Saleh Ashkboos, Alexandra Volkova, Torsten Hoefler, Dan Alistarh

arXiv:2509.23500v2cs.LG

TL;DR

The impact of optimizer choice on LLM performance under PTQ and QAT remains unclear, including whether full-precision optimizer advantages transfer to quantized training. The study trains 50M–1.5B-parameter models with six optimizers, evaluates full-precision training, PTQ, and QAT, and derives optimizer-specific QAT scaling laws. MMR and kurtosis fail to predict PTQ accuracy across optimizers, while Shampoo shows the lowest QAT accuracy degradation and highest parameter efficiency among tested optimizers.

  • Problem

    The impact of optimizer choice on LLM performance under PTQ and QAT remains unclear, including whether full-precision optimizer advantages transfer to quantized training.

  • Method

    The study trains 50M–1.5B-parameter models with six optimizers, evaluates full-precision training, PTQ, and QAT, and derives optimizer-specific QAT scaling laws.

  • Results

    MMR and kurtosis fail to predict PTQ accuracy across optimizers, while Shampoo shows the lowest QAT accuracy degradation and highest parameter efficiency among tested optimizers.

  • Takeaways & Limitations

    Optimizer rankings in full precision need not persist under quantization, making optimizer choice relevant to quantized-model robustness and scaling efficiency.

  • Takeaways & Limitations

    The study leaves additional bitwidths, alternative PTQ schemes, other 4-bit data types, and error propagation in modules such as self-attention for future work.

Abstract

from arXiv · show

As new optimizers gain traction and model quantization becomes standard for efficient deployment, a key question arises: how does the choice of optimizer affect model performance in the presence of quantization? Despite progress in both areas, systematic evidence on optimizer-quantization interactions remains limited. To fill this gap, we study the impact of optimizer choice on model robustness under quantization, considering both post-training quantization (PTQ), and quantization-aware training (QAT). We first train full-precision models, ranging from 50M to 1.5B parameters, with six optimizers, to explore the hyperparameter landscape, and establish well-tuned baselines. We then apply PTQ to evaluate how model performance degrades when trained with different optimizers. We find that outlier-related metrics, such as the max-to-mean ratio (MMR) and Kurtosis, fail to predict the PTQ performance across different optimizers. We show analytically that this is due to the MMR capturing only isolated layer errors, while ignoring how quantization errors accumulate and propagate through the network. To study the QAT degradation, we train quantized models from scratch and compare them to our original-precision baselines. We find that optimizers performing well in the original pretraining setup may not remain optimal under QAT, and that models trained with Shampoo show the lowest accuracy degradation. Finally, we derive scaling laws for quantization-aware training under different optimizers, showing that Shampoo achieves the highest parameter efficiency of all tested optimizers.

1 INTRODUCTION

This paper studies how optimizer choice affects LLM accuracy under PTQ and QAT, addressing limited evidence about optimization–quantization interactions. Across systematic experiments, outlier metrics fail to predict PTQ robustness, while Shampoo is especially robust under QAT and achieves the highest parameter efficiency.

  • The study trains models from 50M to 1.5B parameters with six optimizers and evaluates full-precision, PTQ, and QAT settings.It frames PTQ around equal validation loss and QAT around whether full-precision optimizer advantages transfer under quantization.
  • Muon outperforms the other optimizers across nearly all model sizes in full-precision training.
  • Full-precision accuracy and standard outlier metrics such as MMR and kurtosis do not reliably predict PTQ accuracy across optimizers.For the 760M model, MMR and kurtosis show little to no correlation with accuracy, whereas the proposed metric correlates strongly with zero-shot performance.
  • The paper analyzes quantization-error propagation to explain why isolated outlier statistics can miss errors that accumulate through the network.
  • Shampoo yields the lowest QAT accuracy degradation across almost all model sizes and the highest parameter efficiency under 4-bit QAT.These results indicate that the optimizer best in full precision need not remain best after quantization.

2 BACKGROUND AND EXPERIMENTAL SETUP

The experiments use OLMo2-based models spanning 50M to 1.5B parameters, six optimizers, tuned full-precision baselines, and common evaluation procedures. PTQ uses 4-bit symmetric row-wise quantization, while QAT uses QuEST with quantized forward passes and mostly full-precision backward computation.

  • Model Architecture: Models use OLMo2 with 50M, 125M, 350M, 500M, 760M, and 1.5B parameters, trained on ClimbMix under the Chinchilla regime.
  • Optimizers: The study compares AdamW, PSGD, Shampoo, Muon, Scion, and SOAP, covering baseline, curvature-aware, feature-learning, and gradient-rotation approaches.
  • Experimental Setup and Hyper-parameter Protocol: Each optimizer is tuned through sequential hyperparameter sweeps and eight learning-rate values, with settings selected on smaller models and tested at larger scale.
  • Datasets, Training Steps, and Evaluation Taks: PTQ checkpoints use the common loss across optimizers, while quantized training uses a fixed 20x token-to-parameter ratio and averages accuracy over PIQA, HellaSwag, and ARC-Easy.
  • Quantization: PTQ uses 4-bit symmetric AbsMax row-wise quantization for weights and inputs of linear layers, while other modules remain in original precision.
  • Quantization: QAT uses QuEST, applying Hadamard transforms and MSE-based clipping during the forward pass while masking extreme gradients during backward computation.

3 RESULTS AND ANALYSIS

Across full-precision, PTQ, and QAT experiments, optimizer choice strongly affects quantization robustness, while standard outlier metrics do not reliably predict performance. Error propagation analysis and scaling laws identify Shampoo as especially robust under quantization.

  • 3.1 FULL-PRECISION: Muon achieves the highest full-precision zero-shot accuracy for nearly all model sizes, with its advantage reaching 1.03% at 1.5B parameters.The comparison uses optimized learning rates and excludes the 50M and 125M models from Muon’s consistent advantage.
  • 3.1 FULL-PRECISION: Increasing the learning rate consistently raises MMR across optimizers, while Muon has the lowest MMR in the 760M analysis.MMR is averaged over rows in the input tensor of the last linear layer before the head.
  • 3.2 POST-TRAINING QUANTIZATION: Shampoo consistently minimizes PTQ degradation above 125M parameters, despite having the highest MMR, while low-MMR Muon degrades substantially.PTQ comparisons use models matched at a common validation loss, isolating degradation caused by quantization.
  • 3.2 POST-TRAINING QUANTIZATION: The ABC analysis shows total quantization error is usually dominated by amplified error from preceding modules, not isolated layer error measured by MMR.Across depth, error generally increases with oscillatory behavior, but optimizer-specific profiles range from late spikes to flatter propagation.
  • 3.2 POST-TRAINING QUANTIZATION: Muon has the highest linear-layer gain, potentially explaining its sharp quality degradation despite low MMR, whereas AdamW and Shampoo have the lowest gains.For linear layers, the spectral ratio is near one, so gain is largely determined by the alignment ratio.
  • 3.3 QUANTIZATION-AWARE TRAINING AND SCALING LAW: Under QAT, optimizer rankings differ from full-precision training across nearly all model sizes, and no single optimizer is best for QAT accuracy.Shampoo nevertheless remains the most effective at minimizing accuracy degradation relative to full-precision baselines.
  • 3.3 QUANTIZATION-AWARE TRAINING AND SCALING LAW: Shampoo achieves the highest ρ4bit, indicating the greatest parameter efficiency and resilience to quantization among the evaluated optimizers.The scaling-law parameter-efficiency measure represents the effective parameter capacity retained by a quantized model.

4 CONCLUSION AND FUTURE WORK

The paper systematically studies how optimizer choice affects accuracy under PTQ and QAT, finding that high-precision winners may differ from quantized-setting winners. It also identifies extensions involving additional bitwidths, PTQ schemes, data types, and network modules.

  • The study shows that optimizer choice influences accuracy drop during quantization, with high-precision winners not necessarily remaining optimal under QAT.The analysis covers full-precision baselines, PTQ, QAT, and optimizer-specific QAT scaling laws.
  • Future work includes additional QAT bitwidths, alternative PTQ schemes, other 4-bit data types, and error-propagation analysis for self-attention and other modules.

A.1 ADDITIONAL FULL-PRECISION RESULTS

Additional full-precision results compare average zero-shot accuracy and test loss across model sizes, showing SOAP’s advantage at small sizes and Muon’s advantage at larger sizes.

  • Muon is consistently best for larger full-precision models, while SOAP performs better at small sizes.The table reports average zero-shot accuracy and test loss, which track each other closely.

A.2 ADDITIONAL POST-TRAINING QUANTIZATION RESULTS

Additional PTQ results largely confirm that Shampoo has the mildest degradation after row-wise W4A4 quantization, with a 500M-model exception.

  • Shampoo generally shows the mildest test-loss degradation after row-wise W4A4 PTQ, except at 500M where its degradation is noticeable.Despite the 500M exception in test loss, Shampoo still outperforms other optimizers in downstream tasks.

A.3 ADDITIONAL QUANTIZATION-AWARE TRAINING RESULTS

Additional QAT results show that Shampoo has the smallest degradation relative to full-precision baselines.

  • Shampoo shows the smallest degradation relative to full-precision baselines in W4A4 QAT.

A.4 ABC DECOMPOSITION

The ABC decomposition separates quantization-induced activation changes into accumulated prior error, newly introduced layer error, and their interaction, then analyzes how modules propagate these errors. For linear layers under joint quantization, the propagation gain factors into spectral and alignment terms.

  • Scope and limitation: The framework applies to arbitrary modules and flexible partitionings, including linear layers, convolutions, activations, normalization, attention, CNNs, and Transformers.The authors leave analogous analyses for layer normalization and self-attention to future work.
  • Error decomposition: Quantization changes each layer’s output through both propagated activation error and newly introduced function perturbation.The framework treats these as changes in activation space and function space, respectively.
  • Attribution: Averaging both perturbation orderings follows the Shapley principle and avoids attribution bias between input and function changes.The two valid orderings perturb the function first or the input first.
  • ABC decomposition: The squared relative activation change decomposes exactly into accumulated prior error, new layer error, and their interaction.The terms are denoted A_l, B_l, and C_l, respectively.
  • Error propagation: The gain G_l measures how the l-th module propagates quantization error from the preceding layer into its accumulated-error term.It relates the previous-layer relative error R_l−1 to the propagated component A_l.
  • Linear-layer analysis: For linear layers under joint quantization, G_l factors into a spectral ratio and an alignment ratio.The analysis covers weight and activation quantization and omits B_l and C_l because R_l is dominated by A_l.

A.5 OPTIMIZERS’ MEMORY AND COMPUTATIONAL COMPLEXITIES

This section compares optimizer-step memory and computational complexity for linear layers, including the costs of Newton–Schulz iterations and preconditioner updates.

  • Complexity comparison: Table 8 reports asymptotic memory and per-step computational complexity for optimizers applied to an m×n linear layer.The comparison includes gradients, moments, preconditioners, eigenvector matrices, Newton–Schulz iterations, and update frequency.

1. AdamW

The supplied passages describe optimizer memory and computation costs, scaling-law fitting, and reported training-resource trade-offs rather than AdamW-specific results.

  • Newton–Schulz methods: Newton–Schulz-based updates require O(T(2m^2n + m^3)) time over T iterations.The stated one-step complexity is O(2m^2n + m^3).
  • Optimizer costs: A gradient and first-moment implementation uses 2mn memory, while a two-sided preconditioner uses mn + m^2 + n^2 memory.Applying the preconditioners costs 2m^2n + 2mn^2 operations, with updates adding O(m^3 + n^3).
  • Scaling-law fitting: Scaling-law parameters are fitted by nonlinear least squares on logarithmic loss values using Huber loss and leave-one-out confidence intervals.The fit uses model size, dataset size, and precision, and includes embedding parameters in total parameter count.
Loading 2509.23500v2…