一、论文基本信息论文题目Pruning Attention Heads of Transformer Models Using ASearch: A Novel Approach to Compress Big NLP Architectures*作者Archit Parnami、Rahul Singh、Tarun Joshi发表形式arXiv 2021研究对象BERT-base 中的 self-attention heads任务文本情感分类数据集Amazon Reviews、IMDB Reviews这篇论文研究的是BERT attention head 剪枝。它的核心目标不是重新设计 Transformer也不是压缩 FFN 层而是直接从 BERT 的 self-attention 层中删除冗余 heads。论文指出BERT-base 有 12 个 encoder layers每层 12 个 attention heads总共 144 个 heads作者希望在尽量不损失分类精度的情况下删除其中一部分 heads从而压缩模型。论文摘要中给出的主要结论是该方法最多可以删除 BERT 中约40% attention heads并且不损失精度。这篇论文的核心方法是APruning*。它把 head pruning 看成一个搜索问题在有限精度损失预算内搜索一组可以删除的 attention heads。与简单贪心逐个测试所有 heads 不同A* Pruning 使用启发式搜索来减少需要重新评估的候选 heads 数量。二、论文要解决的问题BERT 这类 Transformer 模型参数量大、推理延迟高、内存占用高不适合直接部署到资源受限设备上。论文在引言中提到大型 Transformer 模型会带来高内存占用、高预测延迟、高功耗、资源受限设备推理性能差、训练困难等问题。(arXiv)传统剪枝方法通常有两类一种是基于权重大小的 magnitude pruning另一种是 Hessian / 二阶信息剪枝。二阶方法更精确但计算昂贵因此实际中更常用 magnitude-based 方法。论文认为对于 BERT 这类迁移学习模型如何在下游任务中有效剪枝还没有被充分研究。(arXiv)这篇论文要解决的问题可以概括为如何在 BERT fine-tuned downstream model 中删除尽可能多的 attention heads同时保证精度下降不超过用户设定的预算这里有两个关键点。第一剪枝对象是完整 attention head不是单个权重。因此它属于结构化剪枝。第二剪枝过程受精度预算约束。也就是说用户可以指定最多允许损失多少 accuracy算法在这个范围内尽量多剪 heads。三、核心思想这篇论文的核心思想可以概括为一句话把 attention head pruning 转化成一个带精度预算约束的搜索问题用 ASearch 寻找低代价的 head 删除序列。*普通 head pruning 可以这样做每次尝试删除一个 head评估模型精度下降多少然后删掉影响最小的那个 head。这个思路直观但非常慢因为每删一个 head 后剩余 heads 的重要性可能会改变需要重新评估所有候选 heads。论文把这个过程理解为搜索每个候选 head 都对应一个可能的剪枝动作。删除某个 head 的代价是模型精度下降。总预算是允许的最大精度损失。目标是在预算内删除尽可能多的 heads。A* Pruning 的作用是不对所有剩余 heads 都反复完整评估而是用启发式规则提前排除那些不太可能满足预算的 heads从而降低搜索开销。四、Local Pruning 和 Global Pruning论文先讨论了两种直接剪枝策略。4.1 Local PruningLocal Pruning 的做法是每次只考虑删除一个具体位置的 head例如第 3 层第 7 个 head。删除后重新评估模型性能选择精度下降最小的 head 永久删除。它的优点是比较细粒度。问题是计算成本很高。BERT-base 有 12 层、每层 12 个 heads总共 144 个 heads。如果每轮都对所有未剪 heads 重新评估搜索次数会非常多。论文指出对于 BERT-baseLocal Pruning 最坏可能需要10,440 次搜索 / 重评估。(arXiv)所以 Local Pruning 的问题是剪得精细但搜索成本太高。4.2 Global Pruning论文中说的 Global Pruning 不是“全模型任意 head 排序剪枝”而是指一次性剪掉所有层中同一个编号的 head。例如把每层的第 5 个 head 都剪掉这样一次删除 12 个 heads。它的优点是搜索次数少。但缺点也很明显它默认不同层中相同编号的 head 具有相似重要性。这个假设并不可靠。第 1 层第 5 个 head 和第 10 层第 5 个 head 可能功能完全不同强行一起剪可能造成较大精度损失。因此Global Pruning 更快但更粗糙Local Pruning 更精确但太慢。论文的 A* Pruning 试图在两者之间取得平衡。(arXiv)五、A* Pruning 的核心机制A* Pruning 的目标是保留 Local Pruning 的细粒度同时减少搜索次数。论文定义了几个关键变量Baseline Accuracy剪枝前模型精度。Budget用户允许损失的最大精度。Accuracy Post Pruning删除某个 head 后的精度。Cost删除某个 head 造成的精度下降。Heuristic估计下一轮删除某个 head 的代价。Search Space当前仍然需要检查的候选 heads。(arXiv)算法每一轮做的事情是先评估当前候选 heads 中每个 head 被删除后的精度。计算每个 head 的删除代价。选择代价最小的 head 删除。更新剩余精度预算。用启发式规则缩小下一轮搜索空间。论文中的启发式假设是随着已经剪掉的 heads 变多再继续剪同一个候选 head 的代价通常不会变小而是会变大。因此当前轮的代价可以作为下一轮代价的低估计。这个估计不会超过真实代价所以可以用来筛掉明显不可能在预算内保留的候选。(arXiv)这就是 A* 在这篇论文里的作用不是用梯度或权重大小直接打分而是用“实际剪掉后精度下降”作为代价再用启发式搜索减少重评估次数。六、它和前面几篇 Head Pruning 论文的区别前面一些 head pruning 论文常用的是重要性打分。例如看删除某个 head 后 loss 变化。看 loss 对 head mask 的梯度敏感度。这篇论文不同它不重点分析 head 的语言功能也不主要用 Taylor gradient importance。它更偏工程搜索直接用剪掉 head 后的 accuracy drop 作为代价。用 ASearch 在精度预算内寻找可剪 heads。所以这篇论文的重点不是解释 Transformer 里哪些 heads 有功能而是给定一个 BERT 下游分类模型如何在精度约束下尽量多删 heads。七、实验设置论文使用BERT-base作为目标模型。BERT-base 有 12 个 encoder layers每层包含 self-attention 层和 dense / FFN 层每个 self-attention 层有 12 个 heads因此总共 144 个 heads。论文只剪 self-attention heads不剪 FFN 层。(arXiv)实验任务是情感分类使用两个二分类数据集数据集任务类型类别数平均长度最大长度训练样本测试样本Amazon ReviewsSentiment22445019200,00025,000IMDB ReviewsSentiment2294354125,00025,000论文说明这两个数据集都是 positive / negative 二分类并且类别平衡。(arXiv)Fine-tuning 设置如下数据集剪枝前测试精度Sequence LengthBatch SizeLearning RateFine-tune EpochsAmazon92.46%64641e-52IMDB88.14%128641e-52论文使用 Hugging Face 的 BERT-base 实现并使用 Adam 优化器进行 fine-tuning。(arXiv)八、实验结果解读8.1 约 40% heads 可以无损删除论文最重要的结果是在 Amazon 和 IMDB 两个情感分类任务上APruning 都能删除约 40% attention heads而不降低测试精度。*具体来说数据集无损剪枝 heads 数量无损压缩比例剪枝后精度Amazon约 57–60 个 heads约 39.58%–41.66%92.46%IMDB约 58–60 个 heads约 40.00%–41.66%88.14%论文解释说一些 heads 被删除后反而会使模型精度提升因此这些 heads 的“cost”被视作 0这也说明 fine-tuned BERT 在情感分类任务上存在明显过参数化。(arXiv)这个结果的含义是对 Amazon / IMDB 情感分类来说BERT-base 中相当一部分 attention heads 不是必要结构。8.2 在允许小幅精度下降时可以剪掉更多 heads当用户允许 1%、2%、3% 这样的 accuracy budget 时A* Pruning 可以剪掉更多 heads。Amazon 结果预算实际使用预算剪掉 heads压缩比例剪枝后精度10.57652.77%91.96%21.598659.72%90.87%32.729666.66%89.74%IMDB 结果预算实际使用预算剪掉 heads压缩比例剪枝后精度10.046651.76%87.28%21.427652.77%86.72%32.608659.72%85.54%论文总结说在更大预算下Amazon 可以剪到约67% headsIMDB 可以剪到约60% heads但精度会有相应下降。(arXiv)这个结果说明APruning 可以提供一条 head 数量和 accuracy 之间的折中曲线。*如果部署场景非常重视精度可以只剪 40% 左右如果部署场景更重视模型压缩可以接受 1–3 个百分点的精度下降剪掉更多 heads。8.3 A* Pruning 明显优于 Random Pruning论文还用 Random Pruning 做对比。随机剪枝的做法是在相同预算下随机选择 heads 删除并重复 100 次。实验显示随机剪枝远不如 A* Pruning。以 Amazon budget3 为例随机剪枝最可能只能删掉约 30 个 heads而 A* Pruning 可以删掉 96 个 heads随机剪超过 80 个 heads 的概率几乎为 0。(arXiv)这个对比很关键不是任意 heads 都可以删。BERT 中确实有冗余 heads但冗余不是均匀分布的。如果随机删除很容易删到重要 heads导致预算迅速耗尽。所以这篇论文真正要证明的是只要剪枝顺序选得好BERT 可以删很多 heads但如果剪枝顺序不好压缩空间会小得多。8.4 剪枝后参数量明显下降论文还展示了 Amazon 数据集上压缩后的参数量变化。原始 BERT 模型参数量约 340M经过 A* Pruning 后参数量可以逐步降低到 199M、161M、137M、114M同时精度从 92.46% 逐步变化到 92.46%、91.96%、90.87%、89.74%。(arXiv)这里要谨慎理解。从标准 BERT-base 结构来看BERT-base 通常约 110M 参数论文图中报告的原始参数量为 340M这一点与常见 BERT-base 参数量说法不一致。既然论文图表如此报告解读时应以论文内部结果为准但复现时需要核对作者具体统计口径。更重要的是方法层面head pruning 能减少模型结构中的 attention 计算分支但 Transformer 的 FFN 层仍然保留。因此如果目标是大幅压缩整个 BERT仅剪 heads 还不够通常还需要结合 FFN pruning、layer pruning、distillation 或 quantization。九、方法优点9.1 剪枝目标很清晰这篇论文的目标非常明确在给定精度损失预算内尽量多删除 attention heads。这比只报告某个固定压缩率更适合实际部署因为不同场景对精度损失的容忍度不同。9.2 使用真实精度下降作为剪枝代价很多剪枝方法使用 proxy score例如权重大小、梯度、attention pattern 等。A* Pruning 直接用剪掉某个 head 后模型 accuracy 的变化作为 cost。这个指标直观而且和最终任务性能直接对应。9.3 比完整 Local Pruning 更省搜索Local Pruning 每一轮都重新评估所有候选 heads成本非常高。A* Pruning 用启发式规则缩小搜索空间减少不必要评估同时仍然保持细粒度 head-level pruning。9.4 证明 BERT 下游任务中 head 冗余明显在 Amazon 和 IMDB 情感分类任务上约 40% heads 可以无损删除这说明 fine-tuned BERT 在这些任务上存在明显过参数化。(arXiv)9.5 比随机剪枝可靠得多随机剪枝结果说明head redundancy 不是随便删都可以利用的。A* Pruning 能找到更好的删除顺序因此在同等预算下能剪掉更多 heads。(arXiv)十、方法局限10.1 只在两个情感分类数据集上验证论文实验只使用 Amazon Reviews 和 IMDB Reviews 两个二分类任务。虽然这能证明方法在情感分类上有效但不能直接说明它在 MNLI、SQuAD、NER、机器翻译、摘要等任务上也同样有效。BERT 的 head 重要性很可能是任务相关的。某些 heads 对情感分类不重要但对句法、推理、问答任务可能重要。10.2 需要大量模型重评估A* Pruning 比完整 Local Pruning 更省但它仍然需要反复删除候选 head 并评估模型精度。表 3 中Amazon / IMDB 不同预算下的搜索次数仍然在数千级别。(arXiv)因此它不属于低成本剪枝方法。相比 Taylor gradient 一次 forward/backward 得到 head importanceA* Pruning 的评估成本更高。10.3 没有深入分析 head 功能这篇论文主要关注剪枝搜索和压缩效果不分析被保留 heads 是否具有句法、位置、语义等功能。所以它回答的是哪些 heads 可以删。但没有深入回答为什么这些 heads 可以删它们在模型中原本做什么。10.4 只剪 attention heads不剪 FFNBERT 中 FFN 层通常占据大量参数和计算。只剪 attention heads 虽然能减少部分计算但整体压缩仍有限。如果目标是实际部署加速仅 head pruning 可能不够需要和其他压缩方法组合。10.5 “精度保证”依赖验证 / 测试评估过程论文说 A* Pruning 在预算内提供 accuracy guarantee本质上是基于反复评估得到的 empirical guarantee。它保证的是在当前数据集评估流程下不超过预算而不是理论上对所有未来样本都保证。如果测试分布变化剪枝后的模型性能仍可能变化。十一、整体评价这篇论文的核心价值在于它把 BERT attention head pruning 明确建模成一个“有精度预算的搜索问题”。相比基于梯度或解释性的 head importance 方法它更加直接删除某个 head 是否值得不看它的权重大小也不看它是否有语言学功能而是看它被删除后对下游任务 accuracy 的实际影响。它的结果说明BERT-base 在情感分类任务中存在大量冗余 attention heads。约 40% heads 可以无损删除。如果允许 1–3 个百分点的精度下降可以删除 50%–67% heads。剪枝顺序非常重要随机删远不如 A搜索。*不过这篇论文更像是一个搜索式 head pruning 方法验证而不是完整的 Transformer 压缩方案。它的实验范围较窄搜索成本较高也没有覆盖 FFN、层数、hidden size 等更大的压缩空间。十二、一句话总结《Pruning Attention Heads of Transformer Models Using ASearch》提出把 BERT attention head 剪枝建模为带精度预算的搜索问题用 ASearch 在有限 accuracy loss 内寻找可删除 heads实验表明在 Amazon 和 IMDB 情感分类任务上BERT-base 中约 40% heads 可以无损删除若允许小幅精度下降则可以删除更多 heads说明下游 BERT 模型中存在明显 attention head 冗余但剪枝顺序必须经过搜索而不能随机选择。**