《论文阅读》AXIAL ATTENTION IN MULTIDIMENSIONAL TRANSFORMERS

留个笔记自用

AXIAL ATTENTION IN MULTIDIMENSIONAL TRANSFORMERS

做什么

点云的概念：点云是在同一空间参考系下表达目标空间分布和目标表面特性的海量点集合，在获取物体表面每个采样点的空间坐标后，得到的是点的集合，称之为“点云”（Point Cloud）。
点包含了丰富的信息，包括三维坐标X，Y，Z、颜色、分类值、强度值、时间等等，不一一列举。
一般的3D点云都是使用深度传感器扫描得到的，可以简单理解为相比2维点，点云是3D的采样

做了什么

这里又是一种transformer的变种，基于轴向注意力，自然地与编码和解码设置中张量的多维度对齐，所提出的层结构允许在解码期间并行计算绝大多数上下文，而不引入任何独立假设

怎么做

首先仍然是摆一个传统transformer来作为回忆

X是输入序列∈R^N×D
显然这里的QK运算的O(N²)是大家优化的重点，这里提出了很多种方法，比如限制了attention的范围，比如堆叠多个稀疏的attention层这样
接下来就开始这篇文章的核心介绍，也就是axial transformer

首先介绍为高维数据张量开发基于自我注意的自回归模型的基本构件，axial attention
这里用Attention_k(x)来代表，这里的意思就是表示这个构件在输入序列x的k轴上执行attention，沿着 k 轴混合信息，同时保持沿着其他轴的信息独立
这里介绍了一种实现的方法，简单来说就是计算Attention_k(x)的时候，通过将除k轴外所有轴转置到批处理轴来实现
这里举了个例子，假如输入x是图像的时候，假设图像大小N = S×S也就是二维正方图时，axial attention的耗时达到了

很好理解，这里相比于常规的transformer的平方耗时，也就是O（N×N），因为常规transformer需要计算每个点之间的关系，对于图像来说，点数量就是N
节约了

简单来说，这里的axial attention的方法就是假设输入是（B，N，H，W）的时候，B是batch，N是Num，H和W是特征图维度，如果传统transformer计算的时候，会计算每个Batch的每个N的H×W的点之间的关系，复杂度就是H×W的平方，这里做法就是把W转换到B上去，也就是留下（B×W，N，H）然后这样进行普通的transformer，复杂度就计算的仅仅是H平方，而这里常见的做法也很多，可以用einsum来实现，而前面节约的时间是根号N也就是S，是因为对于每一个轴它的时间都是S²
了解完了axial attention，接下来就是将其实现成axial transformer

首先定义模型，这里举了个任务的例子是单通道图像x∈R^H×W的分析
然后自然是embedding环节，将H×W的图每个点都embedding，得到H×W×D的特征图，transformer架构的意义就是对D维张量进行分析
这里就提出了两种模型架构
第一种是Inner Decoder内部解码器

首先是第一行，embed（x）就是前面说的对图像x的特征嵌入
然后是shiftright，就是直接将整个输入向右移动一个像素，效果就是将感受野向左移动一个像素，这个操作确保了屏蔽行关注层将当前像素排除在它们的感受野之外
至于PositionEmbeddings就是对位置的嵌入，这里采用的方法是additively factorized，简单来说就是，行是H × 1 × D的嵌入，列是1 × W × D，而位置嵌入就是各自的总和
最后是L_row，这里的意思是被mask的行关注的数量
然后是第二种框架Outer Decoder外部解码器


u表示在当前像素上捕获的上下文，是由正常的transformer block也就是未mask的行和mask的行共同计算的，shiftdown和前面同理，是全部往下移动一个像素，意思就是这样不仅仅能捕捉到当前行的像素，还能捕捉到前面行的像素
然后改进了一种sample的方法

这里的意思就是对于每行都计算一个上下文u，这个u包含了前面所有行包括自己行的信息，对于该行的每一列，都基于u去sample一个点

总结

1.看到后面越来越懵逼，看着看着就完全不懂逻辑在哪了
2.总归主要还是为了理解一下axial attention的作用，在点云处理中感觉完全可以借鉴这种方法，将特征嵌入至位置+特征+法向量等多维空间，然后进行单维的分析