雅可比引导噪声注入：提升大语言模型量化鲁棒性

# 基于雅可比矩阵引导的噪声注入：提升大型语言模型量化鲁棒性
来源：https://arxiv.org/html/2608.20988
Arnav ChavanAffiliation:AmazonNahush LeleAffiliation:AmazonSankalp DayalAffiliation:AmazonDeepak GuptaAffiliation:Amazon

###### 摘要

大型语言模型 \(LLMs\) 的量化常受限于自注意力机制对离散化误差的敏感性。我们发现softmax算子是量化稳定性的瓶颈，因其对异常值敏感且具有状态相关的雅可比矩阵。理论上我们证实抑制该雅可比矩阵的范数有助于约束量化引起的性能下降。基于此，我们提出雅可比矩阵引导的噪声注入（Jacobian-Guided Noise Injection）训练策略，该方法向注意力前逻辑值注入零均值高斯噪声，其方差直接由雅可比矩阵的弗罗贝尼乌斯范数推导得出。与依赖启发式方法或直接惩罚雅可比矩阵的先驱方法不同，本方法基于局部注意力敏感度确定最优噪声方差。我们在SOTA级LLM架构上评估该方法，结果显示其在主流训练后量化（PTQ）方法基础上提升了鲁棒性。实证分析表明，该方法在SigLIP的ImageNet-1K Top-1准确率上相对增益达+37%，在语言模型低比特量化设置下WikiText困惑度相对改善达40%，证明了该方法的有效性。

###### 关键词：

量化、大型语言模型、Softmax、雅可比矩阵正则化、噪声注入

## 1 引言

