Self-Attention Mechanisms in GANs

Convolutional Neural Networks (CNNs), the workhorses of many computer vision tasks including foundational GANs like DCGAN, excel at capturing local spatial hierarchies. The convolutional kernel operates on a small local patch of the input feature map, progressively building up more complex features layer by layer. However, this locality can be a limitation when generating images that require understanding relationships between distant parts of the scene. For example, correctly generating the texture of a large object or ensuring consistency between the left and right sides of a face requires modeling dependencies that span across significant portions of the image. Standard convolutions struggle to capture these long-range dependencies efficiently; information needs to propagate through many layers, potentially getting diluted or distorted.

To address this limitation, researchers adapted the concept of self-attention, originally developed for sequence modeling tasks like machine translation, for use within GAN architectures. The core idea is to allow the network to directly model relationships between features at arbitrary positions within a feature map, regardless of their spatial distance.

How Self-Attention Works in GANs

In the context of image generation within a GAN, a self-attention layer takes a feature map from a preceding layer (e.g., a convolutional layer) as input. For each position (pixel or feature vector) in this input map, the self-attention mechanism calculates an "attention map" that indicates how much focus should be placed on all other positions when computing the output feature at the current position.

For a specific location $i$ in the feature map, the network computes:

1. Query: A representation of the current location $i$, asking "what am I looking for?".
2. Keys: Representations of all other locations $j$, indicating "what information do I hold?".
3. Values: Representations of all other locations $j$, signifying "what information should I provide if attended to?".

The attention weight between location $i$ and location $j$ is typically calculated based on the similarity (often using dot product) between the query at $i$ and the key at $j$. These weights are then normalized (commonly using softmax) across all locations $j$. The final output feature at location $i$ is computed as a weighted sum of the values from all locations $j$, where the weights are the calculated attention scores.

Mathematically, if $x \in \mathbb{R}^{C \times N}$ is the input feature map (where $C$ is the number of channels and $N = H \times W$ is the number of spatial locations), the self-attention output $o_i$ at position $i$ can be formulated as:

$$y_i = \sum_{j} \alpha_{ij} (W_v x_j)$$ $$\alpha_{ij} = \text{softmax}_j (e_{ij})$$ $$e_{ij} = \frac{(W_q x_i)^T (W_k x_j)}{\sqrt{d_k}}$$

Here:

• $x_i, x_j$ are the input feature vectors at positions $i$ and $j$.
• $W_q, W_k, W_v$ are learned weight matrices that transform the input features into queries, keys, and values, respectively.
• $d_k$ is the dimension of the vectors, used for scaling.
• $e_{ij}$ represents the attention energy or score between positions $i$ and $j$.
• $\alpha_{ij}$ is the normalized attention weight, indicating how much position $j$ influences position $i$.
• $y_i$ is the output feature vector at position $i$, computed as a weighted sum of transformed values.

The final output of the self-attention layer, $o$, is typically combined with the original input feature map $x$, often using a learned scaling parameter $\gamma$:

$$o_i = \gamma y_i + x_i$$

This residual connection helps stabilize training and allows the network to easily bypass the attention mechanism if it's not beneficial.

Flow of information in a self-attention mechanism for a single output location, derived from all input locations.

Integration in GANs (SAGAN)

The Self-Attention Generative Adversarial Network (SAGAN) paper demonstrated the effectiveness of incorporating these layers into both the Generator and the Discriminator.

• In the Generator: Self-attention helps synthesize images with better global coherence. For instance, it can ensure that the texture applied to a large object is consistent across its entire surface or that features like eyes in a generated face are correctly positioned relative to each other, even if they are far apart in the pixel space.
• In the Discriminator: Self-attention allows the Discriminator to check for structural consistency across wider regions of the image when determining if it's real or fake. It can verify long-range dependencies that should exist in realistic images but might be missing in generated ones.

Typically, self-attention layers are not used throughout the entire network due to their computational cost. They are often strategically placed in later layers of the Generator (dealing with larger feature maps) and earlier layers of the Discriminator (where feature maps are still relatively large).

Benefits

Integrating self-attention has shown significant improvements:

• Improved Image Quality: Models like SAGAN and BigGAN, which heavily rely on self-attention, demonstrated substantial improvements in metrics like Inception Score (IS) and Fréchet Inception Distance (FID) compared to purely convolutional architectures at the time.
• Modeling Long-Range Dependencies: The primary benefit is the explicit modeling of global context, leading to more structurally coherent and realistic images.
• Training Stability: Some studies suggest that attention can contribute to more stable GAN training dynamics, although stabilization techniques discussed in the next chapter (like spectral normalization, often used alongside attention) also play a major role.

However, there's an important consideration:

• Computational Cost: The computational complexity of the standard self-attention mechanism is quadratic ($O(N^2)$) with respect to the number of spatial locations $N$ in the feature map ($N=H \times W$). This makes it computationally expensive, especially for high-resolution feature maps. Consequently, self-attention is often applied to feature maps of moderate size (e.g., 32x32 or 64x64) within the GAN architecture, rather than at the full output resolution. Research continues on more efficient approximations to attention.

Self-attention mechanisms represent a significant architectural step, enabling GANs to move past local patterns and capture the global structures inherent in complex datasets, advancing image generation quality. Models like BigGAN demonstrated that combining self-attention with other techniques like spectral normalization and architectural adjustments allows for generation at new scale and fidelity.