SparseGPT：大规模语言模型的一次性精确剪枝

Frantar E, Alistarh D. Sparsegpt: Massive language models can be accurately pruned in one-shot[C]//International conference on machine learning. PMLR, 2023: 10323-10337.

第一章：引言与研究动机

大型语言模型（LLMs）特别是生成预训练Transformer（GPT）家族，在广泛的任务上展现了卓越的性能。然而，它们的巨大规模和计算成本使得部署极其困难。以GPT-175B为例，其1750亿参数在半精度（FP16）格式下至少需要320GB存储空间，导致推理时需要至少五块具有80GB内存的A100 GPU。这种资源需求严重限制了这些模型的实际应用。

本文提出了SparseGPT，这是首个能够高效处理100亿参数以上模型的精确一次性剪枝方法。SparseGPT通过将剪枝问题简化为一系列极大规模的稀疏回归实例，然后通过新的近似稀疏回归求解器来解决这些实例。该求解器足够高效，可以在单个GPU上几小时内执行完成对最大的开源GPT模型（1750亿参数）的剪枝。

第二章：问题形式化与现有方法局限性

2.1 层级剪枝问题

后训练压缩通常通过将全模型压缩问题分解为逐层子问题来完成。对于剪枝，Hubara等人（2021a）将这个问题形式化为：对于每一层$\ell$，找到具有特定目标密度的稀疏掩码$M\ell$，以及可能更新的权重$W^c\ell$，使得：

$$\arg\min_{mask\ M_\ell, W^c_\ell} \|W_\ell X_\ell - (M_\ell \odot W^c_\ell)X_\ell\|_2^2$$

其中$W\ell$是未压缩层的权重，$X\ell$是给定的输入，$\odot$表示逐元素乘积。

2.2 掩码选择与权重重构的分离

由于同时优化掩码$M\ell$和剩余权重$W^c\ell$使得问题成为NP困难问题，所有现有方法都采用近似方法。一个特别流行的方法是将问题分离为掩码选择和权重重构两个步骤。具体来说，首先根据某些显著性标准（如权重幅度）选择剪枝掩码$M$，然后在保持掩码不变的情况下优化剩余的未剪枝权重。

一旦掩码固定，公式(1)变成一个线性平方误差问题，可以轻松优化。对于固定的剪枝掩码$M$，矩阵行$w_i$的最优值可以通过求解稀疏重构问题精确计算：

$$w_i^{M_i} = (X_{M_i}X_{M_i}^T)^{-1}X_{M_i}(w_{M_i}X_{M_i})^T$$

其中$X_{M_i}$表示仅包含未被行$i$剪枝的权重对应的输入特征子集。

2.3 扩展到百亿参数的困难

计算最优重构的高计算复杂度主要源于求解每一行需要单独求逆一个$O(d{col} \times d{col})$矩阵。这是因为行掩码$Mi$通常不同，且$(H{Mi})^{-1} \neq (H^{-1}){M_i}$，即掩码Hessian的逆不等于完整逆的掩码版本。

图3展示了行Hessian挑战的可视化：该图说明了行是独立稀疏化的，剪枝的权重以白色显示。每一行需要不同的Hessian逆矩阵，因为它们的掩码模式不同。如果所有行掩码相同，则只需要计算一个共享的逆矩阵，因为$H = XX^T$仅依赖于层输入，对所有行都相同。

第三章：SparseGPT算法详解

3.1 快速近似重构的核心思想

SparseGPT的关键创新在于实现了在具有不同剪枝掩码的行之间重用Hessian的原则性方法。算法首先定义一个递归的$d_{col}$个索引子集序列：

$$U_{j+1} = U_j - \{j\}, \quad \text{其中} \quad U_1 = \{1, ..., d_{col}\}$$

这些子集产生了一系列逆Hessian $(H_{Uj})^{-1} = ((XX^T){Uj})^{-1}$，将在$W$的所有行之间共享。关键的是，更新的逆$(H{U{j+1}})^{-1}$可以通过高斯消元的一步从$B = (H{U_j})^{-1}$高效计算：

$$(H_{U_{j+1}})^{-1} = \left[B - \frac{1}{[B]_{11}} \cdot B_{:,1}B_{1,:}\right]_{2:,2:}$$

3.2 等价的迭代视角与OBS更新

为了激发我们的算法，我们首先从不同的迭代视角看待逐行权重重构，使用经典的OBS更新。假设损失的二次近似，其中当前权重$w$是最优的，OBS更新$\delta_m$提供了剩余权重的最优调整以补偿索引$m$处权重的移除，产生误差$\varepsilon_m$：

$$\delta_m = -\frac{w_m}{[H^{-1}]_{mm}} \cdot H^{-1}_{:,m}, \quad \varepsilon_m = \frac{w_m^2}{[H^{-1}]_{mm}}$$

3.3 最优部分更新

