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

TL;DR

用LRP与Hard Concrete剪枝;WMT英俄仅保留10/48个编码器头,BLEU仅降0.15。

cs.CL 🟡 进阶级 2019-05-23 26 次浏览
Elena Voita David Talbot Fedor Moiseev Rico Sennrich Ivan Titov
Transformer 多头注意力 神经机器翻译 模型剪枝 可解释性

核心发现

方法论

论文先用Layer-wise Relevance Propagation(LRP)衡量每个编码器注意力头对top-1预测logit的贡献,再以最大注意力权重定义置信度,并分析位置、句法依存和罕见词功能。剪枝阶段为每个头加入可学习随机门,采用Louizos等人的Hard Concrete分布近似不可微L0正则,通过交叉熵与门稀疏正则联合微调。

关键结果

  • 重要头高度集中且常具有可解释功能:位置头至少90%指向固定相对位置,最大注意力均值超过0.8;句法头对nsubj、dobj、amod和advmod的识别准确率明显超过位置基线。
  • 在WMT英俄任务中,完整编码器有48个头;剪去38个、仅保留10个时,BLEU只下降0.15。OpenSubtitles英俄中,仅保留4个头,性能损失约0.25 BLEU。
  • 联合剪枝显示编码器自注意力最先被移除,解码器—编码器注意力最重要;WMT基线29.6 BLEU,保留14/31/30个三类头时为29.62,超过半数头可删除且损失不超过0.25。

研究意义

研究改变了“多头越多越有用”的直觉:Transformer的容量并非由所有头平均贡献,而是由少数专门化头承担关键计算。它为注意力可解释性提供了因果性更强的证据——高相关头不仅看起来有规律,也最晚被剪掉。对工业翻译系统而言,这意味着可在保持质量的同时减少推理计算、显存和部署成本。

技术贡献

核心技术是把L0稀疏目标从不可微的非零计数转化为Hard Concrete门的非零概率之和:LC(φ)=Σi(1−P(gi=0|φi)),并优化L=Lxent+λLC。门在每个batch重采样,利用重参数化技巧训练;最终概率趋近0或1,因此测试时可转换为普通稀疏Transformer。作者还固定解码器,仅微调编码器,以避免功能迁移掩盖头的重要性。

新颖性

相较只平均或观察最大注意力权重的分析,论文首次系统结合LRP、置信度、依存关系评测和可学习L0式结构剪枝,直接比较单个头的预测贡献与保留顺序。其新意不只是压缩模型,更是用剪枝结果验证“专门化功能等于真实重要性”。

局限性

  • 实验主要集中于Transformer编码器和英译俄、德、法,结论是否适用于更大模型、其他架构、低资源语言或大型语言模型尚未证明。
  • 句法功能依赖CoreNLP自动依存分析与人工设定阈值,例如比相对位置基线高10%,可能混入解析误差,不能等同于头真正执行完整句法推理。
  • 剪枝从已收敛模型微调;从随机初始化训练同样稀疏结构明显更差,因此方法的压缩收益依赖预训练冗余。

未来方向

作者提出比较其他NMT压缩方法。后续可研究动态、输入相关的门控,跨语言和跨域稳定性,联合剪枝编码器、解码器与前馈层,并评估真实延迟、能耗及不同规模Transformer上的可迁移性。

AI 总览摘要

Transformer翻译系统拥有大量注意力头,但这些头是否同等重要一直不清楚。传统研究常把所有头平均,或只观察最大权重,因此难以区分真正承担翻译任务的组件与冗余组件。Voita等人把问题转化为“谁在推动预测、谁可以被删除”。

研究首先使用Layer-wise Relevance Propagation(LRP)追踪每个编码器头对top-1输出logit的贡献,并以最大注意力权重衡量置信度。高相关头往往具有清晰职能:有的指向相邻词,有的追踪主语、宾语、定中或状语依存关系,还有第一层头偏向句中罕见词。随后,作者为每个头加入Hard Concrete随机门,以可微方式近似L0正则,并从已训练模型出发逐步关闭头。

