【论文笔记】Non-local Neural Networks

Paper：https://arxiv.org/abs/1711.07971v1
Author：Xiaolong Wang, Ross Girshick, Abhinav Gupta, Kaiming He (CMU, FAIR)

0.简述

convolution和recurrent都是对局部区域进行的操作，所以它们是典型的local operations。受计算机视觉中经典的非局部均值（non-local means）的启发，本文提出一种non-local operations用于捕获长距离依赖（long-range dependencies），即如何建立图像上两个有一定距离的像素之间的联系，如何建立视频里两帧的联系，如何建立一段话中不同词的联系等。
non-local operations在计算某个位置的响应时，是考虑所有位置features的加权——所有位置可以是空间的，时间的，时空的。这个结构可以被插入到很多计算机视觉结构中，在视频分类的任务上，non-local模型在Kinetics和Charades上都达到了最好的结果。在图像识别的任务上，non-local模型提高了COCO上物体检测／物体分割／姿态估计等任务的结果。

1. Non-local Neural Networks

1.1 定义

按照非局部均值的定义，我们定义在深度神经网络中的non-local操作如下：
这里写图片描述

其中x表示输入信号（图片，序列，视频等，也可能是它们的features），y表示输出信号，其size和x相同。 $f(x_i,x_j)$ 用来计算i和所有可能关联的位置j之间pairwise的关系，这个关系可以是比如i和j的位置距离越远，f值越小，表示j位置对i影响越小。 $g(x_j)$ 用于计算输入信号在j位置的特征值。C(x)是归一化参数。

从上式中可以看出任意位置上的j，可以是时间的、空间的、时空的任意位置，都可以影响到i的输出值。作为对比，conv的操作是对一个局部邻域内的特征值进行加权求和，比如kernel size=3时， $i-1 \le j \le i+1$ 。 recurrent的操作则是i时刻的值仅基于当前时刻或前一时刻（j=i or i-1）。

另外作为对比，non-local的操作也和fc层不同。公式(1)计算的输出值受到输入值之间的关系的影响（因为要计算pairwise function），而fc则使用学习到的权重计算输入到输出的映射，在fc中 $x_j$ 和 $x_i$ 的关系是不会影响到输出的，这一定程度上损失了位置的相关性。另外，non-local能够接受任意size的输入，并且得到的输出保持和输入size一致。而fc层则只能有固定大小的输入输出。

non-local是一个很灵活的building block，它可以很容易地和conv、recurrent层一起使用，它可以被插入到dnn的浅层位置，不像fc通常要在网络的最后使用。这使得我们可以通过结合non-local以及local的信息构造出更丰富的结构。

1.2 表示形式

接下来我们会讨论f和g的几种形式。有意思的是我们的实验显示不同的表现形式其实对non-local的结果并没有太大影响，表明non-local这个行为才是主要的提升因素。
为了简化，我们只考虑g是线性的情况，即 $g(x_j)=W_gx_j$ ，其中 $W_g$ 是一个可学的权重矩阵，实际中是通过空间域的1*1 conv或时空域的1*1*1 conv实现的。
接下来我们讨论f的不同形式。

Gaussian。从非局部均值和双边滤波器来看，常见的f是高斯函数。在本文中，我们考虑f为如下形式：
这里写图片描述
其中 $x_i^Tx_j$ 是点乘相似度（dot-product similarity）。也可以用欧式距离，但是点乘在深度学习平台上更好实现。此时归一化参数 $C(x)=\sum_{\forall j}f(x_i,x_j)$ 。

Embedded Gaussian。高斯函数的一个简单的变种就是在一个embedding space中去计算相似度，在本文中，我们考虑以下形式：
这里写图片描述
其中 $\theta(x_i)=W_\theta x_i$ 和 $\phi(x_j)=W_\phi x_j$ 是两个embedding。归一化参数和之前一致，为 $C(x)=\sum_{\forall j}f(x_i,x_j)$ 。

我们发现self-attention模块其实就是non-local的embedded Gaussian版本的一种特殊情况。对于给定的i， $\frac {1}{C(x)}f(x_i,x_j)$ 就变成了计算所有j的softmax，即 $y=softmax(x^TW_\theta^TW_\phi x)g(x)$ ，这就是[47]中self-attention的表达形式。这样我们就将self-attention模型和传统的非局部均值联系在了一起，并且将sequential self-attention network推广到了更具一般性的space/spacetime non-local network，可以在图像、视频识别任务中使用。

扫描二维码关注公众号，回复： 3790872 查看本文章

另外和[47]的结论不同，我们发现attention的形式对于我们研究的case来说并不是必须的。为了说明这一点，我们给出另外两种f的表达形式。

Dot product。 f也可以定义成点乘相似度，即：
这里写图片描述
这里我们使用embedded version。在这里，归一化参数设为C(x)=N，其中N是x的位置的数目，而不是f的和，这样可以简化梯度的计算。这种形式的归一化是有必要的，因为输入的size是变化的，所以用x的size作为归一化参数有一定道理。

