【论文速读】Sharing Attention Weights for Fast Transformer

Sharing Attention Weights for Fast Transformer

2019, IJCAI.

这篇论文的主要思想是通过共享注意力权重来加速Transformer。（关于Transformer，网上已有很多优秀的解读了，不了解的可以看看这篇。）

作者在10个公开数据集上进行了实验，对比其他baselines，至少能够提速1.3倍。

Introduction

由于Transformer可并行化计算的特性，使其被广泛运用于自然语言处理任务中。

但是，虽然Transformer在训练阶段速度很快，但是其在测试推理阶段的速度却是大打折扣的。因为在测试阶段，decoder的decoding过程无法并行计算，每一次只能预测出一个单词，然后再利用这个单词预测下一个单词。这种auto-regressive decoding由于无法并行计算，所以时间开销是很大的。当堆叠多层注意力层时，情况会更加糟糕。

近年来，也有学者针对上述问题提出了一些解决的方法。比如，ANN[1]（Average Attention Network）使用加注意力机制和平均注意力模型来替代点积注意力机制；NMT[2]（Non-Autoregressive Transformer）在传统Transformer的基础上提出一种非自回归解码的模型。但是这些方法要么失去了显式建模词与词之间相关性的能力，要么就是模型结构过于复杂难以训练。

与上述方法不同，这篇论文提出的SAN（shared attention network）是通过在层间共享注意力权重来加速Transformer。因为作者观察到，序列中不同位置的注意力分数，在不同的注意力层表现出相似的分布规律。因此在层间共享注意力权重能够减少冗余计算，提升模型的计算速度，同时还能减少内存占用。

SAN的思想很简单，因此很容易实现，并且能够应用于现有的各种Transformer变体中。

Shared Attention Networks

Standard Transformer：

在层间共享注意力权重的前提是，他们各自原本的注意力权重十分相近，否则在两者差距很大的情况下共享注意力权重会大大减弱模型的性能。

令$S_{[i]}$S[i]​表示为注意力权重矩阵的第$i$i列，代表了在获得位置$i$i的编码向量时，其他各个位置的注意力权重。在$S_{[i]}$S[i]​中，$S_{i,j}$Si,j​表示融合位置$j$j到位置$i$i的注意力权重。直觉上看，$S_{[i]}$S[i]​在不同的注意力层中应该不会有较大的变化，因为$S_{[i]}$S[i]​反映的是词与词之间的相关性，而无论堆叠几层注意力层，词与词之间的相关性应该是近似不变的。比如，无论我们堆叠基层网络，一个句子中的主语和动词的相关性都是比较强的。

为了验证这个想法，作者使用JS散度刻画了任意两层网络间注意力权重概率分布的相似性：

颜色越黑，说明注意力权重分布越相似。从上图可以看出，对于self-attention，第2~6层的注意力权重分布都很相近；而对于encoder-decoder attention，相邻的两层或三层网络之间的注意力权重分布比较相似。

基于以上的分析，说明在层间分享注意力权重是具有一定的合理性的。

因此，作者针对Transformer中的self-attention和encoder-decoder attention分别提出了注意力权重共享机制。

1. 对于self-attention，在层间共享注意力权重：
   

其中$S^m$Sm表示第$m$m层网络的注意力权重矩阵，$\pi$π表示共享同一组注意力权重的层数。因此，SAN（Shared Attention Network）的self-attention注意力权重共享策略为，对网络层进行分组，属于同一组的网络共享同一组注意力权重。

2. 而对于encoder-decoder attention，在层间共享加权融合以后的结果：
   

其中$S^m$Sm表示第$m$m层网络的注意力权重矩阵，$V$V表示encoder编码得到的上下文，$\pi$π表示共享同一组注意力权重的层数。因此，SAN（Shared Attention Network）的encoder-decoder attention注意力权重共享策略为，对网络层进行分组，属于同一组的网络共享加权融合以后的上下文结果。

另外，作者还提出了一种共享策略，能够自适应地确定如何对网络层进行分组，使得属于同一组的网络使用相同的注意力权重：

$sim(m,n)$sim(m,n)表示从第$m$m层网络到第$n$n层网络构成的网络组$b(m,n)$b(m,n)中任意两层网络间注意力权重的相似性度量。当$sim(m,n)\geq \theta$sim(m,n)≥θ时，就认为从第$m$m层网络到第$n$n层网络应该共享同一组注意力权重，$\theta$θ为人为设定的一个超参数。

因此，从第1层网络开始，找到满足$sim(1,n)\geq\theta$sim(1,n)≥θ的最大的$n$n，则1~n层共享同一组注意力权重，得到$\pi_1=n$π1​=n，然后再从第$n$n层开始，得到$\pi_2$π2​，直到所有网络层都有考虑到。假设共有$M$M层网络，被分成了$N$N组，则最后的共享策略为$\{\pi_1,\pi_2,...,\pi_N\}${π1​,π2​,...,πN​}，表示自上而下各个分组中包含的网络层数，且$\sum^N_{i=1}\pi_i=M$∑i=1N​πi​=M。

因此SAN训练的过程如下：

Experiments

相关文献

[1] Zhang B, Xiong D, Su J. Accelerating Neural Transformer via an Average Attention Network[C]//Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2018: 1789-1798.

[2] Jiatao Gu, James Bradbury, Caiming Xiong, Victor O.K. Li, and Richard Socher. Non-autoregressive
neural machine translation. In International Conference on Learning Representations, 2018.