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

作者: Soroush Tabesh, Mher Safaryan, Andrei Panferov, Alexandra Volkova, Dan Alistarh

A1 主要贡献

本文旨在解决低比特量化感知训练（QAT）与全精度训练之间的准确性差距问题。现有的QAT方法，尤其是直通估计器（STE），存在收敛缓慢和不稳定的问题，且缺乏形式化的收敛保证。

核心问题与研究目标:
核心问题是标准的QAT方法（如STE）存在性能瓶颈，导致量化后模型精度下降。研究目标是开发一种新的QAT方法，该方法不仅在实践中能提高模型准确性，而且具备坚实的理论基础和收敛保证。

创新点与主要贡献:
本文提出了一种名为CAGE（Curvature-Aware Gradient Estimation，曲率感知梯度估计）的新型QAT方法，其主要贡献如下：
* 提出CAGE方法: 该方法通过一个曲率感知的校正项来增强标准的STE梯度。这个校正项是基于将QAT视为一个多目标优化问题（同时最小化任务损失和量化误差）推导出的帕累托最优性条件，旨在明确地抵消由量化步骤引起的损失增加。它利用了损失景观的局部二阶信息（即Hessian矩阵）。
* 提供理论保证: CAGE在平滑非凸设定下具有强遍历性收敛到帕累托最优点（Pareto-optimal point）的保证，解决了现有工作缺乏理论保证的主要限制。
* 高效且优化器无关的实现: 该方法与优化器无关，但本文提供了一个利用自适应优化器（如Adam）中现有统计量的高效实现。
* 全面的实验验证:
* 在QAT微调任务中，CAGE将相对于先前最佳方法的压缩精度损失减半。
* 在Llama模型的QAT预训练中，使用CAGE进行3比特权重和激活（W3A3）训练达到的准确性，与使用先前最佳方法进行4比特（W4A4）训练的准确性相匹配。
* 实验证明CAGE在多种优化器（包括AdamW, Muon, Shampoo）上均能带来一致的性能提升。

总结而言，CAGE通过引入曲率感知的梯度校正，为缩小低比特量化模型与全精度模型之间的性能差距提供了一个有前景的方案。

A2 先前工作

QAT的核心挑战与STE的局限性。量化感知训练（QAT）的核心挑战在于量化函数不可微，其梯度几乎处处为零。标准方法是直通估计器（STE），最初由Hinton (2012)提出并由Bengio等人(2013b)形式化，它在反向传播过程中绕过了量化算子。然而，STE常导致训练不稳定和次优收敛【Yin et al., Understanding straight-through estimator in training activation quantized neural nets, 2019, ICLR】。

STE的改进方法。因此，大量研究集中于通过可学习的量化参数【Esser et al., Learned step size quantization, 2020, ICLR】、提高梯度本身的准确性【Le et al., AdaSTE: An Adaptive Straight-Through Estimator to Train Binary Neural Networks, 2021, arXiv】或结合近端映射进行损失正则化【Bai et al., Proxquant: Quantized neural networks via proximal operators, 2018, arXiv】来改进STE。例如，AdaSTE【Le et al., AdaSTE: An Adaptive Straight-Through Estimator to Train Binary Neural Networks, 2021, arXiv】从双层优化的角度增强梯度估计。ReSTE【Wu et al., Estimator meets equilibrium perspective: A rectified straight through estimator for binary neural networks training, 2023, arXiv】通过一个修正的估计器来平衡估计误差和梯度稳定性。Liu等人(2023)【Liu et al., Bridging discrete and backpropagation: Straight-through and beyond, 2023, NeurIPS】提出Reinmax，注意到STE是梯度的一阶近似，并提出通过Heun方法构建一个二阶估计器。

基于正则化和近端方法的QAT。此外，ProxQuant【Bai et al., Proxquant: Quantized neural networks via proximal operators, 2018, arXiv】将QAT重新表述为一个通过近端梯度方法解决的正则化学习问题。ProxQuant在随机梯度步骤之间应用一个prox算子来鼓励量化，而不是仅仅依赖STE。与ProxQuant类似，PARQ【Jin et al., Parq: Piecewise-affine regularized quantization, 2025, arXiv】使用凸和分段仿射正则化来引导权重量子化，不同之处在于梯度更新在应用近端算子之前被添加到全精度参数上。

其他QAT理论视角。镜像下降（Mirror Descent）公式将全精度参数解释为量化参数的对偶，并分析了量化算子生成有效镜像映射的情况。特别是，如果我们将softmax投影作为量化算子，则STE的对偶迭代遵循指数梯度下降（EGD）【Ajanthan et al., Mirror Descent View for Neural Network Quantization, 2021, arXiv】。马尔可夫随机场（MRF）和提升概率空间公式【Ajanthan et al., Proximal mean-field for neural network quantization, 2019, arXiv】将模型的每个参数表示为量化值的加权平均，并最小化关于这些概率权重的损失。

