The U-Net Architecture for Diffusion

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. U-Net Structure
5. Time Step Conditioning
6. Skip Connections
7. Self-Attention in U-Net
8. Cross-Attention for Text Conditioning
9. Diffusion Transformer (DiT)
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Describe the encoder-decoder structure with multi-scale processing.
2. Derive the time embedding mechanism (sinusoidal + MLP).
3. Explain why skip connections are critical for diffusion U-Nets.
4. Show how cross-attention enables text-to-image generation.
5. Compare U-Net vs DiT as diffusion backbones.

Notation

• $\mathbf{x}_t \in \mathbb{R}^{H \times W \times C}$ — noisy input at timestep $t$
• $\mathbf{t}_{\text{emb}} \in \mathbb{R}^d$ — time step embedding
• $\mathbf{c} \in \mathbb{R}^{L \times d_c}$ — conditioning (e.g., text embeddings)

Core Intuition

The denoising network $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t, \mathbf{c})$ must predict noise at multiple spatial scales: large-scale structure (is this a face or a landscape?) and fine details (texture, edges). The U-Net processes information at progressively lower resolutions (encoder), captures global context, then reconstructs at full resolution (decoder), combining coarse context with fine-grained features via skip connections.

U-Net Structure

Encoder (downsampling):

• Input: $H \times W \times C$
• Each level: ResBlock → ResBlock → Downsample (stride-2 convolution)
• Feature maps grow in channels: $C \to 2C \to 4C \to 8C$
• Spatial resolution: $H \to H/2 \to H/4 \to H/8$

Bottleneck: ResBlock → SelfAttention → ResBlock at lowest resolution.

Decoder (upsampling):

• Each level: Concatenate skip connection → ResBlock → ResBlock → Upsample
• Mirrors encoder levels in reverse.
• Output: $H \times W \times C$ (same as input).

ResBlock structure: GroupNorm → SiLU → Conv → GroupNorm → SiLU → Conv + skip.

Time Step Conditioning

The network must know which timestep $t$ it's denoising — noise level changes behavior dramatically.

Sinusoidal embedding (same as transformer positional encoding):

MLP projection: $\mathbf{t}_{\text{emb}} = \text{MLP}(\mathbf{t}_{\text{raw}}) \in \mathbb{R}^d$.

Injection into ResBlocks: After the first convolution, add time embedding (broadcast across spatial dimensions):

This modulates the feature maps based on the noise level — the network learns different behavior per timestep.

Skip Connections

Why critical: Without skip connections, the decoder must reconstruct fine details from heavily downsampled (8x smaller) features. Skip connections provide direct access to encoder features at matching resolution.

Concatenation: Decoder features are concatenated with encoder features (channel-wise), then processed by convolutions that learn to combine them.

Information flow:

• Skip: high-frequency details, edges, exact spatial positions.
• Bottleneck: semantic content, global structure, conditioning integration.

Self-Attention in U-Net

Applied at lower resolutions (e.g., $32 \times 32$ or $16 \times 16$) where the sequence length is manageable:

Purpose: Captures long-range spatial dependencies (e.g., symmetric structures, global color consistency) that convolutions miss.

Cost: $O(H^2W^2)$ at the applied resolution — feasible at $32 \times 32$ ($\approx 1M$ entries) but not at $256 \times 256$ ($\approx 4B$ entries).

Cross-Attention for Text Conditioning

In text-to-image models (Stable Diffusion), text information is injected via cross-attention:

• Query: Spatial features from U-Net (flattened to sequence).
• Key/Value: Text encoder output (CLIP or T5 embeddings).

Placement: After each self-attention layer in the U-Net (at selected resolutions).

Effect: Each spatial position "reads" from the text prompt, determining what content to generate at that location.

Diffusion Transformer (DiT)

Alternative backbone (Peebles & Xie, 2023): Replace U-Net entirely with a vision transformer.

Architecture:

• Patchify input: divide noisy image into $p \times p$ patches → sequence of tokens.
• Process with standard transformer blocks (self-attention + FFN).
• Time and class conditioning via adaptive layer norm (adaLN):

Advantages: Simpler architecture, better scaling properties (follows transformer scaling laws), no need for multi-scale design.

Used in: DALL-E 3, Sora, Stable Diffusion 3.

Common Pitfalls

Pitfall 1. Applying self-attention at full resolution. At $256 \times 256$: attention matrix = 4.3 billion entries. Only apply at downsampled resolutions ($\leq 64 \times 64$).

Pitfall 2. Removing skip connections. Without them, the model loses fine-grained spatial information and produces blurry outputs.

Pitfall 3. Using the same architecture for pixel-space and latent-space diffusion. Latent space is lower resolution (typically $64 \times 64$ or $32 \times 32$), allowing self-attention at all levels.

Summary

• U-Net: Multi-scale encoder-decoder with skip connections; standard for image diffusion.
• Time conditioning: Sinusoidal embedding + MLP, injected via addition to ResBlock features.
• Self-attention: Captures global dependencies at low resolutions.
• Cross-attention: Injects text/conditioning information from external encoders.
• DiT: Transformer alternative; simpler, scales better, dominant in latest models.

Exercises

Exercise 1. Compute the total parameters for a U-Net with channel multipliers [1, 2, 4, 8], base channels $C=256$, and attention at the two lowest resolutions.

Exercise 2. Explain why time embedding is added (not concatenated) to spatial features in ResBlocks.

Exercise 3. For input $256 \times 256$: compute the self-attention cost at $32 \times 32$ vs $64 \times 64$ resolution.

Exercise 4. Compare DiT and U-Net in terms of inductive biases for image generation (locality, multi-scale, global context).

Exercise 5. Derive the number of cross-attention FLOPs for text conditioning with text length 77 and spatial resolution $32 \times 32$.