HyperAIHyperAI

Command Palette

Search for a command to run...

7 小时前
LLM

KronQ:基于克罗内克分解海森矩阵的大语言模型量化

Donghyun Lee Yuhang Li Ruokai Yin Priyadarshini Panda

摘要

训练后量化(PTQ)是一种无需重新训练即可压缩大语言模型(LLM)的广泛采用技术。现有的二阶PTQ方法(包括GPTQ)仅根据输入激活统计量构建量化目标,实际上假设所有输出通道对逐层重建目标的贡献相等。我们提出KRONQ,一个通过将梯度协方差引入量化流程来挑战这一假设的PTQ框架。在克罗内克分解海森矩阵近似下,量化损失同时依赖于激活协方差和梯度协方差,KRONQ在两个互补层面上利用这一点:(1)KRONQ引入双向不相干处理,利用梯度协方差将现有的输入侧随机旋转扩展到输出维度,从而降低输入和输出维度上的权重幅度方差;(2)KRONQ基于梯度和激活海森矩阵迹,推导出一种新的灵敏度度量,用于层间混合精度分配。值得注意的是,在LLaMA-3-70B上进行2位权重量化时,GPTQ和GPTAQ出现发散或产生退化量化结果(在WikiText-2上困惑度>2000),而KRONQ实现了7.93的困惑度。

一句话总结

南加州大学和耶鲁大学的研究人员提出 KRONQ,一种训练后量化框架,利用 Kronecker 分解 Hessian 近似下的梯度协方差,实现双向非相干处理,并引入一种新的混合精度分配的灵敏度度量,在 2-bit 仅权重量化的 LLaMA-3-70B 上达到 7.93 困惑度,而 GPTQ 和 GPTAQ 在该设定下发散。

核心贡献

  • KRONQ 是一种训练后量化框架,将梯度协方差纳入量化流程,实现双向非相干处理,将输入侧的随机旋转扩展到输出维度,并降低两个维度上的权重大小方差。
  • 它引入了一种新的层间混合精度灵敏度度量,由梯度和激活 Hessian 迹的乘积计算得出,解决了共享相同输入统计的子层中的歧义,并支持最优的比特宽度分配。
  • 在 LLaMA-2 和 LLaMA-3 模型(7B 到 70B)上,在仅权重和权重与激活两种量化设定下,2/3/4-bit 的实验展示了一致的当前最优性能。最大的改进出现在 2-bit,此时其他方法发散或产生退化的量化,而 KRONQ 在 LLaMA-3-70B 上达到 7.93 困惑度。

引言

参数规模达到数千亿的大型语言模型需要压缩才能实际部署。训练后量化(PTQ)通过将权重和激活转换为低比特表示来减少内存和延迟,在 PTQ 方法中,像 GPTQ 这样的基于补偿的方法依赖于输入激活协方差来近似 Hessian 以修正舍入误差。然而,这些方法忽略了输出侧的梯度结构,隐式地假设输出通道之间重要性相同,从而遗漏了 Kronecker 分解 Hessian 中的梯度协方差项。作者引入 KRONQ,它使用从一次反向传播中获得的梯度协方差 H_G 来增强量化目标。这一添加使得双向非相干处理成为可能,并提供了一个联合 Hessian 迹的灵敏度度量用于层间混合精度,同时在权重更新中代数抵消,使得高效的 GPTAQ 求解器保持不变。KRONQ 在 LLaMA-2 和 LLaMA-3 上 2 到 4-bit 实现了当前最优的准确性,在超低精度下增益最大。

方法

作者提出 KRONQ,一种训练后量化(PTQ)框架,旨在解决现有基于补偿的方法(如 GPTQ 和 GPTAQ)的局限性。标准方法仅依赖输入激活协方差 H_X 作为完整 Hessian 的代理,隐式假设输出侧梯度协方差 H_G 是单位矩阵。然而,输出通道在梯度幅度上表现出显著的变化。如下图所示:

Q、K、V 和 O 投影的梯度协方差 H_G 的归一化对角元素差异达数个数量级,揭示了异构的输出侧灵敏度。为了捕捉这一缺失的因素,作者利用 Kronecker 分解近似(K-FAC)。

对于线性层 y = Wx,每个样本的梯度分解为 ∂L/∂W = g x^⊤,其中 g = ∂L/∂y。完整 Hessian H 的经验 Fisher 近似由下式给出:

H=E[vec(LW)vec(LW)]=E[(xg)(xg)]=E[xxgg].H = \mathbb{E} \left[ \text{vec} \left( \frac{\partial \mathcal{L}}{\partial W} \right) \text{vec} \left( \frac{\partial \mathcal{L}}{\partial W} \right)^\top \right] = \mathbb{E} \left[ (x \otimes g)(x \otimes g)^\top \right] = \mathbb{E} \left[ xx^\top \otimes gg^\top \right].H=E[vec(WL)vec(WL)]=E[(xg)(xg)]=E[xxgg].

