Diffusion U-Net

Problem

Diffusion models train one network ε_θ(x_t, t) to predict the noise that was added to a clean image. Its input is a noisy image and its output is a tensor of the same shape: a guess for every pixel. The same network is called at every one of up to 1000 noise levels, from almost-clean images where it must find faint grain to pure static where it must invent global layout.

That asks a lot of one architecture. It must see the whole image, because deciding where a face goes is a global decision. It must also keep exact pixel positions, because the noise is per-pixel. It must know how noisy its input is. And for text-to-image it must listen to a prompt. A plain stack of convolutions sees only a small window; a classifier throws away position. The question is what shape of network does all four.

Intuition

The U-Net, introduced by Ronneberger et al. for biomedical segmentation, already had the right shape: another image-in, image-out problem. Its down path halves the resolution step by step while widening the channels, so deep layers see the whole picture coarsely. Its up path mirrors it, doubling the resolution back. The two paths drawn side by side make a U.

The trick is the skip connections across the U. Squeezing a 64 × 64 grid to 8 × 8 loses fine detail, and the up path cannot invent it back. So each down-path output is kept and concatenated onto the matching up-path input: the decoder gets both "what the image is about" from below and "exactly where things were" from the side. Think of an editor who reads a summary of the whole book and also keeps the original page open.

Diffusion adds two inputs on top. The timestep becomes a vector, the same sinusoids used for positional encoding, and every residual block adds it to its features, so every layer knows the noise level. The prompt enters through cross-attention layers placed where the grid is small enough for attention to be affordable.

Mechanics

Input and output

Stable Diffusion runs in the latent space of a VAE: a 512 × 512 RGB image becomes a 4 × 64 × 64 latent, and the U-Net denoises that (the latent diffusion idea of Rombach et al.). The U-Net's input and output both have shape [b, 4, 64, 64]. A 3 × 3 convolution, conv_in, lifts the 4 channels to the base width of 320; at the end a GroupNorm, SiLU and 3 × 3 convolution bring 320 channels back to 4, the predicted noise.

Down path, up path and skips

The down path has four levels with channel multipliers [1, 2, 4, 4] on the base width, so 1, 2, 4, 4 × 320 = 320, 640, 1280, 1280 channels. Each level runs 2 residual blocks, then (except the last) a downsample: a 3 × 3 convolution with stride 2 that halves height and width. Every block's output, including conv_in's and each downsample's, is pushed onto a stack: 12 saved tensors.

The up path walks the levels in reverse with 3 residual blocks per level, one more than the down path, so it pops all 12 saved tensors. Each up block concatenates the popped tensor onto its input along the channel axis, torch.cat([h, skip], dim=1), which is why up-block inputs are wide (up to 2560 channels) and the first convolution of each up block narrows them. At the end of each level except the finest, an upsample doubles the resolution: nearest-neighbour interpolation (nn.Upsample(scale_factor=2, mode="nearest")) followed by a 3 × 3 convolution. The original U-Net used a learned transposed convolution (nn.ConvTranspose2d) here instead; resize-then-convolve avoids the checkerboard artifacts transposed convolutions can produce when the kernel size is not divisible by the stride.

skip: concat([h, s], dim = 1)  ·  channels_in = C_h + C_s

Concatenation, not addition, is the U-Net choice: the up block can weigh the skip and the upsampled path however it likes, at the cost of wider convolutions. A ResNet's skip adds; a U-Net's skip stacks.

ResBlock with GroupNorm

The workhorse is a pre-activation residual block. With input x of C_in channels and timestep embedding e:

h = Conv_3×3(SiLU(GN(x))) + Linear(SiLU(e))[:, :, None, None]
y = skip(x) + Conv_3×3(Dropout(SiLU(GN(h))))

The skip is the identity when channels match and a 1 × 1 convolution when they change. The last convolution is initialized to zero, so every block starts as the identity and the untrained network passes its input straight through; the same zero init is used on the output convolution and on each attention block's output projection.

GroupNorm (Wu and He) splits the C channels into G groups and normalizes each group over its channels and all spatial positions, separately for every sample. LayerNorm is the case G = 1 and InstanceNorm the case G = C. Unlike BatchNorm, nothing depends on the other samples in the batch, which matters when a large model fits only a few images per GPU. Stable Diffusion uses G = 32 everywhere, so every channel count must be a multiple of 32; nn.GroupNorm raises a ValueError otherwise. SiLU, also called Swish, is x · σ(x).

Timestep embedding

The integer timestep t becomes sinusoids of geometrically spaced frequencies, the transformer's position encoding with t as the position. With embedding width d and half = d/2:

ω_i = 10000^−i/half,   i = 0 … half − 1  ·  emb(t) = [cos(tω₀), …, cos(tω_half−1), sin(tω₀), …, sin(tω_half−1)]

