bert性能优化之——用另一种方式整合多头注意力_解决多头注意力内存太大的问题

今天我想给大家介绍这样一篇论文：Multi-Head Attention: Collaborate Instead of Concatenate。作者均来自洛桑联邦理工学院

看过我文章的同学肯定知道，我一直在关注bert模型的性能优化相关研究，而这篇论文正好是与transformer的性能优化相关，并且我认为它的方法不需要做太多的适配就能应用在预训练模型上面，实用性较高，因此推荐给大家。

众所周知，经典的transformer架构中采用了multi-head attention机制来引导模型从不同角度学习不同的语义信息，从各种实验对比中也能发现多头机制确实能够提升模型在NLP任务上的精度。然而，随着目前大规模预训练模型的普及，多头注意力机制在带来精度提升的同时，也增加了计算的成本，带来了性能上的限制。

因此最近两年，有些研究人员尝试从不同的维度去探讨是否能从多头机制上去优化transformer的性能。有些工作重点关注了多头中每个头的注意力到底捕捉了哪些语义信息，头与头之间捕捉的信息是否有冗余，例如这篇论文：Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned，提出了一种量化注意力头重要程度的方法。还有一些工作更加激进，提出了多头注意力机制是否有必要的疑问，例如这篇论文：Are sixteen heads really better than one。它对transformer中的每个头都做了消融实验，探讨了每个头在不同下游NLP任务上的作用，最后提出了一种迭代式地剪枝注意力头的方法。

与上述工作不同，本篇论文并非直接对注意力头进行结构性剪枝，而是关注所有注意力头捕捉的通用信息，试图将这些信息提取出来作为sharing weights，每个头各自关注自己独有的工作，从而减少多头注意力计算时的成本。下面我就详细得为大家解读这篇论文的工作。

单个注意力头的减负

在那篇经典的Attention is all you need论文中，对于注意力分数的计算是这样的：

备注一下：论文这里的公式貌似有点问题，最后一项应该是我推导出的项。

因此最后两项计算可以舍弃。又因为前面两项中，不存在  ，因此我们甚至不用去定义这个bias项。

另外，对于上述推导式的第一项，由于其计算了Query和key的相互关系，因此相当于捕捉了上下文的相关信息，而第二项只包含了key的content信息，相当于捕捉了原文内容上的信息。

多头注意力的整合

传统的transformer中，对于不同的注意力采取的整合方式是直接拼接，如下所示：

PCA之前我有文章科普过，相当于是找一个投影空间，让所有高维的点在这个空间上的投影尽量分开，即方差尽量大。因此variance代表了不同主成分上包含原始矩阵的信息量。variance越高，信息量就越多。

实验结果如图所示：

提取通用信息

b模式相当于对不同的head抽取不同维度的矩阵信息；c模式则是让所有head都共享映射矩阵；d模式则是在共享映射矩阵的基础上，进一步压缩最终输出的整合矩阵的维度，达到压缩维度的效果。

实验

论文作者利用上述优化后的transformer架构进行了NMT的实验，主体网络为encoder-decoder，实验结果如下：

可以看到，collabHead在维度缩减到1/4时，仍然能保持跟原始维度相近的效果，说明经过压缩之后，模型只损失了较少的信息。

论文说到这里还没结束，下面要说的才是我比较关注的内容，即如何对预训练后的bert模型应用collabHead，从而提升bert的inference效率。

More collabHead on Bert

通过上述实验，已经能够证明CollabHead模式相比原始的简单拼接模式，在提升性能的同时，只会损失很小的精度。而Bert模型的主体架构也是transformer，因此它也可以利用这个优化达到性能提升的效果。最简单的方式就是在预训练的时候就采用这种架构，而论文也比较推荐这种方式。但是对于我们这种硬件条件有限制的企业来说，从头预训练一个模型似乎不太现实，那么有没有办法直接在finetune过程实施这种优化，从而达到性能提升的效果呢？

答案当然是可以的。论文提出了一种re-parameterize方式，直接对预训练模型中的attention权重进行张量分解，使得分解后得到的矩阵能分别对应上述优化中的各个参数。

该方法的主要核心为Tucker 张量分解，或者更具体得说是CP分解。（CP分解是Tucker分解的一种简化特例）。对于张量分解，可以参考其他博客中的讲解zhuanlan.zhihu.com/p/25 。这里就简单介绍一下。

