Non-local Neural Networks
非局部神经网络

Following the non-local mean [4] and bilateral filters [47], a natural choice of $f$ is the Gaussian function. In this paper we consider:

Here $\mathbf{x}_{i}^{T}\mathbf{x}_{j}$ is dot-product similarity. Euclidean distance as used in [4, 47] is also applicable, but dot product is more implementation-friendly in modern deep learning platforms. The normalization factor is set as $\mathcal{C}(\mathbf{x})=\sum_{\forall j}f(\mathbf{x}_{i},\mathbf{x}_{j})$.

A simple extension of the Gaussian function is to compute similarity in an embedding space. In this paper we consider:

Here $\theta(\mathbf{x}_{i})=W_{\theta}\mathbf{x}_{i}$ and $\phi(\mathbf{x}_{j})=W_{\phi}\mathbf{x}_{j}$ are two embeddings. As above, we set $\mathcal{C}(\mathbf{x})=\sum_{\forall j}f(\mathbf{x}_{i},\mathbf{x}_{j})$.

We note that the self-attention module [49] recently presented for machine translation is a special case of non-local operations in the embedded Gaussian version. This can be seen from the fact that for a given $i$, $\frac{1}{\mathcal{C}(\mathbf{x})}f(\mathbf{x}_{i},\mathbf{x}_{j})$ becomes the softmax computation along the dimension $j$. So we have $\mathbf{y}=\text{\emph{softmax}}(\mathbf{x}^{T}W^{T}_{\theta}W_{\phi}\mathbf{x})g(\mathbf{x})$, which is the self-attention form in [49]. As such, our work provides insight by relating this recent self-attention model to the classic computer vision method of non-local means [4], and extends the sequential self-attention network in [49] to a generic space/spacetime non-local network for image/video recognition in computer vision.

Despite the relation to [49], we show that the attentional behavior (due to softmax) is not essential in the applications we study. To show this, we describe two alternative versions of non-local operations next.

$f$ can be defined as a dot-product similarity:

Here we adopt the embedded version. In this case, we set the normalization factor as $\mathcal{C}(\mathbf{x})=N$, where $N$ is the number of positions in $\mathbf{x}$, rather than the sum of $f$, because it simplifies gradient computation. A normalization like this is necessary because the input can have variable size.

The main difference between the dot product and embedded Gaussian versions is the presence of softmax, which plays the role of an activation function.

Concatenation is used by the pairwise function in Relation Networks [40] for visual reasoning. We also evaluate a concatenation form of $f$:

Here $[\cdot,\cdot]$ denotes concatenation and $\mathbf{w}_{f}$ is a weight vector that projects the concatenated vector to a scalar. As above, we set $\mathcal{C}(\mathbf{x})=N$. In this case, we adopt ReLU [35] in $f$.

The above several variants demonstrate the flexibility of our generic non-local operation. We believe alternative versions are possible and may improve results.

We set the number of channels represented by $W_{g}$, $W_{\theta}$, and $W_{\phi}$ to be half of the number of channels in $\mathbf{x}$. This follows the bottleneck design of [21] and reduces the computation of a block by about a half. The weight matrix $W_{z}$ in Eq.(6) computes a position-wise embedding on $\mathbf{y}_{i}$, matching the number of channels to that of $\mathbf{x}$. See Figure 2.

A subsampling trick can be used to further reduce computation. We modify Eq.(1) as: $\mathbf{y}_{i}=\frac{1}{\mathcal{C}(\mathbf{\hat{x}})}\sum_{\forall j}f(\mathbf{x}_{i},\mathbf{\hat{x}}_{j})g(\mathbf{\hat{x}}_{j})$, where $\mathbf{\hat{x}}$ is a subsampled version of $\mathbf{x}$ (e.g., by pooling). We perform this in the spatial domain, which can reduce the amount of pairwise computation by 1/4. This trick does not alter the non-local behavior, but only makes the computation sparser. This can be done by adding a max pooling layer after $\phi$ and $g$ in Figure 2.

We use these efficient modifications for all non-local blocks studied in this paper.

Our models are pre-trained on ImageNet [39]. Unless specified, we fine-tune our models using 32-frame input clips. These clips are formed by randomly cropping out 64 consecutive frames from the original full-length video and then dropping every other frame. The spatial size is 224$\times$224 pixels, randomly cropped from a scaled video whose shorter side is randomly sampled in $[256,320]$ pixels, following [46]. We train on an 8-GPU machine and each GPU has 8 clips in a mini-batch (so in total with a mini-batch size of 64 clips). We train our models for 400k iterations in total, starting with a learning rate of 0.01 and reducing it by a factor of 10 at every 150k iterations (see also Figure 4). We use a momentum of 0.9 and a weight decay of 0.0001. We adopt dropout [22] after the global pooling layer, with a dropout ratio of 0.5. We fine-tune our models with BatchNorm (BN) [25] enabled when it is applied. This is in contrast to common practice [21] of fine-tuning ResNets, where BN was frozen. We have found that enabling BN in our application reduces overfitting.

We adopt the method in [20] to initialize the weight layers introduced in the non-local blocks. We add a BN layer right after the last 1$\times$1$\times$1 layer that represents $W_{z}$; we do not add BN to other layers in a non-local block. The scale parameter of this BN layer is initialized as zero, following [17]. This ensures that the initial state of the entire non-local block is an identity mapping, so it can be inserted into any pre-trained networks while maintaining its initial behavior.