应用OBS更新$\delta_m$可能会调整所有可用参数的值（在当前掩码$M$中）以补偿$w_m$的移除。但如果我们只更新剩余未剪枝权重中的子集$U \subseteq M$中的权重呢？我们仍然可以从误差补偿中受益，只使用$U$中的权重，同时降低应用OBS的成本。

图4提供了SparseGPT重构算法的可视化：左图展示了给定固定剪枝掩码$M$，我们如何使用一系列Hessian逆$(H_{U_j})^{-1}$增量地剪枝权重矩阵$W$每列中的权重，并更新这些行中位于"右侧"的剩余权重。被剪枝的权重（深蓝色）右侧的权重将被更新以补偿剪枝误差，而未剪枝的权重不会产生更新（浅蓝色）。右图说明了通过迭代分块进行的自适应掩码选择。

3.4 自适应掩码选择机制

SparseGPT采用迭代分块策略，每次为$B_s = 128$列选择剪枝掩码，基于方程(3)中的OBS重构误差$\varepsilon$，使用Hessian序列中的对角值。然后执行下一个$B_s$权重更新，再为下一个块选择掩码，依此类推。这个过程允许每列非均匀选择，特别是也使用相应的Hessian信息，同时也考虑先前的权重更新进行选择。

第四章：实验结果与分析

4.1 模型规模与剪枝难度的关系

图1展示了OPT-175B的稀疏度-困惑度比较：该图比较了SparseGPT与幅度剪枝在不同均匀层级稀疏度下的表现。图中显示，幅度剪枝只能在10%稀疏度之前保持准确性，并在30%稀疏度之后完全崩溃。相比之下，SparseGPT可以达到60%的稀疏度，困惑度仅有轻微增加（从8.35增加到约8.5）。即使在80%稀疏度下，SparseGPT压缩的模型仍能产生合理的困惑度（约16），而幅度剪枝导致完全崩溃（困惑度>100）。

图2展示了整个OPT模型家族的压缩结果：该图显示了使用SparseGPT将整个OPT模型家族（从1.35M到175B参数）压缩到不同稀疏模式时的困惑度与模型大小的关系。图中使用对数尺度显示模型参数量。关键发现是较大的模型更可压缩：它们在固定稀疏度下相对于较小的对应模型损失的准确性显著更少。例如，最大的OPT和BLOOM家族模型可以被稀疏化到50%，困惑度几乎没有增加。

4.2 极限稀疏度测试

对于OPT-175B模型，实验结果表明：

• 50%稀疏度：困惑度从8.35略微降至8.21（这种轻微改善似乎是数据集特定的）
• 60%稀疏度：困惑度保持在可接受的8.5左右
• 80%稀疏度：虽然困惑度增加到约16，但模型仍然功能正常

这意味着SparseGPT可以从这些模型中移除约1000亿个权重，同时保持低精度影响。

4.3 半结构化稀疏模式

半结构化稀疏性（如n:m模式）对硬件加速特别重要。实验结果显示：

• 2:4模式（每4个权重中2个为零）：OPT-175B困惑度仅增加0.39（从8.35到8.74）
• 4:8模式（每8个权重中4个为零）：困惑度增加仅0.1（从8.35到8.45）

4.4 联合稀疏化与量化

图6比较了联合50%稀疏+4比特量化与等效的3比特量化：该图显示了对于≥2.7B参数的OPT家族模型，50%稀疏+4比特的组合始终优于纯3比特量化。对于OPT-175B，50%+4比特达到8.29困惑度，而3比特GPTQ达到8.68。这种改进在所有测试的模型规模上都是一致的，表明稀疏性和量化的组合是一个有前景的方向。

第五章：算法复杂度与实际加速

5.1 计算复杂度分析

SparseGPT的总体成本包括三部分：

1. 初始Hessian计算：$\Theta(n \cdot d_{col}^2)$，其中$n$是输入样本数
2. 遍历逆Hessian序列：$O(d_{col}^3)$
3. 重构/剪枝本身：可以由应用公式(3)到$W$的所有$d{row}$行对所有$d{col}$列的时间上界，即$O(d{col} \cdot d{row} \cdot d_{col})$

总计：$O(d{col}^3 + d{row} \cdot d{col}^2)$。对于Transformer模型，这简化为$O(d{hidden}^3)$，比精确重构快了整整$d_{hidden}$倍。

5.2 实际加速结果

表7和表8展示了实际加速效果：

• CPU加速（使用DeepSparse引擎）：40%稀疏度达到1.57×，50%稀疏度达到1.82×，60%稀疏度达到2.16×
• GPU加速（2:4模式在NVIDIA Ampere上）：Q/K/V/Out层达到1.79×，FC1层达到1.67×，FC2层达到1.54×

附录：数学推导

A. Hessian逆序列的递归计算

给定初始Hessian $H = XX^T + \lambda I$及其逆$H^{-1}$，我们需要高效计算序列$(H_{U_j})^{-1}$。

