Non-local Neural Networks

Standard convolutional layers operate locally, processing information within a confined neighborhood defined by the kernel size. While stacking these layers increases the effective receptive field, efficiently capturing dependencies between spatially distant features (long-range dependencies) remains challenging. For instance, understanding the relationship between a person holding an object and the object itself, even if they appear far apart in the image, requires modeling interactions outside immediate adjacency.

Non-local Neural Networks directly address this limitation by computing the response at a position as a weighted sum of features at all positions in the input feature map. This allows the network to capture global context and model relationships between any two locations, regardless of their spatial distance. Think of it as a generalization of self-attention mechanisms applied to spatial or spatio-temporal data.

The Non-local Operation

The core idea is elegantly captured in a general formula for a non-local operation. Given an input feature map $\mathbf{x}$ (which could be an image, or feature maps from intermediate layers of a CNN), the output feature map $\mathbf{y}$ at a position $i$ is computed as:

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

Let's break down the components:

• $i$: The index of the output position whose response is being calculated.
• $j$: An index that iterates over all possible positions in the input feature map.
• $\mathbf{x}_i$: The feature vector at position $i$ in the input.
• $\mathbf{x}_j$: The feature vector at position $j$ in the input.
• $f(\mathbf{x}_i, \mathbf{x}_j)$: A pairwise function that computes a scalar representing the relationship (e.g., affinity or similarity) between position $i$ and position $j$.
• $g(\mathbf{x}_j)$: A function that computes a representation or transformation of the input feature vector at position $j$. This is often a linear embedding.
• $\mathcal{C}(\mathbf{x})$: A normalization factor, calculated as $\sum_{\forall j} f(\mathbf{x}_i, \mathbf{x}_j)$, ensuring the weights sum appropriately (often to 1, similar to a softmax).

Essentially, the response at position $i$ ($\mathbf{y}_i$) is a weighted average of the transformed features ($g(\mathbf{x}_j)$) from all positions $j$. The weights are determined by the similarity or relationship ($f$) between position $i$ and each position $j$.

Instantiations of the Pairwise and Transformation Functions

The generic non-local operation formula allows for different specific implementations based on the choice of functions $f$ and $g$.

Transformation Function $g$: A common choice for $g$ is a simple linear embedding learned via a $1 \times 1$ convolution: $g(\mathbf{x}_j) = W_g \mathbf{x}_j$ Here, $W_g$ represents the weights of the $1 \times 1$ convolutional layer.

Pairwise Function $f$: Several options exist for the pairwise function $f$, measuring the relationship between $i$ and $j$:

1. Embedded Gaussian: This is perhaps the most common choice and directly relates to self-attention. First, linear embeddings $\theta$ and $\phi$ (again, typically $1 \times 1$ convolutions with weight matrices $W_{\theta}$ and $W_{\phi}$) are applied to the input features. The function is then: $f(\mathbf{x}_i, \mathbf{x}_j) = e^{\theta(\mathbf{x}_i)^T \phi(\mathbf{x}_j)} = e^{(W_{\theta}\mathbf{x}_i)^T (W_{\phi}\mathbf{x}_j)}$ The normalization factor $\mathcal{C}(\mathbf{x})$ in this case makes the weighting $\frac{f(\mathbf{x}_i, \mathbf{x}_j)}{\mathcal{C}(\mathbf{x})}$ equivalent to a softmax function applied over all positions $j$.

2. Dot Product: A simpler version of the Embedded Gaussian, omitting the exponential: $f(\mathbf{x}_i, \mathbf{x}_j) = \theta(\mathbf{x}_i)^T \phi(\mathbf{x}_j)$ Here, the normalization $\mathcal{C}(\mathbf{x})$ might be the number of positions $N$.

3. Concatenation: The embedded features are concatenated, passed through a linear layer (with weight vector $w_f$), and activated (e.g., with ReLU): $f(\mathbf{x}_i, \mathbf{x}_j) = \text{ReLU}(w_f^T [\theta(\mathbf{x}_i), \phi(\mathbf{x}_j)])$ This allows for potentially more complex relationship modeling.

The Embedded Gaussian approach is widely used due to its connection to attention mechanisms and empirical success.

Implementing Non-local Blocks

Non-local operations are typically integrated into existing deep learning architectures as "Non-local Blocks". These blocks often employ a residual connection to facilitate training, similar to ResNet blocks.

