Attention Mechanisms in GANs

While convolutional layers in GANs excel at capturing local patterns, their fixed, local receptive fields limit their ability to model long-range dependencies across an image. For instance, generating a realistic animal might require ensuring consistency between the head orientation and the limb positions, features that could be far apart in the pixel space. Standard convolutions struggle with this, needing many layers to propagate information across large distances, which can dilute signals and increase computational burden without guaranteeing effective modeling of these relationships.

Attention mechanisms provide a powerful alternative, allowing the network to selectively focus on relevant parts of the input, regardless of their spatial distance. In the context of GANs, particularly image generation, self-attention has proven highly effective.

Self-Attention in GANs (SAGAN)

The Self-Attention Generative Adversarial Network (SAGAN) introduced a self-attention module directly into the generator and discriminator architectures. This module allows the network to weigh the importance of features at all other positions when calculating the feature response at a given position.

How it Works

Consider an input feature map $x \in \mathbb{R}^{C \times H \times W}$ from a preceding convolutional layer. The self-attention module computes three representations from this input:

1. Query: $q(x) = W_q x$
2. Key: $k(x) = W_k x$
3. Value: $v(x) = W_v x$

Here, $W_q$, $W_k$, and $W_v$ are learnable weight matrices, typically implemented as $1 \times 1$ convolutions. The dimensions are often reduced for computational efficiency, for example, $W_q, W_k \in \mathbb{R}^{\bar{C} \times C}$ and $W_v \in \mathbb{R}^{C \times C}$, where $\bar{C} = C/8$ is common.

The core idea is to calculate the attention weights based on the similarity between queries and keys. The queries and keys are reshaped into $\mathbb{R}^{\bar{C} \times N}$, where $N = H \times W$ is the number of feature map positions. The attention map $\beta \in \mathbb{R}^{N \times N}$ is computed as:

$$\beta_{j,i} = \frac{\exp(s_{ij})}{\sum_{k=1}^{N} \exp(s_{ik})} \quad \text{where} \quad s_{ij} = q(x_i)^T k(x_j)$$

This equation calculates the dot product similarity $s_{ij}$ between the query at position $i$ and the key at position $j$. A softmax function is then applied row-wise to normalize these similarities into attention weights $\beta_{j,i}$, representing how much the model should attend to position $i$ when synthesizing the output at position $j$.

The value representation $v(x)$ is reshaped to $\mathbb{R}^{C \times N}$. The output of the attention layer $o \in \mathbb{R}^{C \times N}$ is then computed as a weighted sum of the value features, using the attention map:

$$o_j = \sum_{i=1}^{N} \beta_{j,i} v(x_i)$$

This output $o$ is reshaped back to $\mathbb{R}^{C \times H \times W}$. Finally, the module typically adds this attention-derived output back to the original input feature map $x$, often scaled by a learnable parameter $\gamma$ initialized to 0:

$$y = \gamma o + x$$

Initializing $\gamma$ to 0 allows the network to initially rely on local convolutional features and gradually learn to incorporate non-local information as needed during training.

Diagram illustrating the data flow within a self-attention module applied to a feature map. Input features are transformed into Query, Key, and Value representations. Attention weights are computed via matrix multiplication and softmax, then applied to the Value features. The result is scaled and added back to the original input.

Integration in Generator and Discriminator

In SAGAN, these self-attention modules are typically inserted into intermediate layers of both the generator and the discriminator.

• Generator: Adding attention allows the generator to coordinate features across distant spatial locations. This helps in synthesizing objects with complex structures or ensuring global coherence, like correctly forming the shape of an object or rendering backgrounds consistently.
• Discriminator: Attention enables the discriminator to enforce complex geometric constraints or consistency checks on the generated images more effectively. It can verify that relationships between distant parts of the image are plausible according to the real data distribution.

Benefits

• Improved Image Quality: Attention helps GANs capture global dependencies, leading to more realistic and coherent image synthesis, especially for complex scenes or objects with intricate structures.
• Geometric Consistency: Explicitly modeling long-range interactions improves the geometric fidelity of generated samples.
• Computational Cost: The primary drawback is computational cost. The matrix multiplication for the attention map ($s = q^T k$) has a complexity of $O(N^2)$, where $N=H \times W$ is the number of pixels. This quadratic scaling makes it expensive for very high-resolution feature maps. Consequently, attention modules are often applied only to lower or medium-resolution feature maps (e.g., $32 \times 32$ or $64 \times 64$) within the network. Variants like axial attention or patch-based attention have been developed to mitigate this cost for larger images.
• Training Stability: While SAGAN aimed to improve stability through better modeling, incorporating attention can sometimes introduce its own training dynamics that need careful management, often in conjunction with techniques discussed in the next chapter, such as spectral normalization.

"By incorporating attention mechanisms, advanced GAN architectures like SAGAN overcome some fundamental limitations of purely convolutional approaches, enabling the generation of higher-fidelity images with better global structure and consistency. This technique represents a significant step towards capturing the complex, long-range correlations present in data."