Axial Attention,轴向注意力,有行注意力(row-attention)和列注意力(column-attention)之分,一般是组合使用。
原文阅读:https://arxiv.org/pdf/1912.12180v1.pdf
self-attention已经显示出了良好的性能,不过其缺点也是明显的:它的使用使得模型的参数量和计算量开始飙增,特别是应用在nlp的长序列问题和图像的像素点上时。以后者为例,单张图像大小为HW(3),一个像素点需要和其他所有像素点作注意力,即HW,一共又有HW个像素点,则一次self-attention的计算量在O((H*W)2),当图像的分辨率不断增高,这种平方式增长的计算量是不能接受的。
轴向注意力应运而生,简单来说,是分别在图像的竖直方向和水平方向分别进行self-attention,这样计算量被削减为O(2HW)=O(H*W)
显然,单独使用行或列注意力是无法融合全局信息的,只有组合起来才能获得与full attention相似的性能。
以两个维度为例,组合方式可以先行再列,可以先列再行,也可以分别进行行与列注意力,再对结果分别相加。
- 先行再列:x1 = RowAtten(x), out = ColAtten(x1)
- 先列再行:x1 = ColAtten(x), out = RowAtten(x1)
- 分别进行:out = RowAtten(x) + ColAtten(x)
代码展示可以参考我的另一篇博客论文学习——VideoGPT,使用的是对时间和空间共三个维度分别进行axial attention,再对结果相加的形式。
版权归原作者 胖虎干嘛了 所有, 如有侵权,请联系我们删除。