Text Conditioning with Cross-Attention

From Self-Attention to Cross-Attention

Before we can have our model "attend" to a text prompt, let's briefly revisit the concept of attention within the U-Net itself. As you know from your prior work, many U-Net architectures for diffusion use self-attention. In self-attention, each pixel's feature vector in a given layer queries all other pixel features in the same layer. This allows the model to capture long-range spatial dependencies within the image being generated—for example, ensuring that two eyes in a portrait are stylistically consistent. The Query (Q), Key (K), and Value (V) all originate from the same source: the image's intermediate feature map.

Cross-attention adapts this powerful mechanism for conditioning. Instead of the image attending to itself, the image attends to an external source of information—our text prompt. The core idea is:

• The Query (Q) is derived from the image's feature map (the x in the U-Net).
• The Key (K) and Value (V) are derived from the external context, which in our case is the sequence of text embedding vectors from the text encoder.

This allows each spatial location in the nascent image to "look at" the entire text prompt and decide which words or concepts are most relevant to it, pulling in that information to guide its denoising.

To build a strong intuition for this process, let's watch a video that provides an excellent visual breakdown.

The key takeaway is this directional flow of information: the image queries the text, not the other way around. This is what allows the text to guide, or "condition," the image generation process.

Implementing the Cross-Attention Module

With the theory in place, let's move to a practical implementation. A cross-attention mechanism is typically encapsulated in its own nn.Module. We're going to examine a clean, well-structured implementation.

Let's dissect the code you just reviewed. It elegantly captures the concept we discussed.

class CrossAttention(nn.Module):
    def __init__(self, n_heads: int, d_embed: int, d_cross: int, ...):
        super().__init__()
        # Q is projected from the image feature dimension (d_embed)
        self.q_proj = nn.Linear(d_embed, d_embed, ...)
        # K and V are projected from the context dimension (d_cross, e.g., 768 for CLIP)
        self.k_proj = nn.Linear(d_cross, d_embed, ...)
        self.v_proj = nn.Linear(d_cross, d_embed, ...)
        self.out_proj = nn.Linear(d_embed, d_embed, ...)
        # ...

    def forward(self, x, y):
        # x (latent): (Batch_size, Seq_Len_Q, Dim_Q) -> From U-Net
        # y (context): (Batch_size, Seq_Len_KV, Dim_KV) -> From Text Encoder

        input_shape = x.shape
        # ... (multi-head setup) ...

        q = self.q_proj(x)
        k = self.k_proj(y)
        v = self.v_proj(y)

        # ... (transpose for multi-head attention) ...

        weight = q @ k.transpose(-1, -2)
        weight /= math.sqrt(self.d_head)
        weight = F.softmax(weight, dim=-1)
        
        output = weight @ v
        
        # ... (reshape and final projection) ...
        return output

Notice the critical lines:

• In __init__, self.q_proj is initialized to take d_embed (the dimension of the U-Net's internal feature maps) as input.
• self.k_proj and self.v_proj are initialized to take d_cross (the dimension of the text embeddings) as input.
• In forward(self, x, y), q is created from x (the image features), while k and v are created from y (the text context).

This perfectly implements the directional attention flow we need for text-to-image generation.

Integrating Cross-Attention into the U-Net

Now that we have our CrossAttention module, the next step is to integrate it into the U-Net architecture. Typically, this is done by adding a cross-attention layer within the main building blocks (e.g., ResNet or Transformer blocks) of the U-Net.

The following video provides a complete, from-scratch walkthrough of adding class, spatial, and finally, text conditioning to a latent diffusion model. We'll focus on the text conditioning part.

As the video demonstrates, the integration involves a few key steps:

1. Modify the U-Net Blocks: The building blocks of your U-Net (often called ResnetBlock or DownBlock/UpBlock) are modified. Alongside the existing Resnet and SelfAttention layers, you add a CrossAttention layer.
2. Update Forward Passes: The forward method of these blocks is updated to accept an additional context argument. This context is passed to the cross-attention layer.
3. Thread Context Through U-Net: The main Unet's forward method is also updated to accept the context. It then passes this context down to each of its blocks during the forward pass.
4. Prepare Context in Training Loop: Before calling the U-Net, your training script is responsible for taking the text prompts, using the pre-trained text encoder to convert them into embedding vectors, and then passing these vectors as the context argument to the U-Net.

The MinImagen article provides another great visualization of how the conditioning vectors are prepared and used. The diagram under the "Generating Text Conditioning Vectors" section (in resource LINK, which you can refer to in the repo) clearly shows the text embeddings (text_embeds) being processed and concatenated with time tokens to form the final context tensor that is fed into the cross-attention layers.

Conclusion

Today, we've bridged the gap between language and pixels. You have learned how to implement the core mechanism that allows a diffusion model to understand and generate images from text prompts.

Key Takeaways:

• Cross-Attention is the Mechanism for Conditioning: It adapts the self-attention paradigm to allow the model to incorporate external information.
• Directional Information Flow: For text-to-image, the image features (Query) must attend to the text embeddings (Key and Value) to receive guidance.
• Architectural Integration: Cross-attention modules are typically added into the main blocks of the U-Net (e.g., alongside ResNet and self-attention layers).
• Implementation Steps: The process involves creating a CrossAttention module, modifying the U-Net blocks to include it, and updating the forward passes to thread the context (text embeddings) through the network.

Next Lesson Preview:

Now that we can condition our model on a prompt, the next step is to gain finer control over the output. What if we want to explicitly discourage a concept? What if we want to make the model follow the prompt more strongly? In the next lesson, we will explore powerful techniques like negative prompting and prompt re-weighting, which build directly upon the cross-attention framework you've learned today to refine and direct the creative process.