基于梯度缩放的QAT方法。Lee等人(2021a)【Lee et al., Network quantization with element-wise gradient scaling, 2021a, arXiv】提出了逐元素梯度缩放（EWGS）方法，其中每个梯度元素的缩放取决于其符号和相应的量化误差。此外，这些缩放通过从近似Hessian信息中获得的全局缩放因子进行调整。一个相关的方法采用基于位置的缩放梯度，其中梯度估计器仅通过将量化误差乘以一个小的正裕量进行调整【Kim et al., Position-based scaled gradient for model quantization and pruning, 2020, arXiv】。

解决STE中的振荡问题。在基于STE的训练中观察到的另一个现象是权重在相邻量化网格点之间的决策边界附近振荡【Defossez et al., Differentiable model compression via pseudo quantization noise, 2022, arXiv】。为了减轻这些振荡并稳定训练，Nagel等人(2022a)【Nagel et al., Overcoming oscillations in quantizationaware training, 2022a, arXiv】在损失之上引入了一个正则化项来抑制振荡，或者提出了检测振荡权重并在训练期间冻结它们的机制。

CAGE与先前工作的对比。与先前工作相比，CAGE为QAT提供了一个多目标优化的视角。CAGE不是精炼一阶近似或应用标准正则化，而是向梯度中引入一个曲率感知的校正项。

A3 方法细节

3.1 问题描述与动机

QAT的约束优化视角。直观上，量化感知训练（QAT）的关键问题是构建一个梯度估计器，将参数更新到一个“最优”的量化模型，即一个最小化损失的受约束模型。更精确地说，QAT背后的优化问题可以描述为以下约束优化：

$$\min_{x \in \mathcal{Q}} f(x),$$

其中$Q = \{x \in \mathbb{R}^d : x = Q(x)\}$是量化参数的约束集，而$Q : \mathbb{R}^d \to \mathbb{R}^d$是量化算子。看待同一问题的另一种方式是考虑其无约束的重构形式，并相对于$x \in \mathbb{R}^d$最小化目标$f(Q(x))$【Li et al., Training Quantized Nets: A Deeper Understanding, 2017, arXiv】。这里的关键问题是量化算子$Q$的不可微性。因此，Bengio等人(2013a)【Bengio et al., Estimating or propagating gradients through stochastic neurons for conditional computation, 2013a, arXiv】的直通估计器（STE）是近似$f(Q(x))$梯度的默认选择。具体来说，$\nabla_x[f(Q(x))] = J_Q(x)^\top \cdot \nabla f(Q(x))$被近似为$\nabla f(Q(x))$。这实际上是用单位矩阵替换了量化操作的雅可比矩阵【Li et al., Training Quantized Nets: A Deeper Understanding, 2017, arXiv; Hou et al., Analysis of Quantized Models, 2019, ICLR; Chen et al., Quantized Adam with Error Feedback, 2021a, arXiv】。

误差反馈重构。要更深入地了解STE的动态，一种方法是通过误差反馈的视角来审视它。可以很容易地观察到，在迭代$x_t$上使用STE的基于SGD的训练，等价于一个误差反馈的实例【Chen et al., Quantized adam with error feedback, 2021b, https://arxiv.org/abs/2004.14180】，其中量化参数 为$w_t = Q(x_t)$，量化误差为$e_t = x_t - Q(x_t)$。这种等价性在迭代$t \ge 0$中描述如下：

