CoFi剪枝:结构化剪枝结合知识蒸馏,实现Transformer十倍推理加速
2026/9/18 23:32:22 网站建设 项目流程

简介:一份面向自然语言处理研究者与模型部署工程师的ACL'22论文PDF,聚焦陈丹琦团队提出的CoFi结构化剪枝方法。该方法面向Transformer架构,结合粗粒度与细粒度剪枝及逐层蒸馏,在仅损失极少精度的情况下将模型推理加速约10倍,缓解了模型压缩中加速比不高或训练成本过高的难题。资源为单文件PDF,共1个文件,格式为pdf,包体大小约2.68MB,内容完整包含论文原文及图文排版,适合用于精读算法细节、复现实验或梳理剪枝方法演进。已有191人学习/下载,是了解当前NLP模型轻量化前沿思路的实用参考。通过阅读可直接获取CoFi方法对层剪枝、头剪枝、FFN剪枝等不同粒度的处理方案,以及动态层映射的蒸馏策略,有助于快速掌握该方法的动机、公式推导与实验设计。

1. 推理加速的硬约束:为什么剪枝比蒸馏更省成本

模型部署到线上之后,推理延迟直接决定服务成本和用户体验。分布式训练能通过多卡并行把训练时间压下来,但推理阶段面对的是单条请求,多卡并行带来的收益非常有限—— GPU 之间的通信开销甚至会抵消掉并行带来的收益。真正能在推理阶段见效的,是把模型本身变小。这就是模型压缩的用武之地,而模型压缩的两条主线——剪枝和蒸馏——各有各的瓶颈。

剪枝的思路是砍掉贡献度低的参数,常规的结构化剪枝只能做到 2~3 倍加速,再往上精度就会明显下滑;蒸馏可以做到更大的加速比,但通常需要先用大量无标签数据预训练一个学生模型,这个预训练阶段的成本高到让很多团队直接放弃。陈丹琦团队在 ACL'22 提出的 CoFi(Coarse- 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 论文里梳理了剪枝粒度的完整谱系,按粒度和加速比可以整理成下表:

剪枝粒度操作方法典型加速比主要问题
层剪枝直接删掉整个 block(MHA+FFN)删 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 的中间维度通常是 3072(BERT-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, threshold=0.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, temperature=2.0, layer_loss_weight=1.0, output_loss_weight=1.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, reduction="batchmean" ) 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/SQuAD)1 epoch2 epoch 内剩余 epoch额外 20 epoch
小数据集(MRPC/CoLA/STS-B/RTE)4 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_ratemask_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 的优势不是数据增强带来的,而是剪枝和蒸馏联合优化这个机制本身带来的——在处理低资源任务的模型压缩时,这个结论有直接的参考价值。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询