定理A.1：设$H$是对称正定矩阵，$Uj = {j, j+1, ..., d{col}}$。则：

$$(H_{U_{j+1}})^{-1} = (H_{U_j}^{-1})_{2:,2:} - \frac{1}{[H_{U_j}^{-1}]_{11}} \cdot (H_{U_j}^{-1})_{:,1}(H_{U_j}^{-1})_{1,:}$$

证明：使用分块矩阵求逆公式。设$H_{U_j}$的分块形式为：

$$H_{U_j} = \begin{bmatrix} h_{jj} & h_j^T \\ h_j & H_{U_{j+1}} \end{bmatrix}$$

其逆矩阵可写为：

$$(H_{U_j})^{-1} = \begin{bmatrix} \alpha & \beta^T \\ \beta & \Gamma \end{bmatrix}$$

通过分块矩阵求逆公式：

• $\alpha = (h_{jj} - hj^T H{U_{j+1}}^{-1} h_j)^{-1}$
• $\beta = -\alpha H{U{j+1}}^{-1} h_j$
• $\Gamma = H{U{j+1}}^{-1} + \alpha H{U{j+1}}^{-1} h_j hj^T H{U_{j+1}}^{-1}$

通过Schur补的性质，我们得到：

$$H_{U_{j+1}}^{-1} = \Gamma - \frac{\beta \beta^T}{\alpha}$$

这正是高斯消元的一步。□

B. 最优部分更新的推导

考虑移除权重$w_m$并只更新子集$U$中的权重。优化问题变为：

$$\min_{\delta_U} \|W X - (W + \delta_U e_U^T - w_m e_m e_m^T) X\|_2^2$$

其中$e_U$是$U$中索引的指示向量，$e_m$是位置$m$的单位向量。

展开并求导，得到最优更新：

$$\delta_U = -\frac{w_m}{[H_U^{-1}]_{mm}} \cdot H_U^{-1}[:,m]$$

相应的误差为：

$$\varepsilon_m = \frac{w_m^2}{[H_U^{-1}]_{mm}}$$

当$U = M$（所有未剪枝权重）时，这简化为标准OBS公式。

C. 权重冻结的等价性

命题C.1：SparseGPT算法等价于逐列贪婪压缩框架，其中：

• 压缩操作定义为：$\text{compress}(w_j)_i = 0$ 如果 $j \notin M_i$，否则为$w_j^i$
• 每列处理后，相应权重被"冻结"（不再更新）

证明：在迭代$j$，算法：

1. 对所有$i$，如果$j \notin Mi$，设置$w{ij} = 0$（剪枝）
2. 使用$(H_{U_j})^{-1}$更新行中的剩余权重
3. 从未来的更新集合中移除索引$j$

这精确对应于逐列压缩，其中每个权重在处理其列时被固定到最终值。□

D. 自适应掩码选择的误差界

定理D.1：使用基于OBS误差的自适应掩码选择，总重构误差满足：

$$\sum_{j \in \text{pruned}} \varepsilon_j \leq \sum_{j \in \text{pruned}} \frac{w_j^2}{[H^{-1}]_{jj}}$$

其中等号在精确重构时成立，不等号反映了部分更新近似。

证明：对于每个剪枝的权重$j$，部分更新使用$H_{U_j}$而不是完整的$H$。由于$Uj \subseteq {1, ..., d{col}}$，矩阵$H_{U_j}$是$H$的主子矩阵。通过Cauchy交错定理：

$$[H_{U_j}^{-1}]_{jj} \geq [H^{-1}]_{jj}$$

因此：

$$\varepsilon_j = \frac{w_j^2}{[H_{U_j}^{-1}]_{jj}} \leq \frac{w_j^2}{[H^{-1}]_{jj}}$$

对所有剪枝权重求和得到所需的界。

E. 联合稀疏化与量化的误差分析

在联合压缩中，每个权重要么被剪枝（设为0），要么被量化。总误差为：

$$E_{total} = \sum_{j \in \text{pruned}} w_j^2 + \sum_{j \in \text{quantized}} (w_j - \text{quant}(w_j))^2$$

SparseGPT通过OBS更新补偿这两种误差：

$$E_{:,j-i} = \frac{W_{:,j} - M_{:,j} \cdot \text{quant}(W_{:,j})}{[H^{-1}]_{jj}}$$

这种统一处理允许剪枝和量化决策相互影响，产生比顺序应用更好的结果。

结论

SparseGPT代表了大规模语言模型压缩的重要进展。通过巧妙的算法设计和数学优化，它实现了之前认为不可能的目标：在没有任何重训练的情况下，将1750亿参数的模型压缩到50-60%稀疏度，同时保持极高的准确性。这项工作不仅具有重要的实际意义，也揭示了关于大规模神经网络结构的深刻见解——在密集预训练模型的邻域中存在着高度稀疏但仍然准确的模型。