$$\begin{aligned} x_{t+1}=x_{t}-\alpha \widetilde{\nabla} f\left(Q\left(x_{t}\right)\right) \Longleftrightarrow\left\{\begin{array}{l} g_{t}=\alpha \widetilde{\nabla} f\left(w_{t}\right)-e_{t} \\ w_{t+1}=Q\left(w_{t}-g_{t}\right) \\ e_{t+1}=\left(w_{t}-g_{t}\right)-w_{t+1} \end{array}\right. \end{aligned}$$

实际上，右侧的公式表明，STE等价于一个过程，其中在给定步骤中由于权重 quantization 导致的量化误差$e_{t+1}$（即$(w_t - g_t) - w_{t+1}$）被反馈到下一步的梯度中。如果损失函数$f$是一个各向同性的二次函数（即具有单位Hessian矩阵），这个反馈循环将是“精确的”，能够完全校正误差，但在一般情况下，这不是正确的更新方式。

多目标视角。先前关于在非凸设定下公式(1)所述问题的收敛保证的工作，试图找到一个量化点$Q(x^*)$，该点是损失函数$f$的驻点，即$\nabla f(Q(x^*)) = 0$。然而，由于量化不是一个可逆变换，这个目标通常无法实现，并且当前的收敛保证在其速率中具有与量化误差成比例的非消失项【Lin et al., Dynamic model pruning with feedback, 2020, https://arxiv.org/abs/2006.07253; Chen et al., Quantized adam with error feedback, 2021b, https://arxiv.org/abs/2004.14180; Panferov et al., Unified scaling laws for compressed representations, 2025b, https://arxiv.org/abs/2506.01863】 。

多目标视角下的一个实例。为了说明这一点，考虑一个玩具示例，其中我们的标量损失函数是$f(x) = \frac{1}{2}(x - \frac{1}{2})^2$，量化算子$Q(x) = \lfloor x \rfloor$给出一维输入$x \in \mathbb{R}$的整数部分。显然，$\nabla f(Q(x)) = \lfloor x \rfloor - \frac{1}{2}$在任何点都不会消失，其最小绝对值为$\frac{1}{2}$，这来自于量化/舍入误差。

帕累托最优解的定义。相反，我们建议从多目标优化的角度来看待公式(1)中的QAT问题。第一个目标是最小化损失$f(x)$：在非凸设置中，这变成了找到一个驻点，即求解$\nabla f(x) = 0$。第二个目标是满足量化约束$x \in Q$。换句话说，我们的目标是找到一个解$x^*$，使得$\nabla f(x^*) = 0$并且$x^*$与$Q(x^*)$之间的距离最小化。如前所述，很容易看出这两个条件在一般情况下不必同时成立，这意味着对于任何量化点$x \in Q$，$\nabla f(x)$都可能非零。因此，我们定义解为平衡这两个条件的帕累托最优点。

帕累托最优状态的数学表达。设$x \in \mathbb{R}^d$是参数的当前状态。如果我们以足够小的步长朝着$-\nabla f(x)$（或任何与其对齐的方向）更新，那么我们将最小化损失。显然，相反的方向会增加损失。类似地，如果我们朝着$Q(x) - x$更新，那么我们将减少量化误差。为了使这两个目标保持平衡，我们希望我们的迭代收敛到某个帕累托最优状态$x^*$，使得对于某个标量$\lambda > 0$，有$\nabla f(x^*) = \lambda(Q(x^*) - x^*)$。在这种情况下，任何施加于$x^*$的足够小的更新都会损害至少一个目标。换句话说，任何一个目标的局部改进都将以牺牲另一个目标为代价。因此，我们将问题(1)的$\lambda$-帕累托最优解$x^*$定义为：

$$\nabla_{\lambda \mathrm{P}} f\left(Q\left(x^{*}\right)\right):=\nabla f\left(x^{*}\right)+\lambda\left(x^{*}-Q\left(x^{*}\right)\right)=0,$$

其中$\lambda > 0$是平衡两个目标的参数。回想之前提到的玩具示例，对于任何$\lambda > 0$，值$x^*(\lambda) = \frac{1}{2(1+\lambda)}$是$\lambda$-帕累托最优的。请注意，对于相同的$\lambda$，可能存在许多帕累托最优解。在我们的玩具示例中，$x^* = -1/4$是另一个1-帕累托最优值。

CAGE方法的动机。受此多目标视角和公式(2)中新的最优性条件的启发，我们建议将量化误差$e_t = x_t - Q(x_t)$直接整合到训练动态中。通过这样做，我们隐式地对初始任务损失进行正则化，如果量化方案是平滑的，这种正则化可以变得显式（正如我们在理论分析中展示的）。确实，在$Q$的平滑性假设（假设3）下，量化误差可以表示为$x - Q(x) = \nabla\phi(x)$，其中$\phi(x)$是某个平滑的正则化项。值得注意的是，我们的方法自然地避免了处理量化算子的（不存在的）雅可比矩阵$J_Q(x)$，并绕过了对显式“二次”正则化项$\lambda \|x - Q(x)\|_2^2$计算近端步骤的需要，这对于我们实验中使用的动态量化是不可行的【Bai et al., Proxquant: Quantized neural networks via proximal operators, 2018, arXiv】。

误差耦合机制。注意到我们的方法是优化器无关的，我们提出了两种结合量化误差的方式，分别称为耦合校正和解耦校正。区别在于量化误差是在传递给优化器之前加到当前梯度上（耦合校正），还是加在优化器计算出的模型更新之上（解耦校正）。如果基础优化器是SGD，那么这两种校正方法是相同的。我们将为在平滑非凸设置下收敛到帕累托最优点提供理论收敛保证。然而，对于像Adam这样的有状态优化器，这两种校正会产生概念上不同的方案，我们将在下面讨论（将偏置校正项并入学习率$\alpha$）：

耦合与解耦校正在Adam中的公式。
量化误差和mini-batch梯度

优化器状态

优化器更新
$\Delta_t = m_t/(\sqrt{v_t} + \epsilon) + \lambda_t e_t$, (解耦校正)

模型更新

$$x_{t+1}=x_t-\alpha\Delta_t,$$

为了在Adam优化器的背景下比较这两种校正变体，请注意，耦合校正项$\lambda_t e_t$在经过优化器状态处理后变为$(1 - \beta_1)\lambda_t e_t / (\sqrt{v_t} + \epsilon)$。与解耦校正相比，耦合版本引入了一个对角预处理的曲率感知校正项，其中曲率是使用优化器状态内已经维护的Adam统计量来近似的。

与LOTION的关系。同期的工作【Kwun et al., Lotion: Smoothing the optimization landscape for quantized training, 2025, https://arxiv.org/abs/2510.08757】提出了一种名为LOTION的方法，该方法通过无偏随机舍入（RR）来平滑量化损失， 即$f(Q(x)) \approx \mathbb{E}_{\epsilon \sim RR}[f(x + \epsilon)]$，并考虑其二阶展开作为要最小化的代理损失，即$f(x + \epsilon) \approx f(x) + \epsilon^\top \nabla f(x) + \frac{1}{2} \epsilon^\top H(x) \epsilon$。由于舍入的无偏性（即$\mathbb{E}[\epsilon] = 0$），展开式的一阶项在期望上消失，平滑后的损失变为原始损失$f(x)$加上一个额外的曲率感知正则化项$\frac{1}{2} \text{Tr}(H(x) \text{Cov}[\epsilon])$。在这里，精确的Hessian矩阵被其Gauss-Newton近似所取代。如果我们进一步应用经验Fisher近似，并使用Adam统计量来近似Gauss-Newton矩阵的对角线元素，这个正则化项的形式将与我们的耦合校正变体相似（除了平方根）。精确计算正则化损失的梯度估计器需要对正则化项求导，这将涉及损失的三阶导数，是不切实际的。因此，作者建议跳过这一步。

CAGE与LOTION的对比。与LOTION相比，我们建议直接对训练动态进行正则化，而不是对损失本身进行正则化。这种方法的一个优点是，正则化的动态自然地源于我们的帕累托最优性条件，而实际的好处是它避免了对损失正则化项进行微分或计算其近端算子的需要。此外，正如其论文中所述，LOTION需要一次全精度前向传播，并且最终模型是一个仅权重量化的模型。相比之下，CAGE的训练是在权重和激活都量化的情况下完成的，它还可以受益于高性能的低精度GEMM核，从而实现更高效的训练和推理。

3.2 理论分析

SGD下的CAGE迭代。在本节中，我们展示了当基础优化器为SGD时CAGE的主要收敛结果。在这种情况下，耦合和解耦版本重合，迭代过程如下：

$$x_{t+1}=x_{t}-\alpha(\widetilde{\nabla} f(x_{t})+\lambda(x_{t}-Q(x_{t})))$$

对于$t = 0, 1, \dots$。我们对损失的光滑性和随机噪声的性质做出以下假设。

假设1和2。

• 假设1: 损失函数$f: \mathbb{R}^d \to \mathbb{R}$由某个$f^* \in \mathbb{R}$从下方有界，并且是$L_f$-平滑的，即对于任何$x, y \in \mathbb{R}^d$，都有$\|\nabla f(x) - \nabla f(y)\|_2 \le L_f \|x - y\|_2$。
• 假设2: 对于所有迭代$t$，随机梯度$\tilde{\nabla} f(x_t)$是无偏的，即$\mathbb{E}[\tilde{\nabla} f(x_t)] = \nabla f(x_t)$，并且方差是有界的$\mathbb{E}[\|\tilde{\nabla} f(x_t) - \nabla f(x_t)\|_2^2] \le \sigma^2$，其中$\sigma^2 \ge 0$是一个常数。
  这两个假设在优化文献中都相当标准，并且在某种意义上是获得有意义收敛保证的最低要求。此外，为了在分析中处理量化算子$Q$，通常用带有退火超参数的平滑代理来近似实际的非平滑量化算子$Q$。

假设3。存在某个$L_\phi$-平滑的函数$\phi: \mathbb{R}^d \to \mathbb{R}$，使得量化误差$x - Q(x) = \nabla\phi(x)$。
请注意，假设3本质上等同于量化算子$Q$的Lipschitz连续性。使用辅助函数$\phi(x)$来表述此条件的目的是暗示(3)中的迭代旨在最小化正则化损失$f(x) + \lambda\phi(x)$，我们将在分析的后面部分详细阐述。接下来，我们提出收敛结果。

定理1。设假设1、2、3成立，且$L = L_f + \lambda L_\phi$。那么，对于任何$\lambda \ge 0$，步长为$\alpha = \min(\frac{1}{L}, \frac{1}{\sqrt{T}})$的迭代(3)满足：
$\mathbb{E}\|\nabla_{\lambda P} f(Q(\hat{x}))\|^2 \le \frac{2(A_1 + A_2 + A_3)}{\max(1, \sqrt{L})T}$
其中$\hat{x}$是从历史记录$\{x_0, x_1, \dots, x_{T-1}\}$中均匀随机抽取的迭代，且

$$A_{3}=\lambda \max _{x \in\left[x_{0}, x_{T}\right]}\|x-Q(x)\|\left\|x_{0}-x_{T}\right\|.$$

讨论。首先，观察到收敛界显示了对问题（1）的帕累托最优状态的强遍历性收敛，其收敛速率在非量化情况下是最优的。收敛度量是根据我们提出的帕累托梯度（公式(2)）定义的，重要的是，上界中没有因量化误差而产生的非消失项。在实际设置中，距离$\|x_0 - x_T\|$和量化误差$\|x - Q(x)\|$都期望在任何总训练步数$T$下保持有界。因此，这意味着公式(3)中的迭代以$O(\frac{1}{\sqrt{T}})$的速率遍历性收敛到帕累托最优解。

3.3 实践实现

基于AdamW的解耦更新。由于我们的主要应用是训练和微调大型语言模型（LLM），我们在AdamW【Loshchilov & Hutter, Decoupled weight decay regularization, 2019, ICLR】之上实现CAGE。与LOTION【Kwun et al., Lotion: Smoothing the optimization landscape for quantized training, 2025, https://arxiv.org/abs/2510.08757】不同，我们专注于解耦的更新规则，并将校正项保持在预处理路径之外（算法1）。我们在大多数实验中都使用这个选择，因为我们发现它在实践中效果更好。基础优化器AdamW处理从STE获得的mini-batch梯度$g_t$，而CAGE校正作为对参数的轻量级、逐元素的后处理步骤。这种解耦具有多个实际优势：（1）它保留了优化器众所周知的行为，（2）它防止了帕累托平稳方 向$x - Q(x)$被预处理扭曲，这在经验上稳定了低比特训练，以及（3）它通过解耦具有不同行为的项来避免数值错误。

算法 1 CAGE-AdamW (解耦)
要求: 初始参数 x0; 总步数 T; AdamW 超参数 β1, β2, α, ω, ε; 量化函数 Q;
要求: CAGE 系数 λ, 静默比例 s
确保: 参数 xT
1: 初始化 m0 ← 0, v0 ← 0, e0 ← 0
2: for t = 1, 2, . . . , T do
3:   rt ← t/T ▷ 训练进度比例
4:   if rt ≤ s then
5:     λt ← 0
6:   end if
7:   if rt > s then
8:     λt ← λ * (rt - s) / (1 - s)
9:   end if
10:  采样小批量并计算随机梯度 gt (使用量化前向传播)。
11:  xt ← (1 − αω) xt ▷ 解耦的权重衰减
12:  mt ← β1mt−1 + (1 − β1)gt
13:  vt ← β2vt−1 + (1 − β2) gt ⊙ gt
14:  m̂t ← mt/(1 − βt1); v̂t ← vt/(1 − βt2)
15:  x̃t+1 ← xt − α m̂t / (√v̂t + ε)
16:  et ← xt − Q(xt) ▷ 量化误差 (无梯度)
17:  xt+1 ← x̃t+1 − α λt et ▷ 解耦校正
18: end for
19: return xT

解耦更新的具体步骤。给定当前参数集$x_t$和通过反向传播得到的STE梯度估计器$g_t$，AdamW执行以下操作：

$$\tilde{x}_{t+1}=x_t-\alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\varepsilon}, \text{ with } \hat{m}_t=\frac{m_t}{1-\beta_1^t}, \hat{v}_t=\frac{v_t}{1-\beta_2^t}$$

