Non-local neural networks

我来兑现周一的承诺了，今天介绍王小龙(知乎账号: Xiaolong Wang)最新的一篇文章non-local neural networks[1]。因为没有具体地做实验，今天的文章主要捋清non-local是怎么设计的。

Local & non-local

首先我们来看一下文章说的local是什么意思。

Local这个词主要是针对感受野(receptive field)来说的。以卷积操作为例，它的感受野大小就是卷积核大小，而我们一般都选用3*3，5*5之类的卷积核，它们只考虑局部区域，因此都是local的运算。同理，池化(Pooling)也是。相反的，non-local指的就是感受野可以很大，而不是一个局部领域。

那我们碰到过什么non-local的操作吗？有的，全连接就是non-local的，而且是global的。但是全连接带来了大量的参数，给优化带来困难。这也是深度学习(主要指卷积神经网络)近年来流行的原因，考虑局部区域，参数大大减少了，能够训得动了。

那我们为什么还需要non-local？

我们知道，卷积层的堆叠可以增大感受野，但是如果看特定层的卷积核在原图上的感受野，它毕竟是有效的。这是local运算不能避免的。然而有些任务，它们可能需要原图上更多的信息，比如attention。如果在某些层能够引入全局的信息，就能很好地解决local操作无法看清全局的情况，为后面的层带去更丰富的信息。这是我个人的理解。

Non-local block

好了，那我们来看一下文章是怎么设计non-local运算的。为了能够当作一个组件接入到以前的神经网络中，作者设计的non-local操作的输出跟原图大小一致，具体来说，是下面这个公式：

上面的公式中，输入是x，输出是y，i和j分别代表输入的某个空间位置，x_i是一个向量，维数跟x的channel数一样，f是一个计算任意两点相似关系的函数，g是一个映射函数，将一个点映射成一个数，可以看成是计算一个点的特征。也就是说，为了计算输出层的一个点，需要将输入的每个点都考虑一遍，而且考虑的方式很像attention：输出的某个点在原图上的attention，而mask则是相似性给出。参看下图。

Non-local操作可以看成attention

以图像为例，为了简化问题，作者简单地设置g函数为一个1*1的卷积。相似性度量函数f的选择有多种：

Gaussian:

Embedded Gaussian:

Dot Product:

Concatenation:

这相当于embedded的两个点拼接作为带ReLU激活函数全连接层的输入。它在visual reasoning中用的比较多。

这里有两点需要提一下：

后两种选择的归一化系数C(x)选择为x的点数，只是为了简化计算，同时，还能保证对任意尺寸的输入，不会产生数值上的尺度伸缩。

Embedding的实现方式，以图像为例，在文章中都采用1*1的卷积，也就是\theta 和\phi 都是卷积操作。

此外，注意到一个问题，公式(1)中输入x可以是有很多channel的，但是输出y却只有一个channel(f和g的运算结果都是一个数)。为了能让non-local操作作为一个组件，可以直接插入任意的神经网络中，作者把non-local设计成residual block的形式，让non-local操作去学x的residual：

W_z实际上是一个卷积操作，它的输出channel数跟x一致。这样以来，non-local操作就可以作为一个组件，组装到任意卷积神经网络中。

具体实现

如果按照上面的公式，用for循环实现肯定是很慢的。此外，如果在尺寸很大的输入上应用non-local layer，也是计算量很大的。后者的解决方案是，只在高阶语义层中引入non-local layer。还可以通过对embedding(\theta, \phi )的结果加pooling层来进一步地减少计算量。

对于前者，注意到f的计算可以化为矩阵运算，我们实际上可以将整个non-local化为矩阵乘法运算+卷积运算。如下图所示，其中oc为output_channels，卷积操作的输出filter数量。

原文考虑的是T帧的视频为例，这里以一个batch的图像、f选为embedded Gaussian为例，对于其他形式的相似性度量，可以类似地化为矩阵操作。

在tensorflow和pytorch中，batch matrix multiplication可以用matmul函数实现。在keras中，可以用batch_dot函数或者dot layer实现。

具体可以参考：

keras：https://github.com/titu1994/keras-non-local-nets

pytorch：https://github.com/AlexHex7/Non-local_pytorch

跟全连接层的联系

我们知道，non-local block利用两个点的相似性对每个位置的特征做加权，而全连接层则是利用position-related的weight对每个位置做加权。于是，全连接层可以看成non-local block的一个特例：

任意两点的相似性仅跟两点的位置有关，而与两点的具体坐标无关，即f(x_i, x_j) = w_

g是identity函数， g(x)=x

归一化系数为1。归一化系数跟输入无关，全连接层不能处理任意尺寸的输入。

跟Self-attention[2]的联系

这部分在原文中也提到了。Embedding的1*1卷积操作可以看成矩阵乘法：

于是，

这就是那篇名字很6的文章[2]提出来的self-attention。

跟Gram matrix[3]的联系

这部分的insight来自于毛豆大佬(一个喜欢猫的大佬)。

Gram matrix第一次被应用到风格迁移任务中[3]，并在后来成为style loss的标配。

Style loss的gram matrix把一个channel看成一个点(坐标就是整个filter，长度等于一个filter大小H*W)；而公式(1)则是把每个空间位置看成点(坐标是所有filter在该空间位置上的值，长度等于channel数)。两者都是计算任意两个点之间的内积。内积运算也就这两种处理方式了......

也就是说，它们的差别在于沿着filters的不同方向做内积。

基于gram matrix的style loss可以捕捉到纹理信息；从上一节我们知道，non-local层起到attention的作用。而由[4]我们知道，匹配gram matrix相当于最小化feature maps的二次多项式核的MMD距离。而non-local呢？暂时不知道。

参考文献

Wang X, Girshick R, Gupta A, et al. Non-local Neural Networks[J]. arXiv preprint arXiv:1711.07971, 2017.

Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[C]//Advances in Neural Information Processing Systems. 2017: 6000-6010.

Gatys L A, Ecker A S, Bethge M. A neural algorithm of artistic style[J]. arXiv preprint arXiv:1508.06576, 2015.

Li Y, Wang N, Liu J, et al. Demystifying neural style transfer[J]. arXiv preprint arXiv:1701.01036, 2017.