Architecture Modifications for Conditioning

To effectively guide the diffusion model's generation process using conditioning information $y$, we often need to adapt the underlying neural network architecture. The standard U-Net, designed to predict noise $\epsilon_\theta(x_t, t)$ based on the noisy input $x_t$ and timestep $t$, must be modified to also incorporate $y$. The goal is to make the noise prediction conditional, becoming $\epsilon_\theta(x_t, t, y)$.

Several strategies exist for integrating this conditioning information, with the choice often depending on the nature and dimensionality of $y$.

Conditioning on Low-Dimensional Information (e.g., Class Labels)

When the conditioning information $y$ is relatively simple, such as a class label for image generation (e.g., "cat", "dog", corresponding to integer labels 0, 1), we can use straightforward techniques.

1. Embedding and Adding to Timestep Embeddings: A common approach, particularly effective with classifier-free guidance (CFG), is to treat the class label $y$ similarly to the timestep $t$.

   • First, the discrete class label $y$ is converted into a vector embedding, often using a standard embedding layer (torch.nn.Embedding in PyTorch). Let's call this $e_y$.
   • This class embedding $e_y$ is then often added directly to the timestep embedding $e_t$.
   • The combined embedding ($e_t + e_y$) is then projected and added into various blocks of the U-Net, just like the original timestep embedding.

   This approach effectively informs the network about the desired class at multiple levels of feature processing.

   Diagram illustrating the process of combining timestep and class label embeddings before injecting them into U-Net blocks.

2. Concatenation: Another method involves concatenating the embedded conditioning information $e_y$ with the input $x_t$ along the channel dimension. This augmented input is then fed into the U-Net. Alternatively, $e_y$ could be spatially broadcast (tiled to match the spatial dimensions of an intermediate feature map) and concatenated with feature maps deeper within the network. While simple, this might not always effectively propagate the conditioning signal throughout the complex U-Net structure compared to the additive embedding method.

Conditioning on High-Dimensional Information (e.g., Text Embeddings)

When conditioning on richer, higher-dimensional information like text descriptions, simple addition or concatenation is usually insufficient. Text requires capturing sequential dependencies and detailed meanings. For this, cross-attention has become the standard mechanism, forming the foundation of modern text-to-image models like Stable Diffusion.

Cross-Attention Mechanism

Recall that self-attention layers within a U-Net (often in Transformer-style blocks) allow different spatial locations in the image representation to attend to each other. Cross-attention layers work similarly but allow the image representation to attend to the conditioning information $y$.

Here's how it typically works within a U-Net block designed for conditional generation (e.g., text-to-image):

1. Inputs: The layer receives two main inputs:

   • The intermediate image representation from the U-Net (let's denote its features as $z$).
   • The conditioning information, usually a sequence of embedding vectors derived from $y$ (e.g., text embeddings from a pre-trained encoder like CLIP). Let's call this sequence $c$.

2. Query, Value:

   • Query (Q): Derived from the image features $z$ via a linear projection. These represent "what information is needed" at each spatial location of the image feature map.
   • Conditioning (K): Derived from the conditioning embeddings $c$ via a linear projection. These represent "what information is available" in the conditioning sequence.
   • Value (V): Also derived from the conditioning embeddings $c$ via a linear projection. These represent the actual content or features provided by the conditioning information.

3. Attention Calculation: The core operation calculates attention scores based on the similarity between Queries (from image) and Keys (from conditioning). A common way is scaled dot-product attention:

   $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

   Here, $d_k$ is the dimension of the vectors. The softmax ensures the weights sum to 1.

4. Output: The result is a weighted sum of the Value vectors (derived from conditioning $c$), where the weights are determined by how relevant each part of the conditioning sequence is to each part of the image representation. This output is then typically added back to the original image features $z$ (often via a residual connection) or further processed within the U-Net block.

Integration into U-Net

Cross-attention layers are usually inserted into multiple blocks of the U-Net, particularly within the downsampling, bottleneck, and upsampling paths. This allows the conditioning information $y$ to influence the denoising process at various levels of feature abstraction.

Diagram showing the integration of conditioning embeddings ($c$) into image features ($z$) using a cross-attention layer within a U-Net block. Queries are derived from $z$, while Keys and Values come from $c$.

Example: Text-to-Image Generation

In a text-to-image model:

1. A text prompt (e.g., "A photorealistic astronaut riding a horse") is encoded by a text encoder (like CLIP) into a sequence of embedding vectors $c$.
2. During the reverse diffusion process, at each step $t$, the U-Net receives the current noisy image $x_t$, the timestep $t$, and the text embeddings $c$.
3. The cross-attention layers within the U-Net use $x_t$'s features to query the text embeddings $c$, integrating the textual guidance into the predicted noise $\epsilon_\theta(x_t, t, c)$.
4. This guides the denoising process to generate an image consistent with the input text prompt.

By modifying the U-Net architecture to incorporate conditioning information $y$, particularly using mechanisms like additive embeddings or cross-attention, diffusion models gain the ability to perform controlled generation, producing outputs tailored to specific requirements like class labels or detailed text descriptions.