并附带形式为$x_t \leftarrow (1 - \alpha \omega) x_t$的解耦权重衰减正则化。然后CAGE校正项通过瞬时量化误差将参数推向量化支撑集：

$$e_t := x_t - Q(x_t), \quad x_{t+1} = \tilde{x}_{t+1} - \alpha \lambda_t \bar{e}_t$$

预热调度。在实践中，我们发现从训练一开始就激活校正可能会在模型稳定到损失景观之前过度约束动态。与标准的学习率预热调度类似，我们引入了一个静默期比例$s \in [0, 1)$，它在前$sT$步跳过校正，然后线性地将其提升到目标大小：

图 1. 非各向同性Hessian矩阵的二次型函数：在不同条件数下所有方法的最终验证损失（越低越好）。CAGE的损失不包括正则化项。我们发现CAGE持续减少稳态误差。

图 2. 二次型任务的代表性轨迹（投影到前两个主成分PC1和PC2）；CAGE遵循曲率，收敛到最接近最优解的位置。

超参数选择。在实践中，我们发现对于AdamW，参数$s \in [0.8, 0.95]$和$\lambda \in [1, 2.5]$以及线性提升能够稳健地工作。我们在附录C中提供了证明这些选择的扫描实验。这种提升可以防止优化目标的突然转变，并减少训练中的损失尖峰。

