Multi-head Attention中多头的作用

原论文中的说法

Attention is all you need论文中说, 将模型分为多个头, 每个头都是一个子空间, 相当于每个头去关注不同角度的信息, 然后进行融合, 类似于CNN中的每个卷积核捕获不同的纹理信息, 或者理解为是一种ensemble思想.

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.

头和头之间也没有相互之间的制约, 每个空间都是完全独立的, 理论上可以做到子空间关注点的差异化.

在A Multiscale Visualization of Attention in the Transformer Model这篇论文中, 可视化了层中各个头的Attention情况, 如下图:

可以看到:

• 同一层中多数头的关注模型近似, 如第3层中的1, 2, 3, 4头

• 每一层总有几个头与其他头区别较大, 如第2层中的0头

是可以印证不同子空间关注的角度不同, 只是有些头之间的差异程度, 没有我们想象中的这么大.

空间差异探讨

差异是如何产生的. 如果所有头的参数初始化为一模一样的, 由于输入相同, 则输出也是相同的, 因此在梯度更新的时候, 更新的情况也是一样的. 这样无论更新的步数多少, 单层中所有头的参数会一直保持相同, 头和头之间没有差别, 相当于就是一个头.

可以说头和头之间的差别, 起始来源是初始化的差异性.

使用多头可能的原因: 信息冗余

以BERT base为例, 最大长度512, 此时$Q, K \in \mathbb{R}^{512 \times 768}$, 如果直接将两个向量点积, 得到一个(512, 512)大小的Attention矩阵, 去计算其中一个位置的数值时, 需要两个长度为768的向量. 实际上可能不需要向量的全部参数都参与计算, 就能得到高质量的结果了. 这里会产生计算的冗余.

而且依据上面的图, 每个token通常只是注意到有限的若干个token, 说明得到的Attention矩阵是很稀疏的. 这可能是因为使用了Softmax函数对点积结果进行转换的原因, 会将元素间的差距拉大, 比较小的点积值就被转化为接近0的结果. 而这种稀疏性意味着(512, 512)大小的Attention矩阵本身也是冗余的, 里面没有包含这么多信息, 是可以由两个低维矩阵的乘积得到, 相当于可以低秩分解.

从SVD的角度看, 这种稀疏的Attention矩阵的奇异值, 有很多是接近于0的. BERT base中有12个头, 每个头的head size为64, 相当于假设了每个(512, 512)大小的Attention矩阵, 其明显不等于0的奇异值不超过64个, 因此可以用$Q, K \in \mathbb{R}^{512 \times 64}$通过点乘来拟合Attention矩阵.

从这个角度来看, 每个头都可以拟合出单个注意力, 多个头都拟合相当于一种model ensemble, 可以提升整体的效果.

总结

多头可以带来好处:

• 多个子空间从各个角度获取信息, 且每个子空间获取的信息都不差(信息冗余部分), 进行融合, 提升效果

• 计算点积时, 相比于hidden_size * hidden_size的计算量, head_size * head_size会小很多

延伸

这种split, transform, merge的模型设计思想,都能够带来多路信息融合的效果.

参考资料

• 为什么Transformer 需要进行 Multi-head Attention？ - 香侬科技的回答
• BERT中，multi-head 7686412与直接使用768*768矩阵统一计算，有什么区别？ - 苏剑林的回答