Attention Mechanisms for Spatial and External Dependencies

1. Handling Spatial Dependencies with Self-Attention

Convolutional layers are excellent at detecting local patterns, but their receptive field is inherently limited. To model long-range dependencies—for instance, ensuring a person's left eye is stylistically consistent with their right eye, even if they are far apart in the feature map—we need a mechanism that can relate any pixel to any other pixel. This is the role of self-attention.

A Refresher on Self-Attention

You're likely familiar with self-attention from Transformers. In the context of a U-Net, we treat the pixels of a feature map as a sequence. For each pixel, self-attention calculates a weighted average of all other pixels in the feature map, where the weights determine "relevance."

For a thorough, from-scratch refresher on the mechanics of scaled-dot product attention, the following resource by Sebastian Raschka is an excellent read. It builds the entire mechanism in PyTorch, from embeddings to the final context vector.

In short, for each element in our input sequence (i.e., each "pixel" in our flattened feature map x), we generate three vectors:

• Query ($Q$): "What am I looking for?"
• Key ($K$): "What do I contain?"
• Value ($V$): "What information do I offer?"

By taking the dot product of a query from one pixel with the keys from all other pixels, the model computes an attention score, determining how much focus to place on every other pixel. These scores are then used to create a weighted sum of the value vectors, producing an output that enriches the original pixel's representation with context from the entire image.

2. Integrating Self-Attention into the U-Net

Self-attention is computationally intensive ($O(n^2)$ where $n$ is the sequence length, i.e., height × width). Therefore, it's typically applied only at the lower-resolution bottleneck layers of the U-Net, where the feature maps are smaller.

The following video provides a concise visual recap of how self-attention works within the U-Net, before moving on to cross-attention.

3. External Conditioning with Cross-Attention

Self-attention is about internal communication within the image. Cross-attention is about external communication, allowing the U-Net to listen to an outside source, such as a text prompt. This is the core mechanism behind text-to-image models.

The mechanics are almost identical to self-attention, with one crucial difference:

• The Query ($Q$) vector still comes from the image feature map. Each pixel is still asking a question.
• The Key ($K$) and Value ($V$) vectors now come from the external context.

For text-to-image, the external context is the sequence of embeddings from a text encoder like CLIP or a T5 variant. Each word (or token) in the prompt becomes a source for a key and a value.

This way, a pixel's query can "attend" to the words in the prompt. A pixel in a region that should become a "red car" can learn to focus on the "red" and "car" token embeddings to gather the right information for its color and texture.

This video from ExplainingAI gives an excellent explanation of this transition from self- to cross-attention.

4. Implementation and Integration

In practice, self-attention and cross-attention are often implemented within a single, flexible CrossAttention module and combined within a TransformerBlock. This block can then be inserted into the U-Net architecture.

The CrossAttention Module

A well-designed attention module can perform both self- and cross-attention. The logic is simple: if a context tensor is provided, use it to generate Keys and Values (cross-attention). If not, use the input tokens tensor for Keys and Values (self-attention).

The following blog post provides a clean, professional implementation of the components we need.

Putting It All Together in the U-Net

The SpatialTransformer module is the component that gets inserted into the U-Net, typically replacing a standard ResNetBlock at certain resolutions.

1. A ResNet block processes the feature map with convolutions.
2. A SpatialTransformer block then takes that feature map and the text conditioning, refines it using self- and cross-attention, and outputs an updated feature map.
3. The U-Net alternates between these blocks.

The code walkthrough in the ExplainingAI video shows exactly how these new layers are added to the U-Net's downsampling blocks.

def forward(self, x, context = None):
    # Self-attention block
    x = self.attn1(self.norm1(x)) + x
    # Cross-attention block
    x = self.attn2(self.norm2(x), context = context) + x
    # Feed-forward block
    x = self.ff(self.norm3(x)) + x
    return x

Conclusion

Today, we've unlocked two of the most important mechanisms in modern diffusion U-Nets. These transform the model from a simple denoiser into a sophisticated, context-aware generative system.

Key Takeaways:

• Self-Attention is used within the U-Net to model long-range spatial relationships in the image feature map, allowing any pixel to attend to any other pixel.
• Cross-Attention is the key to external conditioning. It allows the image features (Queries) to attend to an external context like text embeddings (Keys and Values), infusing the generation process with semantic guidance.
• These are implemented practically using flexible Attention modules combined into TransformerBlocks, which are then placed at lower-resolution stages of the U-Net.
• The overall flow within such a block is typically: Self-Attention -> Cross-Attention -> Feed-Forward Network, with normalizations and residual connections throughout.

Next Lesson Preview:

We've now seen two powerful forms of conditioning: timestep embeddings and cross-attention on text. In the next lesson, we will look at a simpler, yet very effective, conditioning method: implementing class conditioning using learned embeddings. This is a foundational technique used in models like ADM (Abe et al., 2021) and is a great building block before we tackle even more advanced guidance methods.