结果显示,模型依赖的是少数“重型工人”。WMT英俄模型的48个编码器头可剪去38个,BLEU仅下降0.15;OpenSubtitles中只保留4个头,损失约0.25 BLEU。扩展到全部注意力后,编码器自注意力最容易删除,解码器—编码器注意力最关键。研究说明,专门化头既具有可解释性,也具有不可替代性,为高效部署和结构化理解Transformer提供了直接证据。

深度分析

研究背景

Transformer由Vaswani等人提出,以编码器—解码器、多头注意力和前馈网络取代循环结构。8头模型曾比同规模单头模型高近1 BLEU,但注意力可视化通常平均所有头或只看最大权重,无法衡量不同头的真实贡献。论文聚焦WMT神经机器翻译中的编码器自注意力。

核心问题

核心问题是:翻译质量是否依赖大量独立注意力头?重要头是否学习稳定、可解释的语言功能?编码器、解码器自注意力和交叉注意力的敏感度有何不同?难点在于注意力权重不等于因果贡献,而直接使用L0计数又不可微。

核心创新

论文有三项创新。第一,用LRP直接衡量头对预测logit的相对贡献。第二,将头分类为相对位置头、句法依存头和罕见词头,并用CoreNLP依存结构检验。第三,引入Hard Concrete随机门和L0式正则,在保留翻译能力的同时学习结构稀疏性,从而把解释与压缩统一起来。

方法详解

  • �� 注意力计算为Attention(Q,K,V)=softmax(QKᵀ/√dk)V,多头输出为Concat(gi·headi)WO。
  • �� LRP按层传播top-1 logit相关性,得到每层头的排序;置信度是排除EOS后最大权重的平均值。
  • �� 位置头要求至少90%指向固定相对位置;句法头比较nsubj、dobj、amod、advmod准确率与位置基线。
  • �� Hard Concrete门gi∈[0,1],以Σ(1−P(gi=0))近似L0,并优化交叉熵加λ稀疏项。
  • �� 从收敛模型微调,测试时按概率较大的一端将门离散为0或1。

实验设计

数据包括WMT英俄、英德、英法,每个方向使用250万句对;另用600万句OpenSubtitles2018英俄数据检验领域影响。测试集为50,000句。模型采用相同参数规模的Transformer,指标为BLEU。实验比较完整模型、剪枝模型和从头训练的同稀疏结构,并分别剪编码器、全部三类注意力及不同解码器层。

结果分析

WMT英俄完整模型为29.6 BLEU;编码器48头剪至10头时仅降0.15。OpenSubtitles中48头剪至4头,约降0.25。全部注意力剪枝时,WMT保留14/31/30个编码器、解码器自注意力、交叉注意力头,得29.62 BLEU;从头训练同结构仅29.47。交叉注意力最重要,较长句子的WMT更依赖解码器自注意力。

应用场景

该方法可用于机器翻译推理压缩:先训练完整模型,再用Hard Concrete识别冗余头,导出稀疏结构,降低显存和矩阵运算。它也适合模型审计,通过LRP和依存分析检查模型是否学习相邻词、语法关系或罕见词处理功能。实际部署仍需测量硬件上的真实延迟,而非只看头数量。

局限与展望

研究规模和语言范围有限,且依存标签来自CoreNLP,功能判定受解析质量和阈值影响。固定解码器的剪枝设置有助于归因,却未覆盖全模型联合重组。稀疏模型从头训练性能较差,说明剪枝利用了已训练网络中的冗余。未来应比较更多压缩算法,研究动态门控、其他语言和更大模型,并报告端到端速度与能耗。

通俗解读 非专业人士也能看懂

把翻译模型想成一家大型厨房。每个注意力头都是一名厨师:有人专门找相邻食材,有人负责确认主料和配料的关系,还有人专门寻找菜单中最少见的食材。传统做法只看全厨房的平均动作,无法知道谁真正重要。

研究者先检查每名厨师的动作是否影响最终菜品,再给每人装一个开关。训练时,系统会逐渐关闭不太有用的人,同时尽量保持菜品评分不变。结果发现,48名厨师中大多数可以离场,英俄翻译评分只下降0.15 BLEU;有明确专长的人通常最后才被关闭。

