[论文阅读]Losing Heads in the Lottery: Pruning Transformer Attention in Neural Machine Translation

前言

论文名：Losing Heads in the Lottery: Pruning Transformer Attention in Neural Machine Translation
论文作者：Abigail See et.al.
期刊/会议名：EMNLP 2020
本文作者：XMU_MIAO
日期：2020/11/20

摘要

一、Introduction

  注意力机制是Transformer架构的重要组成部分。近期的研究表明大部分的attention heads对其做的决定不自信并且在训练后能够被删去。但在训练前对其进行修剪会造成质量低。在本文中，在训练早期，我们利用彩票假设对heads进行剪枝而不是在模型完全收敛的时候。 我们在Turkish->English机器翻译任务上的实验表明对transformer-big模型上heads剪枝3/4而平均BLUE损失为0.1是有可能的。剪枝模型的推理速度是原来的1.5倍，尽管代价是更长的训练时间。这种剪枝方法是对其他方法发补充，如teacher-student 训练，我们English->German student在75%编码器注意力稀疏度的情况下损失0.2BLEU。

二、Contribution

• 彩票假说可以应用于整个参数块的剪枝，而不仅仅是对单个参数
• 大部分attention heads可以在训练时及早删除并不会对质量造成重大损失。最激进的attention剪枝损失在80%-90%块稀疏的情况下损失1个BLEU点
• 与从头开始训练具有相同结构的模型相比，彩票假说获得更好的结果
• 模型剪枝显示了层中head数量的分布模式，例如，随着层次的增加，context attention变得越发重要，decoder只需要在第一层进行self-attention，其与层的self-attention是冗余的，可以被剪枝

三、Methodology

本文的方法是在训练早期对Transformer模型中heads进行剪枝，剪枝的流程遵循"彩票方法"(lottery ticket approach)，剪枝依据是启发式度量"注意力自信度"(attention importance)。 以下对二者以及本文方法的基本流程进行介绍。

3.1 彩票方法

彩票假说：一个随机初始化的稠密神经网络包含一个子网络，该子网络被初始化后（初始化为checkpoint点对应参数）进行单独训练能够在最多训练与原网络相同迭代次数内达到相似的测试精度。，将彩票假说应用于模型剪枝时，得到“中奖彩票”可以按照以下步骤进行：
（1） 初始化一个稠密网络并迭代训练$j$次，并保存在$checkpoint$
（2） 对模型按照某种准则进行剪枝
（3） 并将剩余参数初始化为$checkpoint$中对应参数值，得到中奖彩票

3.2 注意力自信度

  一个注意力head被定义为自信当其给一个句子中的一个词赋予了很大的权重。这样的注意力head能够实现强对齐，应当被保留在模型中。
  一个注意力head的softmax层计算其关注的词的概率分布。我们将这个概率分布的最大值作为“自信度”，形式化表示如下：$$confidence=maxattention(s)$$其中$s$表示源句词构成的序列，一个attention head自信度值取一段时间内该attention head的自信度的平均值。

3.3 本文方法基本流程

  利用启发式方法一次性对attnetion heads剪枝太多不太可行，本文采用迭代剪枝的方法，具体流程如上图所示：
（1）初始化网络，并迭代训练$x$次，将结果保存在$checkpoint$中
（2） 迭代训练$y$次，并利用attention自信度进行剪枝操作，并将剩余参数设置为$checkpoint$中对应参数值
（3） 重复步骤（2）直至达到预期剪枝量
（4） 对剪枝模型继续训练直至收敛

四、Experiments

待更新

五、Analysis

5.1 Head Distribution

---

表3和表4展示了在整个迭代剪枝过程中注意力分布的变化。在Turkish->English任务中，decoder中的heads随着层数的变深被剪枝得越多。第一层似乎是最重要的，其余层的都不重要。context attention与decoder中的self-attention相互关联，每一层的重要性都高于前一层。两种语言中，encoder的中间层attnetion似乎都没有第一层和最后一层attention重要。

5.2 Architecture or Initialisation?

  为了检验彩票假说在我们论文中的正确性，我们重新初始化经过修剪的模型，同时保持其结构（即彩票假说中参数随机初始化，而不是设置为checkpoint中对应参数值）。表4和表5对比了经彩票假说初始化模型和随机初始化模型之间的平均BLEU差异。
  从表中可以看出随着稀疏度的变大，经彩票假说初始化的模型和随机初始化的模型的差距在拉大，这在一定程度上也证实了彩票假说。

总结

  本文提出一种Transformer中对于attention head的块剪枝方案，剪枝流程遵循“彩票假说”，剪枝依据采用启发式的“注意力自信度”，一个高自信度的attention head对源句的词会分配更大的权重。该剪枝方案在训练前期进行剪枝操作，并迭代剪枝至期望值，并对最终剪枝模型进行训练收敛。实验表明在NMT任务上能够对attention head剪枝50%-72%而不对翻译质量造成太严重的损失并且能够与其他方法互补使用（如teacher-student training）。
  本文只是研究了对attention head进行块修剪，对Transformer中其他部分进行修剪是继续研究的一个方向，并且可以尝试采用其他启发式剪枝方法以及是否要采用彩票假设。