什么是轴向注意力（Axial Attention）机制

Axial Attention，轴向注意力，有行注意力（row-attention）和列注意力（column-attention）之分，一般是组合使用。

原文阅读：https://arxiv.org/pdf/1912.12180v1.pdf

self-attention已经显示出了良好的性能，不过其缺点也是明显的：它的使用使得模型的参数量和计算量开始飙增，特别是应用在nlp的长序列问题和图像的像素点上时。以后者为例，单张图像大小为H*W(*3)，一个像素点需要和其他所有像素点作注意力，即H*W，一共又有H*W个像素点，则一次self-attention的计算量在O((H*W)²)，当图像的分辨率不断增高，这种平方式增长的计算量是不能接受的。

轴向注意力应运而生，简单来说，是分别在图像的竖直方向和水平方向分别进行self-attention，这样计算量被削减为O(2*H*W)=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，再对结果相加的形式。