一、论文基本信息论文题目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模型维度为 512FFN 维度为 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 PropagationLRP。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 翻译模型里有三类 attentionencoder self-attentiondecoder self-attentiondecoder-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不涉及 FFNTransformer 中 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 往往最后被剪掉说明它们确实承担了模型预测中的主要工作。