论文阅读|对Non-local 的轻量化改变ANLNet

Introduction

《ANLNet：Asymmetric Non-local Neural Networks for Semantic Segmentation》发表于ICCV 2019，是基于Non-Local的思路往轻量化方向做改进。**Non-Local模块是一种效果显著的技术，但同时也受限于过大计算量而难以很好地嵌入网络中应用。为了解决以上问题，ANLNet基于Non-Local结构并融入了金字塔采样模块，在不牺牲性能的情况下极大减少计算量和内存负担，得到APNB模块，在充分考虑了长距离依赖的前提下，对APNB做修改，融入了不同层次的特征，得到AFNB，从而在保持性能的同时极大地减少计算量。

code：https://github.com/MendelXu/ANN.git.

• 结合SPP和Non-Local，前者从不同大小区域内提取出关键点，后者对这些关键点建模长距离依赖；
• Non-Local 需要计算每一个像素点，可以看作是点对点建模；ANL-Net只计算通过SPP提取出的关键点，可以看作是点对区域建模；
• AFNB包含两个输入，一个高级特征图，一个低级特征图，分别对应图中的Stage5和Stage4。其中高级特征图作为Query，低级特征图作为Key和Value；
• APNB的建模方式与AFNB类似。

Non-local 的不足：一个不同的non-local块如上图a所示。该块首先要计算所有位置之间的相似性，这需要通过矩阵相乘运算得到，给定一个尺寸为C×H×W的输入特征图，其计算复杂度为$O(C{H}^2{W}^2)$。之后它还需要gather所有位置对它们的影响，这一步的计算复杂度也是$O(C{H}^2{W}^2)$。显然，这种方式的很不高效，高昂的计算代价使得它很难在实际中使用。于是及需要一个不牺牲其性能，同时更加高效的work。

作者观察到只要key分支和value分支的输出有相同的尺寸，non-local block最终的输出尺寸就不会变化。所以如果key分支和value分支值取少量的表示性点，即将图1 a中比较大的N值变成b中的一个小很多的S，这样就可能实现不牺牲性能的情况下极大地减小复杂度，。这一思想就是该work中的核心创新思想。

具体到如何选择更少的表示性点，这里将spatial pyramid pooling嵌入到non-local blocks中来实现（APNB）。APNB和 标准non-local block的cost比较见图2。

此外，作者在APNB的基础上融合了不同stages的特征，得到AFNB，旨在计算low-level和high-level 特征图之间每个像素的关系，从而生成了一个带有long range interactions的融合特征。

一些work表明，充分利用 long range dependencies能够提高性能，但是由于增加一个卷积层感受野也不足以覆盖相关区域，而且加深网络深度或选择一个更大的kernel来增加感受野的方式会带来很大的计算代价，很不高效，所以只是依赖卷积来捕获long range dependencies是受限的

Asymmetric Non-local Neural Network

论文中受限回顾了non-local的详细设计细节，然后介绍Asymmetric Pyramid Non-local Block，Asymmetric Fusion Non-local Block。

Revisiting Non-local Block

如图1所示，假设输入特征 $X\in R^{C\times H\times W}$，然后使用3个1×1卷积 $W_{\varphi }(X),W_{\Theta },W_{\gamma }$来将X变成不同的embeddings ：

$$\varphi =W_{\varphi }(X),\theta =W_{\theta }(X),\gamma =W_{\gamma }(X),这里\varphi ,\theta ,\gamma \in R^{\hat{C}\times H\times W}。$$

$\hat{C}$是生成的新embeddings的通道数量（一般比C小，这里降维了），接下来，三个embeddings被flatten到 $\hat{C}\times N，N=H\cdot W$。然后通过矩阵相乘 $V={\varphi }^T\times \theta$得到一个相似性矩阵 $V\in R^{N\times N}$.

之后，还会再V上做一个正则化操作得到一个统一的相似性矩阵， $\overrightarrow{V}$这里正则化函数选择的是softmax，等价于self-attention机制，再许多任务中都适用。

对于γ中的每个位置，attention层的输出是： $O=\overrightarrow{V}\times {\gamma }^T$ 。这的 $O\in R^{N\times {\gamma }^T}$。最后的输出为：

$$Y=W_o(O^T)+XorY=cat(W_0(O^T),X).$$

