企业官网建设流程全解析

一、论文基本信息

论文题目：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 Pruning

Local Pruning 的做法是：每次只考虑删除一个具体位置的 head，例如第 3 层第 7 个 head。删除后重新评估模型性能，选择精度下降最小的 head 永久删除。

它的优点是比较细粒度。

问题是计算成本很高。BERT-base 有 12 层、每层 12 个 heads，总共 144 个 heads。如果每轮都对所有未剪 heads 重新评估，搜索次数会非常多。论文指出，对于 BERT-base，Local 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)

实验任务是情感分类，使用两个二分类数据集：

论文说明，这两个数据集都是 positive / negative 二分类，并且类别平衡。(arXiv)

Fine-tuning 设置如下：

论文使用 Hugging Face 的 BERT-base 实现，并使用 Adam 优化器进行 fine-tuning。(arXiv)

八、实验结果解读

8.1 约 40% heads 可以无损删除

论文最重要的结果是：在 Amazon 和 IMDB 两个情感分类任务上，APruning 都能删除约 40% attention heads，而不降低测试精度。*

具体来说：

论文解释说，一些 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 结果：

IMDB 结果：

论文总结说，在更大预算下，Amazon 可以剪到约67% heads，IMDB 可以剪到约60% heads，但精度会有相应下降。(arXiv)

这个结果说明：

APruning 可以提供一条 head 数量和 accuracy 之间的折中曲线。*

如果部署场景非常重视精度，可以只剪 40% 左右；如果部署场景更重视模型压缩，可以接受 1–3 个百分点的精度下降，剪掉更多 heads。

8.3 A* Pruning 明显优于 Random Pruning

论文还用 Random Pruning 做对比。随机剪枝的做法是：在相同预算下随机选择 heads 删除，并重复 100 次。

实验显示，随机剪枝远不如 A* Pruning。以 Amazon budget=3 为例，随机剪枝最可能只能删掉约 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 importance，A* Pruning 的评估成本更高。

10.3 没有深入分析 head 功能

这篇论文主要关注剪枝搜索和压缩效果，不分析被保留 heads 是否具有句法、位置、语义等功能。

所以它回答的是：哪些 heads 可以删。

但没有深入回答：为什么这些 heads 可以删，它们在模型中原本做什么。

10.4 只剪 attention heads，不剪 FFN

BERT 中 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 冗余，但剪枝顺序必须经过搜索而不能随机选择。**