应用 K-FAC 独立性假设 x ⊥⊥ g,期望分解为 H ≈ H_X ⊗ H_G,其中 H_X = E[xx^⊤] 且 H_G = E[gg^⊤]。这将复杂度从 O(d_in^2 d_out^2) 降低到 O(d_in^2 + d_out^2)。将其代入标准的逐层 PTQ 目标,得到 Kronecker 分解量化目标:

minW^tr[HGΔWHXΔW],\min_{\widehat{W}} \text{tr} \left[ H_G \Delta W H_X \Delta W^\top \right],Wmintr[HGΔWHXΔW],

其中 ΔW = W - Ŵ。基于 GPTAQ,作者通过引入非对称校准进一步修正输入漂移,得到最终目标:

minW^tr[HG(ΔWHXΔWWΔXXΔW)].\min_{\widehat{W}} \text{tr} \left[ H_G \left( \Delta W H_X \Delta W^\top - W \Delta X X^\top \Delta W^\top \right) \right].Wmintr[HG(ΔWHXΔWWΔXXΔW)].

在应用于该目标的逐列 Optimal Brain Surgeon (OBS) 更新下,量化第 p 列后的权重补偿简化为:

ΔW:,p+1:=δp[HX1]p,p+1:+W:,p[P]p,p+1:,\Delta W_{:, p+1:} = -\delta_p [H_X^{-1}]_{p, p+1:} + W_{:, p} \cdot [P]_{p, p+1:},ΔW:,p+1:=δp[HX1]p,p+1:+W:,p[P]p,p+1:,

其中 δ_p 是缩放的量化误差,P 是 GPTAQ 非对称校正矩阵。

当归一化 Hessian 是非相干的,即其特征向量不与坐标轴对齐时,量化误差减小。虽然之前的方法对 H_X 应用了非相干处理,但作者研究了 H_G 是否也是相干的。这促使了双向非相干处理(BiIP)。如下图所示:

处理前,非相干度量 μ(H_G)/√d_out 高达 0.99,处理后降至 0.10 以下,证实了输出侧旋转有效地使 H_G 非相干。将输入侧的对角线缩放扩展到列和行两个方向,作者应用:

WSGWSX,SX=diag([HX]jjW:,j2)1/4,SG=diag([HG]iiWi,:2)1/4,W \leftarrow S_G W S_X, \quad S_X = \text{diag} \left( \frac{[H_X]_{jj}}{\|W_{:, j}\|^2} \right)^{1/4}, \quad S_G = \text{diag} \left( \frac{[H_G]_{ii}}{\|W_{i, :}\|^2} \right)^{1/4},WSGWSX,SX=diag(W:,j2[HX]jj)1/4,SG=diag(Wi,:2[HG]ii)1/4,

其中输出侧的 S_G 项是新颖的。随后,应用正交变换:

WUWV,HXVHXV,HGUHGU,ΔXXVΔXXV,W \leftarrow U W V^\top, \quad H_X \leftarrow V H_X V^\top, \quad H_G \leftarrow U H_G U^\top, \quad \Delta X X^\top \leftarrow V \Delta X X^\top V^\top,WUWV,HXVHXV,HGUHGU,ΔXXVΔXXV,

将 U 和 V 实例化为随机 Hadamard 变换。这确保了 H_G 和 H_X 都变得非相干。图示表明,仅使 H_X 非相干几乎不会改变输出通道的变异系数(CV_out),而 BiIP 同时显著降低了 CV_in 和 CV_out。

为了在子层间分配比特预算,作者根据灵敏度度量对它们进行排序。在二阶近似下,第 ℓ 层的期望量化损失正比于 ε_ℓ^2 · tr(H^(ℓ))。利用 Kronecker 近似 H ≈ H_X ⊗ H_G,这自然扩展到两个因子:

E[L]tr(HG())tr(HX()).\mathbb{E}[\mathcal{L}_\ell] \propto \text{tr}(H_G^{(\ell)}) \cdot \text{tr}(H_X^{(\ell)}).E[L]tr(HG())tr(HX()).

作者将 KRONQ 灵敏度得分定义为 s_ℓ = tr(H_G^(ℓ)) · tr(H_X^(ℓ)),在 BiIP 之后计算,并为具有较大 s_ℓ 的子层分配更高的比特宽度。如下图所示:

在 tr(H_G) · tr(H_X) 下的子层排序与仅使用 tr(H_X) 的排序存在显著差异。在 Transformer 注意力块中,查询(Q)、键(K)和值(V)投影共享相同的输入,因此 H_X^(Q) = H_X^(K) = H_X^(V)。联合得分通过 H_G 打破了这种退化,因为 Q、K 和 V 接收不同的下游梯度,从而产生严格更好的困惑度和比特宽度权衡。

实验

