论文阅读——Non-Local Neural Networks

2020 年 6 月 14 日
AI

一、摘要

当前深度学习中用于捕获长距离依赖关系的主要方式——堆叠卷积或者使用循环模块，均为局部邻域处理组件。本文提出了另一种捕获长距离依赖关系的通用组件——非局部处理模块。
【待完善】

二、Non-Local Neural Network

2.1 回顾Non local means

Non local means filter，即非局部均值滤波算法，是图像去噪领域一个非常有名的算法。我读研一时有个大四的学妹在做毕设，问我说她想找个传统去噪算法作为她跑的DL代码的对比方法，我就给她推荐了NLM（当时我也没听过几个传统去噪算法哈哈，所以应该算是大名鼎鼎了）。

NLM的思路非常直观，首先其出发点在于：在去噪任务中，对图像的每个小块，如果能用和其相同性质的其他小块加权平均（去噪任务主流的方法是在不改变原始输入信号的情况下通过加权平均，或者说滤波，来去除噪声），会对当前小块噪声抑制起到较好的效果。这样相同性质的小块越多，加权平均之后的效果就越好。

因此，NLM算法简单来说就是两个步骤：一是确定当前小块与图像所有位置的小块之间的接近程度，即所有小块与当前小块加权平均时的权值；二是将各个小块的像素值与当前小块按照该权值，加权平均之后作为去噪的结果。算法示意图如Fig. 1：

Fig. 1 非局部均值滤波算法

Fig. 2展示了对一副大图，以中心像素为基准，图像所有像素的权重可视化结果。可以看出与中心像素所在邻域性质相似的区域获得的权重明显高一些。

Fig. 2 NLM算法中各像素与中心像素的接近程度可视化

NLM的模型如下：
$NL[v](i) = \sum_{j\in{I}}w(i,j)v(j) \tag{1}$ ，其中 $v$ 代表图像， $i,j$ 代表位置下标； $w$ 表示权重。非局部均值滤波算法中具体使用的 $w(i,j)$ 为：
$w(i,j)=\frac{1}{Z(i)}e^{-\frac{||v(N_i)-v(N_j)||^2_{2,a}}{h^2}} \tag{2}$ ， $Z(i)$ 为归一化因子，其值为当前小块和图像所有小块权重之和：
$Z(i)=\sum_je^{-\frac{||v(N_i)-v(N_j)||^2_{2,a}}{h^2}} \tag{3}$

这里着重记录一下 $w(i,j)$ 。求 $i,j$ 两位置的权重 $w$ 也就是衡量图像在这两个位置的相似度，最简单的方法是直接评估两位置对应像素灰度值的差距。考虑到使用邻域比单独的像素可靠一些，邻域相似度高则认为这两个像素的相似度高，故论文中用比较两个像素所在邻域来代替像素值直接比较。

衡量两个图像块相似度最常用的方法是计算他们之间的欧氏距离，不过在求欧式距离的时候，不同位置的像素的权重应当是不一样的，距离块的中心越近的像素，权重越大，距离中心越远，权重越小。这个权重可以看作一个符合高斯分布的kernel。（不过实际计算中考虑到计算量的问题，常常采用均匀分布的权重。另外原始NLM算法速度很慢，后续又出了很多优化版本，这里不做讨论。）

联想：在harris特征匹配算法中，计算像素之间相似程度也用到邻域之间的相似度，不过计算的是归一化互相关系数；sift匹配则是对邻域提取到的特征，利用余弦相似定理（归一化的dot product）来计算特征相似程度。

通过下面一段matlab代码，可以看出NLM具体计算过程：

 for r=rmin:1:rmax
     for s=smin:1:smax                                
         if(r==i1 && s==j1) continue; end;
         W2= input2(r-f:r+f , s-f:s+f);                
         d = sum(sum(kernel.*(W1-W2).*(W1-W2)));
         w=exp(-d/h);                 
         if w>wmax                
             wmax=w;                   
         end
         sweight = sweight + w;
         average = average + w*input2(r,s);                           
     end 
 end

其中kernel由以下代码生成：

function [kernel] = make_kernel(f)              
kernel=zeros(2*f+1,2*f+1);   
for d=1:f    
  value= 1 / (2*d+1)^2 ;    
  for i=-d:d
    for j=-d:d
        kernel(f+1-i,f+1-j)= kernel(f+1-i,f+1-j) + value ;
    end
  end
end
kernel = kernel ./ f;

2.2 从NLM到NLNN

上面一部分回顾了经典的NLM算法。另外之所以去细看NLM的论文，实际上就是因为在看Non-local Neural Networks论文时对Fig. 3这张论文截图有两个疑问：

