Self-Attention for Images: Complete Guide — Non-Local Networks

Summary

Self-attention for images, introduced by Wang et al. in their pioneering work on non-local neural networks, represents a major advancement in computer vision network architecture. Unlike classical convolutions that operate only on local neighborhoods, self-attention establishes connections between all pairs of positions in an image, enabling the capture of long-range dependencies essential for many vision tasks.

This comprehensive guide explores in depth the self-attention mechanism as applied to images, from its mathematical foundations to its practical implementation with PyTorch. We will see how non-local layers complement traditional CNNs, why this mechanism has become an indispensable component of modern vision architectures, and how to effectively integrate it into your own models. Whether you are a deep learning researcher or an engineer looking to improve the accuracy of your segmentation models, this guide will provide you with all the keys necessary to master visual self-attention in your machine learning projects.

Mathematical Principle

The non-local formulation

The core of self-attention for images lies in the non-local layer, a fundamental operation that revolutionizes the way networks process spatial information. A non-local layer computes for each position i in the image:

y_i = (1 / C(x)) · Σ_j f(x_i, x_j) · g(x_j)

Let’s break down this essential formula component by component:

• f(x_i, x_j) = exp(x_i^T · W_theta^T · W_phi · x_j) is the similarity function between positions i and j. It measures how semantically related two regions of the image are. The exponential guarantees that similarity scores remain strictly positive, and the projected matrix product allows the network to learn which dimensions are important for comparison.
• g(x_j) = W_g · x_j is the value transformation, which projects each position into a representation space optimized for aggregation.
• C(x) = Σ_j f(x_i, x_j) is the normalization factor, which ensures that the summed contributions of all positions form a consistent weighted average.

This formulation radically differs from classical convolutional operations. In a standard convolution, each pixel receives information only from its immediate neighbors — typically a 3×3 or 5×5 neighborhood. Self-attention, on the other hand, connects every pixel to every other pixel in the image, capturing global relationships that would otherwise be inaccessible to convolutional filters.

Formal comparison with convolution

To properly understand the complementarity between these two operations, consider their theoretical receptive fields:

Where k is the kernel size, C the number of channels, and N the total number of positions (height × width). Self-attention sacrifices local computational efficiency for a global and complete receptive field.

The Q, K, V Projection Matrices

In the practical implementation, the weights W_theta, W_phi, and W_g correspond respectively to the projection matrices:

• Query (Q = W_theta · x): what position i is looking for in the rest of the image
• Key (K = W_phi · x): what each position j offers as information
• Value (V = W_g · x): the actual information to aggregate

The attention score between two positions is computed by softmax(Q · K^T / √d_k), where d_k is the key dimension. This formulation, popularized by the transformer attention mechanism, is directly applicable to images by treating each spatial position as a token.

Intuition: The Magnifying Glass and the Wide View

To intuitively understand the difference between convolution and self-attention, imagine looking at a photograph in two different ways:

Convolution is like looking at an image through a magnifying glass. As you move a magnifying glass across a photograph, you see each detail in great precision — the texture of a cat’s fur, the veins on a leaf, the outlines of a building. But your field of view is limited: you cannot simultaneously see the top-left corner and the bottom-right corner of the image. Convolution works exactly the same way. Its filters meticulously scan each local region, progressively learning increasingly complex patterns as layers are stacked. However, even in a deep network of a hundred layers, connecting two distant pixels requires a long, indirect path through all the intermediate layers.

Self-attention is like stepping back to see the whole picture. When you step back from a painting, your eye instantly perceives relationships between elements that are very far apart: you notice that the main subject is in the left third according to the rule of thirds, that the light comes from an upper-right source creating consistent shadows, or that two characters in the scene are looking at each other across an empty space. Self-attention performs exactly this operation: it connects the top-left corner directly to the bottom-right corner, without going through intermediaries.

Both approaches are necessary and complementary. The magnifying glass is indispensable for fine textures, precise edges, repetitive patterns — everything that pertains to local detail. The wide view is essential for global composition, semantic relationships between distant objects, and the structural coherence of the scene. A network using only convolutions would lack global vision; a network using only self-attention would waste enormous computational resources on details that 3×3 filters would easily capture. It is precisely this complementarity that makes non-local networks so effective: they integrate self-attention within existing convolutional architectures, without replacing them.

Python Implementation with PyTorch