The structure of a common Non-local Block (using the Embedded Gaussian version) is as follows:

1. Input: Feature map $\mathbf{x}$ (e.g., dimensions $H \times W \times C$).
2. Embeddings: Compute $\theta(\mathbf{x})$, $\phi(\mathbf{x})$, and $g(\mathbf{x})$ using separate $1 \times 1$ convolutions. $\theta$ and $\phi$ typically reduce the channel dimension (e.g., to $C/2$) to save computation, while $g$ might maintain or reduce it. Let the outputs be $\mathbf{\Theta}$, $\mathbf{\Phi}$, $\mathbf{G}$.
3. Affinity Calculation: Reshape $\mathbf{\Theta}$ and $\mathbf{\Phi}$ to matrices where one dimension represents the spatial locations ($H \times W$) and the other represents the channels ($C/2$). Calculate the dot product $\mathbf{\Theta}^T \mathbf{\Phi}$ (after appropriate transpositions). This gives a matrix of size $(HW) \times (HW)$ representing affinities between all pairs of locations.
4. Normalization: Apply softmax row-wise (or column-wise, depending on matrix setup) to the affinity matrix to get normalized weights $\mathbf{A}$.
5. Weighted Sum: Reshape $\mathbf{G}$ similarly to $\mathbf{\Theta}$ and $\mathbf{\Phi}$ (dimension $HW \times C'$). Compute the weighted sum by matrix multiplying the normalized affinities $\mathbf{A}$ with $\mathbf{G}$. Let the result be $\mathbf{Y}'$.
6. Reshape and Project: Reshape $\mathbf{Y}'$ back to the spatial dimensions $H \times W \times C'$. Apply a final $1 \times 1$ convolution (with weights $W_z$) to project it back to the original channel dimension $C$. Let this be $\mathbf{Y}$. The weights $W_z$ are often initialized to zero, so the block initially behaves like an identity function.
7. Residual Connection: Add the output $\mathbf{Y}$ back to the original input $\mathbf{x}$: $\mathbf{z} = \mathbf{Y} + \mathbf{x}$.

This structure allows the network to learn long-range dependencies while retaining the original information flow through the residual connection.

Data flow within a Non-local Block using Embedded Gaussian affinity and a residual connection. The $\theta$, $\phi$, and $g$ transformations are typically implemented using $1 \times 1$ convolutions.

Connection to Self-Attention

If you are familiar with the Transformer architecture, you'll recognize the Embedded Gaussian non-local operation as being equivalent to the scaled dot-product self-attention mechanism.

• $\theta(\mathbf{x}_i)$ corresponds to the "Query".
• $\phi(\mathbf{x}_j)$ corresponds to the "Key".
• $g(\mathbf{x}_j)$ corresponds to the "Value".

The non-local block computes the attention weights between each query position $i$ and all positions $j$, then uses these weights to form a weighted sum of the values. Non-local networks essentially introduced self-attention to the computer vision domain, applying it directly to feature maps rather than sequences of word embeddings.

Applications and Considerations

Non-local blocks have proven effective in tasks where long-range interactions are significant:

• Video Analysis: Capturing temporal dependencies between frames for action recognition.
• Object Detection & Segmentation: Modeling context, such as relationships between objects or between object parts and the whole object. Mask R-CNN, for example, benefits from incorporating contextual information.
• Image Generation: Improving the global coherence and structure of generated images.

However, a significant consideration is computational cost. Calculating the pairwise affinities involves comparing every position with every other position. If the input feature map has $N = H \times W$ spatial locations, the complexity is $O(N^2)$, which can be demanding for high-resolution feature maps.

Mitigation Strategies: To manage this cost, implementations often:

• Reduce channel dimensions via the $1 \times 1$ convolutions in $\theta$, $\phi$, and $g$.
• Apply spatial subsampling (e.g., max-pooling) to the ($\phi$) and value ($g$) inputs before the pairwise computations. This reduces the number of positions $j$ involved in the sum.
• Insert Non-local blocks only in deeper layers of the network where feature maps are smaller, or sparingly throughout the architecture.

Non-local networks provide a powerful and general mechanism for incorporating non-local interactions within deep learning models for vision. They represent an important step towards models that can reason about global image structure and context, complementing the local feature extraction capabilities of standard convolutions and leading to architectures like Vision Transformers.