一、论文基本信息论文题目Movement Pruning: Adaptive Sparsity by Fine-Tuning作者Victor Sanh、Thomas Wolf、Alexander M. Rush发表会议NeurIPS 2020方法名称Movement Pruning研究对象BERT-base 等预训练语言模型的任务级剪枝相关实现Hugging Face 提供了 movement pruning / PruneBERT 相关代码和模型说明。官方相关仓库说明这个方法用于对预训练模型如 BERT 做 fine-pruning并强调它与 magnitude pruning 的区别magnitude pruning 保留“远离 0 的权重”movement pruning 保留“在微调过程中远离 0 的权重”。(GitHub)这篇论文的核心问题是在预训练模型迁移到下游任务时传统 magnitude pruning 是否仍然合适论文的回答是不够合适。原因是BERT 这类模型的权重主要来自大规模预训练下游 fine-tuning 通常只让权重发生很小变化。如果只看权重绝对值大小那么剪枝决策很大程度上已经被预训练阶段决定了而不是由当前下游任务决定。Movement Pruning 的思路是不要只看权重当前有多大而要看它在 fine-tuning 过程中是朝着“更重要”方向移动还是朝着“可删除”方向移动。论文摘要明确指出movement pruning 是一种简单、确定性的 first-order weight pruning 方法更适合预训练模型的 fine-tuning 场景。二、论文要解决的问题传统 magnitude pruning 的逻辑非常直接权重绝对值大说明重要。权重绝对值小说明不重要。这种方法在从零开始训练的普通监督学习模型中很常见也经常有效。因为模型权重完全由当前任务数据训练出来权重大小和任务重要性之间通常存在一定关系。但 BERT 这类预训练模型不一样。预训练模型先在大规模通用语料上训练再在下游任务上 fine-tune。Fine-tuning 阶段的数据通常比预训练数据小得多权重不会发生剧烈变化。论文指出在 transfer learning 场景中权重值大多由原始预训练模型决定fine-tuning 只做较小调整因此 magnitude pruning 很难真正根据下游任务学习剪枝决策。这就导致一个问题一个在预训练阶段绝对值很大的权重不一定对当前下游任务重要。一个原本绝对值较小的权重也可能在 fine-tuning 中变得对当前任务重要。所以 Movement Pruning 要解决的问题是如何让剪枝标准适应下游任务 fine-tuning而不是被预训练权重大小锁死三、核心思想Movement Pruning 的核心思想可以概括为保留那些在 fine-tuning 过程中远离 0 的权重剪掉那些在 fine-tuning 过程中趋向 0 的权重。这和 magnitude pruning 的判断标准完全不同。Magnitude pruning 关心的是这个权重现在离 0 有多远Movement pruning 关心的是这个权重在当前任务训练过程中是在远离 0还是在靠近 0如果一个权重虽然当前绝对值很大但 fine-tuning 过程中一直往 0 靠近说明当前任务可能不需要它。如果一个权重虽然当前绝对值不大但 fine-tuning 过程中持续远离 0说明当前任务正在强化它。论文中明确说movement pruning 与 magnitude pruning 的区别在于低值和高值权重都可能被剪掉只要它们在训练中收缩而被保留的是那些在训练中远离 0 的连接。这个思想非常适合 BERT 这种预训练模型因为它真正利用了 fine-tuning 过程中产生的任务特定信号。四、它是不是结构化剪枝严格来说原始 Movement Pruning 主要是非结构化权重剪枝。它剪的是单个 weight而不是完整 attention head、FFN neuron、Transformer layer 或 channel。也就是说它产生的是稀疏矩阵某些权重变成 0。矩阵整体形状不变。Transformer 层数、hidden size、attention head 数量都不变。所以它和 Head Pruning、LayerDrop、DynaBERT 这类结构化方法不同。不过Movement Pruning 的思想后来可以扩展到 block pruning 或结构化 pruning 中。例如后续工作会把 movement pruning 的思想用于更大的结构块。但这篇原始论文的核心方法是adaptive unstructured sparsity by fine-tuning。这点很重要因为它直接影响部署非结构化剪枝可以大幅减少非零参数量和存储。但如果硬件或推理框架不能高效利用稀疏矩阵实际推理速度不一定明显提升。论文也明确指出当前硬件对稀疏模型推理支持有限从推理速度角度看同等大小的小型 dense model 有时更理想。五、方法细节Score-Based PruningMovement Pruning 使用了一套score-based pruning框架。简单说每个权重 (W) 旁边都会有一个对应的重要性分数 (S)。训练时模型同时学习两类东西原始权重 (W)。剪枝分数 (S)。最后根据分数 (S) 生成二值 maskmask 1表示保留该权重。mask 0表示剪掉该权重。这和 magnitude pruning 的区别在于magnitude pruning 的分数直接来自权重绝对值而 movement pruning 的分数是在 fine-tuning 中学习出来的。可以理解为Magnitude pruning权重自己决定自己是否重要。Movement pruning额外学习一个“是否保留”的评分系统。六、Hard Movement Pruning论文提出了 hard movement pruning。它的做法是根据分数 (S) 保留 top-(v)% 的权重其余权重 mask 为 0。这里 (v) 表示保留比例。例如保留 10% 权重表示 90% 权重被剪掉。保留 3% 权重表示 97% 权重被剪掉。问题是top-k / top-v 操作本身不可导。也就是说不能直接通过普通反向传播更新分数 (S)。论文采用straight-through estimator直通估计器。简单说前向传播时用硬 mask反向传播时假装 mask 操作可以传梯度让梯度直接传给分数 (S)。论文明确说明由于 Top-v 的梯度几乎处处为 0它们在 backward pass 中忽略 Top-v把梯度 straight-through 到 (S)。这种设计的含义是前向时模型真的按照稀疏 mask 运行。反向时分数仍然可以学习。这就是 hard movement pruning 的核心。七、为什么它能反映“movement”Movement Pruning 的分数更新和权重梯度有关。直观上如果一个权重的更新方向让它远离 0那么对应分数会增加更可能被保留。如果一个权重的更新方向让它靠近 0那么对应分数会下降更可能被剪掉。论文解释得很清楚当权重在训练中远离 0 时score 会增加当权重向 0 收缩时score 会降低。因此score 可以被看作一种 movement accumulator也就是对权重移动方向的累积记录。这就是这篇论文最核心的思想。它不是看某一时刻的权重大小而是看 fine-tuning 过程中权重的运动趋势。所以 Movement Pruning 更像是在问当前任务到底在强化哪些连接当前任务又在弱化哪些连接八、Soft Movement Pruning除了 hard movement pruning论文还提出了soft movement pruning。Hard 版本用 top-v 保留固定比例权重soft 版本则使用连续 mask并加入稀疏正则让模型自己学出稀疏结构。可以这样理解Hard movement pruning直接规定最后保留多少权重。Soft movement pruning通过正则项鼓励分数变稀疏再根据阈值得到 mask。Soft 版本通常表现更好尤其在高稀疏率下。论文结果中soft movement pruning 在 SQuAD、MNLI、QQP 上高稀疏区间普遍优于 magnitude pruning、L0 regularization 和 hard movement pruning。原因也比较直观Hard top-v 对保留比例控制更直接但训练中更硬、更不平滑。Soft 版本让 mask 学习更连续配合正则更容易优化。所以论文最终实验中最强的通常是Soft Movement Pruning。九、渐进式剪枝不是一开始就剪到很稀疏Movement Pruning 不是一开始就把模型剪到 3% 或 10% 权重。它采用类似 automated gradual pruning 的方式训练前期先保持较高保留率随着 fine-tuning 进行逐渐提高稀疏率最后达到目标保留比例。论文遵循 Zhu and Gupta 的 cubic sparsity scheduling并且发现剪枝结束后加入一定 cool-down 步骤可以提升高稀疏率下表现。这个设计很重要。如果一开始就强行剪到极高稀疏度模型还没适应任务性能很容易崩。渐进式剪枝的好处是先让模型学习任务。再逐步收紧稀疏约束。最后用 cool-down 让稀疏模型稳定下来。十、实验设置论文主要实验对象是BERT-base-uncased。作者做的是task-specific pruning也就是针对每个下游任务分别 fine-prune。论文说明实验使用 BERT-base-uncased并冻结 embedding modules只 fine-tune Transformer layers 和任务特定 head所有稀疏率都相对于 BERT-base 统计对应模型大小。实验任务包括三个英文 benchmark任务类型数据量指标SQuAD v1.1抽取式问答约 8K 训练样本EM / F1MNLI自然语言推理约 393K 训练样本matched / mismatched accuracyQQP句子对相似 / 重复问题判断约 364K 训练样本accuracy / F1论文明确说明SQuAD 是 span extractionMNLI 和 QQP 是 paired sentence classification。对比方法包括Magnitude Pruning。L0 Regularization。Reweighted Proximal Pruning。LayerDrop。mini-BERT。论文还测试了 movement pruning 与知识蒸馏结合后的效果。十一、实验结果解读11.1 高稀疏率下Movement Pruning 明显优于 Magnitude Pruning论文最重要的结果是在高稀疏率下movement pruning 明显优于 magnitude pruning。当保留权重较多时例如剩余 70% 以上magnitude pruning 表现很好甚至通常优于 movement pruning。但随着稀疏率提高magnitude pruning 开始快速下降而 movement pruning 更稳定。论文在结果部分明确指出低稀疏区间 magnitude pruning 表现更好但高稀疏区间 movement pruning 明显胜出。这正好符合论文的理论判断低稀疏率时剪掉少量小权重就够了。高稀疏率时必须根据任务重新判断哪些连接真正重要。Magnitude pruning 只看权重大小不能充分适应下游任务movement pruning 利用 fine-tuning 中的任务信号因此高稀疏率更强。11.2 只保留 10% 权重时Soft Movement Pruning 表现最好论文 Table 2 报告了 10% 和 3% 剩余权重下的结果。在SQuAD Dev上BERT-base fine-tuned 是80.4 / 88.1 EM/F1。当只保留10% 权重时Magnitude Pruning 为67.7 / 78.5L0 Regularization 为69.9 / 80.0Movement Pruning 为71.9 / 81.7Soft Movement Pruning 为71.3 / 81.5。当只保留3% 权重时Soft Movement Pruning 达到69.5 / 79.9明显优于 Magnitude Pruning 的40.1 / 54.5。在MNLI Dev上10% 权重下 Soft Movement Pruning 达到80.7 / 81.13% 权重下达到79.0 / 79.6明显优于 Magnitude Pruning。在QQP Dev上10% 权重下 Soft Movement Pruning 达到90.5 / 87.13% 权重下达到89.3 / 85.6也显著优于 Magnitude Pruning。这个结果很强因为它说明即使只保留极少量 encoder 权重BERT 仍然可以保留相当多的下游任务性能。11.3 加入蒸馏后性能进一步提升论文进一步把 pruning 和 knowledge distillation 结合起来。Teacher 使用 fine-tuned BERT-base训练目标是任务损失和输出分布蒸馏损失的组合。论文指出蒸馏会提升所有剪枝方法在所有数据集和稀疏率下都有帮助。加入蒸馏后在SQuAD上10% 权重时 Soft Movement Pruning 达到76.6 / 84.93% 权重时达到72.7 / 82.3。在MNLI上10% 权重时达到81.2 / 81.83% 权重时达到79.5 / 80.1。在QQP上10% 权重时达到90.2 / 86.83% 权重时达到89.1 / 85.5。这里的重点是Movement Pruning 本身已经很强但和蒸馏结合后在极高稀疏率下仍能保持更高性能。论文摘要也指出结合蒸馏后模型可以在只剩3% 参数的情况下仍保持较小精度损失。(arXiv)11.4 Movement Pruning 与 Magnitude Pruning 得到的稀疏模型差异很大论文分析发现movement pruning 和 magnitude pruning 最终保留下来的权重分布差异明显。Magnitude pruning 会直接删除接近 0 的权重因此剩余权重分布呈现明显“两边大、中间空”的形态。Movement pruning 则不同。它不仅可能保留小权重也可能剪掉大权重。论文观察到movement pruning 的 score 和权重绝对值之间没有简单关系高绝对值和低绝对值权重都可能重要也都可能被剪。这说明Movement Pruning 学到的不是“权重大就保留”而是“对当前任务有用就保留”。这正是它在 transfer learning 场景中优于 magnitude pruning 的根本原因。11.5 全局剪枝与局部剪枝差距不大但高稀疏下全局略有优势论文还比较了 local selection 和 global selection。Local selection 是在每个矩阵内部选出 top-v 权重。Global selection 是在全网络所有矩阵中统一选出 top-v 权重。直觉上global selection 可以让不同层、不同矩阵获得不同稀疏率因此可能更好。但论文发现二者整体差距不大只有在高稀疏率下global movement pruning 有一定优势例如在 SQuAD 只保留 3% 权重时差距约 2.3 F1。论文还发现全局一阶方法倾向于给低层保留更多权重而高层被剪得更重。这个现象很有意思BERT 低层可能承担更通用、基础的表示功能因此在极高稀疏下更需要保留。十二、方法优点12.1 专门适合预训练模型 fine-tuning 场景Movement Pruning 最大的贡献就是指出预训练模型剪枝不能只看权重大小必须考虑 fine-tuning 过程中的任务适应方向。这比传统 magnitude pruning 更符合 BERT 这类迁移学习模型的特点。12.2 高稀疏率下效果明显更好论文最有说服力的地方是在 10%、3% 剩余权重这种极端稀疏下movement pruning 仍能保持较好性能而 magnitude pruning 大幅下降。这说明它真正适合extreme sparsity场景。12.3 方法简单不需要二阶信息Movement Pruning 使用一阶梯度信号不需要计算 Hessian 或复杂二阶近似。所以它比二阶 pruning 方法更容易实现也更适合大模型。论文也明确指出它不需要昂贵的二阶导数计算importance scores 可以作为标准 fine-tuning 的副产物获得。12.4 可以和蒸馏结合它本身不依赖 teacher但如果加入知识蒸馏性能还能进一步提升。论文结果显示distillation 会增强所有剪枝方法而 soft movement pruning distillation 表现最强。这说明它和蒸馏、量化等压缩方法是互补的。12.5 可以得到极小存储模型Hugging Face 的实现说明中提到movement pruning 可以获得极高稀疏模型例如只保留约 5% 权重时仍保留约 95% dense performance也展示了通过标准稀疏存储把 BERT encoder 大幅压缩的可能性。(GitHub)这说明它对模型存储压缩很有价值。十三、方法局限13.1 原始方法是非结构化稀疏不一定带来真实加速这是 Movement Pruning 最大的局限。它剪掉的是单个权重得到的是稀疏矩阵。如果底层硬件和推理框架不能高效执行稀疏矩阵乘法那么实际推理速度可能不会明显提升。Hugging Face 相关说明也指出标准 PyTorch 推理下movement pruning 得到的极稀疏模型并没有明显 inference speed 提升需要专门面向稀疏架构的推理设置。(GitHub)所以它更直接的收益是模型存储变小。而不是一定端到端推理马上变快。13.2 需要任务级 fine-pruningMovement Pruning 是在下游任务 fine-tuning 阶段学习稀疏结构。这意味着每个任务通常要单独 fine-prune 一次。如果希望得到一个通用稀疏 BERT然后直接迁移到所有任务这篇论文没有完全解决。13.3 对低稀疏率不一定优于 magnitude pruning论文结果显示在低稀疏率、保留权重较多时magnitude pruning 往往表现更好或更稳定。所以 Movement Pruning 的优势主要在高稀疏率。如果只剪掉少量权重传统 magnitude pruning 可能已经足够。13.4 训练和调参比简单 magnitude pruning 更复杂Magnitude pruning 很简单按权重大小排序剪掉即可。Movement Pruning 需要额外 score 参数、mask 学习、稀疏率调度、warm-up、cool-down、soft/hard 选择等。所以它的工程复杂度更高。13.5 没有直接解决结构化压缩问题如果目标是减少 heads、减少 layers、减少 FFN blocks、减少 hidden sizeMovement Pruning 原始形式不是直接答案。它可以作为更大结构块剪枝的思想基础但原论文主要做的是 weight-level pruning。十四、整体评价Movement Pruning 的核心价值在于它重新定义了预训练模型剪枝时“重要性”的含义。在普通监督学习中权重大小往往能反映任务重要性但在 BERT 这样的预训练模型中权重大小更多反映预训练阶段的通用知识。下游任务 fine-tuning 时真正重要的是哪些连接被当前任务进一步强化哪些连接正在被弱化。因此Movement Pruning 的判断标准非常自然保留那些在 fine-tuning 中远离 0 的权重。剪掉那些在 fine-tuning 中趋向 0 的权重。这让它比 magnitude pruning 更适合 task-specific transfer learning 场景。实验也证明在只保留 10% 甚至 3% encoder 权重的极端稀疏条件下movement pruning 尤其是 soft movement pruning明显优于 magnitude pruning 和 L0 regularization。不过它的核心局限也很清楚它主要产生非结构化稀疏。如果没有稀疏硬件或专门推理库参数少不等于速度快。因此它更适合存储压缩、稀疏模型研究或者作为后续 block pruning / structured sparse pruning 的基础思想。十五、一句话总结《Movement Pruning: Adaptive Sparsity by Fine-Tuning》提出一种面向预训练语言模型的任务自适应剪枝方法不再根据权重绝对值大小决定保留或删除而是根据 fine-tuning 过程中权重是否远离 0 来学习剪枝分数实验表明在 BERT 的 SQuAD、MNLI 和 QQP 任务上Movement Pruning 尤其是 Soft Movement Pruning 在高稀疏率下明显优于 Magnitude Pruning并且与蒸馏结合后可以在只保留极少参数时仍保持较强性能。