Non-local Neural Networks

2023-11-02 19:17:16

1. 摘要

卷积和循环神经网络中的操作都是一次处理一个局部邻域，在这篇文章中，作者提出了一个非局部的操作来作为捕获远程依赖的通用模块。

受计算机视觉中经典的非局部均值方法启发，我们的非局部操作计算某一位置的响应为所有位置特征的加权和。而且，这个模块可以插入到许多计算机视觉网络架构中去。

2. 介绍

在深度神经网络中，捕获远程依赖非常重要。卷积神经网络依靠大的感知野来对远程依赖建模，这是通过重复叠加卷积块来实现的。但同时，它也有一些限制。首先，它在计算上效率低下。其次，它会导致需要仔细解决的优化难题。最后，这些挑战使得多跳依赖性建模非常困难，例如，当需要在远程位置之间来回传递消息时。

假设有一个 3×3 的卷积核，那我们的每一个激活值就只与这 9 个点的值有关，所以也就是局部（local）的。所谓非局部（non-local），也就是说，每一个激活值可能不至与它相邻近的点有关，还可能与离它比较远的点也有关。比如下面第一帧的球，就与后两帧的球相关。

因此，作者提出的非局部操作计算某一位置的响应为所有位置特征的加权和，这些位置可以是空间的、时间的或者时空的。

使用非局部操作有几个优点：（a）与循环和卷积操作的渐进行为相反，非本地操作通过计算任意两个位置之间的相互关系直接捕获远程依赖性，而不管它们的位置距离如何；（b）非局部操作非常有效，即使只有几层（例如5）也能达到最佳效果；（c）最后，非本地操作保持可变的输入大小，并且可以很容易地与其他操作组合（例如卷积）。

3. 网络结构

3.1. 定义

我们要计算输出位置 i 的响应，j 枚举了所有的位置，函数 f 则计算 i 和所有 j 的关系，最后再进行一个归一化。因为 j 考虑了所有可能的位置，所以这个操作是非局部的。而卷积操作只考虑某一个区域，循环操作只考虑当前和最近的时间序列，它们都是局部的。至于全连接，它的权重是学习到的，不是 i 和 j 的一个函数关系。此外，上面的公式可以接受不同的输入大小，并且能保持输出大小和输入一致，而全连接则必须固定输入和输出大小。

3.2. 实例

当然了，f 和 g 可以有不同的选择。但是实验发现，模型对 f 和 g 的选择并不敏感，也就是说通用的非局部操作才是改进的关键。

为了简化，g 只考虑线性形式：

W 是要学习的权重，这也就可以通过一个 1×1 的卷积来实现。

函数 f 则有多种选择：

Gaussian

高斯函数是最常见的，欧氏距离也同样适用，不过点积在深度学习平台实现则更友好。归一化则采取：

Embedded Gaussian

另一种则是在嵌入空间使用高斯函数：

采取两个映射：

归一化则与上面保持一样。给定 i，上面的公式实际上也就是沿着 j 方向的一个 Softmax：

为了说明 Softmax 并不是必须的，作者还提出了另外两个选择，它们的归一化用位置 j 的数量 N 来实现。

Dot product

Concatenation

3.3. 非局部块

作者进而将上面的方程包装到一个非局部块，从而可以方便地插入到现有的框架中去，一个局部块定义如下：

’+‘ 是引入了残差连接，这样将非局部块插入到现有训练好的模型中后，如果权重系数初始化为零，就可以维持原状态不变。

整个流程如上图所示，为了减小计算，引入了瓶颈结构，先将通道数减为一半。除此之外，还可以对 x 进行空间下采样，也就是在上图中的 φ 和 g 后面加入最大池化，这样可以进一步提高效率。其一个 TensorFlow 实现如下：

def NonLocalBlock(input, subsample=True):
    """
    @Non-local Neural Networks
    Non-local Block
    """
    
    _, height, width, channel = input.get_shape().as_list()     # (B, H, W, C)

    theta = tf.layers.conv2d(input, channel // 2, 1)    # (B, H, W, C // 2)
    theta = tf.reshape(theta, [-1, height*width, channel // 2])  # (B, H*W, C // 2)

    phi = tf.layers.conv2d(input, channel // 2, 1)  # (B, H, W, C // 2)
    if subsample:
        phi = tf.layers.max_pooling2d(phi, 2, 2)  # (B, H / 2, W / 2, C // 2)
        phi = tf.reshape(phi, [-1, height * width // 4, channel // 2])  # (B, H * W / 4, C // 2)
    else:
        phi = tf.reshape(phi, [-1, height * width, channel // 2])   # (B, H*W, C // 2)
    phi = tf.transpose(phi, [0, 2, 1])  # (B, C // 2, H*W)

    f = tf.matmul(theta, phi)   # (B, H*W, H*W)
    f = tf.nn.softmax(f)    # (B, H*W, H*W)

    g = tf.layers.conv2d(input, channel // 2, 1)  # (B, H, W, C // 2)
    if subsample:
        g = tf.layers.max_pooling2d(g, 2, 2)  # (B, H / 2, W / 2, C // 2)
        g = tf.reshape(g, [-1, height * width // 4, channel // 2])  # (B, H*W, C // 2)
    else:
        g = tf.reshape(g, [-1, height * width, channel // 2])  # (B, H*W, C // 2)

    y = tf.matmul(f, g)     # (B, H*W, C // 2)
    y = tf.reshape(y, [-1, height, width, channel // 2])    # (B, H, W, C // 2)
    y = tf.layers.conv2d(y, channel, 1)  # (B, H, W, C)

    y = tf.add(input, y)    # (B, W, H, C)

    return y