A4 实验环境与结果

实验环境总结

• 量化流程: 实验使用标准的行式矩阵量化器，基于QuEST【Panferov et al., Quest: Stable training of llms with 1-bit weights and activations, 2025a, ICML】。具体步骤包括：

  1. 通过Hadamard变换对张量进行“旋转”以缓解异常值。
  2. 根据目标比特宽度$b$，使用预计算的MSE最优高斯裁剪因子$k_b$来设置量化尺度。
  3. 应用对称量化、重新缩放并反转Hadamard变换。
  4. 反向传播使用QuEST的信任掩码STE（trust-masked STE）。

• 实验设置:

  • 合成损失实验: 在一个二次型目标函数$f(x) = \frac{1}{2} x^\top Ax - b^\top x$上进行优化，其中矩阵$A$的条件数$\kappa(A)$分别为{1, 10, 100}，使用4比特量化。

  • 量化感知微调:

    • 模型: Llama-3.2-3B。
    • 数据集: Tulu-SFT。
    • 任务: 权重和激活都量化到MXFP4格式进行QAT微调。在GSM8K、HellaSwag和WinoGrande上评估零样本或少样本性能。

  • 预训练实验:

    • 模型: Llama风格的Transformer，参数量从30M到800M不等。
    • 数据集: C4数据集。
    • 任务: 从零开始进行QAT预训练，权重/激活比特宽度为{2, 3, 4}。训练预算为每个参数100个token。

  • 优化器通用性实验:

    • 模型: 50M参数的OLMo2家族模型。
    • 数据集: ClimbMix。
    • 优化器: AdamW, Muon, Shampoo, SOAP。

  • 与LOTION的对比实验:

    • 模型: 100M参数的Llama风格模型。
    • 数据集: C4。
    • 任务: 仅权重量化（W4A16）的QAT，使用AbsMax量化器。

