[论文评析] ArXiv,2021, Focal Self Attention技术分析

动机

通过整合fine-grained local attention和coarse-grained golbal attention,来克服各自的问题。

Focal self Attention

所谓的Focal self attention简单来说就是对距离Query越近的区域进行细粒度fine granulity的attention, 对距离Query越远的区域进行粗粒度的attention，通过调整粒度的级别最终会得到一个层次化的feature maps, 然后把这些feature map经过Flatten、Concatenation等操作转化为Vector, 然后对Vector分别进行Linear projection得到 Key和Value, 然后进行常规的Scaled Dot-Product Attention (SDPA)即可。

基本概念

为了控制粒度，作者引入了三个基本定义，分别为：
1.Focal level $L$: 用于进行Focal self attention的粒度级别的数量；
2.Focal Window size $s_w^l$: 第$l$个粒度级别上用于提取tokens的sub-window的尺寸；
3.Focal Region size $s_r^l$: 第$l$个粒度级别上水平/垂直方向上包含的sub-window的数量；

流程

1.网格化
假设Feature map的原始尺寸为$[H,W,d]$, 常规做法是对每个token都要进行一次attention, time和memory开销太大， 因此通过网格化将feature map分成一个个的sub windows， 然后逐window进行attention, 即所谓的window-wise self attention。
假设sub window的尺寸为$[s_p,s_p]$, 则总的windows数量为$\frac{HW}{s_p{s}_p}$, 记网格化后的尺寸为$[sp,sp,H/s_p\cdot W/s_p\cdot d]$.
注：通过设置不同的大小的$s_p$(实际上就是$s_w^l$), 可以得到不同粒度级别的surroundings.
2. 子窗口池化
假设输入：$X\in R^{H\cdot M\cdot d}$,
对于第$l$个粒度级别， 其网格化操作形式化定义为：
$\hat{X}=Reshape(X)\in R^{s_w^l\mathrm{x}{s}_w^l(H/s_w^l\cdot W/s_w^l\cdot d)}$,
然后池化操作形式化为：
$x^l=f_p^l\in R^{H/s_w^l\cdot W/s_w^l\cdot d}$.
然后遍历$l,l\in 1,2,...,h$， 最终得到输入feature map的层次化surroundings： { x l } l = 1 h \{x^{l}\}_{l=1}^{h} {xl}l=1h​.
3.聚合操作
$x^l$为feature map，首先通过Flatten操作将其转化为Vector, 然后再把$h$个这样的Vector进行Concat操作， 最终得到的tokens向量为 $x_{total}\in R^{s\mathrm{x}d}$, 其中$s={\sum }_{l=1}^h(s_w^l{)}^2$.

4.计算Query, Key和Value
通过分别进行Linear projection即可得到Query, Key和Value，形式化定义如下：
$Q=f_q(x^1)$
$K=f_k(x_{total})$
$V=f_v(x_{total})$

5.Attention计算
对于位于第$i$个sub-window $Q_i\in R^{s_p\cdot s_p\cdot d}$的Query， 其对应的Key和Value分别记为$K_i\in R^{s\cdot d}$, $V_i\in R^{s\cdot d}$.
计算公式如下：
$Attention(Q_i,K_i,V_i)=softmax(\frac{Q_i{K}_i^T}{\sqrt{d}+B})V_i$
最终输出的尺寸依然为$R^{s_p\cdot s_p\cdot d}$
然后遍历 i , i ∈ { 1 , 2 , . . , H / s p ⋅ W / s p } i, i \in\{1,2,.., H/s_{p} \cdot W/s_{p}\} i,i∈{1,2,..,H/sp​⋅W/sp​}重复上述操作即可。

时间复杂度分析

网格化实际上就是对数据进行Reshape， 假设原始输入$X\in R^{M\mathrm{x}N\mathrm{x}d}$
,Reshape之后第$l$层$\hat{X}\in R^{s_w^l\mathrm{x}{s}_w^l\mathrm{x}(\frac{M}{s_w^l}\cdot \frac{N}{s_w^l}\cdot d)}$
对其池化的时间复杂度为$O(s_w^l\cdot s_w^l\frac{M}{s_w^l}\cdot \frac{N}{s_w^l}\cdot d)=O(M\cdot N\cdot d)$, 对所有L个层池化的时间复杂度为$O(O(L\cdot M\cdot N\cdot d))$.

