Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned

📅 2026/7/18 4:58:58 👁️ 阅读次数 📝 编程学习
Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned

一、论文基本信息

论文题目:Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned

作者:Elena Voita、David Talbot、Fedor Moiseev、Rico Sennrich、Ivan Titov

发表会议:ACL 2019

研究对象:Transformer 机器翻译模型中的 attention heads

官方代码:lena-voita/the-story-of-heads

这篇论文研究的是 Transformer 多头注意力中的一个核心问题:不同 attention heads 对模型预测的贡献是否一样?如果不一样,哪些 heads 真正重要,哪些 heads 可以被剪掉?论文最终发现,只有少数 heads 承担了明显重要且可解释的功能,而大量 heads 可以被剪掉,对翻译质量影响很小。ACL Anthology 摘要中给出的典型结果是:在 English–Russian WMT 数据集上,剪掉 encoder 中 48 个 heads 里的 38 个,BLEU 只下降 0.15。


二、论文要解决的问题

Transformer 的 multi-head attention 通常被认为可以让模型从不同子空间、不同关系角度理解序列。直觉上,每个 head 都应该有自己的作用。

但这篇论文提出一个更细的问题:Transformer 中的每个 attention head 是否真的都重要?

以前很多 attention 可解释性分析会看所有 heads 的平均 attention,或者只看最大 attention 权重。但这样会掩盖一个问题:不同 heads 的重要性可能差别很大。有些 head 可能承担了关键功能,有些 head 可能只是冗余结构。论文明确指出,简单平均 heads 或只看最大 attention,都没有显式考虑不同 heads 的重要性差异。(ACL Anthology)

所以这篇论文主要回答四个问题:

第一,翻译质量到底依赖哪些 encoder heads?

第二,重要 heads 是否具有稳定且可解释的语言功能?

第三,不同 attention 类型对 head 数量的敏感性是否不同?

第四,能否大量剪掉 attention heads,同时保持翻译质量?


三、核心结论

这篇论文最重要的结论是:Transformer 中只有少数 attention heads 真正承担关键功能,大多数 heads 对最终翻译质量不是不可替代的。

更具体地说,作者发现重要 heads 往往具有比较明确的功能,例如:

关注相邻词的位置型 heads。

关注特定句法依存关系的 syntactic heads。

关注句子中低频词的 rare-word heads。

而在剪枝实验中,这些具有明确功能的 heads 往往最后才被剪掉。这说明它们确实比普通 heads 更重要,不只是可视化上看起来有规律。论文结论中也明确说,specialized heads 是最后被剪掉的,从而直接验证了它们的重要性。(ACL Anthology)

这篇论文的题目“Specialized Heads Do the Heavy Lifting”其实就是在说:真正重要的是少数专门化 heads,它们承担了主要工作;其余大量 heads 可以被剪掉。


四、实验数据和模型设置

论文主要研究机器翻译任务。源语言固定为 English,目标语言包括 Russian、German 和 French。作者为了控制训练数据规模,对每个语言对都使用相同数量的 WMT 句对,数量为2.5M sentence pairs;此外,English–Russian 还额外使用 OpenSubtitles2018 语料做跨领域分析。

模型方面,论文采用Transformer-base设置:encoder 和 decoder 都是 6 层,每个 multi-head attention 有 8 个 heads,模型维度为 512,FFN 维度为 2048。因此 encoder self-attention 一共有:

6 层 × 8 heads = 48 个 encoder heads

这也是论文中经常说“48 个 encoder heads”的来源。(ACL Anthology)

这点要注意:这篇论文不是用 Transformer-large,而是用 Transformer-base。


五、Head 重要性是怎么判断的

论文用了两种互补思路来判断 head 重要性。

第一种是confidence,置信度

一个 head 如果经常把大部分 attention mass 放在某一个 token 上,说明它的注意力分布很集中。论文把 head confidence 定义为:在开发集上,对每个 token 取该 head 最大 attention 权重,再求平均,并且排除 EOS。直观理解就是:这个 head 平时是不是很明确地关注某个 token。

但作者没有只依赖 confidence,因为 attention 分布集中不一定等于对最终预测重要。所以论文又用了第二种方法:Layer-wise Relevance Propagation,LRP。