• 硬件配置: 8xH100 GPUs。

• 软件配置: 基于QuEST的QAT实现，主权重和优化器状态为FP32，计算使用bfloat16。

实验结果总结

• 合成损失实验 (图1, 图2):

  • 结果: 在所有条件数下，CAGE都显著降低了最终的稳态误差。轨迹图显示，CAGE找到的量化解比其他方法更接近真实最优解。
  • 结论: 实验验证了CAGE能够有效利用曲率信息，在量化域中找到更优的解。

• 量化感知微调 (图3):

  • 结果: 在Llama-3.2-3B模型的微调中，CAGE在GSM8K、HellaSwag和WinoGrande等基准上持续提高了准确率。
  • 结论: CAGE将模型量化到MXFP4格式时产生的量化误差大致减半，优于SOTA的QAT基线。

  
  图 3. Llama-3.2-3B (Tulu-SFT)上的QAT准确率，比较CAGE与最先进的MXFP4基线（使用QuEST）。CAGE在GSM8K/HellaSwag/WinoGrande上的结果持续改进，基本上将量化误差减半。两种情况都应用了Hadamard变换以减轻异常值。

• 预训练实验 (表1, 图5, 图7):

  • 结果:

    • 损失对比 (表1): 在所有模型尺寸的W4A4训练中，CAGE的最终验证困惑度始终显著优于QuEST基线。
    • 帕累托前沿 (图5): 在W2A2, W3A3, W4A4三种精度下，CAGE相对于QuEST在验证损失与模型大小的帕累托前沿上均表现出优越性。值得注意的是，W3A3 CAGE的性能优于W4A4 QuEST。
    • 损失与正则化项演变 (图7): 基线方法的量化误差在训练后期增加，而CAGE在激活正则化后，量化误差显著且持续下降，验证了其设计动机。

  • 结论: CAGE在从头预训练中能稳定地提升模型性能，并能以更低的比特数达到甚至超越基线方法在更高比特数下的性能。

  
  图 5. 使用W2A2、W3A3和W4A4进行权重和激活量化的验证损失与模型大小（字节）的关系，比较了CAGE+QuEST、基线QuEST和BF16。可以观察到CAGE始终产生更低的损失。

  Table 1. W4A4预训练结果（最终验证困惑度；越低越好）。HT = Hadamard变换。BF16是全精度参考。
  

  
  图 7. 使用QuEST基线端到端训练50M Llama模型在W4A4设置下的验证损失（左）、聚焦于最后阶段的步数（中）以及量化误差（正则化项）值（右）的图示。观察到，虽然整体训练行为相似，但CAGE在训练的最后阶段显著降低了量化模型的损失。