dot product和embeded gaussian的版本的主要区别在于是否做softmax，softmax在这里的作用相当于是一个激活函数。

Concatenation。 Concat是在Relation Networks [38]中用到的pairwise function。我们也给出了一个concat形式的f，如下：
这里写图片描述
这里[.,.]表示的是concat， $w_f$ 是能够将concat的向量转换成一个标量的权重向量。这里设置C(x)=N。

以上我们定义了多种变种的pairwise function f，这说明了我们的non-local的灵活性，我们相信也会有别的变种能够提升性能。

1.3 Non-local Block

我们将(1)式中的non-local操作变形成一个non-local block，以便其可以被插入到已有的结构中。
我们定义一个non-local block为：
这里写图片描述
其中 $y_i$ 已经在(1)式中给出了， $+x_i$ 则表示的是一个residual connection。residual connection的结构使得我们可以在任意的pretrain模型中插入一个新的non-local block而不需要改变其原有的结构（如 $W_z=0$ 作为初始化则完全和原始结构一致）。
一个示例的non-local block如图所示，我们看到式(2),(3),(4)中的这些pairwise function可以通过矩阵乘法来进行，(5)则可以直接concat得到。
这里写图片描述
non-local block的pairwise的计算可以是非常lightweight的，如果它用在高层级，较小的feature map上的话。比如，图2上的典型值是T=4，H=W=14 or 7。通过矩阵运算来计算parwise function的值就和计算一个conv layer的计算量类似。另外我们还通过以下方式使其更高效。

Non-local Blocks的高效策略。我们设置 $W_g$ , $W_\theta$ , $W_\phi$ 的channel的数目为x的channel数目的一半，这样就形成了一个bottleneck，能够减少一半的计算量。 $W_z$ 再重新放大到x的channel数目，保证输入输出维度一致。
还有一个subsampling的trick可以进一步使用，就是将(1)式变为： $y_i=\frac {1}{C(\hat x)}\sum_{\forall j}f(x_i,\hat x_j)g(\hat x_j)$ ，其中 $\hat x$ 是 $x$ 下采样得到的（比如通过pooling），我们将这个方式在空间域上使用，可以减小1/4的pairwise function的计算量。这个trick并不会改变non-local的行为，而是使计算更加稀疏了。这个可以通过在图2中的 $\phi$ 和 $g$ 后面增加一个max pooling层实现。
我们在本文中的所有non-local模块中都使用了上述的高效策略。

2. 视频分类模型

为了理解non-local networks的操作，我们在视频分类任务上进行了一系列的ablation experiments。
2D ConvNet baseline （C2D）。为了独立开non-local nets中时间维度的影响vs 3D ConvNets，我们构造了一个简单的2D baseline结构。
Table 1给出了ResNet-50 C2D backbone。输入的video clip是32帧，大小为224*224。Table 1中的所有卷积都是用的2D的kernel，即逐帧对输入视频进行计算。唯一和temporal有关的计算就是pooling，也就是说这个baseline模型简单地在时间维度上做了一个聚合的操作。
这里写图片描述

Inflated 3D ConvNet （I3D）。 Table 1中的C2D模型可以通过inflate的操作转换成一个3D卷积的结构。具体地，一个2D k*k大小的kernel可以inflate成3D t*k*k大小的kernel，只要将其权重重复t次，再缩小t倍即可。

我们讨论2种inflate的方式。一种是将residual block中的3*3的kernel inflate成3*3*3的，另一种是将residual block中的1*1的kernel inflate成3*1*1的。这两种形式我们分别用 $I3D_{3*3*3}$ 和 $I3D_{3*1*1}$ 表示。因为3D conv的计算量很大，我们只对每2个residual blocks中的1个kernel做inflate。对更多的kernel做inflate发现效果反而变差了。另外conv1层我们inflate成5*7*7。

Non-local network。 我们将non-local block插入到C2D或I3D中，就得到了non-local nets。我们研究了插入1，5，10个non-local blocks的情况，实现细节将在后面给出。

2.1 Non-local Network实现细节

Training。 我们的模型是在ImageNet上pretrain的，没有特殊说明的话我们使用32帧的输入。32帧是通过从原始长度的视频中随机选择1个位置取出64个连续帧，然后每隔1帧取1帧得到的最终的32帧。spatial size是224*224大小，是将原始视频rescale到短边为[256,320]区间的随机值，然后再random crop 224*224大小。我们在8卡GPU上进行训练，每卡上有8 clips（也就是说总的batchsize是64 clips）。我们一共迭代了400k iterations，初始lr为0.01，然后每150k iterations lr下降1/10。momentum设为0.9，weight decay设为0.0001。dropout在global pooling层后面使用，dropout ratio设为0.5。
我们finetune模型的时候 BN是打开的，这和常见的finetune ResNet的操作不同，它们通常是frozen BN。我们发现在我们的实验中enable BN有利于减少过拟合。
在最后一个1*1*1 conv层（表示 $W_z$ ）的后面我们加了一个BN层，其他位置我们没有增加BN。这个BN层的scale参数初始化为0，这是为了保证整个non-local block的初始状态相当于一个identity mapping，这样插入到任何预训练网络中在一开始都能保持其原来的表现。

