Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CAGE: Curvature-Aware Gradient Estimation For Accurate Quantization-Aware Training

Soroush Tabesh, Mher Safaryan, Andrei Panferov, Alexandra Volkova, Dan Alistarh

TL;DR

CAGE tackles the remaining gap between low-bit quantization-aware training and full-precision training by injecting a curvature-aware correction into the gradient, derived from a Pareto-optimality view that balances task loss and quantization constraints. The method yields strong theoretical guarantees (ergodic convergence to a Pareto-optimal point in the smooth non-convex setting) and an optimizer-agnostic, efficient implementation built on AdamW, leveraging existing optimizer statistics. Empirically, CAGE improves QAT performance in both fine-tuning and large-scale pretraining of Llama-style models, achieving lower loss and higher effective capacity than prior baselines (e.g., QuEST), with gains across multiple optimizers and precisions (W2A2, W3A3, W4A4). This work suggests that incorporating local curvature information and explicit quantization-error regularization can substantially close the gap to full-precision models, enabling more efficient training and deployment of large language models.

Abstract

Despite significant work on low-bit quantization-aware training (QAT), there is still an accuracy gap between such techniques and native training. To address this, we introduce CAGE (Curvature-Aware Gradient Estimation), a new QAT method that augments the straight-through estimator (STE) gradient with a curvature-aware correction designed to counteract the loss increase induced by quantization. CAGE is derived from a multi-objective view of QAT that balances loss minimization with the quantization constraints, yielding a principled correction term that depends on local curvature information. On the theoretical side, we introduce the notion of Pareto-optimal solutions for quantized optimization, and establish that CAGE yields strong convergence guarantees in the smooth non-convex setting. In terms of implementation, our approach is optimizer-agnostic, but we provide a highly-efficient implementation that leverages Adam statistics. CAGE significantly improves upon the prior state-of-the-art methods in terms of accuracy, for similar computational cost: for QAT fine-tuning, it halves the compression accuracy loss relative to the prior best method, while for QAT pre-training of Llama models, its accuracy for 3-bit weights-and-activations (W3A3) matches the accuracy achieved at 4-bits (W4A4) with the prior best method. The official implementation can be found over https://github.com/IST-DASLab/CAGE .

CAGE: Curvature-Aware Gradient Estimation For Accurate Quantization-Aware Training

TL;DR

CAGE tackles the remaining gap between low-bit quantization-aware training and full-precision training by injecting a curvature-aware correction into the gradient, derived from a Pareto-optimality view that balances task loss and quantization constraints. The method yields strong theoretical guarantees (ergodic convergence to a Pareto-optimal point in the smooth non-convex setting) and an optimizer-agnostic, efficient implementation built on AdamW, leveraging existing optimizer statistics. Empirically, CAGE improves QAT performance in both fine-tuning and large-scale pretraining of Llama-style models, achieving lower loss and higher effective capacity than prior baselines (e.g., QuEST), with gains across multiple optimizers and precisions (W2A2, W3A3, W4A4). This work suggests that incorporating local curvature information and explicit quantization-error regularization can substantially close the gap to full-precision models, enabling more efficient training and deployment of large language models.

Abstract

Despite significant work on low-bit quantization-aware training (QAT), there is still an accuracy gap between such techniques and native training. To address this, we introduce CAGE (Curvature-Aware Gradient Estimation), a new QAT method that augments the straight-through estimator (STE) gradient with a curvature-aware correction designed to counteract the loss increase induced by quantization. CAGE is derived from a multi-objective view of QAT that balances loss minimization with the quantization constraints, yielding a principled correction term that depends on local curvature information. On the theoretical side, we introduce the notion of Pareto-optimal solutions for quantized optimization, and establish that CAGE yields strong convergence guarantees in the smooth non-convex setting. In terms of implementation, our approach is optimizer-agnostic, but we provide a highly-efficient implementation that leverages Adam statistics. CAGE significantly improves upon the prior state-of-the-art methods in terms of accuracy, for similar computational cost: for QAT fine-tuning, it halves the compression accuracy loss relative to the prior best method, while for QAT pre-training of Llama models, its accuracy for 3-bit weights-and-activations (W3A3) matches the accuracy achieved at 4-bits (W4A4) with the prior best method. The official implementation can be found over https://github.com/IST-DASLab/CAGE .
Paper Structure (27 sections, 1 theorem, 21 equations, 9 figures, 2 tables, 1 algorithm)

This paper contains 27 sections, 1 theorem, 21 equations, 9 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Prior Work
  4. CAGE: Curvature-Aware Gradient Estimation
  5. Problem Description and Motivation
  6. Error coupling.
  7. Relationship to LOTION.
  8. Theoretical Analysis
  9. Discussion.
  10. Practical Implementation
  11. Decoupled update.
  12. Warmup schedule.
  13. Experiments
  14. Quantization pipeline.
  15. Synthetic Loss Experiments
  16. ...and 12 more sections

Key Result

Theorem 1

Let Assumptions asm:smooth, asm:unbiased, asm:quant-field hold and $L = L_f + \lambda L_{\phi}$. Then, for any $\lambda\ge0$, the iterates (iter-1) with step-size $\alpha = \min(\frac{1}{L}, \frac{1}{\sqrt{T}})$ satisfy where $\hat{x}$ is a random iterate drawn from the history $\{x_0,x_1,\dots,x_{T-1}\}$ uniformly at random and

Figures (9)

  • Figure 1: Quadratic with non-isotropic Hessian: final validation loss for all methods various condition numbers (lower is better). The loss for CAGE does not include the regularizer term. We find that CAGE consistently reduces stationary error.
  • Figure 2: Representative trajectories (projected to the top two principal components PC1 and PC2) for the quadratic task; CAGE follows curvature and converges closest to the optimum.
  • Figure 3: QAT accuracy on Llama-3.2-3B (Tulu-SFT) for CAGE vs. the state-of-the-art MXFP4 baseline, using QuEST panferov2025quest. CAGE consistently improves GSM8K/HellaSwag/WinoGrande results, essentially halving the error due to quantization. Hadamard transform is applied in both cases for outlier mitigation.
  • Figure 4: Validation loss during QAT (W4A4) for AdamW, Shampoo, SOAP, and Muon, for standard QAT and CAGE. Solid lines denote training with CAGE, dashed -- without. Across all optimizers, the method yields lower validation loss.
  • Figure 5: Validation loss versus model size (bytes) for weight and activation quantization using W2A2, W3A3 and W4A4, comparing CAGE+QuEST to baseline QuEST and BF16. Observe that CAGE yields consistently lower loss.
  • ...and 4 more figures

Theorems & Definitions (1)

  • Theorem 1