That is the latent-diffusion and DiT convention. The original DDPM code divides by half − 1 instead of half and puts sin before cos; the two are equivalent up to a permutation and a slight change of frequencies, and any one model only ever sees its own. Stable Diffusion embeds t at width 320 and passes it through Linear(320, 1280) → SiLU → Linear(1280, 1280). Each ResBlock then projects the 1280-wide vector to its own channel count and adds it to every spatial position: a per-channel bias that depends on the noise level. Guided-diffusion's ADM variant instead produces a scale and a shift and applies them after the second GroupNorm (adaptive GroupNorm), which Dhariwal and Nichol found improved sample quality.

Self- and cross-attention at low resolution

Convolutions are local; attention lets every position see every other. Attention compares all pairs of positions, so its cost grows with the square of the token count, and the U-Net places it only where the grid is small. DDPM used self-attention at 16 × 16 only; Stable Diffusion v1 places a transformer block after every ResBlock at 64, 32 and 16 and in the middle block at 8 (configured as attention_resolutions [4, 2, 1], which are downsampling factors, not sizes).

Each block first applies GroupNorm and a 1 × 1 convolution, then flattens the grid to a sequence, [b, C, H, W] → [b, HW, C], and runs a pre-norm transformer layer: self-attention over the HW positions, cross-attention whose queries come from the image and whose keys and values come from the 77 × 768 CLIP text-encoder output, and a GEGLU feed-forward (Linear to 8C, split in half, one half times GELU of the other, Linear 4C → C). A zero-initialized 1 × 1 convolution maps back and the result is added to the block input. Cross-attention is how the prompt reaches the pixels: each position asks which words concern it. With 8 heads at every level, the head width is 40, 80, 160 as the channels grow.

DiT: the transformer alternative

Peebles and Xie asked whether the U-Net's convolutional inductive bias is needed at all. A Diffusion Transformer (DiT) treats the latent the way a Vision Transformer treats an image: cut the 4 × 32 × 32 latent into p × p patches (a Conv2d with kernel and stride p), add fixed sin-cos position embeddings, and run a stack of plain transformer blocks at one constant width. There is no down path, no up path and no skip; a final linear layer turns each token back into a p × p × 8 patch (4 channels of noise and 4 of learned variance), and the patches are reassembled.

Conditioning changes too. The timestep (embedded with the same sinusoids at width 256, then an MLP) and the class label (a learned embedding table with one extra row for "no class", used for classifier-free guidance) are summed into one vector c. Instead of being added to features, c drives adaLN-Zero: a SiLU and Linear(D, 6D) produce a shift, a scale and a gate for each of the attention and MLP branches. The LayerNorms have no learned affine; the shift and scale come from c, and each branch output is multiplied by its gate before the residual add. The Linear is initialized to zero, so every gate starts at zero and every block starts as the identity, the same trick as the U-Net's zero-initialized convolutions. Among the conditioning schemes the paper compared (in-context tokens, cross-attention, adaLN and adaLN-Zero), adaLN-Zero gave the best FID.

x ← x + α₁ · Attn(LN(x)(1 + γ₁) + β₁)  ·  x ← x + α₂ · MLP(LN(x)(1 + γ₂) + β₂),   (β, γ, α)_1,2 = Linear(SiLU(c))

The paper's finding was that DiTs scale cleanly: more compute per forward pass, whether from a larger model or from smaller patches (more tokens), gave better samples, and DiT-XL/2 set a new state of the art on class-conditional ImageNet at 256 × 256. Later large text-to-image models moved to transformer backbones in this spirit.

Worked example

All numbers below come from tested helpers in ml_models_docs.py (unet_walk, dit_params, timestep_embedding_ldm, timestep_embedding_ddpm, unet_group_norm), written in plain Python. The Stable Diffusion count reproduces the 859.52M the latent diffusion code reports for this configuration, and the DiT-XL/2 count matches the paper's 675M.

Implementation

A two-level U-Net with timestep conditioning in plain PyTorch. It keeps every idea from Stable Diffusion's network at toy size: GroupNorm-SiLU-conv ResBlocks, the embedding added per channel, a stride-2 downsample, nearest upsample plus conv, concatenated skips, and zero-initialized last convolutions.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


def timestep_embedding(t, dim, max_period=10000):
    half = dim // 2
    freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32) / half)
    args = t[:, None].float() * freqs[None]
    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)


