简介一份面向自然语言处理研究者与模型部署工程师的ACL22论文PDF聚焦陈丹琦团队提出的CoFi结构化剪枝方法。该方法面向Transformer架构结合粗粒度与细粒度剪枝及逐层蒸馏在仅损失极少精度的情况下将模型推理加速约10倍缓解了模型压缩中加速比不高或训练成本过高的难题。资源为单文件PDF共1个文件格式为pdf包体大小约2.68MB内容完整包含论文原文及图文排版适合用于精读算法细节、复现实验或梳理剪枝方法演进。已有191人学习/下载是了解当前NLP模型轻量化前沿思路的实用参考。通过阅读可直接获取CoFi方法对层剪枝、头剪枝、FFN剪枝等不同粒度的处理方案以及动态层映射的蒸馏策略有助于快速掌握该方法的动机、公式推导与实验设计。1. 推理加速的硬约束为什么剪枝比蒸馏更省成本模型部署到线上之后推理延迟直接决定服务成本和用户体验。分布式训练能通过多卡并行把训练时间压下来但推理阶段面对的是单条请求多卡并行带来的收益非常有限—— GPU 之间的通信开销甚至会抵消掉并行带来的收益。真正能在推理阶段见效的是把模型本身变小。这就是模型压缩的用武之地而模型压缩的两条主线——剪枝和蒸馏——各有各的瓶颈。剪枝的思路是砍掉贡献度低的参数常规的结构化剪枝只能做到 2~3 倍加速再往上精度就会明显下滑蒸馏可以做到更大的加速比但通常需要先用大量无标签数据预训练一个学生模型这个预训练阶段的成本高到让很多团队直接放弃。陈丹琦团队在 ACL22 提出的 CoFiCoarse- and Fine-grained Pruning把这两条路合并成一条用一个统一的掩码框架同时做粗粒度和细粒度的剪枝再配合动态层映射的逐层蒸馏把加速比推到 10 倍的同时精度几乎无损而且绕开了成本高昂的学生模型预训练阶段。这篇论文的完整标题是《Structured Pruning Learns Compact and Accurate Models》代码开源在 princeton-nlp/CoFiPruning。如果你正在负责 transformer 模型的线上推理优化或者想在 GLUE、SQuAD 这类任务上把模型压到原来的十分之一还能保住精度这篇值得读完。2. 剪枝粒度谱系层、头、FFN 维度的加速比天花板2.1 为什么推理加速不能照搬分布式训练训练阶段可以把一个大 batch 切到多张卡上并行计算梯度聚合的通信开销被大量计算掩盖。推理阶段面对的是在线请求batch 通常很小而且延迟敏感多卡通信反而会把 p99 延迟拉高。单卡上唯一可行的方向就是让模型本身变得更快——也就是说减少需要参与计算的参数和计算路径。这里要明确一个概念模型压缩并不等于推理加速。有些压缩方法比如非结构化剪枝确实减少了参数量但生成的稀疏矩阵在 GPU 上无法被高效利用推理时间几乎不变。所以本文讨论的剪枝默认是结构化剪枝——剪完之后模型的张量维度变小计算图可以真正收缩的剪枝。2.2 从 transformer 结构看剪枝落点CoFi 针对的是 transformer 架构的教师模型。transformer 的基本组成单元是 block每个 block 包含一个多头自注意力MHA和两个前馈神经网络FFN。从参数规模看MHA 和 FFN 的参数量比大约是 1:2也就是说 FFN 占了模型参数的大头。需要注意的一个工程细节是在 GPU 上MHA 和 FFN 的推理时间基本相当但在 CPU 上FFN 会耗费更多的推理时间。这意味着剪枝策略不能只盯着参数量还要看目标部署硬件。如果你的模型部署在 CPU 上FFN 层的剪枝收益会比 MHA 更明显如果在 GPU 上两边的剪枝收益都值得挖掘。2.3 四类剪枝粒度的加速比对比CoFi 论文里梳理了剪枝粒度的完整谱系按粒度和加速比可以整理成下表剪枝粒度操作方法典型加速比主要问题层剪枝直接删掉整个 blockMHAFFN删 50% 层约 2X粒度太粗容易伤精度头剪枝用 mask 矩阵保留部分 attention head只留 1 个 head 仅 1.4X加速比不够head 计算路径仍存在FFN 剪枝删掉 FFN 层或 FFN 中间维度视修剪比例而定单独使用无法兼顾 MHA 的冗余细粒度/非结构化剪枝删除 MHA 和 FFN 中更小的块或单个权重难以获得实际加速稀疏矩阵难优化训练难度大2.4 每种粒度的边界在哪里层剪枝是当时研究最多的方向因为实现最简单——把 transformer 的第几层直接拿掉剩下的层照样前向传播。但层剪枝的问题在于粒度太粗一个 block 里有 MHA 也有 FFN可能 MHA 还有冗余而 FFN 已经快被剪光了一刀切必然造成浪费。头剪枝的加速比数据很反直觉即使把所有 head 剪到只剩 1 个加速比也只有 1.4X。原因是 MHA 的计算路径不只是 head 内部的矩阵乘法还有 concatenation 之后的输出投影层这个投影层的计算量并不会因为 head 变少而显著下降。FFN 剪枝针对的是 FFN 的中间维度intermediate size这是 transformer 里参数量最大的部分。但单独剪 FFN 会忽略 MHA 中的冗余而且如果所有层的 FFN 都被剪到同样比例模型结构仍然不够紧凑。2.5 为什么需要粗粒度与细粒度联合从 2.3 的表格可以看出一条规律粒度越粗加速越明显但精度损失越大粒度越细精度控制越精准但加速越难。单独使用任何一种粒度都摸不到 10 倍加速的门槛。CoFi 的核心洞察是把不同粒度的剪枝放进同一个框架里联合优化——用层级别的掩码决定「这一层 MHA/FFN 是否保留」用 head 级别的掩码决定「这一层哪些 head 保留」用维度级别的掩码决定「FFN 的哪些中间维度和隐藏层维度保留」。三个粒度的掩码是联合训练的而不是分别剪完再叠加。这样模型可以自适应地为每一层选择最合适的剪枝策略某一层可能整个 FFN 被剪掉另一层可能只剪几个 head 和部分 FFN 维度。3. CoFi 掩码机制粗粒度与细粒度剪枝的统一框架3.1 从 mask 矩阵到层级别掩码CoFi 的掩码设计是这套方案最精巧的部分。先回顾一下常规的头剪枝做法通过一个由 0/1 组成的 mask 矩阵来保留一部分 head。但 CoFi 的作者注意到一个致命问题——当 mask 矩阵全为 0也就是某一层的所有 head 都被剪掉时模型优化会变得非常困难。梯度通过全零的 mask 回传时这一整层的参数都得不到更新信号后续训练基本就废了。CoFi 的解决方案是为每一层的 MHA 和 FFN 分别引入两个层级别的掩码变量。MHA 的输出可以写成output z_L * MultiHeadAttention(x)其中z_L是层级的掩码变量决定这一层 MHA 是否参与计算。同理FFN 层也有一份自己的层掩码。在训练过程中这些掩码元素的值介于 0 和 1 之间保持可导推理阶段才把低于阈值的掩码变量映射为 0得到最终的剪枝模型。这样设计的一个直接好处是模型可以在训练过程中直接学到「某些层的 MHA 冗余度太高不如整层去掉」而不需要在训练结束后再去做层筛选。3.2 head 维度和 FFN 维度的细粒度掩码层掩码负责「要不要这一层」head 掩码和维度掩码负责「这一层里留哪些部分」。具体来说CoFi 引入了两组细粒度掩码变量head 掩码控制 MHA 内部哪些 head 需要保留。每个 head 对应一个可学习的标量掩码。FFN 维度掩码控制 FFN 中间层的哪些维度需要保留。FFN 的中间维度通常是 3072BERT-base级别每个维度对应一个可学习的标量。这里还有一个容易被忽略的设计CoFi 不仅对 MHA 的 head 和 FFN 的中间维度做剪枝还对 MHA 和 FFN 的输出在隐藏层维度上做剪枝。思路是把维度掩码应用到模型的所有权重矩阵上并且这个掩码是跨层共享的。跨层共享的原因是残差连接的存在——残差结构使得隐向量中的每个维度都能连接到下一层对应的维度所以「哪些维度保留」这个决策应该在所有层间保持一致否则第 3 层保留了第 5 个维度第 4 层却把第 5 个维度剪掉了信息传递会断裂。3.3 预期稀疏度从全局目标反推逐层阈值掩码变量训练好之后怎么确定哪些掩码该置零CoFi 的做法是引入一个全局的预期稀疏度。预期稀疏度定义了剪枝后模型保留的参数比例或剪掉的参数比例它由整个模型的参数量、block 层数、隐藏层维度、MHA 每个 head 的维度、FFN 的中间维度共同决定。在模型训练阶段所有掩码元素都保持连续值参与前向和反向传播。推理阶段阈值由每个权重矩阵的预期稀疏度确定低于阈值的掩码被映射为 0。这个阈值不是拍脑袋定的而是让每一层在满足全局稀疏度约束的前提下保留各自最重要的 head 和维度。一个值得注意的点CoFi 的稀疏度是「预期」的不是「精确」的。因为不同层的冗余程度不同训练出来的掩码分布天然不均匀所以最终每层的剪枝比例会出现差异——这正是作者想要的效果让每层剪到最合适的位置。3.4 掩码参数的初始化与训练流程初次跑 CoFi 的时候最容易犯的错误是把所有掩码初始化为 1。这样做的问题是每一层都认为自己应该被完整保留剪枝过程会变得非常缓慢。常见做法是把掩码初始化在接近 1 的位置但配合一个较大的稀疏度惩罚让模型在训练过程中快速学会「放弃」。下面是一个参考实现框架import torch import torch.nn as nn class CoFiMask(nn.Module): CoFi 的层级掩码定义 - layer_mask: 每层 MHA 和 FFN 是否保留 - head_mask: 每层每个 head 是否保留 - intermediate_mask: 每层 FFN 中间维度是否保留 - hidden_mask: 跨层共享的隐藏维度是否保留 def __init__(self, num_layers, num_heads, intermediate_size, hidden_size): super().__init__() # 初始化时所有掩码接近 1但保留学习空间 self.layer_mask nn.Parameter(torch.ones(num_layers, 2)) self.head_mask nn.Parameter(torch.ones(num_layers, num_heads)) self.intermediate_mask nn.Parameter( torch.ones(num_layers, intermediate_size) ) self.hidden_mask nn.Parameter(torch.ones(num_layers, hidden_size)) def forward(self, threshold0.0): # 训练阶段直接返回连续掩码推理阶段做二值化 if threshold 0: return ( (self.layer_mask threshold).float(), (self.head_mask threshold).float(), (self.intermediate_mask threshold).float(), (self.hidden_mask threshold).float(), ) return self.layer_mask, self.head_mask, self.intermediate_mask, self.hidden_mask掩码的更新取决于蒸馏损失函数——掩码参数和模型参数是同时训练的模型在每一步更新中同时优化「当前剪枝状态下还能不能保留教师模型的知识」这件事。训练过程中稀疏度不是一步到位的而是在一定步数内从 0 逐步增长到目标值即预期的稀疏度这个调度逻辑会在第 5 章详细展开。提示实际运行时注意把掩码参数放入优化器的 param_group 中并且可以考虑给掩码设置独立的学习率。常见做法是掩码学习率比模型参数学习率小一个数量级避免稀疏结构震荡太剧烈。4. 动态层映射蒸馏把逐层蒸馏从死板中解放出来4.1 固定层映射的局限传统的蒸馏方法中教师网络和学生网络的层映射往往是预先定义好的。比如 DistilBERT 的做法是每隔一层从教师模型取一次隐藏状态来对齐学生模型这个映射关系是固定的——学生模型的第 2 层对应教师模型的第 4 层从训练开始到结束都不变。这个固定映射的问题在于剪枝后的学生模型每一层的功能可能已经发生了偏移。比如剪枝后学生模型的第 5 层实际上融合了教师模型第 7 层和第 8 层的功能这时候还硬要拿教师模型的第 5 层去对齐学生模型的第 5 层蒸馏损失就会给出互相矛盾的梯度信号。4.2 动态搜索层映射CoFi 提出了一种动态层映射机制。假设教师模型有一组备选的层集合需要从中选出哪些层参与向学生模型的知识蒸馏同时还要确定教师层的第 i 层映射到学生层的第 j 层。这个映射不是训练前就写死的而是在蒸馏过程中动态搜索的。具体做法是定义一个层映射函数表示教师模型的第 i 层映射到学生模型的第 j 层。隐藏层的蒸馏损失可以定义为学生模型第 j 层的隐藏表示经过一个线性变换后与教师模型第 i 层的隐藏表示的 MSE 距离。这个线性变换的作用是把学生模型的隐藏维度映射到教师模型的维度空间因为剪枝后的学生模型隐藏维度通常比教师模型小。关键细节是这个层映射函数是使蒸馏损失最小化的方向动态选择的——计算两层集合之间的距离时MSE 操作可以并行执行。这样每一轮训练中蒸馏目标始终指向「当前状态下学生模型最需要对齐的教师层」而不是机械地按固定索引对齐。4.3 蒸馏损失函数的完整形式CoFi 的最终损失函数由两部分组成逐层蒸馏损失和输出蒸馏损失。输出蒸馏就是常规的 soft label 蒸馏——学生模型的输出概率分布与教师模型的输出概率分布之间的 KL 散度。总损失是两者的加权和其中权重的设置会显著影响剪枝效果。这里给出一个参考实现思路def cofi_distillation_loss( student_hidden_states, # 学生模型各层隐藏状态 teacher_hidden_states, # 教师模型各层隐藏状态 student_logits, teacher_logits, temperature2.0, layer_loss_weight1.0, output_loss_weight1.0, ): # 1. 动态层映射为学生模型的每一层找最匹配的教师层 layer_loss 0.0 matched_teacher_layers set() for student_layer_idx, student_hidden in enumerate(student_hidden_states): best_loss float(inf) best_teacher_idx -1 for teacher_idx, teacher_hidden in enumerate(teacher_hidden_states): if teacher_idx in matched_teacher_layers: continue # 线性变换后计算 MSE这里简化省略了变换矩阵 mse torch.nn.functional.mse_loss( student_hidden, teacher_hidden ) if mse best_loss: best_loss mse best_teacher_idx teacher_idx matched_teacher_layers.add(best_teacher_idx) layer_loss best_loss # 2. 输出层蒸馏KL 散度 student_log_probs torch.nn.functional.log_softmax( student_logits / temperature, dim-1 ) teacher_probs torch.nn.functional.softmax( teacher_logits / temperature, dim-1 ) output_loss torch.nn.functional.kl_div( student_log_probs, teacher_probs, reductionbatchmean ) return layer_loss_weight * layer_loss output_loss_weight * output_loss这段代码的要点是层映射的搜索是贪心的——从学生模型的第一层开始逐层找当前未匹配的教师层中 MSE 最小的那一层。虽然论文中描述为动态搜索实际实现中贪心匹配是常见且有效的简化方式。4.4 为什么能绕开无标签数据预训练DistilBERT 这一类蒸馏方法能获得高精度很大程度上是因为他们用大量无标签数据先预训练了学生模型。这个预训练阶段通常需要几天甚至几周的 GPU 时间对工业团队来说几乎不可接受。CoFi 不需要预训练学生模型因为剪枝本身就是从原模型出发的——剪出来的模型在初始时就继承教师模型的全部参数剪枝过程中每一层都在和教师模型的对应层做对齐相当于把预训练和蒸馏融合在了一个阶段里。论文中的实验数据显示CoFi 在训练时间上相比 DistilBERT 有大幅缩减。关键在于 CoFi 的掩码是连续的梯度可以在掩码参数上正常回传模型在剪枝过程中始终能接收到「如何逼近教师模型」的信号不存在预训练阶段的知识断层。提示如果你的目标任务数据量很小建议增大输出蒸馏的权重。逐层蒸馏在小数据集上容易过拟合到隐藏状态的细节而输出蒸馏提供的是更鲁棒的类别分布信号。5. CoFi 训练全流程稀疏度调度与关键超参设置5.1 不同规模数据集的训练策略差异CoFi 在 GLUE 和 SQuAD v1.1 上做了完整的实验验证但它在不同规模数据集上的训练调度策略并不相同。这个差异特别值得注意——如果你直接把小数据集的策略套到大模型上训练时间会浪费很多。对于 MNLI、QNLI、SST-2、QQP 这四个相对较大的 GLUE 数据集以及 SQuAD 数据集CoFi 的训练流程是先训练 20 个 epoch再额外微调 20 个 epoch。前 20 个 epoch 内先用蒸馏目标函数微调模型 1 个 epoch然后在接下来的 2 个 epoch 内让模型达到目标稀疏度之后持续剪枝搜索更优结构。对于 MRPC、CoLA、STS-B、RTE 这四个较小的 GLUE 数据集训练总时长增加到 100 个 epoch 并微调 20 个 epoch。前 4 个 epoch 用蒸馏目标函数微调接下来的 20 个 epoch 逐步剪到目标稀疏度。小数据集上的剪枝需要更多步数来稳定搜索因为数据量少、梯度噪声大剪枝收敛天然更慢。5.2 稀疏度调度的时间表把上面的策略抽象成一张表方便对照数据集规模蒸馏微调达到目标稀疏度继续剪枝搜索最终微调大数据集MNLI/QNLI/SST-2/QQP/SQuAD1 epoch2 epoch 内剩余 epoch额外 20 epoch小数据集MRPC/CoLA/STS-B/RTE4 epoch20 epoch 内剩余 epoch额外 20 epoch稀疏度从 0 逐步增长到目标值这个增长过程中模型有足够的时间去探索哪些 head、哪些维度可以被移除。达到目标稀疏度后的「继续剪枝」阶段是 CoFi 的一个独特设计——目标稀疏度只是一个下限模型可以继续朝着更稀疏的方向搜索只要精度不崩。在代码层面稀疏度调度的实现可以参考下面的伪代码def get_sparsity(step, total_steps, warmup_steps, target_sparsity): 稀疏度从 0 线性增长到 target_sparsity warmup_steps 内保持稀疏度为 0让模型先适配蒸馏目标 if step warmup_steps: return 0.0 progress (step - warmup_steps) / max(total_steps - warmup_steps, 1) return min(target_sparsity, target_sparsity * progress)这里的warmup_steps要对应蒸馏微调的 epoch 数——论文里的「先用蒸馏目标函数微调 1 个 epoch」本质上就是让稀疏度保持为 0模型先学习如何匹配教师模型的输出分布。5.3 训练结束后的微调精度恢复的关键一步论文里反复强调了一个实验发现训练结束后的微调阶段对最终精度影响极大。原因在于剪枝过程中模型始终带着稀疏度约束在优化被裁剪的维度依然会产生梯度信号而微调阶段不再施加稀疏度惩罚模型可以在当前稀疏结构下充分收敛到更优的参数空间。实际操作中微调阶段应该把掩码参数冻结只更新模型本身的权重。否则掩码继续变化模型结构还在变微调就失去了「稳定结构、优化参数」的意义。5.4 完整训练流程的参考命令CoFi 的开源实现基于 HuggingFace Transformers训练流程大致如下# 1. 在目标数据集上微调教师模型如果直接使用预训练模型可跳过 python run_glue.py \ --model_name_or_path bert-base-uncased \ --task_name mnli \ --do_train \ --do_eval \ --output_dir ./teacher_model # 2. 用 CoFi 进行剪枝蒸馏 python run_glue.py \ --model_name_or_path ./teacher_model \ --task_name mnli \ --do_train \ --do_eval \ --output_dir ./cofi_pruned_model \ --distill_teacher ./teacher_model \ --pruning_type cofi \ --target_sparsity 0.8 \ --sparsity_warmup_epoch 1 \ --sparsity_adjust_epoch 2 \ --num_train_epochs 20 \ --per_device_train_batch_size 32 \ --learning_rate 3e-5 \ --mask_learning_rate 3e-6 # 3. 剪枝完成后微调 python run_glue.py \ --model_name_or_path ./cofi_pruned_model \ --task_name mnli \ --do_train \ --do_eval \ --output_dir ./cofi_pruned_model_finetuned \ --num_train_epochs 20 \ --learning_rate 3e-5命令中几个关键参数的理解target_sparsity是目标稀疏度论文中对比了不同稀疏度下的加速比和精度80% 的稀疏度大致对应把模型压到原来的五分之一。sparsity_warmup_epoch对应蒸馏微调的 epoch 数sparsity_adjust_epoch对应达到目标稀疏度所需的时间窗口。5.5 训练成本对比论文的实验结果中CoFi 在相同加速比和模型大小下取得了比现有方法更高的准确率。在与 DistilBERT 的对比中CoFi 能达到基本一致的加速比和略高的准确率但训练时间大幅缩减——因为 DistilBERT 需要额外的大规模无标签数据预训练阶段这部分成本在 CoFi 中直接省掉了。对比结果如下方法加速比精度是否需要无标签数据预训练DistilBERT与 CoFi 基本一致略低于 CoFi需要成本极高CoFi最高 10X几乎无损不需要提示learning_rate和mask_learning_rate的比例很关键。如果掩码学习率过大稀疏结构会在训练早期就固化后面的蒸馏阶段变成「在错误结构上死磕」如果过小稀疏度增长跟不上调度表训练结束后还没有收敛到目标稀疏度。6. 复现实验中的三个发现FFN 冗余、上层偏好与微调效应6.1 消融实验如何定位每个模块的贡献论文做了一系列消融实验来验证 CoFi 各个模块的有效性。第一组消融是剪枝单元——分别去掉层剪枝、head 剪枝、FFN 维度剪枝观察精度变化。实验结果表明缺少任何一层粒度的剪枝都会带来明显的精度下降说明粗粒度和细粒度剪枝是互补关系而不是可以互相替代的关系。第二组消融针对蒸馏目标函数。分别去掉逐层蒸馏、改成固定层映射、去掉输出蒸馏对比最终剪枝效果。动态层映射的收益在低稀疏度下不明显但在高稀疏度比如 90%下差异显著——结构被剪得很厉害时固定的层映射会产生大量无效对齐而动态映射始终能为每一层找到「当前最值得学习」的教师层。6.2 剪枝后结构分析FFN 冗余最多CoFi 的一个很有趣的实验结论是剪枝后的模型结构在不同数据集上呈现出明显差异但存在两个共同规律。第一几乎在所有的稀疏度设置下FFN 层的剪枝比例都明显高于 MHA 层。这说明 FFN 层承载了更多的冗余信息——这和参数规模的直觉一致FFN 的参数量是 MHA 的两倍中间维度有大量冗余。从部署的角度看如果你的模型运行在 CPU 上优先剪 FFN 会收获更大的推理加速比。第二CoFi 倾向于更多地剪掉上层的网络结构。这个现象和学术界观察到的高层语义冗余是一致的——靠近输出的层更多地在适配特定任务而底层保留的是更通用的语言表示。这个发现对实际操作有直接指导意义如果不想跑完整的 CoFi 剪枝流程一个低成本的手工替代方案是优先剪掉后 1/3 层的 FFN 维度。6.3 微调阶段对精度恢复的具体量化实验数据表明训练结束后的微调对模型精度保持起着关键作用。对比「剪枝后不微调直接评估」和「微调 20 个 epoch 后评估」大部分数据集上微调能带来不小的精度恢复。原因可以拆成两层剪枝过程的目标函数是「蒸馏损失 稀疏度惩罚」的组合此时的模型参数是为了满足这个复合目标而优化的去掉稀疏度惩罚、只优化任务损失之后模型参数的分布可以重新适应任务本身的难度分布。这里有一个实用的验证技巧。剪枝训练完成后不要急着微调先做一步操作把每层每个 head 的 mask 值打印出来统计剩余数量。这个分布会告诉你模型认为哪些层应该保留、哪些层应该剪掉。如果某个数据集的剪枝结果中某一层的 MHA 被完全剪掉但在验证集上精度并没有明显下降说明这个 head 集合确实是冗余的。import torch def analyze_masks(model): 打印 CoFi 剪枝后每层的 mask 分布 for name, param in model.named_parameters(): if mask not in name: continue mask_binary (param 0).float() kept_ratio mask_binary.mean().item() print( f{name}: 保留比例 {kept_ratio:.2%}, f绝对数量 {int(mask_binary.sum().item())} )数据增强后的对比实验同样值得参考作者对 CoFi 和 DistilBERT 分别做了相同的数据增强后CoFi 依然保持了更高的性能。这说明 CoFi 的优势不是数据增强带来的而是剪枝和蒸馏联合优化这个机制本身带来的——在处理低资源任务的模型压缩时这个结论有直接的参考价值。本文还有配套的精品资源点击获取