Following [46] we perform spatially fully-convolutional inference on videos whose shorter side is rescaled to 256. For the temporal domain, in our practice we sample 10 clips evenly from a full-length video and compute the softmax scores on them individually. The final prediction is the averaged softmax scores of all clips.

Table 2a compares different types of a single non-local block added to the C2D baseline (right before the last residual block of res₄). Even adding one non-local block can lead to $\scriptstyle\sim$1% improvement over the baseline.

Interestingly, the embedded Gaussian, dot-product, and concatenation versions perform similarly, up to some random variations (72.7 to 72.9). As discussed in Sec. 3.2, the non-local operations with Gaussian kernels become similar to the self-attention module [49]. However, our experiments show that the attentional (softmax) behavior of this module is not the key to the improvement in our applications; instead, it is more likely that the non-local behavior is important, and it is insensitive to the instantiations.

In the rest of this paper, we use the embedded Gaussian version by default. This version is easier to visualize as its softmax scores are in the range of $[0,1]$.

Table 2b compares a single non-local block added to different stages of ResNet. The block is added to right before the last residual block of a stage. The improvement of a non-local block on res₂, res₃, or res₄ is similar, and on res₅ is slightly smaller. One possible explanation is that res₅ has a small spatial size (7$\times$7) and it is insufficient to provide precise spatial information. More evidence of a non-local block exploiting spatial information will be investigated in Table 2d.

Table 2c shows the results of more non-local blocks. We add 1 block (to res₄), 5 blocks (3 to res₄ and 2 to res₃, to every other residual block), and 10 blocks (to every residual block in res₃ and res₄) in ResNet-50; in ResNet-101 we add them to the corresponding residual blocks. Table 2c shows that more non-local blocks in general lead to better results. We argue that multiple non-local blocks can perform long-range multi-hop communication. Messages can be delivered back and forth between distant positions in spacetime, which is hard to do via local models.

It is noteworthy that the improvement of non-local blocks is not just because they add depth to the baseline model. To see this, we note that in Table 2c the non-local 5-block ResNet-50 model has 73.8 accuracy, higher than the deeper ResNet-101 baseline’s 73.1. However, the 5-block ResNet-50 has only $\scriptstyle\sim$70% parameters and $\scriptstyle\sim$80% FLOPs of the ResNet-101 baseline, and is also shallower. This comparison shows that the improvement due to non-local blocks is complementary to going deeper in standard ways.

We have also tried to add standard residual blocks, instead of non-local blocks, to the baseline models. The accuracy is not increased. This again shows that the improvement of non-local blocks is not just because they add depth.

Our method can naturally handle spacetime signals. This is a nice property: related objects in a video can present at distant space and long-term time interval, and their dependency can be captured by our model.

In Table 2d we study the effect of non-local blocks applied along space, time, or spacetime. For example, in the space-only version, the non-local dependency only happens within the same frame: i.e., in Eq.(1) it only sums over the index $j$ in the same frame of the index $i$. The time-only version can be set up similarly. Table 2d shows that both the space-only and time-only versions improve over the C2D baseline, but are inferior to the spacetime version.

Table 2e compares our non-local C2D version with the inflated 3D ConvNets. Non-local operations and 3D convolutions can be seen as two ways of extending C2D to the temporal dimensions.

Table 2e also compares the number of parameters and FLOPs, relative to the baseline. Our non-local C2D model is more accurate than the I3D counterpart (e.g., 75.1 vs. 74.4), while having a smaller number of FLOPs (1.2$\times$ vs. 1.5$\times$). This comparison shows that our method can be more effective than 3D convolutions when used alone.

Despite the above comparison, non-local operations and 3D convolutions can model different aspects of the problem: 3D convolutions can capture local dependency. Table 2f shows the results of inserting 5 non-local blocks into the I3D_3×1×1 models. These non-local I3D (NL I3D) models improve over their I3D counterparts (+1.6 point accuracy), showing that non-local operations and 3D convolutions are complementary.

Finally we investigate the generality of our models on longer input videos. We use input clips consisting of 128 consecutive frames without subsampling. The sequences throughout all layers in the networks are thus 4$\times$ longer compared to the 32-frame counterparts. To fit this model into memory, we reduce the mini-batch size to 2 clips per GPU. As a result of using small mini-batches, we freeze all BN layers in this case. We initialize this model from the corresponding models trained with 32-frame inputs. We fine-tune on 128-frame inputs using the same number of iterations as the 32-frame case (though the mini-batch size is now smaller), starting with a learning rate of 0.0025. Other implementation details are the same as before.

Table 2g shows the results of 128-frame clips. Comparing with the 32-frame counterparts in Table 2f, all models have better results on longer inputs. We also find that our NL I3D can maintain its gain over the I3D counterparts, showing that our models work well on longer sequences.

Table 3 shows the results from the I3D authors [7] and from the Kinetics 2017 competition winner [3]. We note that these are comparisons of systems which can differ in many aspects. Nevertheless, our method surpasses all the existing RGB or RGB + flow based methods by a good margin. Without using optical flow and without any bells and whistles, our method is on par with the heavily engineered results of the 2017 competition winner.