为什么作者在选择 $f(x_i, x_j)$ 时说“一个自然的选择是使用高斯函数”
作者说使用”gaussian function”但是给出的式子只是一个矩阵乘法之后的指数函数，如何代表高斯？

Fig. 3 论文中提到的gaussian function可能并不严谨？

我自己先尝试对这2个问题进行解答：

首先使用exponential function并不是gaussian的体现，按照NLM论文的说法，指数函数的作用其实就是为了让欧式距离较大位置的权重能够快速衰减到0附近，也即”fast decay”。（值得注意的是，后续还出现了不少针对nlm中的指数kernel进行优化的论文。）

然后真正体现gaussian的是上面代码段中在两个邻域计算欧式距离时加入的gaussian kernel衰减矩阵，其作用在于，按照邻域中像素距离中心的远近来控制求邻域间欧式距离时的像素权重，即 $||v(N_i)-v(N_j)||^2_{2,a}$ 代表gaussian版本的欧式距离。

从这个角度来说, NLNN(Non Local Neural Network)论文中提到的gaussian function貌似并不是真的用了gaussian衰减，因为NLNN在比较 $i,j$ 两个位置时，甚至并未通过比较任何形式的邻域 $N(i), N(j)$ 来确定该 $j$ 位置与 $i$ 的接近程度，而仅仅是直接比较输入信号在 $i,j$ 两个位置对应向量或者向量的embedding的距离测度。

具体来说，对于NLM，输入信号为2D图像，通过计算以 $i,j$ 像素为中心的邻域相似度来决定位置 $j$ 处的权重；对于NLNN，若输入信号为3D特征图，则每个空间位置 $j$ 对应一个向量，故论文计算的是 $x_i$ 和 $x_j$ 两个1D向量的点积再对得到的标量过一个指数函数。

接下来回到NLNN论文中，首先看一下作者给出的通用non-local operation模型:
$y_i=\frac{1}{C(x)}\Sigma_j f(x_i, x_j)g(x_j) \tag{4}$

其中， $i$ 代表当前待计算的output position的index， $x$ 代表输入信号（图像、序列、视频或者他们的特征）， $y$ 为输出信号，通常和 $x$ 保持size相同。二元函数 $f$ 输出一个标量，如两个位置 $i,j$ 的某种接近程度的metric。一元函数负责给出输入信号在 $j$ 位置处的一个表示。计算出的输出信号在 $i$ 处的响应最后会经过一个归一化因子 $C(x)$ 以进行归一化。

上式从形式上看与NLM的模型，即式(1)没有差异。接下来作者在”Instantiations”部分为该通用模版举了几个具体的例子，使之更易理解，同时也用于说明non local block的通用性。

$g(x_j)$ 代表位置 $j$ 处输入信号的一个表示，为了简洁起见，作者就直接把其形式固定为原始输入的一个embedding： $g(x_j)=W_gx_j$ ， $W_j$ 为 $x_j$ 对应的权重矩阵，如最简单的，对于CNN中的feature， $W_j$ 可以是 $1×1$ 卷积（或者进一步扩展到多个 $1\times1$ 卷积，即由一个卷积层提取之后的新特征？）

然后作者讨论了对 $f(x_i, x_j)$ 函数的选择。这里有必要再对NLM和NLNN的区别做一个强调：

输入信号形式： NLM输入信号形式固定，为2D图像，NLNN输入信号形式多样，可以是2D(第一层输入特征)、3D（中间层输入特征）甚至4D（加入时序的中间层特征）
计算相似程度时有无使用邻域信息： $f(x_i, x_j)$ 需要输出一个标量代表 $i,j$ 两个位置的相似程度。由于NLM输入形式简单， $x_i, x_j$ 均只代表一个灰度数值，故求相似程度时利用邻域信息提高鲁棒性。NLM在使用邻域信息时为了突出中心像素坐标，所以要用gaussian kernel。而对于NLNN来说 $x_i,x_j$ 各自代表一个向量，求两向量之间的相似程度来表征 $x_i, x_j$ 两个位置的相似程度本身就比较鲁棒（当然或许可以把NLNN也扩展为求$K_i,K_j两个向量邻域对相似程度？这样从思想上来说感觉更接近NLM），所以NLNN未利用邻域信息。

所以NLNN论文中的gaussian与embedded gaussian或许改成exponential与embedded exponential更贴切？毕竟后面还把式子转化成softmax，softmax中的exp与gaussian有隐含关系吗？

2.3 NLNN block公式推导

2.3.1 权重与softmax

