Non-local neural networks

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操作的输出跟原图大小一致,具体来说,是下面这个公式:

\mathbb{y}_i = \frac{1}{\mathcal{C}(\mathbb{x})}\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j) g(\mathbb{x}_j) (1)

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

Non-local操作可以看成attention

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

  • Gaussian: f(\mathbb{x}_i, \mathbb{x}_j)=e^{\mathbb{x}_i^T \cdot \mathbb{x}_j}, \mathcal{C}(x)=\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j)
  • Embedded Gaussian: f(\mathbb{x}_i, \mathbb{x}_j)=e^{\theta(\mathbb{x}_i)^T \cdot \phi(\mathbb{x}_j)}, \mathcal{C}(x)=\sum_{\forall j} f(\mathbb{x}_i, \mathbb{x}_j)
  • Dot Product: f(\mathbb{x}_i, \mathbb{x}_j)=\theta(\mathbb{x}_i)^T \cdot \phi(\mathbb{x}_j), \mathcal{C}(x)=|\{i|i \text{ is a valid index of } \mathbb{x}\}|
  • Concatenation: f(\mathbb{x}_i, \mathbb{x}_j)=\text{ReLU}(\mathbb{w}_f^T \cdot [\theta(\mathbb{x}_i), \phi(\mathbb{x}_j)]), \mathcal{C}(x)=|\{i|i \text{ is a valid index of } \mathbb{x}\}| ,这相当于embedded的两个点拼接作为带ReLU激活函数全连接层的输入。它在visual reasoning中用的比较多。

这里有两点需要提一下:

  1. 后两种选择的归一化系数C(x)选择为x的点数,只是为了简化计算,同时,还能保证对任意尺寸的输入,不会产生数值上的尺度伸缩。
  2. Embedding的实现方式,以图像为例,在文章中都采用1*1的卷积,也就是 \theta\phi 都是卷积操作。

为了能让non-local操作作为一个组件,可以直接插入任意的神经网络中,作者把non-local设计成residual block的形式,让non-local操作去学x的residual:

\mathbb{z}_i =W_z \cdot \mathbb{y}_i + \mathbb{x}_i

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:

titu1994/keras-non-local-netsgithub.com图标

pytorch:

AlexHex7/Non-local_pytorchgithub.com图标


跟全连接层的联系

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

  • 任意两点的相似性仅跟两点的位置有关,而与两点的具体坐标无关,即 f(\mathbb{x}_i, \mathbb{x}_j)=w_{ij}
  • g是identity函数, g(\mathbb{x}_i)=\mathbb{x}_i
  • 归一化系数为1。归一化系数跟输入无关,全连接层不能处理任意尺寸的输入。


跟Self-attention[2]的联系

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

\theta(\mathbb{x}_i)= W_{\theta} \cdot \mathbb{x}_i; \quad \phi(\mathbb{x}_j)=W_{\phi} \cdot \mathbb{x}_j \quad \Rightarrow \quad \theta(\mathbb{x})= W_{\theta} \cdot \mathbb{x}; \quad \phi(\mathbb{x})= W_{\phi} \cdot \mathbb{x}

于是,

y=\text{softmax}(\mathbb{x}^T \cdot W_{\theta}^T \cdot W_\phi \cdot \mathbb{x}) \cdot g(\mathbb{x})

这就是那篇名字很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呢?暂时不知道。


参考文献

  1. Wang X, Girshick R, Gupta A, et al. Non-local Neural Networks[J]. arXiv preprint arXiv:1711.07971, 2017.
  2. Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[C]//Advances in Neural Information Processing Systems. 2017: 6000-6010.
  3. Gatys L A, Ecker A S, Bethge M. A neural algorithm of artistic style[J]. arXiv preprint arXiv:1508.06576, 2015.
  4. Li Y, Wang N, Liu J, et al. Demystifying neural style transfer[J]. arXiv preprint arXiv:1701.01036, 2017.
编辑于 2018-01-30

文章被以下专栏收录