• 优化器通用性 (图4):

  • 结果: 对于AdamW, Shampoo, SOAP和Muon等多种优化器，启用CAGE后，验证损失均一致下降。
  • 结论: CAGE的增益不依赖于特定的优化器更新机制，是一种与优化器无关的技术。

  
  图 4. AdamW、Shampoo、SOAP和Muon在标准QAT和CAGE下的QAT（W4A4）验证损失。实线表示使用CAGE进行训练，虚线表示不使用。在所有优化器中，该方法都产生了更低的验证损失。

• 精度缩放定律分析 (图6):

  • 结果: 通过拟合数据/模型/精度缩放定律，发现CAGE在不同比特率下均提高了模型的有效容量因子eff(P)。在4比特时参数效率提升超过10%，在2比特时提升20%。
  • 结论: CAGE能更有效地利用模型参数，从而在给定的参数预算下获得更好的性能。

  
  图 6. 拟合的参数效率eff(P)，相对于标准16位进行归一化，跨不同精度。CAGE比QuEST提高了有效模型容量，在4位附近达到峰值。

• 与LOTION的对比 (图8):

  • 结果: 在作者复现的LOTION版本上进行对比，CAGE在训练后期取得了比LOTION更低的最终验证损失，而LOTION的性能趋于平稳。
  • 结论: CAGE在仅权重量化场景下同样表现优越。

  
  图 8. CAGE与LOTION在100M参数Llama风格模型上进行预训练的损失比较，该模型在C4上使用AbsMax W4A16量化进行仅权重量化QAT训练。两种方法在训练早期轨迹相似，但CAGE持续改进并获得了显著更低的最终损失，而LOTION则趋于平稳。

A5 结论

本文介绍了CAGE（曲率感知梯度估计），一种新颖的量化感知训练（QAT）方法，旨在减轻低比特量化带来的精度下降。CAGE通过一个有原则的、曲率感知的校正项来增强梯度，从而解决了直通估计器（STE）的局限性。该方法源于将QAT重构为一个多目标优化问题，旨在平衡任务损失最小化与量化约束的满足。

理论贡献在于为量化优化定义了帕累托最优解，并且CAGE在平滑非凸设定下提供了强遍历性收敛到帕累托最优点的保证，为通常基于启发式的STE方法提供了理论基础。

实践贡献方面，虽然该方法与优化器无关，但本文开发了一种基于AdamW的高效解耦实现，利用现有的优化器统计信息来提高稳定性和性能。本文的公式也推广了同期的工作【Kwun et al., Lotion: Smoothing the optimization landscape for quantized training, 2025, https://arxiv.org/abs/2510.08757】 。

实验验证表明，在预训练（Llama风格模型高达800M参数）和后训练QAT中，CAGE都表现出色，持续优于采用现代异常值缓解技术的强基线。特别是，缩放定律分析证实，与现有的QAT方法相比，CAGE提高了模型的有效容量。

CAGE的良好性能表明，结合局部曲率信息和正则化对于QAT性能是有益的。未来工作可能探索CAGE在极低比特领域的应用，以及它与其他压缩技术（如稀疏化和向量量化）的结合。

A6 附录

A 收敛性分析：定理1的证明

正则化损失的平滑性。设$F(x) = f(x) + \lambda\phi(x)$为整体正则化损失，其平滑性常数为$L = L_f + \lambda L_\phi$。根据平滑性不等式，我们有：

期望的推导。应用条件期望$E_t[\cdot] = E[\cdot | x_t]$，利用随机梯度的无偏性和方差的有界性，我们得到：