大型语言模型 \(LLMs\) 在众多自然语言任务中取得了显著成功（40 (https://arxiv.org/html/2608.20988#bib.bib25)；41 (https://arxiv.org/html/2608.20988#bib.bib26)；39 (https://arxiv.org/html/2608.20988#bib.bib28)）。然而，由于其巨大的计算和内存需求，高效部署这些模型仍具挑战性（12 (https://arxiv.org/html/2608.20988#bib.bib44)；37 (https://arxiv.org/html/2608.20988#bib.bib43)）。量化通过降低权重和激活的精度，为高效部署提供了有前景的路径（18 (https://arxiv.org/html/2608.20988#bib.bib18)）。例如，与FP16模型相比，4比特权重量化可将内存占用减少4倍，推理速度提升超过3倍，甚至能在移动GPU上部署700亿参数的模型（24 (https://arxiv.org/html/2608.20988#bib.bib13)）。

然而，简单的量化往往导致显著性能下降（10 (https://arxiv.org/html/2608.20988#bib.bib12)；24 (https://arxiv.org/html/2608.20988#bib.bib13)），这源于低精度下的近似噪声和舍入误差扰动了中间计算（32 (https://arxiv.org/html/2608.20988#bib.bib35)）。在较低比特宽度（如4比特或以下）时，由于表征能力降低和激活异常值的存在，性能下降更为严重。近期研究表明，量化误差对数学推理、多步规划和长上下文任务的影响尤为突出，这推动了鲁棒的量化感知和激活感知方法的发展（22 (https://arxiv.org/html/2608.20988#bib.bib1)；23 (https://arxiv.org/html/2608.20988#bib.bib3)；46 (https://arxiv.org/html/2608.20988#bib.bib4)）

参见标题图 1：使用ERQ和RepQViT在ImageNet-1K Top-1准确率上的PTQ结果，基于SigLIP base 16-384。一个关键观察是，Transformer（42 (https://arxiv.org/html/2608.20988#bib.bib23)）中的自注意力机制对量化误差特别敏感，这是由于softmax、归一化等高度敏感操作的存在。我们在量化模型部署中观察到了这一现象（图2 (https://arxiv.org/html/2608.20988#S2.F2)），其中误差传播在注意力层尤为显著，导致与预期激活值产生重大偏离。与线性层中误差传播被固定权重矩阵所约束不同，注意力中的Softmax运算符表现出状态相关敏感性，这种敏感性根据输入分布剧烈变化（20 (https://arxiv.org/html/2608.20988#bib.bib19)；26 (https://arxiv.org/html/2608.20988#bib.bib20)）。当逻辑值包含异常值或幅值较大时，问题尤为严重，因为这些值的量化误差会通过softmax门的指数非线性被指数级放大（45 (https://arxiv.org/html/2608.20988#bib.bib14)；6 (https://arxiv.org/html/2608.20988#bib.bib15)）。对较大逻辑值的微小扰动会导致输出概率产生不成比例的巨大变化，引发灾难性的误差传播。这造成了不可预测的误差放大，而标准量化技术无法解决。

本文中，我们分析了通过softmax运算符的量化误差传播，并推导出使用softmax雅可比矩阵范数约束该误差的条件。此外，我们提出雅可比矩阵引导的噪声注入，一种旨在提升量化模型下游性能的训练策略（图1 (https://arxiv.org/html/2608.20988#S1.F1)）。我们的主要贡献是：
1. 1\. 我们分析了softmax雅可比矩阵的谱范数与量化误差放大之间的关系，并证明最小化逻辑值扰动下的期望损失能隐式地正则化该范数。
2. 2\. 我们推导出近似雅可比矩阵弗罗贝尼乌斯范数的表达式，并利用它校准噪声注入方差，为启发式方法提供了更简单的替代方案。
3. 3\. 我们证明了雅可比矩阵引导的噪声注入方法在多个LLM架构和量化设置下提升了量化鲁棒性，在W4A4设置下恢复了高达37%的准确率且零推理开销。

## 2 方法论

本节中，我们将Transformer（42 (https://arxiv.org/html/2608.20988#bib.bib23)）中自注意力Softmax算子识别为量化稳定性的瓶颈。我们从理论上分析了约束该量化误差所需的条件，并观察到这些条件可通过隐式Hessian正则化（1 (https://arxiv.org/html/2608.20988#bib.bib11)）来满足。最后，为在实际场景中应用此正则化，我们提出了一种微调框架以实现鲁棒量化。

### 2.1 Softmax雅可比矩阵的敏感性

在标准Transformer中，自注意力机制根据预激活逻辑值计算注意力概率。为分析误差传播，令 \(z\in\mathbb{R}^N\) 表示给定查询 \(i\) 的预激活逻辑值矩阵的单行向量（即 \(z_i=q^Tk_i/\sqrt{d}\)）。对应的注意力概率向量 \(a\in\mathbb{R}^N\) 通过Softmax函数计算：
\[ a=S(z)=\frac{e^{z_i}}{\sum_{j=1}^N e^{z_j}} \tag{1} \]
当模型以量化格式部署时，权重和激活的离散化会向逻辑值引入有界扰动 \(\delta\in\mathbb{R}^N\)，使得量化后的预激活值为 \(z_q=z+\delta\)。传播到注意力分布的误差为：
\[ \Delta a=S(z+\delta)-S(z) \tag{2} \]
使用一阶泰勒展开，我们通过雅可比矩阵 \(J_S(z)\in\mathbb{R}^{N\times N}\) 近似该误差：
\[ \Delta a\approx J_S(z)\delta \implies \|\Delta a\|_2\leq \|J_S(z)\|_2\|\delta\|_2 \tag{3} \]
关键漏洞在于softmax雅可比矩阵 \(J_S(z)\) 的构造：
\[ J_S(z)_{i,j}=\partial a_i/\partial z_j=a_i(\mathbf{1}_{i=j}-a_j) \tag{4} \]
与线性层中输入雅可比矩阵是常数权重矩阵（例如 \(\nabla_X(XW)=W^T\)）不同，\(J_S(z)\) 是稠密且严格状态相关的。雅可比矩阵的结构揭示了softmax输出对其输入逻辑值的敏感性。当注意力集中于单个token时，会产生安全的饱和区域，此时雅可比矩阵范数趋近于零；而当注意力在多个token间均匀分配时，则会产生高度敏感的区域，此时量化误差被急剧放大（范数随着质量分布在更多token上而减小）。标准微调目标未考虑这一点。因此，无正则化的模型可能学习到使预激活值停留在这些敏感区域的参数，最大化了 \(\|J_S(z)\|_2\)。在此类情况下，即使最小的量化误差 \(\|\delta\|_2\) 也会引发不可预测的、指数级的 \(\Delta a\) 放大，导致灾难性任务退化。所提出的噪声注入策略旨在通过使注入噪声成为雅可比矩阵状态的函数，在训练期间隐式地诱导这种正则化。

参见标题图 2：Siglip-base-384中注意力输出的激活余弦相似度。
### 2.2 约束Softmax雅可比矩阵

为严格约束量化误差 \(\|\Delta a\|_2\)，我们必须限制雅可比矩阵的谱范数 \(\|J_S(z)\|_2\)。令 \(\mathcal{L}(z)=(l\circ S)(z)\) 表示作为预激活逻辑值函数的端到端损失。直接惩罚 \(\|J_S(z)\|_2\) 在计算上代价过高，因为它不是静态参数，而是状态相关的矩阵。我们可以通过对逻辑值Hessian矩阵 \(\nabla^2_z\mathcal{L}(z)\) 进行正则化来建立关于雅可比矩阵的理论边界。应用多元链式法则，逻辑值Hessian矩阵可分解为：
\[ \nabla^2_z\mathcal{L}(z)=J_S(z)^T\nabla^2_a l(a)J_S(z)+\sum_{k=1}^N\frac{\partial l}{\partial a_k}\nabla^2_z S_k(z) \tag{5} \]
令 \(H_a=\nabla^2_a l(a)\) 表示激活Hessian矩阵。利用Gauss-Newton近似并忽略二阶残差项，关系简化为：
\[ \nabla^2_z\mathcal{L}(z)\approx J_S(z)^T H_a J_S(z) \tag{6} \]
在局部最优点附近，损失景观是局部凸的，这意味着 \(H_a\) 是半正定的（\(H_a\succeq 0\)）。令 \(\lambda_{\min}>0\) 表示 \(H_a\) 的最小正特征值。根据半正定矩阵的性质，我们可以约束Hessian矩阵的迹：
\[ \text{Tr}\left(J_S(z)^T H_a J_S(z)\right)\geq\lambda_{\min}\text{Tr}\left(J_S(z)^T J_S(z)\right) \tag{7} \]
算法 1 雅可比矩阵引导的噪声注入
0:模型 \(\mathcal{M}\)，更新间隔 \(T\)，缩放因子 \(\alpha\)
1: 对每个训练步 \(t\) 执行
2:    如果 \(t\mod T=0\) 则
3:       执行前向传播以计算每层的注意力 \(P^{(\ell)}\)
4:       对每层 \(\ell\) 执行
5:          使用公式12 (https://arxiv.org/html/2608.20988#S2.E12) 计算 \(\|J_S^{(\ell)}\|_F^2\)
6:          更新 \(\sigma_i^{(\ell)}\leftarrow\alpha\cdot\sqrt{\mathbb{E}_{b,h}[\|J_S\|_{F,i}^2]}\)
7:          钳制：\(\sigma\leftarrow\text{clamp}(\sigma,\sigma_{\min},\sigma_{\max})\)
8:       结束循环
9:    结束如果
10:   计算逻辑值：\(Z=QK^T/\sqrt{d}\)
11:   采样噪声：\(\epsilon_i\sim\mathcal{N}(0,\sigma_i^2)\)
12:   扰动：\(\tilde{Z}=Z+\epsilon\)
13:   应用softmax：\(\tilde{A}=\text{Softmax}(\tilde{Z})\)
14:   计算损失 \(\mathcal{L}(\tilde{A})\) 并反向传播
15: 结束循环

根据弗罗贝尼乌斯范数的定义，\(\text{Tr}\left(J_S(z)^T J_S(z)\right)=\|J_S(z)\|_F^2\)。由于谱范数满足 \(\|J_S(z)\|_2^2\leq \|J_S(z)\|_F^2\)，我们得到：
\[ \|J_S(z)\|_2\leq\sqrt{\frac{\text{Tr}\left(\nabla^2_z\mathcal{L}(z)\right)}{\lambda_{\min}}} \tag{8} \]
公式8 (https://arxiv.org/html/2608.20988#S2.E8) 直接约束了公式3 (https://arxiv.org/html/2608.20988#S2.E3) 中定义的误差放大。因此，我们理想的正则化目标应当惩罚逻辑值Hessian矩阵的迹：
\[ \mathcal{L}_{\text{ideal}}(z)=\mathcal{L}(z)+\lambda\cdot\text{Tr}\left(\nabla^2_z\mathcal{L}(z)\right) \tag{9} \]
其中 \(\lambda>0\) 控制正则化强度。然而，计算此目标需要二阶导数的连续反向传播，对于大型Transformer而言计算上难以实现。表1 (https://arxiv.org/html/2608.20988#S2.T1) 展示了如何仅使用一阶梯度高效近似此惩罚项。

### 2.3 随机扰动

我们现在推导 \(\mathcal{L}_{\text{ideal}}\) 的一阶近似。考虑修改目标以最小化在直接应用于逻辑值的连续、零均值扰动 \(\epsilon\in\mathbb{R}^N\) 下的期望损失：\(\mathbb{E}_{\epsilon}[\mathcal{L}(z+\epsilon)]\)。我们通过围绕未扰动逻辑值 \(z\) 的扰动损失的二阶泰勒级数展开来分析该目标的行为：
\[ \mathcal{L}(z+\epsilon)=\mathcal{L}(z)+\nabla_z\mathcal{L}(z)^T\epsilon+\frac{1}{2}\epsilon^T\nabla^2_z\mathcal{L}(z)\epsilon+\mathcal{O}(\|\epsilon\|^3) \tag{10} \]
我们显式定义扰动 \(\epsilon\) 为从 \(\mathcal{N}(0,\sigma^2 I)\) 采样的各向同性高斯噪声。利用其统计特性（\(\mathbb{E}[\epsilon]=0\) 和 \(\mathbb{E}[\epsilon\epsilon^T]=\sigma^2 I\)），对泰勒展开取期望导致一阶梯度项完全消失：
\[ \mathbb{E}_{\epsilon\sim\mathcal{N}(0,\sigma^2 I)}[\mathcal{L}(z+\epsilon)]\approx\mathcal{L}(z)+\frac{\sigma^2}{2}\text{Tr}\left(\nabla^2_z\mathcal{L}(z)\right) \tag{11} \]
此推导得出了一个关键相似性：最小化高斯逻辑值扰动下的期望损失与我们第2.2节 (https://arxiv.org/html/2608.20988#S2.SS2) 推导的要求相符。因此，扰动方差 \(\sigma^2\) 充当了公式9 (https://arxiv.org/html/2608.20988#S2.E9) 中的隐式正则化强度 \(\lambda\)，对应关系为 \(\lambda=\sigma^2/2\)。

### 2.4 雅可比矩阵引导的噪声注入

表 1：Llama-3.2-3B 和 Qwen2.5-3B 的零样本PTQ结果。Acc是7个基准测试的平均准确率（↑）。表1 (https://arxiv.org/html/2608.20988#S2.T1) 确立了方差为 \(\sigma^2\) 的高斯噪声注入隐式地以强度 \(\lambda=\sigma^2/2\) 正则化了Hessian矩阵的迹。然而，固定的全局 \(\sigma\) 对所有层和位置一视同仁，忽略了雅可比矩阵范数（以及量化敏感性）在网络中剧烈变化的事实。位于敏感区域（第2.1节 (https://arxiv.org/html/2608.20988#S2.SS1)）的位置需要比