KRONQ 在 LLaMA 模型上,在仅权重、分组以及权重与激活量化设定下,针对 GPTQ 和 GPTAQ 进行评估,使用困惑度和推理基准上的零样本准确率。该方法始终达到最低困惑度和最高准确率,在超低比特宽度下增益最大,此时仅依赖激活协方差的方法失败,并且它成功量化了 LLaMA-3-70B,而基线方法发散。这些改进泛化到更新的模型和更难的基准,而使用联合输入输出灵敏度得分的混合精度分配进一步优化了困惑度-比特权衡。消融实验证实了输出侧曲率、对角线缩放和双向非相干处理都是必不可少的,并且该方法提供了显著的内存节省和解码加速。

在 4-bit 每通道仅权重量化中,像 OSTQuant 和 SpinQuant 这样的几种方法在 LLaMA 模型上的困惑度接近 FP16 基线,而文中报告 KRONQ 在所有比特宽度下都达到最低困惑度,在 2 和 3 bits 时增益最大,其他方法则严重退化。零样本准确率的改进与困惑度趋势一致,KRONQ 在所有推理基准上始终优于基线。在列出的 4-bit 方法中,OSTQuant 在 LLaMA-2-7B 上记录了最低的 WikiText-2 困惑度(5.64)以及在 PiQA、ArcC 和 WG 上的最高平均零样本准确率(64.0)。KRONQ 在 W2 和 W3 时提供了最显著的困惑度降低,此时仅依赖激活协方差的方法失败,并且是唯一能够在这些比特宽度下对 LLaMA-3-70B 产生有效量化的方法。

在 2-bit 仅权重分组量化下,KRONQ 在 LLaMA-2-7B 上实现了最低困惑度和最高零样本准确率,困惑度达到 7.61,而 GPTQ 为 274.0,OmniQuant 为 11.06,平均准确率为 53.7。文中指出,KRONQ 在 LLaMA-2-13B 上也保持了对 GPTQ 和 GPTAQ 的领先。在 LLaMA-2-7B 上,GPTQ 在分组量化下困惑度崩溃至 274.0,而 KRONQ 达到 7.61,优于 OmniQuant(11.06)和 GPTAQ(23.19)。KRONQ 在 LLaMA-2-7B 上达到平均零样本准确率 53.7,超越 OmniQuant(48.3)和 GPTQ(39.3),在所有列出的推理基准上均有提升。

在权重与激活量化 W2A4 设定下,KRONQ 在 LLaMA-2 和 LLaMA-3 模型上始终实现最低困惑度和最高平均零样本准确率。最大的困惑度降低发生在 LLaMA-2-7B 上,KRONQ 将困惑度从 GPTQ 的 36.74 降至 9.38。准确率提升与困惑度改善一致,KRONQ 在所有基准上均优于 GPTQ 和 GPTAQ。KRONQ 将 LLaMA-2-7B 的困惑度从 36.74(GPTQ)降至 9.38,相对下降 74%。在 LLaMA-3-8B 上,KRONQ 的平均零样本准确率达到 51.0,比 GPTAQ 高出 8.8 个百分点。

在四个近期模型家族中,KRONQ 在 W4 和 W2 下均达到最低 WikiText-2 困惑度。在 W2 时差距急剧扩大,基线方法往往灾难性退化,而 KRONQ 保持稳健,在其中一个案例中,GPTAQ 甚至表现不如 GPTQ,而 KRONQ 依然出色。在更近期模型上,KRONQ 相对于 GPTQ 和 GPTAQ 的困惑度优势在 2-bit 量化时急剧扩大,增益常常超过一个数量级。在 DeepSeek-R1-Distill-Llama-8B 的 W2 下,GPTAQ 的困惑度比 GPTQ 更差,但 KRONQ 实现了远低得多的困惑度,突显了其输出侧校正的互补优势。

在 DeepSeek-R1-Distill-Llama-8B 和 Gemma-3-12B-IT 上,于仅权重 4-bit 量化下,KRONQ 在四个硬推理基准上始终优于 GPTQ 和 GPTAQ。最显著的增益出现在生成式编码任务上,KRONQ 在两个模型上的得分几乎翻倍于次优方法。在数学和知识基准上的大幅改进证实了该方法泛化能力远超困惑度。KRONQ 在每个模型-基准组合上都实现了最高准确率。在 LiveCodeBench 上,KRONQ 对于 DeepSeek-R1-Distill-Llama-8B 和 Gemma-3-12B-IT 的得分几乎翻倍于次优方法。在 AIME-2024 数学基准上,KRONQ 在 Gemma-3-12B-IT 上的准确率是 GPTQ 和 GPTAQ 的两倍以上。KRONQ 相对于 GPTAQ 的优势在需要复杂推理和生成的任务(如代码和数学)上最大。

评估涵盖 LLaMA、DeepSeek 和 Gemma 模型的仅权重、分组以及权重与激活量化,测量困惑度和推理基准上的零样本准确率。KRONQ 始终达到最低困惑度和最高准确率,在 2- 和 3-bit 宽度下优势最为显著,此时竞争方法如 GPTQ 和 GPTAQ 往往崩溃或严重退化。该方法在困惑度之外的泛化良好,在包括代码生成和数学的硬推理任务上带来显著增益,并且在基线方法失败的新模型家族上仍保持稳健。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供