Principle:Zai org CogVideo UNet Backbone
| Knowledge Sources | |
|---|---|
| Domains | Video_Generation, Diffusion_Models, Neural_Architecture |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
The UNet backbone is the core denoising network architecture for diffusion models, using an encoder-decoder structure with skip connections, timestep conditioning, and optional cross-attention to iteratively denoise latent representations.
Description
The UNet backbone serves as the learned denoising function in diffusion-based generative models. During training, it learns to predict the noise added to a sample (or equivalently, predict the clean sample) at any given diffusion timestep. During inference, it is applied iteratively across a schedule of decreasing noise levels to transform pure noise into a coherent output.
The architecture derives from the original U-Net design for image segmentation, adapted for generative modeling with several critical additions:
Timestep conditioning injects information about the current noise level into every residual block of the network. The timestep is first encoded using sinusoidal positional embeddings (following the transformer literature), then projected through a two-layer MLP to produce a conditioning vector. This vector is injected into residual blocks either additively or via a FiLM-style scale-and-shift mechanism applied after normalization.
Skip connections between corresponding encoder and decoder levels preserve fine-grained spatial information that would otherwise be lost through progressive downsampling. Encoder features are concatenated (not added) with decoder features, doubling the channel count at the concatenation point, which a subsequent convolution then reduces.
Cross-attention enables conditioning on external signals such as text embeddings. At designated resolution levels, spatial transformer blocks interleave self-attention (over spatial positions) with cross-attention (between spatial positions and conditioning tokens). This is the mechanism by which text prompts guide image or video generation.
Multi-head attention is used throughout, with configurable head counts. Two QKV attention orderings are supported: the legacy format (split heads before splitting QKV) and the newer format (split QKV before splitting heads), both producing equivalent results.
Gradient checkpointing allows trading computation time for memory by recomputing intermediate activations during the backward pass rather than storing them, which is essential for training large models on limited GPU memory.
Usage
Apply this architecture as the core denoising network in any diffusion-based generation pipeline. It is suitable for both image and video generation when combined with appropriate attention mechanisms (spatial transformers for images, spatiotemporal transformers for video).
Theoretical Basis
The diffusion model learns to reverse a forward noising process. Given a clean sample x_0, the forward process adds Gaussian noise at each timestep:
x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
where epsilon ~ N(0, I) and alpha_bar_t is the cumulative noise schedule.
The UNet epsilon_theta is trained to predict the noise:
L = E[|| epsilon - epsilon_theta(x_t, t, c) ||^2]
where c is optional conditioning (text, class labels, etc.).
Timestep embedding uses sinusoidal encoding:
emb_i = sin(t / 10000^(2i/d)) for even i emb_i = cos(t / 10000^(2i/d)) for odd i t_emb = MLP(emb) = Linear(SiLU(Linear(emb)))
Residual block with timestep conditioning:
# Additive mode: h = Conv(SiLU(GroupNorm(x))) h = h + Linear(SiLU(t_emb)) # broadcast over spatial dims h = Conv(Dropout(SiLU(GroupNorm(h)))) output = SkipConnection(x) + h # FiLM (scale-shift) mode: h = Conv(SiLU(GroupNorm(x))) scale, shift = chunk(Linear(SiLU(t_emb)), 2) h = GroupNorm(h) * (1 + scale) + shift h = Conv(Dropout(SiLU(h))) output = SkipConnection(x) + h
Cross-attention for text conditioning:
# In SpatialTransformer block: h = reshape(x, "b c h w -> b (h w) c") # Self-attention Q_s, K_s, V_s = Linear(LayerNorm(h)) h = h + Attention(Q_s, K_s, V_s) # Cross-attention with text context c Q_c = Linear(LayerNorm(h)) K_c, V_c = Linear(c) # c is text embedding (b, seq_len, dim) h = h + Attention(Q_c, K_c, V_c) output = reshape(h, "b (h w) c -> b c h w")
Skip connections in the U-Net:
# Encoder: store features at each level
encoder_features = []
for block in input_blocks:
h = block(h, t_emb, context)
encoder_features.append(h)
# Bottleneck
h = middle_block(h, t_emb, context)
# Decoder: concatenate with stored encoder features
for block in output_blocks:
h = concat([h, encoder_features.pop()], dim=channels)
h = block(h, t_emb, context)
The total downsampling factor determines the latent resolution at the bottleneck:
bottleneck_resolution = input_resolution / 2^(num_levels - 1)