这说明大型模型里存在很多备用劳动力。它们可能帮助训练更稳定,却不一定在最终工作中不可替代。这个方法因此同时回答两个问题:模型怎样分工,以及怎样在不明显降低质量的情况下减小模型。

简单解释 像给14岁少年讲一样

想象你在玩一个超强翻译游戏。游戏里有48个小助手,每个助手都盯着句子里的不同词,再把线索交给系统。有的助手总看前后相邻的词,有的专门找“谁做了什么”,还有的会注意特别生僻的词。

问题是:是不是助手越多越强?研究者给每个助手装上一个智能开关,然后一个个尝试关掉。如果翻译分数几乎不变,就说明这个助手可能只是备用。结果很惊人:英俄任务里关掉38个助手,只留下10个,BLEU只少0.15分!

研究者还发现,真正有专长的助手最难被关掉,就像球队里的守门员或关键传球手。相反,很多普通助手做的事情互相重叠。注意力头不是平均分工,而是少数成员承担关键工作。

这对手机翻译和服务器很有用:如果能删掉不用的助手,模型就更省内存、更快。不过,不能简单认为所有模型都能这样删;这项研究主要测试了几种英语翻译任务,换语言、换模型或换场景后还要重新检查。

术语表

Multi-Head Self-Attention(多头自注意力)

模型让多个独立头同时计算词与词之间的关系,再拼接结果。每个头可学习不同的信息模式。

论文主要分析编码器自注意力头。

Layer-wise Relevance Propagation(逐层相关性传播,LRP)

一种解释方法,把输出预测的相关性逐层传回网络组件。相关性越高,表示该组件对指定预测贡献越大。

用于给每层注意力头排序。

Hard Concrete

带有0和1概率质量的随机门分布,可近似离散开关,同时保持训练可微。它由Concrete分布拉伸和截断得到。

用于学习哪些注意力头应被关闭。

L0 Regularization(L0正则化)

惩罚非零参数或组件的数量,而不是数值大小。目标是直接减少模型中活跃组件。

论文用非零门概率之和近似L0。

BLEU

机器翻译常用自动评价指标,通过比较生成译文与参考译文的n-gram重合度评分。分数越高通常表示翻译更接近参考答案。

用于比较剪枝前后的翻译质量。

开放问题 这项研究留下的未解疑问

  • 1 专门化头是否在更大Transformer、低资源语言和大型语言模型中同样稳定?现有实验语言和规模有限,需要跨架构验证。
  • 2 注意力头的句法标签是否代表真正推理,还是数据偏差与位置规律的组合?需要干预实验和更可靠的因果评估。
  • 3 头数量减少是否必然带来真实硬件加速?还需报告稀疏算子、批量大小和设备相关的端到端延迟。

应用场景

近期应用

机器翻译模型压缩

翻译服务商可先训练完整Transformer,再用Hard Concrete门剪除冗余头并导出稀疏模型。WMT英俄结果表明,编码器仅保留10/48个头仍接近原始BLEU,但部署前需在目标硬件上验证速度。

模型行为审计

研究团队可用LRP定位高贡献头,再用相对位置和CoreNLP依存关系检查其功能。这能帮助发现模型是否依赖相邻词、主谓关系或罕见词线索,并辅助诊断跨域性能变化。

远期愿景

可解释的稀疏Transformer

未来可构建动态、输入相关的注意力路由,让不同句子只激活必要头,从而同时提升效率和可解释性。主要障碍是门控稳定性、硬件支持以及跨语言迁移。

原文摘要

Multi-head self-attention is a key component of the Transformer, a state-of-the-art architecture for neural machine translation. In this work we evaluate the contribution made by individual attention heads in the encoder to the overall performance of the model and analyze the roles played by them. We find that the most important and confident heads play consistent and often linguistically-interpretable roles. When pruning heads using a method based on stochastic gates and a differentiable relaxation of the L0 penalty, we observe that specialized heads are last to be pruned. Our novel pruning method removes the vast majority of heads without seriously affecting performance. For example, on the English-Russian WMT dataset, pruning 38 out of 48 encoder heads results in a drop of only 0.15 BLEU.

cs.CL