一、论文基本信息论文题目Are Sixteen Heads Really Better than One?作者Paul Michel、Omer Levy、Graham Neubig发表会议NeurIPS 2019研究方向Transformer 模型压缩、Attention Head Pruning、多头注意力结构分析论文链接NeurIPS / arXiv代码作者公开了 head pruning 实验代码。(NeurIPSarXivGitHub)这篇论文提出了一个非常直接但影响深远的问题Transformer 中的 Multi-Head Attention 真的需要这么多 Head 吗例如 BERT-base 每层有 12 个 attention headsTransformer 大模型甚至可能有 16、32、64 个 heads。传统观点认为不同 head 学习不同类型的信息关系因此 head 越多表达能力越强。但是这篇论文通过大量实验发现很多 attention heads 在推理阶段其实是冗余的删除大量 heads 后模型性能几乎不下降甚至某些层只保留一个 head 也可以正常工作。论文研究对象包括 Transformer 翻译模型和 BERT 类模型并分析了 attention head pruning 对速度、内存和精度的影响。(NeurIPS)1. 论文核心问题很直接Transformer 真的需要这么多 attention heads 吗Transformer、BERT 这类模型通常每层都有多个 heads。论文质疑的是这些 heads 是否都在发挥不可替代的作用还是其中很多只是冗余结构。2. 论文的主要结论是很多 attention heads 在推理阶段是冗余的。作者发现训练好的 Transformer / BERT 中有相当一部分 heads 可以在测试阶段被移除而模型性能不会明显下降。有些层甚至可以只保留一个 head。(NeurIPS 会议录)3. 多头注意力在训练阶段有价值但推理阶段不一定都必要。论文不是否定 multi-head attention 的训练价值而是指出模型训练完成后并不是每个 head 都对最终预测有明显贡献。4. 判断 head 是否冗余最直接的方法是“删掉它看性能是否下降”。如果某个 head 被 mask 掉后BLEU 或 accuracy 基本不变说明其他 heads 或后续层可以弥补它的作用那么这个 head 就可以认为是冗余的。5. 论文还使用了基于 loss sensitivity 的 head importance 估计。简单说就是看模型 loss 对某个 head 的开关是否敏感。如果关掉这个 head 后 loss 变化很小说明它不重要如果 loss 变化大说明它重要。6. 剪枝对象是完整 attention head因此这是结构化剪枝。它不是删除零散权重而是删除一个完整 head 对应的注意力计算分支。所以它属于 Transformer 中的attention head-level structured pruning。7. 论文采用的是层内归一化后的全局排序剪枝。不同层的梯度尺度可能不同因此先在每层内部对 head importance 做归一化再把所有 heads 放到一起排序优先剪掉全局最不重要的 heads。8. 不同类型的 attention 对 head 数量的依赖程度不同。在机器翻译模型中encoder self-attention 和 decoder self-attention 更容易剪encoder-decoder attention 更敏感剪太多会导致性能严重下降。论文报告encoder-decoder attention 剪掉超过 60% heads 会出现明显性能崩塌而 self-attention 在保留较少 heads 时仍能产生相对合理的翻译结果。(arXiv)9. BERT 中也存在明显 head 冗余。论文在 BERT-base 上做了自然语言理解实验发现很多 heads 对下游任务影响很小可以被剪掉而不明显损害性能。10. 这篇论文的重要意义不是提出复杂剪枝算法而是证明了 Transformer 内部存在大量 head redundancy。它告诉我们Transformer 的计算冗余不仅存在于 FFN、隐藏维度和层数中也存在于 multi-head attention 的 head 维度中。11. 这篇论文对后续 Transformer 压缩影响很大。后续很多 BERT 压缩、Transformer 剪枝、LLM 高效化方法都沿用了类似思路先评估结构单元的重要性再删除冗余结构。12. 它的局限也很明显。它主要是训练后分析和剪枝没有从训练阶段直接约束 head 冗余而且如果只是用 mask 关闭 head并不会真正加速必须物理删除对应参数和计算分支才能获得实际部署收益。二、论文背景为什么要研究 Attention HeadTransformer 的核心结构 Multi-Head Self-AttentionMulti-head attention 会把输入表示拆成多个子空间然后其中(h)head 数量每个 head 有独立的 (Q,K,V) 参数例如BERT-base12 层 Transformer每层 12 个 heads总共有个 attention heads。大型 TransformerGPTBERT-largeT5head 数量更多。问题来了如果很多 head 学习的信息高度重复那么这些 head 是否只是增加计算量这正是论文研究的问题。三、核心发现论文最大的发现Multi-head attention 在训练阶段非常重要但是推理阶段存在大量冗余。作者首先做了一个非常简单的实验训练完成一个 Transformer。然后逐个关闭 attention head。观察模型性能变化。结果非常令人意外很多 head 被删除后模型性能几乎不变。甚至部分 Transformer 层只保留一个 head性能仍然可以维持。论文摘要明确指出大量 attention heads 可以在测试阶段删除而不会显著影响性能部分层甚至可以减少到单 head。这说明Multi-head attention 的实际推理需求远低于训练时的设计规模。四、Attention Head 为什么可以被剪理解这个问题需要先理解 head 的作用。理论上每个 head 应该学习不同关系例如一个 head 学习主语和谓语关系另一个 head 学习长距离依赖另一个 head 学习位置关系另一个 head 学习语义对应关系但是论文实验发现现实中的 Transformer 并没有让所有 head 都承担独立功能。很多 head学到类似模式对最终输出贡献很小可以被其他 head 替代因此Multi-head attention 存在大量容量冗余。这和 CNN 中 filter redundancy 非常类似。CNN冗余 filter → 删除 filterTransformer冗余 head → 删除 attention head五、Head Pruning 方法论文首先提出如果直接删除某些 head需要一个重要性评价指标。作者采用类似剪枝中的 Taylor importance 思想。对于一个 attention head引入 mask其中表示保留 head。表示删除 head。Multi-head attention 可以表示为然后定义 head importance表示删除该 head 后loss 增加多少。如果很小。说明删除这个 head 对模型影响很小。因此优先删除。作者进一步使用一阶 Taylor 近似这样不需要真的删除每个 head 再重新测试而是一次 forward/backward 就可以估计所有 head 重要性。这和 CNN 剪枝中的 Taylor Filter Pruning 思路非常接近。六、 Head 删除实验论文首先做一个非常直观的实验。对于每个 attention layer只保留一个 head。观察性能。结果一些层中单个 head 可以达到接近原始 multi-head attention 的效果。这说明并不是所有层都需要大量 heads。特别是某些 encoder self-attention 层某些中间层head redundancy 非常明显。但是decoder attention 和 encoder-decoder attention 对 multi-head 依赖更强。论文分析指出不同 attention 类型对 head 数量的依赖程度不同。(NeurIPS)单独删除 head 只能说明“某些 head 不重要”。但真正部署需要一次删除多个 head。因此论文设计了系统 pruning流程第一步计算所有 head importance。第二步按照 importance 排序。第三步删除 importance 最低的 heads。第四步测试模型性能。例如原始12 heads/layer剪掉20%40%60%heads观察accuracy / BLEU变化。实验发现很多任务可以删除大量 heads而性能下降非常有限。论文指出在不同任务和模型中可以删除大量 attention heads而不会明显损失性能同时还能获得速度和内存收益。(NeurIPS)七 实验这篇论文的实验主要分成两组主实验和一组补充实验主实验用WMT2014 英法机器翻译 Transformer-large和MultiNLI 自然语言推理 BERT-base补充实验又加入了SST-2、CoLA、MRPC、IWSLT2014 德英翻译。论文在实验设置中明确说明WMT 模型是 Vaswani 等人的Transformer large6 层、每层 16 个 headsBERT 使用bert-base-uncased12 层、每层 12 个 heads并在 MultiNLI 上 fine-tune。实验类别数据集 / 任务使用模型模型结构细节评价指标主要用途主实验 1WMT2014 English-to-French机器翻译Transformer-large6 层 Transformer每层16 个 attention heads包含 encoder self-attention、decoder self-attention、encoder-decoder attention 三类注意力BLEU分析翻译模型中 attention head 是否冗余逐个 mask head只保留单个 head按重要性迭代剪枝主实验 1 的测试集newstest2013同上使用 Ott et al. 的预训练 WMT En-Fr 模型BLEU主要报告测试结果计算 head ablation 和 pruning 后性能泛化测试newstest2014同上在 newstest2013 上选 head在 newstest2014 上验证BLEU验证“重要 head”是否能跨测试集保持一致域外测试MTNT English-to-French同上同一个 WMT Transformer-largeBLEU检查 head 重要性是否能泛化到 noisy / out-of-domain 翻译数据主实验 2MultiNLI / MNLI自然语言推理BERT-base-uncased12 层 Transformer encoder每层12 个 attention heads只包含 self-attentionAccuracy分析 BERT 中 attention head 是否冗余逐 head mask单层只保留一个 head整体 head pruning主实验 2 的测试集MNLI matched validation setFine-tuned BERT-base使用 BERT-base-uncased在 MNLI 上 fine-tuneAccuracy主要报告 BERT head pruning 结果域外测试MNLI mismatched validation setFine-tuned BERT-base同一个 MNLI fine-tuned BERTAccuracy检查 head 重要性是否能跨 matched / mismatched 分布保持一致补充实验SST-2情感分类Fine-tuned BERTGLUE 版本 Stanford Sentiment TreebankAccuracy附录中验证 BERT head pruning 在其他分类任务上的效果补充实验CoLA语法可接受性判断Fine-tuned BERTGLUE 版本 Corpus of Linguistic AcceptabilityMatthews correlation附录中验证 head pruning 在语言可接受性任务上的效果补充实验MRPC句子复述识别Fine-tuned BERTGLUE 版本 Microsoft Research Paraphrase CorpusF1附录中验证 head pruning 在句对匹配任务上的效果训练动态实验 / 补充实验IWSLT2014 German-to-English机器翻译较小 Transformer6 层每层8 个 headsBLEU分析训练过程中 head 重要性如何形成以及 pruning 曲线随 epoch 如何变化论文正文明确提到主实验只考虑两个训练好的模型WMT Transformer-large和BERT-base-uncased。其中 WMT 模型在WMT2014 English-to-French上训练并在newstest2013上报告 BLEUBERT 在MultiNLI上 fine-tune并在matched validation set上报告 accuracy。补充实验在附录 B 中给出额外使用了SST-2、CoLA、MRPC 和 IWSLT2014 German-to-English。其中 SST-2、CoLA、MRPC 都使用 fine-tuned BERTIWSLT 使用的是论文第 6 节描述的较小 Transformer 模型。论文还指出在 SST-2 上最多约60% heads、在 CoLA 和 MRPC 上约50% heads可以被剪掉而没有明显性能影响。简单说主线实验可以概括为机器翻译WMT2014 En-Fr / Transformer-large / BLEU自然语言推理MultiNLI / BERT-base / Accuracy补充 NLP 任务SST-2、CoLA、MRPC / BERT-base / Accuracy、Matthews correlation、F1训练动态分析IWSLT2014 De-En / 小型 Transformer / BLEU八、哪些 Head 更重要论文进一步分析是不是所有 head 都一样答案不是。不同位置的 head 重要性不同。作者发现Encoder Self-Attention通常更容易剪。原因多个 encoder head 存在较强冗余。Encoder-Decoder Attention更加敏感。原因decoder 需要依赖 encoder 输出进行生成。因此attention alignment 更依赖多个 head。论文摘要指出self-attention layers 可以明显剪枝而 encoder-decoder layers 对 multi-headedness 更依赖。(NeurIPS)九、和 CNN Filter Pruning 的关系这篇论文实际上把 CNN 剪枝思想迁移到了 Transformer。CNN结构Conv Filter重要性Filter importance剪删除 filterTransformer结构Attention Head重要性Head importance剪删除 head对应关系CNN filter pruningTransformer head pruning所以这篇论文是 Transformer 结构化剪枝的重要开端。十、和传统 Attention Dropout 的区别很多人会问head pruning 和 attention dropout 有什么区别区别非常大。Attention Dropout发生在训练阶段。作用防止过拟合。推理所有 head 都存在。Head Pruning发生在训练完成之后。作用减少模型结构。推理head 永久删除。所以Dropout 是正则化。Pruning 是模型压缩。十一、实验意义这篇论文的重要价值不是提出一个复杂算法而是提出一个非常关键的问题Transformer 的多头设计是否真的被充分利用它证明Transformer 参数规模增长过程中存在大量结构冗余。尤其是Attention heads。这为后续大量 Transformer 压缩方法奠定基础例如Head pruningBERT compressionDistilBERTTinyBERTMobileBERTEfficient Transformer十二、优点12.1 简单有效方法非常简单计算 head importance。排序。删除。无需重新设计 Transformer。12.2 结构化剪枝删除的是完整 attention head。不是删除单个参数。因此可以真正减少计算量。12.3 适用于 Transformer这是早期系统研究 Transformer 冗余的重要工作。它证明Transformer 并不是越大越好。12.4 提供可解释分析论文不仅剪枝还分析哪些 head 重要。哪些 layer 依赖多 head。帮助理解 Transformer 内部机制。十三、局限性13.1 Head importance 是后验分析论文主要研究训练完成后删除 head。但是它没有解决训练过程中如何避免产生冗余 head。13.2 删除 head 后结构重新训练不足大量实验采用prune → evaluate如果剪得非常激进通常需要 fine-tuning。13.3 只关注 Attention HeadTransformer 还有大量其他结构例如FFN neuronshidden dimensionlayer depthembedding size这些同样存在冗余。后续工作进一步研究Transformer block pruning。13.4 不一定等价真实加速删除 head 可以减少理论 FLOPs。但真实速度还受到GPU kernelmemory accessbatch size影响。因此需要硬件测试。十四、整体评价Are Sixteen Heads Really Better than One? 是 Transformer 结构化剪枝领域非常重要的一篇论文。它提出了一个简单但深刻的问题Multi-head Attention 中的 head 数量是不是过度设计通过实验它证明大量 attention heads 在推理阶段是冗余的可以被安全删除。它的贡献主要有三个第一首次系统分析 Transformer attention head redundancy。第二提出基于 loss sensitivity 的 head importance estimation。第三证明 Transformer 可以通过 head pruning 获得速度和内存收益。从模型压缩角度看这篇论文对应 CNN 时代的 filter pruning。它告诉我们Transformer 的“宽度”不仅存在于 hidden dimension也存在于 attention heads而 heads 同样可以被结构化剪枝。十五、一句话总结《Are Sixteen Heads Really Better than One?》证明了 Transformer 的 Multi-Head Attention 存在大量冗余 head通过基于损失敏感度的 head importance 估计可以删除大量 attention heads 而保持模型性能从而降低 Transformer 推理成本这篇论文开启了 Transformer Attention Head Pruning 方向为后续 BERT、LLM 高效压缩提供了重要基础。