Basic Non-Local Block

Let’s start by implementing the fundamental component: the NonLocalBlock with an efficient embedded version (embedded Gaussian), which is the most commonly used variant in practice.

import torch
import torch.nn as nn
import torch.nn.functional as F


class NonLocalBlock(nn.Module):
    """
    Non-Local Block (spatial Self-Attention) for images.

    Implements the 'embedded Gaussian' variant described by
    Wang et al. (CVPR 2018). This module is typically placed
    after a residual block in a ResNet, capturing long-range
    dependencies that convolutions cannot directly model.

    Args:
        in_channels: number of input channels
        reduction: reduction factor for the internal dimension
                   (default 2, so internal_dim = in_channels // 2)
        sub_sample: spatial subsampling factor (1, 2, or 4)
                    Reduces the size of the attention matrix
    """

    def __init__(self, in_channels, reduction=2, sub_sample=1):
        super(NonLocalBlock, self).__init__()
        self.sub_sample = sub_sample
        inter_channels = max(in_channels // reduction, 1)

        # Three 1×1 linear transformations for Q, K, V
        self.theta = nn.Conv2d(
            in_channels, inter_channels,
            kernel_size=1, stride=1, padding=0, bias=False
        )
        self.phi = nn.Conv2d(
            in_channels, inter_channels,
            kernel_size=1, stride=1, padding=0, bias=False
        )
        self.g = nn.Conv2d(
            in_channels, inter_channels,
            kernel_size=1, stride=1, padding=0, bias=False
        )

        # Output 1×1 convolution to project back to original space
        self.W_z = nn.Conv2d(
            inter_channels, in_channels,
            kernel_size=1, stride=1, padding=0, bias=False
        )
        # Zero initialization to ensure the block starts
        # as an identity operation
        nn.init.zeros_(self.W_z.weight)

        # Batch normalization
        self.bn = nn.BatchNorm2d(in_channels)

        # Pooling for subsampling
        if sub_sample > 1:
            self.phi_pool = nn.MaxPool2d(kernel_size=sub_sample)
            self.g_pool = nn.MaxPool2d(kernel_size=sub_sample)
        else:
            self.phi_pool = None
            self.g_pool = None

    def forward(self, x):
        """
        Forward pass of the non-local block.

        Args:
            x: input tensor of shape (B, C, H, W)

        Returns:
            output tensor of same shape as input
        """
        batch_size = x.size(0)

        # --- Query: projection of all positions ---
        theta_out = self.theta(x)
        theta_out = theta_out.view(batch_size, theta_out.size(1), -1)
        theta_out = theta_out.permute(0, 2, 1)  # (B, N, C_inter)

        # --- Key: projection with optional subsampling ---
        phi_input = self.phi_pool(x) if self.phi_pool else x
        phi_out = self.phi(phi_input)
        phi_out = phi_out.view(batch_size, phi_out.size(1), -1)

        # --- Value: projection with optional subsampling ---
        g_input = self.g_pool(x) if self.g_pool else x
        g_out = self.g(g_input)
        g_out = g_out.view(batch_size, g_out.size(1), -1)
        g_out = g_out.permute(0, 2, 1)  # (B, N', C_inter)

        # --- Compute attention matrix ---
        # Attention scores: (B, N, C) @ (B, C, N') → (B, N, N')
        attention = torch.bmm(theta_out, phi_out)
        # Softmax normalization along N' (source positions)
        attention = F.softmax(attention, dim=-1)

        # --- Aggregation of weighted values ---
        # (B, N, N') @ (B, N', C) → (B, N, C)
        y = torch.bmm(attention, g_out)
        # (B, N, C) → (B, C, N) → (B, C, H, W)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, -1, x.size(2), x.size(3))

        # --- Output projection and residual connection ---
        y = self.W_z(y)
        y = self.bn(y)
        output = y + x

        return output

Integration into a ResNet

The main advantage of non-local networks is that they integrate into existing architectures. Here’s how to insert a non-local block into a standard ResNet-50, typically at stage 4 (lowest resolution, where the computational cost of attention is most manageable):

class ResNet50NonLocal(nn.Module):
    """
    ResNet-50 enhanced with a non-local block at stage 4.
    This combination leverages convolutions for local details
    and self-attention for long-range dependencies.
    """

    def __init__(self, num_classes=1000):
        super(ResNet50NonLocal, self).__init__()

        resnet = torch.hub.load('pytorch/vision:v0.13.0',
                                'resnet50', pretrained=True)

        # The first three stages remain convolutional
        self.conv1 = resnet.conv1
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool
        self.layer1 = resnet.layer1
        self.layer2 = resnet.layer2
        self.layer3 = resnet.layer3

        # Stage 4: self-attention added here
        self.layer4 = resnet.layer4
        self.nonlocal_block = NonLocalBlock(
            in_channels=2048,
            reduction=2,
            sub_sample=1
        )

        self.avgpool = resnet.avgpool
        self.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        # Self-attention after stage 4 convolutions
        x = self.nonlocal_block(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

Comparison with/without self-attention in semantic segmentation

To concretely illustrate the contribution of self-attention, consider a semantic segmentation task on the Cityscapes dataset. The goal is to classify each pixel of an urban image (road, pedestrian, vehicle, building, etc.).

class SegmentationHeadWithAttention(nn.Module):
    """
    Segmentation head comparing performance with and without
    a non-local block to demonstrate the contribution of
    self-attention.
    """

    def __init__(self, in_channels=2048, num_classes=19,
                 use_attention=True):
        super(SegmentationHeadWithAttention, self).__init__()
        self.use_attention = use_attention

        # Channel reduction
        self.reduce = nn.Sequential(
            nn.Conv2d(in_channels, 512, kernel_size=3,
                      padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )

        # Optional non-local block for capturing global context
        if use_attention:
            self.nonlocal_block = NonLocalBlock(
                in_channels=512,
                reduction=2,
                sub_sample=2
            )

        # Final prediction convolution
        self.classifier = nn.Conv2d(512, num_classes,
                                    kernel_size=1)

    def forward(self, x):
        x = self.reduce(x)
        if self.use_attention:
            x = self.nonlocal_block(x)
        x = self.classifier(x)
        x = F.interpolate(x, scale_factor=8,
                          mode='bilinear',
                          align_corners=True)
        return x

Analyzing the Contribution of Self-Attention

On Cityscapes, models equipped with non-local blocks typically report an improvement of +1.5 to +2.5 mIoU points (mean Intersection over Union) compared to their purely convolutional counterparts. The gains are particularly notable for:

• Large structures: roads, sky, buildings — objects that extend over large portions of the image
• Fragmented objects: fences, poles, vegetation — where semantic coherence between distant parts is crucial
• Cluttered scenes: crowded streets where understanding the global context helps distinguish pedestrians from vehicles

Key Hyperparameters

The non-local block has several hyperparameters that significantly impact its performance and computational cost:

reduction (reduction factor)

This parameter controls the internal dimension of the non-local block, i.e., the number of intermediate channels used for the Q, K, and V projections. A value of reduction=2 is the standard choice and represents a good trade-off between expressiveness and memory.

• reduction=1: no reduction. Maximum expressiveness but quadratic memory cost. Can be prohibitive with many channels.
• reduction=2: recommended for most applications. Cuts channels in half, significantly reducing the memory required for the attention matrix.
• reduction=4 or more: for very wide models or limited resources. Gradual loss of modeling capacity.
• reduction=8: usually too aggressive. Severe reduction decreases the network’s ability to distinguish subtle relationships between positions.

sub_sample (spatial subsampling)

Subsampling reduces the spatial resolution at which keys and values are computed. With sub_sample=2, the number of positions is divided by 4 (divided by 2 in height and width).

• sub_sample=1: no reduction. The attention matrix is of size (H×W) × (H×W). Accurate but costly.
• sub_sample=2: MaxPool2d(2). 75% reduction in the number of positions. Often the best trade-off.
• sub_sample=4: MaxPool2d(4). For high resolutions where the full cost would be prohibitive.

Stage Selection for Integration

The intuition behind this choice is simple: the lower the resolution, the fewer pairs of positions to compare. At 7×7 (stage 4 of a ResNet-50), there are only 49 positions to compare pairwise, i.e., 2,401 pairs — perfectly manageable on a GPU even with a large batch size.

Advantages and Limitations

Major Advantages

1. Native global receptive field: Unlike convolutions that require deep stacking of layers to progressively widen their receptive field, self-attention achieves global reach in a single operation. This is a decisive advantage for tasks where contextual understanding is paramount.
2. Explicit modeling of long-range relationships: In an urban scene, understanding that a parked car on the left and a pedestrian crossing on the right belong to the same traffic scene requires a global view. Self-attention directly establishes these semantic connections without needing to pass through intermediate layers.
3. Natural complementarity with CNNs: Non-local networks do not replace convolutions; they complement them. Convolutions excel at extracting local features (edges, textures, patterns), while self-attention captures global structure. Together, they form a rich, multi-scale visual representation.
4. Integration flexibility: The non-local block can be inserted anywhere in an existing network, simply as an additional module with a residual connection. No deep architectural modification is necessary.
5. Documented improvements: The original work by Wang et al. demonstrates significant improvements in image classification (ImageNet), object detection (COCO), and video segmentation (YouTube-VOS).

Limitations and Challenges

1. Quadratic computational cost: The O(N²) complexity in terms of the number of positions N is the main drawback. For an image of 224×224 pixels, N = 50,176, which would yield an attention matrix of over 2.5 billion elements. Non-local networks are therefore almost always applied to feature maps at reduced resolution.
2. High memory consumption: The complete attention matrix must be stored in memory for gradient computation during training. For large intermediate feature maps, this can quickly saturate GPU memory.
3. Absence of spatial inductive bias: Unlike convolutions, which naturally incorporate a bias toward spatial locality and translation invariance, self-attention treats all pairs of positions identically. This can be a disadvantage when local spatial structure is informative.
4. Training stability: Adding non-local layers can sometimes make training unstable, requiring adjustment of the learning rate and careful initialization of output weights (hence the zero initialization in our implementation).
5. Partial interpretability: Although attention maps can be visualized to understand which positions the network considers related, interpretation remains difficult in deeper layers where representations become increasingly abstract.

4 Concrete Use Cases

1. Semantic Segmentation of Medical Images

In medical imaging (MRI, CT scans, radiography), anatomical structures often extend over large regions of an image. A non-local block allows the network to understand that a bright region in the upper-left corner and another in the lower-right corner belong to the same organ. The ability to establish coherent global connections significantly improves segmentation accuracy, especially for tumors whose boundaries can be blurry and dispersed.

Recommended configuration: NonLocalBlock with reduction=2, sub_sample=2, integrated after the last convolutional block of a U-Net’s encoder. On the BraTS dataset for brain tumor segmentation, an improvement of +2% in Dice score is typically observed.

2. Object Detection in Satellite Images

Satellite images present unique characteristics: extremely high resolution, objects of very varying sizes (from tiny vehicles to vast agricultural zones), and the need to understand large-scale spatial relationships (a road crossing a landscape, a seaport connected to the sea). Self-attention excels in this context because it allows the detector to connect distant visual cues — for example, associating a tire track on a path with vehicles located several kilometers further away.

Recommended configuration: Multiple non-local blocks (sub_sample=2) inserted in the FPN (Feature Pyramid Network) at different scales, enabling multi-scale attention that captures both relationships local to dense urban areas and global relationships in rural areas.

3. Image Super-Resolution

Super-resolution consists of reconstructing a high-resolution image from a low-resolution version. Convolution-based methods face a fundamental limitation: to reconstruct precise detail, they must infer missing information solely from the local neighborhood. Self-attention breaks this limitation by allowing the model to search for similar texture cues anywhere in the image.

Recommended configuration: NonLocalBlock with reduction=4 to limit cost (super-resolution networks often have many layers), integrated into an RCAN or EDSR network. On the DIV2K benchmark, adding a non-local block improves PSNR by an average of +0.15 dB.

4. Action Recognition in Videos

Videos introduce an additional dimension: time. A non-local block extended to the three spatial-temporal dimensions (H×W×T) allows capturing complex spatio-temporal dependencies. For example, to recognize the action “opening a door,” the network must connect the hand approaching the handle (t=1) to the door opening (t=5). Self-attention directly establishes these temporal connections without requiring a recurrent network or a deep stack of 3D convolutions.

Recommended configuration: NonLocalBlock3D with sub_sample=2 on the temporal dimension as well, integrated into an I3D or SlowFast network. On the Kinetics-400 dataset, networks equipped with non-local blocks achieve approximately +1.5% top-1 accuracy.

See Also

• Mastering Reflexive Position in Python: Complete Guide for Developers
• Euclidean Algorithm in Python: Computing the Greatest Common Divisor Efficiently