Stand-Alone Self-Attention in Vision Models阅读笔记

NeurIPS 2019
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, Jonathon Shlens

一、简介

受限于感受域的大小设定，卷积很难获取长距离的像素关系，而在序列模型中，已经能很好地用attention来解决这个问题。

将内容交互的模块（注意力机制）作为视觉模型的主要单元。为此，我们提出了一个简单的局部自我关注层（local self-attention layer），该层既可以用于小输入，也可以用于大输入。

二、背景

2.1 卷积

卷积神经网络利用较小的领域（即核大小）来鼓励网络学习特定层内的局部相关结构。

如图是一个3×3的核，通过与对应的像素相乘加权求和的方式得到中心点像素的结果。

CNN采用权重共享的方式，强调学习表示中的平移等价性，因此将卷积的参数计数与输入大小分离。

2.2 自注意力机制

注意力机制的作用是学习一个背景下的重要区域。独立自注意力层定义了三个概念query,key, value。

self-attention直接模拟远距离交互的能力及并行性。

如图是Single-headed attention（单头注意力机制）。

输入是一个特征图，分别通过卷积层得到querys、keys和values。然后将每一个query单元与keys中的各个key单元相乘求和，再经过softmax得到参数，该参数与values对应像素位置相乘后加权得到最终的输出单元。

公式如下：

该注意力机制中缺少编码位置信息，这使得他的排列是等变的，从而限制了视觉任务的表现力。

已经有实验证明，相对位置要优于绝对位置，绝对位置要由于没有加位置。

如上图所示，使用二维相对位置编码(relative attention)，将行偏移和列偏移的编码进行concatenate。该二维相对位置编码是可以被训练的。

上式就是同时考虑query和key的内容间的相似性以及相对位置。

由于考虑了相对位置，self-attention也拥有了类似卷积的平移不变性。

三、Fully Attentional Vision Models

构建一个完全注意力的架构。

Replacing Spatial Convolutions：

空间卷积为核大于1×1的卷积，文中把所有的空间卷积替换成self-attention layer，若需要下采样，则在层后接一个stride为2的2×2的平均池化。

self-attention layer：

CNN的初始层，又称stem，在学习局部特征（如边缘）发挥着关键的作用，后面的层用来识别全局对象。并且由于输入图像通常比较大，初始层的任务还有下采样。

stem中的内容包含RGB像素，这些像素是高度空间相关的，独立起来则失去了意义，没有丰富的content信息，使用第二部分中的self-attention效果不佳。

文中提出了一个空间感知(spatially aware)自注意力层来解决标准的自注意力层无法包含空间信息的问题。

将输入经过m个1×1的卷积核得到m个V。emb是位置矩阵，分row和col，该矩阵是可以学习的。

p（a, b, m）表示V中注入位置的信息，由于不同相对位置的emb不同，所以同一个像素点在不同的相对距离下就有不同的值，类似与卷积的属性，有利于边缘特征学习。

x是由Q和K得出的参数，与self-attention layer得参数求的方法一致。

四、实验