【论文阅读】非局部神经网络浅读（简介+代码）

卷积和循环神经网络中每个时刻只能捕捉相邻位置元素的关系，而作者受到非局部均值滤波去噪算法的启发，设计了一个非局部网络块用来建模任意位置元素的关系，是自注意力应用于视觉任务的扩展，本质也可以看作“加权平均求和”。作者提出的通用的非局部网络块可以灵活地处理图像、序列、视频等多种任务，且可以放在网络的任意位置，在多个视觉任务中均使得模型效果提升。

卷积网络通过卷积核对相邻位置的元素进行建模（局部操作），如果需要建模远距离位置的话需要堆叠多个卷积核；而循环神经网络需要堆叠RNN块，依赖上一时刻的隐状态来预测当前时刻的输出（局部操作），从而对相邻位置的元素进行建模。这两种网络的“局部”操作带来的问题主要有：

• 计算低效，当两个元素的位置很远时，需要堆叠很多操作才能建立两者之间的联系；
• 优化困难；
• 多跳依赖建模（multihop dependency modeling）困难，例如在两个远距离位置来回传递消息非常困难。

而非局部操作的优势体现在：

• 不论距离远近，可以直接对两个位置的元素建模；
• 实验证明，非局部操作高效，且可以使用较少的层实现最佳效果；
• 灵活适应不同的输入尺寸，且与其他操作可以直接组合（“即插即用”）

有趣的是，作者提到这篇非局部网络模块与Transformer中的自注意力模块非常相似，self-attention通过关注所有位置（查询、加权求和）来预测某个位置的输出，而这可以看作非局部均值的一种形式，因此可以将这篇工作看到将机器翻译任务中提出的self-attention扩展到更通用的任务中。

笔者认为这篇非局部神经网络的贡献主要有两点：

• 将自注意力扩展到更加通用的任务中；
• 解释了自注意力有效的关键原因是非局部计算机制，而非QKV的具体计算设计、softmax等。

本文是笔者的个人理解与总结，如果有误欢迎指出交流！

非局部均值滤波

在介绍这篇非局部网络之前，我们先来简单看一下“非局部均值滤波”的思路。（下面我只做一些浅显的介绍，如果读者需要深入了解，可以参考 这篇博客）

传统的图像处理算子受限于算子的大小，每次计算只能关注相邻位置的信息，例如3*3的均值滤波计算时只能考虑附近的9个元素：

$$\begin{array}{c}\left[\begin{array}{ccc}1/3& 1/3& 1/3\\ 1/3& 1/3& 1/3\\ 1/3& 1/3& 1/3\end{array}\right]\end{array}$$

而当图像中距离较远的像素存在强相关性时，普通的算子难以建模这种关系，因此出现了“非局部均值滤波”，其核心思想是：当前像素的估计值由图像中与它具有相似邻域结构的像素加权平均得到，可以简单理解为当前像素的估计值是与图像中所有像素的加权平均。（其实到这里稍微思考一下就会觉得这也是Attention的主要思想，key和query计算的是注意力权重，再与value进行计算也就是加权求和）

基本原理

下面来定义一下网络中非局部操作的计算，当前计算元素的索引为$i$，$j$代表所有位置的索引，$x$为输入数据（可以是图像、序列、视频及其特征），$y$是输出（与输出同尺寸）。

$$\begin{array}{c}{y}_i=\frac{1}{C(x)}\sum _{\forall j}f(x_i,x_j)g(x_j)\end{array}$$

函数$f$计算两个$x$的关系，得到权重（标量）；函数$g$用于提取$j$位置元素的表示。（这里也可以类比Attention机制，$f(x_i,x_j)$相当于 key和 query计算，而$g(x_j)$相当于 value）另外为了数值稳定性，在求和之后使用$C(x)={\sum }_{\forall j}f(x_i,x_j$