LRP 用来衡量某个 head 对模型最终预测的贡献。论文具体做法是:计算每个 head 输出中的神经元对模型 top-1 logit 的贡献,把一个 head 内所有神经元的 relevance 加起来,并在同层 heads 之间归一化;最后在开发集所有生成步上取平均,得到这个 head 的总体 relevance。

因此,这篇论文中的 head 重要性不是简单看 attention 图,而是更接近:这个 head 对模型最终预测到底贡献了多少。


六、重要 heads 有哪些功能

论文最有价值的部分之一,是它不仅说“哪些 heads 重要”,还进一步分析“这些 heads 在做什么”。

作者总结出三类主要功能。

6.1 Positional Heads:关注相邻词

第一类是positional heads

这类 head 会稳定地关注相邻位置,比如前一个词或后一个词。论文把一个 head 认定为 positional head 的标准是:至少 90% 的时间,它的最大 attention 权重都落在某个固定相对位置上,实际中通常是 -1 或 +1,也就是相邻 token。

这种 head 的作用可以理解为:在 Transformer 没有显式递归结构的情况下,帮助模型捕捉局部顺序信息。

论文还发现,positional heads 通常也是高 confidence heads,并且往往被 LRP 排得很重要;所有语言对中的 positional heads,其平均最大 attention 权重都超过 0.8。

这说明位置型 head 不是偶然出现的注意力模式,而是模型稳定使用的一类功能结构。


6.2 Syntactic Heads:关注句法依存关系

第二类是syntactic heads

作者检查了几个依存关系,包括 nominal subject、direct object、adjectival modifier 和 adverbial modifier。方法是先用 CoreNLP 对 held-out sentences 生成依存结构,然后看每个 head 最大 attention 是否指向与当前 token 存在特定依存关系的 token。

如果某个 head 在某类依存关系上的命中率,比“只看最常见相对位置”的 baseline 高至少 10%,作者就认为这个 head 具有 syntactic function。

论文发现,确实有一些 heads 学会了明显高于位置 baseline 的句法关系检测能力。例如在 English–Russian 实验中,direct object 关系上,最好的 head 能明显超过基于位置的 baseline。

这说明:Transformer encoder 的一部分 attention heads 不是随便看词,而是在捕捉某些句法结构。

这点对理解 Transformer 很重要,因为它说明 self-attention 不是完全黑箱的,某些 heads 具有清晰的语言学功能。


6.3 Rare Words Heads:关注低频词

第三类是rare-word heads

论文发现,在所有模型中,第一层都有一个被 LRP 判断为特别重要的 head。进一步分析后发现,这个 head 会关注句子中最不常见的词。对于 OpenSubtitles 模型,当句子中最稀有的 token 不在 top-500 高频词中时,这个 head 有 66% 的情况会指向最稀有 token,有 83% 的情况会指向两个最稀有 token 之一;在 WMT 模型中,它指向两个最稀有 token 之一的比例也超过 50%。

这个发现很有意思,因为低频词往往更难翻译,也更容易影响句子含义。

所以 rare-word head 可以理解为:帮助模型特别关注可能翻译困难或信息量较高的词。


七、剪枝方法:随机门控 + L0 正则松弛

论文不仅分析 heads,还提出了一个 pruning 方法,用来验证哪些 heads 真正重要。

它的做法是:给每个 attention head 乘上一个标量 gate。如果 gate 接近 1,这个 head 保留;如果 gate 接近 0,这个 head 被关闭。不同于普通输入相关 gate,这里的 gate 是head-specific的参数,不随句子变化。

作者希望直接用 L0 正则来惩罚非零 gates,也就是鼓励模型尽量关闭不重要 heads。但 L0 不可导,不能直接优化。因此论文采用Hard Concrete distribution作为可微随机松弛。这个分布在 0 和 1 上都有概率质量,训练时可以通过重参数化技巧反向传播。

训练目标可以直观理解为两部分:

一部分是翻译交叉熵损失,保证模型还能翻译。

另一部分是 head 数量惩罚,鼓励模型关闭更多 heads。

通过调节正则系数,作者可以得到不同稀疏程度的模型。官方代码 README 也说明,训练时从已收敛的 Transformer 模型开始,加入 head gates 和 L0 relaxation 后继续训练;不同正则系数会产生不同数量的保留 heads。(GitHub)

