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

TL;DR

CAGE introduces a curvature-aware gradient correction for quantization-aware training, significantly reducing accuracy loss and matching higher-bit performance at lower bit-widths, with strong theoretical guarantees and efficient implementation.

Contribution

It proposes a novel curvature-aware gradient estimator for QAT, providing theoretical convergence guarantees and practical improvements over prior methods.

Findings

Halves accuracy loss in QAT fine-tuning.

Achieves 3-bit W3A3 accuracy comparable to 4-bit W4A4.

Provides a theoretically grounded, optimizer-agnostic method.

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 .