模型剪枝方法全解

目录写在前面一、为什么需要剪枝过参数化是个普遍现象二、剪枝的基本流程三、非结构化剪枝Unstructured Pruning3.1 幅值剪枝Magnitude Pruning3.2 非结构化剪枝的硬件问题四、结构化剪枝Structured Pruning4.1 卷积网络的 Filter / Channel 剪枝4.2 Transformer 的结构化剪枝4.3 结构化 vs 非结构化的权衡五、重要性判据判断谁该被剪是剪枝的核心5.1 基于梯度的方法5.2 基于 Taylor 展开的方法5.3 基于激活值的方法六、剪枝的时机什么时候剪6.1 训练后剪枝Post-Training Pruning6.2 训练中剪枝Pruning During Training6.3 稀疏化训练Training with Sparsity Regularization6.4 初始化时剪枝Pruning at Initialization6.5 迭代剪枝Iterative Pruning七、2:4 结构化稀疏NVIDIA 的硬件友好方案八、LLM 时代的剪枝为什么旧方法不够用九、SparseGPT用二阶信息做训练后剪枝十、Wanda更简单、同样有效十一、剪枝与其他压缩技术的关系十二、一些值得深想的地方写在最后做模型压缩剪枝往往是最先被提到的技术但也是最容易被讲浅的一个。很多教程只说把权重小的去掉然后贴一段torch.nn.utils.prune的代码就算完事了。这篇文章希望把剪枝讲透——从动机、数学直觉、各类方法一直到 LLM 时代的新挑战和新方案SparseGPT、Wanda——把整条脉络梳理清楚。一、需要剪枝的原因过参数化深度学习有一个几乎反直觉的特点过参数化over-parameterization对训练有好处但对推理是负担。训练一个大模型参数越多越容易收敛、越不容易陷入糟糕的局部极值。然而训练完成后这些参数里有相当一部分实际上对模型的输出贡献很小甚至几乎为零。把它们留着不过是在推理时浪费计算和内存。有多浪费Han et al. (2015) 的经典工作《Learning both Weights and Connections for Efficient Neural Networks》里发现AlexNet 和 VGG 这些早期大模型有高达 80% 以上的权重可以被移除而性能下降几乎可以忽略不计。这个观察启发了整个剪枝研究领域。更有理论支撑的是 Frankle 和 Carlin 在 2019 年提出的彩票假说Lottery Ticket Hypothesis一个大的随机初始化网络里存在一个小得多的子网络中奖票如果用相同的初始权重单独训练这个子网络它能达到和完整网络相当的性能。剪枝的本质就是找到这张中奖票。剪枝的实际价值体现在三个方向减小模型体积存储、加速推理计算量、降低内存占用带宽。在边缘部署、移动端、实时推理等场景下这三点都直接关系到能不能用。二、剪枝的基本流程在深入各种具体方法之前先把标准的剪枝流程讲清楚后面所有内容都是在这个框架里做变体。训练 → 剪枝 → 微调Fine-tune是最经典的三阶段范式第一步在完整数据集上正常训练一个大模型让它充分收敛。第二步按照某种重要性判据importance criterion把不重要的参数移除或置为零得到一个稀疏化的模型。这个步骤会导致精度下降。第三步对剪枝后的模型做微调也叫 fine-tune 或 recovery training在数据集上继续训练若干 epoch让剩余的权重重新调整适应把精度恢复回来。在 PyTorch 里剪枝通常不是真的删掉参数而是给参数打一个二值 maskmask 为 1 的位置权重保留mask 为 0 的位置权重被遮盖等效于置零。这样做的好处是剪枝可逆可以随时改变 mask 重新选择被剪的参数适合迭代式剪枝。import torch import torch.nn as nn import torch.nn.utils.prune as prune # 一个简单的全连接层 layer nn.Linear(100, 50) # 对这一层的 weight 做 L1 非结构化剪枝去掉幅值最小的 30% prune.l1_unstructured(layer, nameweight, amount0.3) # 查看稀疏度 sparsity float(torch.sum(layer.weight 0)) / layer.weight.nelement() print(fSparsity: {sparsity * 100:.1f}%) # 约 30% # 把 mask 永久固化真正删除被剪的连接 prune.remove(layer, weight)PyTorch 的实现细节值得说一下调用prune.l1_unstructured之后模块会多出两个属性——weight_orig原始权重和weight_mask二值掩码。在 forward pass 里实际用的是weight weight_orig * weight_mask。只有调用prune.remove之后才会把这个遮盖操作固化真正修改weight参数。三、非结构化剪枝Unstructured Pruning非结构化剪枝是最细粒度的剪枝方式独立地看待每一个权重根据它的重要性决定是否保留不考虑它在矩阵里的位置关系。3.1 幅值剪枝Magnitude Pruning最简单也是最广泛使用的重要性判据权重的绝对值越小说明它对输出的贡献越小应当优先被剪掉。直觉上这很合理——如果一个权重接近于零把它真的置为零前向传播的输出几乎不会变化。形式上对于权重矩阵重要性分数就是然后找一个阈值把所有的权重置零。# 全局幅值剪枝对整个模型的所有层统一设置稀疏度 parameters_to_prune [ (model.fc1, weight), (model.fc2, weight), (model.fc3, weight), ] prune.global_unstructured( parameters_to_prune, pruning_methodprune.L1Unstructured, amount0.5, # 全局移除 50% 的权重 )注意这里用的是global剪枝而不是每层独立地剪 50%。全局剪枝的意思是把所有层的权重放在一起排序找到全局的第 50% 分位数作为阈值统一决定哪些权重被剪。这比每层独立剪枝local pruning效果通常更好因为不同层对精度的贡献差异很大——靠近输出的层往往更敏感不应该和靠近输入的层一视同仁地各剪 50%。3.2 非结构化剪枝的硬件问题非结构化剪枝有一个根本性的工程困境它产生的稀疏矩阵在标准硬件上并不能直接加速。CUDA 的矩阵乘法 kernel 是专门针对稠密矩阵优化的。一个 70% 稀疏的矩阵在 GPU 上做矩阵乘法实际消耗的时间和一个稠密矩阵几乎没有区别——因为 GPU 仍然要把那些零值的位置也处理完而不是自动跳过它们。要真正利用稀疏性加速有两条路一是用专门的稀疏矩阵运算库如 NVIDIA 的 cuSPARSE在稀疏度足够高通常 90%时才有实际加速效果二是使用 NVIDIA 在 Ampere 架构引入的2:4 结构化稀疏N:M sparsity——后面会讲到。这个问题是非结构化剪枝最大的局限参数数量确实降下去了但延迟往往没有同比例改善。四、结构化剪枝Structured Pruning既然非结构化剪枝产生的稀疏矩阵不友好更激进的方案是直接删除整个结构单元——结构化剪枝。结构化剪枝不是移除单个权重而是移除完整的 filter卷积核、channel通道、attention head、甚至整层。一旦某个结构被移除对应的矩阵维度直接减小模型天然变成一个更小的稠密网络标准硬件可以直接加速不需要任何稀疏矩阵支持。4.1 卷积网络的 Filter / Channel 剪枝对于卷积网络最常见的结构化剪枝单位是 filter输出 channel。每个卷积层有若干个 filter每个 filter 负责提取一种特征。如果某个 filter 的权重整体都比较小或者它的输出激活值绝大多数时候接近零那这个 filter 的贡献很有限可以整体移除。判据一L1/L2 Norm计算每个 filter 的权重向量的 L2 范数作为重要性分数去掉范数最小的若干个 filter。# 对卷积层做 L2 结构化剪枝按 dim0输出通道维度去掉 25% 的 filter prune.ln_structured(conv_layer, nameweight, amount0.25, n2, dim0)判据二基于 Batch Normalization 的 scaling 因子BN 层对每个通道学习一个缩放因子 γgamma。如果某个通道的 γ 很小说明这个通道在整个网络里被压制了相当于这个 filter 的输出被 BN 主动抑制。这个方法不需要单独计算重要性直接复用 BN 的 γ 参数作为剪枝信号非常自然。判据三基于激活值的 APoZAPoZAverage Percentage of Zero activations统计一个 filter 的输出激活值中有多少比例接近零通常是过 ReLU 之后为零的比例。一个经常输出全零的 filter 实际上什么也没做可以安全移除。4.2 Transformer 的结构化剪枝对于 Transformer 系模型BERT、LLaMA 等结构化剪枝的目标变成 attention head、FFN 的神经元、甚至整个 Transformer 层。Attention Head 剪枝每个多头注意力层有若干个 head可以通过测量每个 head 对模型输出的影响比如对 loss 的梯度信号或者直接把某个 head 掩掉看精度下降多少来判断它的重要性去掉不重要的 head。微软的 FAR 和 Michel et al. 2019 的工作都在这个方向做了系统研究发现大量 Transformer 模型里很多 head 几乎可以被移除而不影响性能。FFN 维度剪枝Transformer 的前馈网络有一个先扩展通常 4x再压缩的结构中间层的神经元也可以被结构化地剪掉直接减小 FFN 的中间维度。层剪枝Layer Pruning最粗粒度的结构化剪枝直接去掉整个 Transformer block。研究发现深层模型如 BERT-Large 的 24 层里靠近输出的若干层对最终性能贡献相对较小可以被整体移除。这是目前对 LLM 做结构化压缩最直接的手段之一。4.3 结构化 vs 非结构化的权衡这两种剪枝方式各有优劣选择哪一种主要取决于部署场景非结构化剪枝的优点是灵活可以达到非常高的稀疏度80%而不显著损失精度因为它以最精细的粒度保留了最重要的权重缺点是产生的稀疏矩阵在标准 GPU 上难以加速。结构化剪枝的优点是产物是一个更小的稠密模型可以在任何硬件上直接加速工程友好缺点是粒度粗在相同精度损失下能达到的压缩率通常低于非结构化剪枝。五、重要性判据判断谁该被剪是剪枝的核心剪枝的质量在很大程度上取决于如何衡量参数的重要性。上面提到的幅值、L2 norm 都是比较朴素的方法研究者们提出了越来越精细的判据。5.1 基于梯度的方法参数的重要性不只看它的值还要看改变它对损失函数的影响。一阶梯度 ∂L/∂w 反映的是将某个权重增大一点点时 loss 的变化率可以用来衡量权重的重要程度。Optimal Brain DamageOBDLeCun et al. 1990是这个方向的经典工作利用损失函数对权重的二阶偏导Hessian 矩阵对角元素来更准确地估计删除某个权重对 loss 的影响。其后继工作 Optimal Brain SurgeonOBS进一步考虑了 Hessian 的非对角元素并在删除某个权重后对其他权重做补偿更新理论上更精确。这类方法的精度更高但计算 Hessian 或其近似的代价相当高限制了它在大模型上的直接应用——直到后来被 SparseGPT 以近似的方式带回了 LLM 剪枝领域后面详述。5.2 基于 Taylor 展开的方法Molchanov et al. (2017) 提出了一个实用的近似方法用 Taylor 展开来估计移除参数 w_i 对损失的影响这个乘积同时考虑了权重的幅值和梯度比单纯用幅值更准确。直觉上一个权重很大但梯度为零说明当前 loss 对它不敏感可以剪一个权重很小但梯度很大说明改变它会显著影响 loss应该保留。5.3 基于激活值的方法另一条思路是不看权重本身而是看权重对应的输出激活值。如果一个神经元的输出在大量样本上都接近零说明它在实际推理中几乎没有被激活可以被移除。这种方法不依赖梯度只需要跑一遍前向传播统计激活分布计算非常高效。六、剪枝的时机什么时候剪除了剪什么什么时候剪也是影响最终效果的重要因素。6.1 训练后剪枝Post-Training Pruning最直接模型训练完了直接在上面做剪枝。优点是不需要修改训练流程可以对任何现成的预训练模型应用缺点是训练时模型不知道自己将来要被剪剪完之后精度下降可能比较大需要微调来恢复。6.2 训练中剪枝Pruning During Training在训练过程中动态地剪枝让模型在稀疏状态下继续学习。最经典的是Gradual Magnitude PruningGMPZhu and Gupta, 2018从一个小的初始稀疏度开始随着训练的进行按照一个调度策略通常是三次函数曲线逐渐提高稀疏度最终达到目标稀疏度。这个过程让模型有时间适应稀疏化比一次性剪掉效果好很多。6.3 稀疏化训练Training with Sparsity Regularization在 loss 函数里加入 L1 正则项来惩罚权重的绝对值之和L1 正则化天然会把很多权重推向零让模型主动学会稀疏化。这和剪枝有相通之处但更柔和——它不强制置零而是让模型在训练中自然形成稀疏结构。6.4 初始化时剪枝Pruning at InitializationSNIPLee et al. 2019等工作尝试在训练一开始就判断哪些连接是重要的直接从一个稀疏子网络开始训练。这条路最理想但判断还没训练的模型哪些参数是重要的本身就很困难实际效果不如先训练后剪枝稳定。6.5 迭代剪枝Iterative Pruning不论在哪个时间点剪迭代剪枝往往比一次性剪one-shot pruning效果更好。具体做法每次只剪一小部分比如剩余权重的 10%然后微调恢复精度再剪再微调……反复迭代直到达到目标稀疏度。这符合直觉每次只做小幅度的改变给模型时间适应避免精度断崖式下跌。彩票假说的实验里用的就是迭代幅值剪枝Iterative Magnitude Pruning, IMP在高稀疏度下效果远好于一次性剪枝。七、2:4 结构化稀疏NVIDIA 的硬件友好方案非结构化剪枝的硬件困境促使了一种折中方案的出现N:M 稀疏N:M sparsity特别是 2:4 稀疏。NVIDIA 在 Ampere 架构A100中专门引入了对 2:4 稀疏的硬件加速支持。2:4 稀疏的规则是在每连续 4 个权重里最多保留 2 个非零值另外 2 个必须为零。这不是完全的结构化剪枝没有删整行或整列但也不是完全随机的非结构化剪枝——它有一个局部的规律性硬件可以针对这个规律设计专门的 sparse tensor core。2:4 稀疏的好处是理论上能达到 2x 的矩阵乘法加速同时因为仍然有一定的权重选择灵活性精度损失比完全结构化剪枝小。Wanda 等方法也支持 2:4 稀疏模式的生成。八、LLM 时代的剪枝为什么旧方法不够用当模型规模从 BERT 的 1 亿参数膨胀到 LLaMA 的 700 亿甚至更大传统剪枝方法遇到了一个严峻的新挑战没法微调。传统剪枝的三段式流程训练 → 剪枝 → 微调在小模型上可行是因为微调的成本可以接受。但对一个 700 亿参数的模型哪怕是在少量数据上做几个 epoch 的微调也需要数十甚至数百张高端 GPU对绝大多数团队和个人来说是完全不现实的。更糟糕的是Frantar 和 Alistarh 在 SparseGPT 的工作里发现对 LLM 直接做朴素的幅值剪枝效果很差——把 LLaMA-7B 剪到 50% 稀疏度幅值剪枝会让困惑度perplexity从 6.3 暴涨到 17.3几乎不可用。这和小模型上幅值剪枝可以轻松剪掉 90% 而精度基本不损失的经验完全相反。为什么 LLM 对幅值剪枝这么脆弱原因之一来自一个有趣的现象LLM 的激活值里会出现大量级特征large magnitude features。在某些特定维度channel上激活值会比其他维度大出 10 到 100 倍。LLM 严重依赖这些大幅值特征来做预测而大幅值特征对应的权重往往也比较大——所以纯粹基于权重幅值排序可能会保留一些权重值不小但对应的激活值微小的权重同时删掉一些权重值适中但连接到大激活特征的关键权重。这种判断失误在小模型上后果有限但在 LLM 里会造成精度急剧下滑。九、SparseGPT用二阶信息做训练后剪枝SparseGPTFrantar 和 Alistarh, 2023是第一个实现了对十亿级 LLM 做高质量剪枝的方法它把 OBS 系列的二阶思路带回了现代大模型。核心思路不只是把不重要的权重置零在置零之后还对剩余权重做补偿更新让剩余权重调整来弥补被删权重留下的坑。可参考我的模型量化相关文章有所介绍形式上对于一层权重矩阵 W 和对应的输入 XSparseGPT 把剪枝问题建模为一个最小化重建误差的优化问题即找到一个满足稀疏约束的使得它的输出和原来的的输出尽可能相同。这个问题的求解利用了 Hessian 矩阵的近似逆矩阵来进行高效的逐列求解。SparseGPT 的结果令人印象深刻LLaMA-7B 在 50% 非结构化稀疏度下困惑度只从 6.3 升至约 6.9保住了绝大多数性能——而幅值剪枝在同等稀疏度下会到 17.3。对 OPT-175B 和 BLOOM-176B 这样的更大模型SparseGPT 可以在单台服务器上不到 4.5 小时完成剪枝。不过 SparseGPT 也有代价需要计算 Hessian 近似有一定的内存和计算开销。十、Wanda更简单、同样有效WandaSun et al. 2023arxiv:2306.11695在 SparseGPT 之后提出了一个更轻量级的方案切入点正是那个大量级激活特征的现象。核心思路权重的重要性分数 权重幅值 × 对应输入激活值的 L2 范数。其中是第个输入神经元的激活值在一批校准数据上统计的。这个分数同时考虑了两件事一是权重本身够不够大二是这个权重对应的输入通道是不是那种大量级特征。只有两个条件都满足才值得保留——一个权重本身很大但它接收的输入通道激活值很小贡献也不大反之一个权重值适中但它连接的是大激活特征通道那它其实很重要。Wanda 的 PyTorch 实现极为简洁def prune_wanda(W, X, sparsity_ratio): # W: 权重矩阵 (C_out, C_in) # X: 校准数据的输入激活 (N*L, C_in) # 计算每个输入神经元的激活 L2 范数 input_norms X.norm(p2, dim0) # (C_in,) # 重要性分数 权重幅值 × 对应输入范数 score W.abs() * input_norms.unsqueeze(0) # (C_out, C_in) # 每行独立排序per-output basis _, sorted_idx torch.sort(score, dim1) # 把分数最低的 sparsity_ratio 比例的权重置零 pruned_idx sorted_idx[:, :int(W.shape[1] * sparsity_ratio)] W.scatter_(dim1, indexpruned_idx, srctorch.zeros_like(W)) return WWanda 有几个值得注意的特点无需权重更新剪完直接用不需要像 SparseGPT 那样对剩余权重做补偿更新。论文里给出了一个有趣的解释——他们的实验发现当用 Wanda 的度量找到最优稀疏子网络时额外的权重更新并不能显著改善结果这暗示 LLM 本身就存在精确的稀疏子网络Wanda 能直接找到它。Per-output 粒度每行独立地找最不重要的权重来剪而不是全局排序。这避免了某些输出神经元被过度剪枝而另一些几乎不被剪枝的不平衡情况。效果对比LLaMA-7B50% 非结构化稀疏Wanda 的困惑度是 7.26而纯幅值剪枝是 17.29SparseGPT 是约 6.9。Wanda 在几乎不引入额外计算开销的情况下效果接近 SparseGPT远好于幅值剪枝。Wanda 还支持 2:4 N:M 结构化稀疏模式只需要修改排序的方式在每 4 个连续权重里找分数最低的 2 个剪掉就能生成 NVIDIA 硬件友好的稀疏格式。十一、剪枝与其他压缩技术的关系剪枝不是孤立的在实际工程里通常和其他压缩手段组合使用。剪枝 量化Quantization剪枝减少权重数量量化降低每个权重的精度比如从 FP32 → INT8。两者正交可以叠加。Han et al. 2016 的 Deep Compression 工作就是把剪枝、量化、Huffman 编码三管齐下实现了对 AlexNet 约 40× 的压缩。SparseGPT 的后续工作也探索了在稀疏化的同时对剩余权重做量化。剪枝 知识蒸馏Knowledge Distillation用大模型作为教师在剪枝后的小模型微调阶段用教师模型的 soft logits 作为训练信号能比直接用 hard label 微调恢复更多精度。这在 BERT 压缩系列工作里被广泛使用。剪枝 vs 低秩分解Low-Rank Factorization低秩分解把权重矩阵分解为两个更小矩阵的乘积LoRA 是其在微调里的应用。两种方法都在减少有效参数量但思路不同剪枝是稀疏化保留部分权重置零其余低秩分解是压缩用更少的信息来近似整个矩阵。十二、细节讨论高稀疏度下大稀疏模型 vs 小稠密模型谁更好Wanda 的论文里有一组有趣的对比50% 稀疏的 LLaMA-65B在零样本精度上优于同参数量的稠密 LLaMA-30B。这个结论并不是放之四海皆准——对于结构化稀疏结论相反无微调情况下小稠密模型更好。这暗示非结构化稀疏在大模型上存在规模红利大模型本身携带的信息更丰富稀疏化只是丢掉了冗余保留了密度更高的信息核心。为什么不同层的稀疏度应该不同全局剪枝给所有层统一设置相同稀疏度是一个粗放的策略。实际上不同层对最终性能的敏感度差异非常大。通常第一层和最后一层比中间层更敏感注意力层中的某些权重矩阵如 Query、Key比 Value 更敏感FFN 层通常比注意力层更能承受高稀疏度。让不同层用不同的稀疏度是进一步提升剪枝效果的重要方向。彩票假说的真正含义是什么彩票假说意味着大网络里存在天然的小子网络这个子网络用原始初始化权重训练就能达到大网络的精度。这个假说如果成立意味着我们不需要先训练大网络、再剪枝——理论上应该可以找到那张中奖票直接从头训练。但实践中找到中奖票本身就依赖大网络的训练结果否则你怎么知道哪些权重是重要的所以绕不过去。它更多的意义是从理论上解释了为什么神经网络可以被高度稀疏化而精度不降。最后模型剪枝从 1990 年代 OBD/OBS 的理论探索到 2015 年 Song Han 的幅值剪枝实验到 2019 年彩票假说的理论洞见再到 2023 年 SparseGPT 和 Wanda 为 LLM 量身定制的训练后方案走了将近三十年。每一代的进步都在解决上一代的局限从只能剪小模型到能剪百亿参数的 LLM从需要微调到训练后一步到位从忽略硬件到与 NVIDIA 稀疏加速硬件协同设计。学剪枝最值得记住的几个核心认知非结构化剪枝灵活但硬件不友好结构化剪枝硬件友好但精度有代价幅值在小模型上够用但 LLM 需要考虑激活值迭代剪枝好于一次性剪枝剪完通常需要微调除非你用了 SparseGPT/Wanda 这类专门为免微调设计的方法。参考资料Learning both Weights and Connections for Efficient Neural NetworksHan et al., NeurIPS 2015Deep CompressionHan et al., ICLR 2016The Lottery Ticket HypothesisFrankle Carlin, ICLR 2019SparseGPT: Massive Language Models Can be Accurately Pruned in One ShotFrantar Alistarh, ICML 2023A Simple and Effective Pruning Approach for Large Language Models (Wanda)Sun et al., ICLR 2024PyTorch Pruning Tutorial: https://pytorch.org/tutorials/intermediate/pruning_tutorial.html