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

arXiv: https://arxiv.org/abs/2510.18784 · HTML: https://mlsanity.com/paper/2510.18784/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 46 references that are in the map.
- First posted on arXiv 21 Oct 2025.

## Builds on (most cited first, 30 of 31)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Learned Step Size Quantization](https://mlsanity.com/paper/1902.08153/) (2019)
- [Quantizing deep convolutional networks for efficient inference: A whitepaper](https://mlsanity.com/paper/1806.08342/) (2018)
- [2 OLMo 2 Furious](https://mlsanity.com/paper/2501.00656/) (2024)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [LSQ+: Improving low-bit quantization through learnable offsets and better initialization](https://mlsanity.com/paper/2004.09576/) (2020)
- [Dynamic Model Pruning with Feedback](https://mlsanity.com/paper/2006.07253/) (2020)
- [Overcoming Oscillations in Quantization-Aware Training](https://mlsanity.com/paper/2203.11086/) (2022)
- [Scaling Laws for Precision](https://mlsanity.com/paper/2411.04330/) (2024)
- [Network Quantization with Element-wise Gradient Scaling](https://mlsanity.com/paper/2104.00903/) (2021)
- [Training Quantized Nets: A Deeper Understanding](https://mlsanity.com/paper/1706.02379/) (2017)
- [ProxQuant: Quantized Neural Networks via Proximal Operators](https://mlsanity.com/paper/1810.00861/) (2018)
- [Differentiable Model Compression via Pseudo Quantization Noise](https://mlsanity.com/paper/2104.09987/) (2021)
- [Bridging Discrete and Backpropagation: Straight-Through and Beyond](https://mlsanity.com/paper/2304.08612/) (2023)
- [Quantized Adam with Error Feedback](https://mlsanity.com/paper/2004.14180/) (2020)
- [QuEST: Stable Training of LLMs with 1-Bit Weights and Activations](https://mlsanity.com/paper/2502.05003/) (2025)
- [Estimator Meets Equilibrium Perspective: A Rectified Straight Through Estimator for Binary Neural Networks Training](https://mlsanity.com/paper/2308.06689/) (2023)
- [Mirror Descent View for Neural Network Quantization](https://mlsanity.com/paper/1910.08237/) (2019)
- [AdaSTE: An Adaptive Straight-Through Estimator to Train Binary Neural Networks](https://mlsanity.com/paper/2112.02880/) (2021)
- [Proximal Mean-field for Neural Network Quantization](https://mlsanity.com/paper/1812.04353/) (2018)
- [LOTION: Smoothing the Optimization Landscape for Quantized Training](https://mlsanity.com/paper/2510.08757/) (2025)

## Built on by (most cited first, 4 of 4)

- [WinQ: Accelerating Quantization-Aware Training of Language Models Around Saddle Points](https://mlsanity.com/paper/2605.17471/) (2026)
- [Zero-Shot Quantization via Weight-Space Arithmetic](https://mlsanity.com/paper/2604.03420/) (2026)
- [HESTIA: A Hessian-Guided Differentiable Quantization-Aware Training Framework for Extremely Low-Bit LLMs](https://mlsanity.com/paper/2601.20745/) (2026)
- [TinyNeRV: Compact Neural Video Representations via Capacity Scaling, Distillation, and Low-Precision Inference](https://mlsanity.com/paper/2604.09220/) (2026)

## Abstract (arXiv)

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 .