[论文解读] Alleviating the Inequality of Attention Heads for Neural Machine Translation

Zewei Sun, Shujian Huang+2 more|arXiv (Cornell University)|Sep 21, 2020

Multimodal Machine Learning Applications参考文献 21被引用 6

一句话总结

本文提出 HeadMask，一种通过在反向传播过程中随机或有策略地掩码注意力头来缓解基于 Transformer 的神经机器翻译中注意力头不平等的训练方法。实验表明，该方法在多个语种对上均提升了翻译性能，减少了注意力头重要性的方差，并增强了模型鲁棒性，证实了均衡的头训练可提升模型容量与泛化能力。

ABSTRACT

Recent studies show that the attention heads in Transformer are not equal. We relate this phenomenon to the imbalance training of multi-head attention and the model dependence on specific heads. To tackle this problem, we propose a simple masking method: HeadMask, in two specific ways. Experiments show that translation improvements are achieved on multiple language pairs. Subsequent empirical analyses also support our assumption and confirm the effectiveness of the method.

研究动机与目标

• 研究并解决 Transformer 模型中注意力头重要性不均的问题。
• 验证假设：不平衡的训练会导致少数关键头占据主导，而其他头被低估。
• 设计一种确保所有注意力头实现更均衡学习的训练方法。
• 通过实证验证：降低对特定头的依赖可提升模型鲁棒性与翻译质量。

提出的方法

• 提出 HeadMask，一种基于掩码的训练技术，即在每个训练批次中随机选择并掩码部分注意力头。
• 引入一种变体，优先掩码最重要的头，以将训练负担重新分配给表现较弱的头。
• 通过反向传播中的梯度估计（不更新参数）来识别关键头的重要性。
• 在前向与反向传播过程中应用头掩码，仅更新未被掩码的头，以强制实现均衡学习。
• 采用两阶段流程：第一阶段，计算损失与梯度但不更新参数；第二阶段，在最终参数更新前掩码激活最高的头。
• 使用泰勒近似法评估头的重要性，以衡量损失对头移除的敏感度。

实验结果

研究问题

• RQ1 在翻译过程中，Transformer 中的注意力头在重要性上有多大的不均衡？
• RQ2 注意力头训练不均衡是否会导致模型过度依赖少数关键头？
• RQ3 在训练过程中采用掩码策略能否降低头重要性的方差并提升整体性能？
• RQ4 降低对特定头的依赖是否能增强模型鲁棒性与泛化能力？
• RQ5 在随机掩码与基于重要性的掩码策略中，哪种能带来更好的翻译质量与头重要性平衡？

主要发现

• 与基线相比，Random-18 在 Zh→En 上提升 +0.42 BLEU，Ro→En 上提升 +0.87，Tr→En 上提升 +0.79。
• Impt-18 在 Ro→En 上取得最佳结果，提升 +0.78 BLEU，但在其他语种对上表现不如 Random-18。
• 两种 HeadMask 变体均显著降低了头重要性的方差，其中 Impt-18 在 Zh→En 上将方差降至 9.13（基线为 77.28）。
• 所有语种对的平均头重要性均下降，表明模型对单个头的依赖性降低。
• 采用 HeadMask 训练的模型表现出更强的鲁棒性，即使在推理时掩码关键头，性能仍能保持稳定。
• 实证分析证实：在训练中掩码重要头可使头重要性分布更平坦，实现更均衡的模型行为。