接下来就是这一章最重要的内容：将上述公式（以embedded gaussian为例子）简化为softmax的形式。为了便于理解，这里可以以我们比较熟悉的CNN某中间层特征来作为输入信号，其shape为HxWx1024。对于 $i,j$ 两个位置，计算其embedded gausssian:
$z_i = \theta(x_i) = W_ix_i$

$z_j = \phi(x_j) = W_jx_j$

对于输入信号的某个位置 $i$ ，因为i确定是C为常数，故可写成：
$y_i=\sum_{\forall{j}}\frac{1}{C}e^{z_i^Tz_j}g(x_j)$

然后把C代入进来：
$y_i=\sum_{\forall{j}}\frac{e^{z_i^Tz_j}}{\Sigma_je^{z_i^Tz_j}} g(x_j)$

可以看出对于固定的 $i$ ，中间那一块符合softmax的定义，因此有：
$y_i=\sum_{\forall{j}} softmax(z_i^Tz_j)g(x_j)$

至此我们推导了对于单个位置 $i$ ，输出信号的响应 $y(i)$ 。接下来需要对这个公式进行向量化。大家在写反向传播相关练习时都知道，单个样本的公式比较好推导，但是如果要在代码中实现，还是需要向量化之后的公式。

2.3.2 公式向量化

首先保持 $i$ 固定，去掉 $j$ ：
$y_i = sum(softmax(z_i^TZ)*g(X))$

注意，这里的 $Z$ 和 $X$ 本身是3维向量 $C\times H \times W$ ，但是如果推导时直接用3维的shape，后面会非常麻烦。故这里的一个技巧是将 $g(X)$ 和 $Z$ 均看成 $C×HW$ 的矩阵（假设 $g(X)$ 使用的embedding也是C个kernel）。这样sum就是沿着列方向进行。

可以先检查下各个向量的维度： $z_i^T$ 为 $1×C$ ， $Z$ 为 $C×HW$ ，softmax相当于是activation不改变维度，故经过softmax之后的维度为 $1×HW$ ，然后 $g(X)$ 为 $C×HW$ ，故对于每个 $i$ ，通过广播得到 $C×HW$ 的矩阵，然后对列求和得到输出响应。

然后尝试去除 $i$ ,得到最终的向量化公式：
为了实现对样本维度的向量化，一种个人比较喜欢的做法是先将 $i$ 可能的取值代进去，并排写到一起观察规律，比如下面的草稿：

Fig. 4 将yi展开观察规律

看起来要求 $Y=[y_1, y_2, ..., y_{hw}]^T$ ，需要做的就是对每一个 $i$ 的 $x_i^TX$ 向量（注意这张图中的 $x$ 相当于上面用的 $z$ 记号)分别与 $g(X)$ 广播相乘并把结果矩阵堆到一起，再对列求和即可，也就是下图的第一行：

Fig. 5 第一行矩阵运算结果等效于第二行的矩阵乘法（A、B、C代表维度）

（图示：第一个矩阵的每一行分别与第二个矩阵相乘（利用广播），得到的A个CxB的矩阵concat到一起，再沿着列方向求和）
这里直接给出结论，第一行的结果和第二行的结果等价。感兴趣的同学可以证明下？（我挑了几个位置对比了下是一样的hh）

到这里，我们可以写出去掉 $i$ 之后的向量化公式：
$Y = g(X)·[softmax(Z^T·Z)]^T$

其中”·”代表dot product，即矩阵乘法，其在向量化过程中起到关键的一步，并且将让人头疼的sum去掉了。这样我们就可以写出NLNN关键组件的代码，同时也可以看懂论文中的Figure2，尤其是为什么右上那个乘法用的是矩阵乘法。（注意，看该图时可以忽略掉维度T，并将embedding所用权重简化为 $1×1$ ，这样会稍微容易理解一些）

Fig. 6 NLNN论文中的figure 2

严格来说，根据输出响应的定义，softmax之后那个矩阵乘法的结果就是当 $f(x_i, x_j)$ 采取embedded gaussia时的输出信号。但是这张图后面还有个通过1024个1×1 conv将输出映射回emedding之前的维度，然后加上原始输入信号。个人理解这里最左边那条通路以及逐元素相加运算应该算是构成了NLNN版“卷积”的shortcut。

【待续】

image.png

四、记录

对长距离依赖关系的建模方式：

堆叠卷积层
使用循环网络组件
graphic models，如条件随机场、GNN

词语记录：

showcase：v. display/exhibit
amenable：合适的，易顺服的
affinity：n. 接近，亲和力

参考文献：
//developers.google.com/machine-learning/clustering/similarity/measuring-similarity
//blog.csdn.net/piaoxuezhong/java/article/details/78345929