这个方法的优势是:它不是手工设定每层剪多少 heads,而是让模型自己学习哪些 heads 应该关闭。


八、剪枝实验结果

8.1 只剪 encoder self-attention heads

论文首先只对 encoder self-attention 做剪枝。原因是作者主要想分析 encoder 中哪些 heads 重要。

结果非常明显:

在 OpenSubtitles English–Russian 上,48 个 encoder heads 最后只保留 4 个,BLEU 只下降 0.25。

在更复杂的 WMT English–Russian 上,只保留 10 个 encoder heads,就可以与完整模型只差 0.15 BLEU。

论文在第 6.2 节明确报告了这两个结果。

这说明:encoder self-attention 中绝大多数 heads 对最终翻译质量不是不可替代的。

换句话说,encoder 中真正承担关键功能的 heads 只占少数。


8.2 重要功能的 heads 最后被剪掉

剪枝过程中作者观察到,那些之前被识别出来的 positional heads、syntactic heads 和 rare-word heads 通常会保留下来,直到剪枝非常激进时才被删除。论文提到,即使剪掉 2/3 的 heads,模型仍保留了前面识别出的所有功能类型。

这点非常关键。

如果一个 head 只是“看起来有解释性”,但实际对模型没用,那么剪枝时它应该很早被删掉。可实验结果相反:这些 specialized heads 往往最后被剪掉。

所以这篇论文通过剪枝实验证明:

可解释的 specialized heads 确实是重要 heads,而不是可视化幻觉。


8.3 剪所有 attention 类型

Transformer 翻译模型里有三类 attention:

encoder self-attention

decoder self-attention

decoder-encoder attention

论文进一步对所有 attention heads 都加 gates,观察哪些类型更容易被剪掉。结果显示,对于 WMT 数据,模型可以剪掉接近 3/4 的 encoder heads,以及超过 1/3 的 decoder self-attention 和 decoder-encoder attention heads,而翻译质量没有明显下降;甚至剪掉超过一半所有 heads 时,BLEU 下降也不超过 0.25。

不过不同 attention 类型敏感性不同。论文发现,模型倾向于优先剪 encoder self-attention heads,而 decoder-encoder attention heads 最重要。原因很直观:decoder-encoder attention 负责让译文生成过程读取源句信息,没有它翻译无法正常进行。

这说明:并不是所有 attention heads 都同等冗余。Encoder self-attention 最容易剪,decoder-encoder attention 最敏感。


九、从头训练小 head 模型不如先训练大模型再剪

论文还有一个很重要的发现:如果直接从头训练一个 head 数量很少的模型,效果通常不如先训练完整模型,再通过 pruning 得到相同 head 配置。

表 2 中右侧列给出了“from scratch”的结果。作者发现,从随机初始化开始训练相同稀疏结构时,性能下降更明显;他们认为,这和模型压缩中的常见现象一致:通过剪枝学到的稀疏结构,不一定能从头训练到相同效果。

这点很重要,因为它说明:

多头结构可能在训练阶段帮助模型发现有用功能;训练完成后,很多 heads 才变得可剪。

也就是说,这篇论文并不是简单地说“Transformer 一开始就不需要这么多 heads”,而是更精确地说:

完整多头结构对训练可能有帮助,但训练完成后,推理阶段只需要少数关键 heads。


十、和 “Are Sixteen Heads Really Better than One?” 的区别

这两篇论文都研究 attention head redundancy,但侧重点不完全一样。

Are Sixteen Heads Really Better than One?更关注:删除 heads 后模型性能如何变化,并使用 loss sensitivity / Taylor-style importance 来做 head pruning。

这篇Specialized Heads Do the Heavy Lifting更关注:哪些 heads 重要,它们是否有可解释功能,以及这些功能性 heads 是否真的会在剪枝中被保留。

所以它的贡献不只是“head 可以剪”,而是进一步说明:重要 heads 往往是专门化 heads,它们承担位置、句法、低频词等具体语言功能;剩余大量 heads 才是主要冗余来源。

这使它比单纯剪枝论文多了一层解释性分析。


十一、方法优点

11.1 不只做剪枝,还解释了为什么能剪

很多剪枝论文只告诉你哪些结构可以删,但不解释被保留下来的结构在做什么。