每个Query所属sub-window $Q_i\in R^{s_p\mathrm{x}{s}_p\mathrm{x}d}$, 对应$K_i,V_i\in R^{s\mathrm{x}d}$, 因此计算Attention的时间复杂度为$O((s_p{)}^2\cdot d\cdot s)$, 共有$\frac{M}{s_p}\cdot \frac{N}{s_p}$个这样的sub-window, 因此进行Attention总的时间复杂度为$M\cdot N\cdot d\cdot ({\sum }_1^L(s_r^l{)}^2)$。

因此，总的时间复杂度为：$O(M\cdot N\cdot d\cdot (L+{\sum }_{l=1}^L(s_r^l{)}^2))$.

点评

1.fine-grained local attention不能抓住global的信息，coarse-grained golbal attention虽然能抓住global information但因为粒度粗，因此，这两种方式实际上都不能发挥出NLP中原始Transformer中attention的建模能力。
focal self attention 实际上正是对这两种attention的整合，因为它既能抓住closet surroundings的fine-grained local信息，同时又能抓住far surroundings的coarse-grained global信息。具体来说， Key和Value的来源不是某个单一粒度级别的tokens, 而是多个粒度级别的层次化的tokens的融合。
2.所谓的focal self attention最核心的东西可以看成是常规scaled dot-product attention的一个前置操作。

与Transformer中MHA的比较

1.Transformer中的MHA处理的是序列，对于长度为n，维度为d的序列$X\in R^{n\mathrm{x}{d}_{model}}$, $Q\in R^{n\mathrm{x}{d}_k}$, 对应$K\in R^{m\mathrm{x}{d}_k},V\in R^{m\mathrm{x}{d}_v}$, $d_k,d_v$完全取决于Linear projection输出的维度，$m$为K-V对的数量；而在focal self attention这里，$Q_i\in R^{s_p\cdot s_p\cdot d}$, 其中$d$为Figure的channel, $s_p$取决于网格化时设置的每个sub-window的宽度，对应$K_i,V_i\in R^{s\cdot d}$, 其中s为所有L层池化结果经Flatten后再Concat的长度。
2.原始Transformer中为MHA, 通过多次Linear projection得到多个子空间，分别进行SDPA再聚合，论文中提到这既可以节省memory，提高efficancy,同时提升泛化能力； focal self attention中用的是网格化方法，Attention是在Window级别，而不是每个Query position, 这减少了memory, 节省了time.

几个疑点

1.论文中Fig.4 与正文中描述略不同， 图中是不同level的tokens先concatenation后进行linear projection产生$K_i,V_i$, 文本描述部分是对每个level的分别进行linear projection (即 $f_k,f_v$)， 然后再concatenation.

2.Fig. 4中Query position所在区域恰好位于Feature map的中心，因此每个level的tokens都是方形， 然而当Query position位于边缘时，怎么办？ 需要做填充吗？ 从论文中没有看到类似信息。

3.$s_p$, $s_w^l$与$s_r^l$之间应该满足的关系：
这一点论文中并没有明确说明。
很明显， { s w l } l = 1 L \{s_{w}^{l}\}_{l=1}^{L} {swl​}l=1L​应该都能被$s_p$整除， 特别的，当$s_w^l=1$时，粒度级别最细； 当$s_w^l=s_p$时， 粒度级别最粗。

三者之间的关系， Fig.4中Query所在的蓝色区域周围上下左右包含的尺寸为$s_w^l\mathrm{x}{s}_w^l$的方格的数量都为2, 因此可以很容易得出结论：$s_r^l=4+\frac{s_p}{s_w^l}$.

Focal Transformer

如图， Focal Transformer包含多个Stage： { S t a g e i } i = 1 4 \{\rm Stage i\}_{i=1}^{4} {Stagei}i=14​, 每个$\mathrm{S}\mathrm{t}\mathrm{a}\mathrm{g}\mathrm{e}\mathrm{i}$里面通过Stack相同的building block组成,数量为$N_i$，每个building block主要由Focal Self-Attention和Multi-Layer Perceptron 组成， 最核心的就是上面提到的Focal self-attention.

可以看到，在每个Stage之前都有一个Patch Embedding (PE)， PE实际上就是一个Convolution层， 其中第一层的PE作用是将$X\in R^{M\mathrm{x}N\mathrm{x}3}$映射到$R^{M\mathrm{x}N\mathrm{x}d}$, 然接下来，每个Stage开始之前都会有一个PE将上一个Stage输出的Feature map的Spatial dimension减小为原来的1/2， Channel dimension增大为原来的2倍。

总结

Reference

1.ArXiv, 2021, Focal Self-attention for Local-Global Interactions in Vision Transformers.