$$\begin{aligned} \begin{aligned} \mathbb{E}_{t}\left[F\left(x_{t+1}\right)-F\left(x_{t}\right)\right] & \leq-\alpha\left\|\nabla f\left(x_{t}\right)+\lambda\left(x_{t}-Q\left(x_{t}\right)\right)\right\|^{2}+\frac{L \alpha^{2}}{2} \mathbb{E}_{t}\left[\left\|\widetilde{\nabla} f\left(x_{t}\right)+\lambda\left(x_{t}-Q\left(x_{t}\right)\right)\right\|^{2}\right] \\ & =-\alpha\left\|\nabla_{\lambda \mathrm{P}} f\left(Q\left(x_{t}\right)\right)\right\|^{2}+\frac{L \alpha^{2}}{2}\left(\left\|\nabla_{\lambda \mathrm{P}} f\left(Q\left(x_{t}\right)\right)\right\|^{2}+\sigma^{2}\right) \\ & =-\alpha\left(1-\frac{L \alpha}{2}\right)\left\|\nabla_{\lambda \mathrm{P}} f\left(Q\left(x_{t}\right)\right)\right\|^{2}+\frac{L \alpha^{2}}{2} \sigma^{2} \\ & \leq-\frac{\alpha}{2}\left\|\nabla_{\lambda \mathrm{P}} f\left(Q\left(x_{t}\right)\right)\right\|^{2}+\frac{L \alpha^{2}}{2} \sigma^{2}, \end{aligned} \end{aligned}$$

这里我们强制步长限制为$\alpha \le \frac{1}{L}$。

累加不等式。应用全期望并对从$t = 0, \dots, T-1$得到的不等式求和，我们得到：

其中$\alpha = \min(\frac{1}{L}, \frac{1}{\sqrt{T}})$。剩下的任务是界定源于量化的项$\phi(x_0) - \phi(x_T)$。

量化误差项的界定。由于关于量化的假设3，我们可以将标量函数$\phi$表示为与路径无关的线积分：

$$\phi(x)=\int_{x_{0}}^{x}(I-Q) \cdot d \mathbf{r},$$

其中$I$是恒等映射，即对于任何$x \in \mathbb{R}^d$，$I(x) = x$。由于$I - Q$是一个梯度场（即$x - Q(x) = \nabla\phi(x)$），上述线积分不依赖于端点$x_0$和$x$是如何连接的。因此，我们选择直接路径$r(t) = x_0 + (x - x_0)t$，对于$t \in [0, 1]$，并将积分简化为通常的黎曼积分：

最终边界。利用这个关系，我们可以如下界定该项：

$$\begin{aligned} \begin{aligned} \phi(x_0) - \phi(x_T) & \leq \left| \int_0^1 \langle r(t) - Q(r(t)), r'(t) \rangle \, dt \right| \\ & \leq \int_0^1 \|r(t) - Q(r(t))\| \cdot \|r'(t)\| \, dt \leq \max_{x \in [x_0, x_T]} \|x - Q(x)\| \cdot \|x_0 - x_T\|. \end{aligned} \end{aligned}$$

B 超参数和可复现性

共享设置。序列长度$L=512$；批大小为64，梯度累积为8（有效tokens/step在不同尺寸间保持固定）。优化器为AdamW，参数为$\beta_1 = 0.9, \beta_2 = 0.95, \epsilon = 10^{-8}$，权重衰减为0.1，梯度裁剪为1.0。使用余弦学习率调度，预热期为总步数的10%。主权重和优化器状态为FP32；计算使用bfloat16。硬件：8× H100（启用NCCL和编译）。

量化。对权重和激活使用QuEST行式Hadamard量化器；比特宽度$b \in \{2, 3, 4\}$（对称，MSE最优高斯裁剪）。在变换域中使用信任掩码STE。

CAGE设置。在所有实验中，除非另有说明，我们使用解耦变体，静默比率$s = 0.9$，正则化参数$\lambda = 2.0$。

Table 2. Llama风格预训练运行的超参数（按模型大小）。Token预算遵循每个参数100个token（TPP）的规则。

C CAGE超参数的消融实验

二维超参数扫描。我们对CAGE的超参数进行了二维扫描：静默比率$s$和正则化系数$\lambda$。图9报告了在50M Llama模型上使用W4A4量化在C4上训练的验证困惑度。我们发现$s \in [0.8, 0.95]$和$\lambda \in [1, 2.5]$构成了一个可靠的操作区域。在此范围内，CAGE始终优于基线。

图 9. 在50M Llama模型上使用W4A4量化时，CAGE在不同静默比率s和正则化系数λ选择下的验证困惑度。