class ResBlock(nn.Module):
    def __init__(self, cin, cout, emb_dim, groups=32):
        super().__init__()
        self.norm1 = nn.GroupNorm(groups, cin)
        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.emb = nn.Linear(emb_dim, cout)
        self.norm2 = nn.GroupNorm(groups, cout)
        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
        nn.init.zeros_(self.conv2.weight)
        nn.init.zeros_(self.conv2.bias)
        self.skip = nn.Identity() if cin == cout else nn.Conv2d(cin, cout, 1)

    def forward(self, x, emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.emb(F.silu(emb))[:, :, None, None]
        h = self.conv2(F.silu(self.norm2(h)))
        return self.skip(x) + h


class TinyUNet(nn.Module):
    def __init__(self, cin=4, ch=64, emb_dim=256):
        super().__init__()
        self.ch = ch
        self.time = nn.Sequential(nn.Linear(ch, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim))
        self.conv_in = nn.Conv2d(cin, ch, 3, padding=1)
        self.down1 = ResBlock(ch, ch, emb_dim)
        self.down = nn.Conv2d(ch, ch, 3, stride=2, padding=1)
        self.down2 = ResBlock(ch, 2 * ch, emb_dim)
        self.mid = ResBlock(2 * ch, 2 * ch, emb_dim)
        self.up2 = ResBlock(4 * ch, 2 * ch, emb_dim)
        self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(2 * ch, 2 * ch, 3, padding=1))
        self.up1 = ResBlock(3 * ch, ch, emb_dim)
        self.out = nn.Sequential(nn.GroupNorm(32, ch), nn.SiLU(), nn.Conv2d(ch, cin, 3, padding=1))
        nn.init.zeros_(self.out[-1].weight)
        nn.init.zeros_(self.out[-1].bias)

    def forward(self, x, t):
        emb = self.time(timestep_embedding(t, self.ch))
        h1 = self.down1(self.conv_in(x), emb)               # [b, 64, 32, 32], saved
        h2 = self.down2(self.down(h1), emb)                 # [b, 128, 16, 16], saved
        h = self.mid(h2, emb)
        h = self.up2(torch.cat([h, h2], dim=1), emb)         # 128 + 128 -> 128
        h = self.up1(torch.cat([self.up(h), h1], dim=1), emb)  # 128 + 64 -> 64, at 32 x 32
        return self.out(h)


model = TinyUNet()
x = torch.randn(8, 4, 32, 32)
t = torch.randint(0, 1000, (8,))
assert model(x, t).shape == x.shape

Attention slots in after a ResBlock at a small resolution. A simplified block (Stable Diffusion's adds the 1 × 1 projections and a GEGLU feed-forward) using nn.MultiheadAttention, whose kdim and vdim let keys and values come from a text encoder of a different width:

class SpatialAttention(nn.Module):
    def __init__(self, ch, ctx_dim, heads=8):
        super().__init__()
        self.norm = nn.GroupNorm(32, ch)
        self.ln1, self.ln2 = nn.LayerNorm(ch), nn.LayerNorm(ch)
        self.self_attn = nn.MultiheadAttention(ch, heads, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(ch, heads, kdim=ctx_dim, vdim=ctx_dim, batch_first=True)

    def forward(self, x, ctx):
        b, c, height, width = x.shape
        h = self.norm(x).flatten(2).transpose(1, 2)        # [b, H*W, C]
        q = self.ln1(h)
        h = h + self.self_attn(q, q, q, need_weights=False)[0]
        q = self.ln2(h)
        h = h + self.cross_attn(q, ctx, ctx, need_weights=False)[0]
        return x + h.transpose(1, 2).reshape(b, c, height, width)

Tradeoffs

• U-Net versus DiT. The U-Net's multi-resolution convolutions are a strong prior for images and work well at moderate scale; its design has many knobs (levels, multipliers, where attention goes). DiT has one width, one block and a patch size, scales predictably with compute, and reuses transformer tooling, but it is less data-efficient and its token count, and cost, grows quickly as patches shrink.
• Where attention goes. Attention at fine resolutions improves global coherence and prompt following but its cost grows with the fourth power of the side. Stable Diffusion v1 pays for it at 64 × 64; DDPM kept it at 16 × 16 only.
• Concatenated skips. They preserve detail but make up-path convolutions wide; most of the parameters sit in the 1280-channel levels.
• Conditioning. Adding the embedding is cheapest; adaptive normalization (scale and shift, as in ADM's AdaGN or DiT's adaLN-Zero) is more expressive for global signals like timestep and class; cross-attention is needed for sequences like prompts.
• Zero initialization. Starting every residual branch at zero makes deep networks stable from step one, at no cost.
• Fixed sizes. The latent side must survive repeated halving (a multiple of 8 for three downsamples), and a network trained at one resolution drifts at others without fine-tuning.

Connections

References

• Ronneberger, Fischer and Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. MICCAI 2015. arXiv:1505.04597
• Ho, Jain and Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS 2020. arXiv:2006.11239
• Dhariwal and Nichol. Diffusion Models Beat GANs on Image Synthesis. NeurIPS 2021. arXiv:2105.05233
• Rombach et al. High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022. arXiv:2112.10752
• Peebles and Xie. Scalable Diffusion Models with Transformers. ICCV 2023. arXiv:2212.09748
• Wu and He. Group Normalization. ECCV 2018. arXiv:1803.08494
• Ho and Salimans. Classifier-Free Diffusion Guidance. 2022. arXiv:2207.12598
• Shazeer. GLU Variants Improve Transformer. 2020. arXiv:2002.05202
• Odena, Dumoulin and Olah. Deconvolution and Checkerboard Artifacts. Distill, 2016. distill.pub/2016/deconv-checkerboard
• PyTorch v2.15.0: Conv2d, ConvTranspose2d, Upsample, GroupNorm, SiLU, MultiheadAttention.