$W_o$由一个1×1卷积实现，作为一个权重参数来调整 non-local 操作相对于原始输入X的重要性，更重要的是讲通道从 $\hat{C}$恢复到原始输入X的通道数C。

Asymmetric Pyramid Non-local Block

前面提到了，为了减少Non-local的计算复杂度和内存消耗，将key分支和value分支值取少量的表示性点，即由原来的N到S。这里具体的数据流公式表示如下面：

将N变成S后：

Solution。所以，关键就变成了如何选取S个点。这里是在θ和γ之后分别增加采样模块 $P_{\theta }$来采样几个稀疏的anchor点，表示为 ${\theta }_P\in R^{\hat{c}\times S}，{\gamma }_P\in R^{\hat{c}\times S}$。这里的S就是要采样的点数。数学表示如下：

$${\theta }_P=P_{\theta }(\theta ),{\gamma }_P=P_{\gamma }(\gamma ).$$

$\varphi 和anchor点{\theta }_p$之间的相似性矩阵 $V_P$的可以由下计算：

$$V_P={\varphi }^T\times {\theta }_p$$

这里计算得到的相似性矩阵 $V_P$是一个非对称矩阵，size为N×S。随后的步骤和标准的Non-local一样，对 $V_P$做正则化，然后和 ${\gamma }_P^T$相乘得到最终的attention输出 $O=\overrightarrow{V_P}\times {\gamma }_P^T$。最后最终的输出为

$$Y_P=cat(W_0(O_P^T),X).$$

这样相比原来 $O(\hat{C}{N}^2)的复杂度，现在复杂度为O(\hat{C}NS),计算复杂度为原来的\frac{S}{N}$，理想情况下S要足够小才能极大的减小计算复杂度，但是通道的减少可能会找出性能的下降，所以这里利用SPP空间金字塔池化来选点，这样即能够减少点， 也能够利用SPP的全局表示能力保证性能不下降，而且SPP是parameter-free且高效的。

于是得到了图3中的设计：

其中关键的改变就是θ和γ后都分别增加了一个SPP模块来采样表示性点，采样过程如图4. 使用了几个池化层，然后池化结果flatten并concatenate起来作为下层的输入。这里四个池化层的宽度和高度都一样，大小为{1，3，6，8}，所以总共有S=1+3×3+6×6+8×8=110个anchor点。（两个分支的SPP模块都一样，因为要保证这两个分支的维度一样）

这里SPP提供了足够的全局统计数据来弥补通道降维可能导致的性能下降。假设输入特征图宽和高都96，则N=W×H，总共节约计算为96×96/110=84倍！

Asymmetric Fusion Non-local Block

通常的特征融合都是addition/concatenation，是以pixel-wise和局部的方式执行的，而这里借助了APNB的非对称性，通过APNB来融合多level的特征来捕获long range dependencies，这个模块叫做Fusion Non-local Block.

FNB有两个输入：一个是high -level的特征图 $X_h\in R^{C_h\times N_h}$,还有一个是low-level 的特征图$X_l\in R^{C_l\times N_l}$, $N_h,N_l$分别是$X_h,X_l$的空间位置数量，$C_h,C_l$则分别是通道数，这里首先对他们使用一个1×1卷积降维到相同的通道$\hat{C}$，得到两个embedding ${\epsilon }_h\in R^{\hat{C}\times N_h}和{\epsilon }_l\in R^{\hat{C}\times N_l}$。然后他们相乘得到相似性矩阵$V_F={\epsilon }_h^T\times {\epsilon }_l.$ $V_F\in R^{N_h\times N_l}$，然后对$V_F$运用正则化，然后再和前面一样聚集所有位置对他们的影响$O_F=\overrightarrow{V_F}\times {\epsilon }_l^T$，输出的注意力$O_F\in R^{N_h\times \hat{C}}$，表示的是${\epsilon }_l$对${\epsilon }_h$的奖励，是从${\epsilon }_l$精心选择的。之后$O_F$送入1×1卷积恢复通道$C_h$。于是得到的输出如下:

$$Y_F=cat(W_0(O_F^T),X_h).$$

整个流程可以见图3. AFNB继承自APNB。

Experiments

本人非学习语义分割的，所以只关注其对Non-local的改进，更多的实验信息请见论文。