这篇论文的价值在于,它把 head importance、head function 和 pruning 结果联系起来。

重要 heads 有可解释功能,剪枝时也最晚被删除。

这个闭环非常有说服力。


11.2 剪枝粒度是完整 head,属于结构化剪枝

论文剪掉的是完整 attention head,而不是零散参数。这样剪枝后的模型可以在结构上变小,理论上可以减少 attention 计算。

因此它属于 Transformer 中比较典型的attention head structured pruning


11.3 L0 gate 方法比手工排序更灵活

这篇论文不是简单按一个指标排序后直接剪,而是用带 L0 松弛的 gate 训练,让模型自动学会哪些 heads 应该关闭。

这种方法可以同时考虑模型适应过程,而不是只做一次性打分。


11.4 证明了 encoder heads 的高度冗余

在 encoder self-attention 中,大量 heads 可以被删除而 BLEU 几乎不变。这说明 Transformer encoder 的多头注意力存在明显冗余。

这个结论对后续 Transformer 压缩非常有启发。


11.5 解释了不同 attention 类型的重要性差异

论文不仅剪 encoder,还分析 decoder self-attention 和 decoder-encoder attention。结果表明 decoder-encoder attention 更难剪,这符合机器翻译任务的结构直觉。

这让我们在做 head pruning 时不能简单“一刀切”。


十二、方法局限

12.1 主要实验集中在机器翻译

论文实验主要是 Transformer 机器翻译,包括 WMT 和 OpenSubtitles。它没有像后续 BERT 剪枝论文那样覆盖大量自然语言理解任务。

因此,它的结论最直接适用于encoder-decoder Transformer 翻译模型


12.2 主要分析的是 attention heads,不涉及 FFN

Transformer 中 FFN 通常占据大量参数和计算量。即使剪掉很多 heads,整个模型的总加速也可能有限。

所以,如果目标是部署压缩,仅剪 heads 可能不够,还需要结合 FFN 剪枝、层剪枝、低秩分解或蒸馏。


12.3 L0 gate pruning 需要额外 fine-tuning

论文的剪枝不是直接一键删除,而是在已训练模型基础上加入 gates 并继续训练正则目标。

这比简单 Taylor importance 剪枝更复杂,训练成本也更高。


12.4 可解释性仍然有限

论文识别出 positional、syntactic、rare-word 三类 heads,但并不是所有重要 heads 都能被完全解释。某些 heads 可能具有复杂功能,不能简单归入这几类。

因此,论文证明了“部分重要 heads 可解释”,但不是说所有重要 heads 都完全可解释。


12.5 真实加速依赖实现

如果只是把 gate 置零,而底层计算仍然执行所有 heads,那么不会真正加速。要获得实际部署收益,需要物理删除对应 Q/K/V 和输出投影中的 head 参数,并使用更小的 attention 计算图。


十三、整体评价

这篇论文的核心贡献不是简单证明 attention heads 可以被剪,而是说明“哪些 heads 真正在工作”。

它先用 LRP 找出重要 heads,再分析这些 heads 的功能,发现它们往往具有明确语言角色:关注相邻词、关注句法依存关系、关注低频词。随后,论文用 L0 stochastic gates 做 head pruning,发现这些 specialized heads 往往最后被剪掉,从而反过来验证它们确实承担了关键功能。

这使论文形成了一个非常清晰的逻辑闭环:

少数 heads 重要。

重要 heads 通常具有可解释功能。

不重要 heads 可以被大量剪掉。

剪枝过程中,专门化 heads 会被保留下来。

这篇论文对后续 Transformer 剪枝的启发很直接:做 head pruning 时,不应该只看 head 数量,而应该关注 head 的功能差异和 attention 类型差异。尤其在机器翻译中,encoder self-attention heads 冗余很大,而 decoder-encoder attention heads 更敏感。


十四、一句话总结

《Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned》证明了 Transformer 中真正重要的是少数具有明确功能的 specialized heads,例如位置、句法和低频词 heads;通过基于随机 gates 和 L0 松弛的结构化剪枝,大量冗余 heads 可以被删除而几乎不损失 BLEU,其中专门化 heads 往往最后被剪掉,说明它们确实承担了模型预测中的主要工作。