矩阵A，B，C通常称为因子矩阵，在一定程度上包含了原始X中各个维度上的主要信息，而矩阵G通常称为核张量，用于表征不同因子矩阵之间互相关联的程度。

回到我们的多头注意力优化问题上。我可以尝试对多头拼接后的key/query映射权重参数进行张量分解： 

下图为其矩阵分解图例，看完应该会对上述分解有个直观的理解：图参考自sandia.gov/~tgkolda/pub

备注：张量分解其实也是一个优化问题，一般通过计算分解前和分解后矩阵对应元素的mse loss来进行优化学习。因此分解过程也是需要一定的时间成本。

当然整体的re-parameterize操作除了上述步骤外，还有一些应用的小技巧。在实际应用时，通常按如下步骤进行：

1、使用原始的bert模型对下游任务进行finetune，得到一个finetune-bert-task模型。

2、对finetune-bert-task进行re-parameterize操作，得到压缩后的re-finetune-bert-task模型。

3、使用re-finetune-bert-task模型对下游任务再进行少量迭代的finetune。

经过实验验证，步骤3对于最后模型的精度还是很有帮助的，建议保留。

实验验证

我本来是想将该方法在我的NER任务中进行实际验证，但是发现好多张量分解的工具在tensorflow的静态图（尤其是estimator模式）下不太适配。如果有的同学对这方面实现兴趣，可以看一下tensorly框架，它支持tensorflow2.0的动态图模式、pytorch以及MXNET，在github上的一个issue上搜到其貌似也支持静态图，配置如下：

import tensorly as tl

tl.set_backend("tensorflow_graph")

另外还有一个框架tensorD,我试了一下也有bug。好在论文作者非常良心得放出了基于pytorch的开源代码，链接如下：github.com/epfml/collab。核心内容主要在两个文件：collaborative_attention.py 以及swap.py。前者主要定义了CollabHead的优化逻辑。后者则是定义了如何将预训练模型中的权重参数进行re-parameterize。下面主要看一下re-parameterize的内容。代码如下：

new_layer = CollaborativeAttention(

            dim_input=layer.dim_input,

            dim_value_all=layer.dim_value_all,

            dim_key_query_all=dim_shared_query_key,

            dim_output=layer.dim_output,

            num_attention_heads=layer.num_attention_heads,

            output_attentions=False,

            attention_probs_dropout_prob=layer.attention_probs_dropout_prob,

            use_dense_layer=layer.use_dense_layer,

            use_layer_norm=layer.use_layer_norm,

            mixing_initialization=MixingMatrixInit.CONCATENATE,

        )

_, factors = parafac(

                WQWKT_per_head.detach(), dim_shared_query_key, init="random", tol=tol

            )

WQ_shared, mixing, WK_shared = factors

new_layer.key.weight.data.copy_(WK_shared.transpose(0, 1))

new_layer.query.weight.data.copy_(WQ_shared.transpose(0, 1))

new_layer.mixing.data.copy_(mixing)

首先，定义出带CollabHead结构的attention模型，然后调用了tensorly的CP分解parafac方法，得到三个因子矩阵。最后分别将三个矩阵权重参数赋给计算图中对应的变量中。

bq_per_head = layer.bQ.reshape([layer.num_attention_heads, -1])

content_bias = bq_per_head.unsqueeze(1) @ WK_per_head

content_bias = content_bias.squeeze(1)

new_layer.content_bias.weight.data.copy_(content_bias)

另外，根据论文所说，该方法还可以与其他模型压缩方法一起使用，例如和albert还有DistilBert等联合使用，效果也还不错，感兴趣的同学可以进一步验证。

小结

本次解读的论文主要通过分析transformer中注意力头之间的冗余信息，并设计了一种优化后的多注意力头整合方法，将通用的信息提取出来共享于所有注意力头，让每个注意力头可以专注于捕捉独有的信息。此方法经过验证，可以在提升性能的同时只会损失极小的精度。另外，本论文还提出对预训练的bert模型中的attention权重进行张量分解，使得本文的多注意力头整合方法同样可以适用于预训练模型，拓展了本方法的适用范围，个人认为该方法值得在实际业务中进行尝试。

由于本人的数学功底不深，因此对于张量分解这块内容的解读比较浅，还有很多问题亟待解答，例如是不是可以通过其他分解方式更高效得去分解bert模型权重？这个就需要感兴趣的同学继续深入探索了。