Inference。 推理时，在我们将视频rescale到短边256进行推理。时域上我们从整个视频中平均采样10个clip，然后分别计算他们的softmax scores，最后做平均得到整个视频的score。

3. 实验

关于视频分类的实验，我们在Kinetics上进行全面的实验，另外也给出了Charades上的实验结果，显示出我们的模型的泛化性。这里只给出Kinetics上的结果，更多的请看原文。
Table 2给出了ablation results。
这里写图片描述

f的表现形式的影响。表2a比较了不同的non-local block的形式插入到C2D得到的结果（插入位置在res4的最后一个residual block之前）。发现即使只加一个non-local block都能得到~1%的提高。
有意思的是不同的non-local block的形式效果差不多，说明是non-local block的结构在起作用，而对具体的表达方式不敏感。本文后面都采用embedded Gaussian进行实验，因为这个版本有softmax，可以直接给出[0,1]之间的scores。

哪个阶段加入non-local blocks？表2b比较了一个non-local block加在resnet的不同stage的效果，具体加在不同stage的最后一个residual block之前。发现在res2，res3，res4层上加non-local block效果类似，加在res5上效果稍差。这个的可能原因是res5的spatial size比较小，只有7*7，可能无法提供精确的spatial信息了。

加入更多的non-local blocks。表2c给出了加入更多non-local block的结果，我们在resnet-50上加1 block（在res4），加5 block（3个在res4，2个在res3，每隔1个residual block加1个non-local block），加10 block（在res3和res4每个residual block都加non-local block）。在resnet101的相同位置加block。发现更多non-local block通常有更好的结果。我们认为这是因为更多的non-local block能够捕获长距离多次转接的依赖。信息可以在时空域上距离较远的位置上进行来回传递，这是通过local models无法实现的。
另外需要提到的是增加non-local block得到的性能提升并不只是因为它给base model增加了深度。为了说明这一点，表2c中resnet50 5blocks能够达到73.8的acc，而resnet101 baseline是73.1，同时resnet50 5block只有resnet101的约70%的参数量和80%的FLOPs。说明non-local block得到的性能提升并不只是因为它增加了深度。

时空域上做non-local。我们的方法也可以处理时空域的信息，这一特性非常好，在视频中相关的物体可能出现在较远的空间和较长的时间，它们的相关性也可以被我们的模型捕获。表2d给出了在时间维度，空间维度和时空维度分别做non-local的结果。仅在空间维度上做就相当于non-local的依赖仅在单帧图像内部发生，也就是说在式(1)上仅对index i的相同帧的index j做累加。仅在时间维度上做也类似。表2d显示只做时间维度或者只做空间维度的non-local，都比C2D baseline要好，但是没有同时做时空维度的效果好。

Non-local net vs. 3D ConvNet。表2e比较了我们的non-local C2D版本和inflated 3D ConvNets的性能。Non-local的操作和3D conv的操作可以看成是将C2D推广到时间维度的两种方式。
表2e也比较了param的数量，FLOPs等。我们的non-local C2D模型比I3D更加精确（75.1 vs 74.4），并且有更小的FLOPs（1.2x vs 1.5x）。说明单独使用时non-local比3D conv更高效。

Non-local 3D ConvNet. 不管上面的比较，其实non-local操作和3D conv各有各的优点：3D conv可以对局部依赖进行建模。表2f给出了在 $I3D_{3*1*1}$ 上插入5个non-local blocks的结果。发现NL I3D都能够在I3D的基础上提升1.6个点的acc，说明了non-local和3D conv是可以相互补充的。

更长的输入序列。 最后我们也实验了更长输入序列的情况下模型的泛化性。输入clip包含128帧连续帧，没有做下采样，是一般情况下取的32帧的4倍长度。为了将这个模型放入显存中，每个GPU上只能放下2 clips。因为这么小的batchsize的原因，我们freeze所有的BN层。我们从32帧训练得到的模型作为初始化模型，然后用128帧进行finetune，使用相同的iterations数目（虽然batchsize减小了），初始lr为0.0025，其他设置和之前保持一致。
表2g给出了128帧的实验结果，和表2f的32帧的结果相比，所有模型都表现得更好，说明我们的模型在长序列上的效果也很好。

和state-of-the-art的比较。表3给出了Kinetics上各个方法的结果。
这里写图片描述

Done

最后想说，这篇文章除了整个故事非常完整以外，实验部分做的非常好，面面俱到地分析了各种情况，是以后写文章做实验可以学习的地方。