Diffusion Model Fundamentals

Beyond VAE and GAN architectures, Diffusion Models represent another powerful paradigm in generative AI. This article explores the core concepts of Diffusion Models, along with Conditional variants and Latent Diffusion Models, accompanied by practical code implementations.

Core Mechanism

The diffusion model operates through two complementary processes: a forward process that progressively adds noise to an image, and a reverse process that reconstructs the original image from heavily corupted data. The architecture employs a noise scheduler to control the noise addition at each timestep.

The forward diffusion process transforms clean data toward a standard Gaussian distribution, while the reverse process learns to denoise by modeling the conditional probability distributions. Using \\(x_0\\) to denote the original image and \\(t\\) to represent the diffusion timestep, the forward process follows:

Forward Process Derivation

By recursively substituting the noise equations, we can derive a closed-form expression for the noisy image at any timestep:

Reverse Process

The reverse process attempts to reconstruct the original data by learning:

Training Objective

Training a diffusion model involves maximizing the marginal likelihood of the data. Starting from the log-likelihood objective:

Practical Implementation

The core training loop involves adding noise to images at random timesteps and training the model to predict that noise. Here's a simplified implementation using popular libraries:

import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler

# Initialize the noise scheduler with 1000 diffusion steps
noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

for batch in training_dataloader:
    images = batch['image']  # Shape: (B, C, H, W)
    
    # Sample random timesteps for each image in the batch
    batch_timesteps = torch.randint(
        0, 
        noise_scheduler.config.num_train_timesteps,
        (images.shape[0],),
        device=images.device,
        dtype=torch.int64
    )
    
    # Generate random noise for the forward process
    random_noise = torch.randn(images.shape, device=images.device)
    
    # Create noisy images using the scheduler
    noisy_images = noise_scheduler.add_noise(
        images, random_noise, batch_timesteps
    )
    
    # Predict the noise using the model
    noise_prediction = denoising_network(noisy_images, batch_timesteps)
    
    # Compute MSE loss between predicted and actual noise
    loss = F.mse_loss(noise_prediction, random_noise)
    
    loss.backward()
    optimizer.step()

Conditional Diffusion Models

Conditional variants extend the basic diffusion framework by incorporating additional conditioning information such as text descriptions, class labels, or reference images. This guidance significantly improves generation quality and control.

Conditioning can be integrated through several approaches:

• Additive Conditioning: Encoded conditions are added directly to the noise representation or timestep embeddings
• Attention-based Fusion: Condition embeddings are fused through cross-attention mechanisms, as implemented in Stable Diffusion

The attention mechanism in conditional models follows:

Latent Diffusion Models

Latent Diffusion Models (LDM) address the computational efficiency challenge of operating in high-dimensional pixel space. The key innovation involves first compressing images into a lower-dimensional latent space using an autoencoder, then performing diffusion in this compressed representation.

The architecture consists of three main components:

1. Autoencoder: Encodes images \\(x \in \mathbb{R}^{3 \times H \times W}\\) to latents \\(z = \varepsilon(x) \in \mathbb{R}^{d \times h \times w}\\)
2. Diffusion UNet: Operates on the latent space to learn the denoising process
3. Conditioning Encoder: Encodes external conditioning (text, images) into representations \\(\tau_\theta(y)\\)

This compression dramatically reduces the computational cost while preserving essential image features. For example, encoding a 256×256 image to a 32×32 latent reduces the computation by approximately 64× in the diffusion process.

Generation Strategies

DDPM Sampling

The original DDPM sampling formula recursively applies the learned denoising function:

DDIM Acceleration

Denoising Diffusion Implicit Models (DDIM) relax the Markov chain assumption, enabling much faster sampling through non-Markovian transitions:

Network Architectures

DiT (Diffusion Transformer)

DiT adapts the Transformer architecture for diffusion tasks by first compressing images into latent patches, then applying transformer blocks with conditioning from timesteps and labels.

The input pipeline transforms latents into a sequence of patch tokens through embedding, then processes them through transformer blocks that incorporate adaptive layer normalization (adaLN):

# DiT model initialization
# Configuration: depth=12, hidden_size=384, patch_size=4, num_heads=6
batch_size = 16
latents = torch.randn(batch_size, 4, 32, 32).to(device)
time_indices = torch.randint(0, 1000, (batch_size,)).to(device)
class_labels = torch.randint(0, 1000, (batch_size,)).to(device)

# Patch embedding produces: (batch, num_patches, hidden_size)
# For 32x32 latents with patch_size=4: 8x8=64 patches
# Each patch: 4x4x4 = 64 values, mapped to 384 dimensions

# Time and label embeddings: (batch, hidden_size)
time_emb = time_embedding(time_indices)
label_emb = label_embedding(class_labels)

The DiT block applies adaptive conditioning to both attention and feedforward layers:

def dit_block_forward(x, conditioning):
    # Generate adaptive modulation parameters
    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
        adaptive_modulation(conditioning).chunk(6, dim=1)
    
    # Adaptive layer norm with learned affine transformations
    modulated_attn = modulate(normalize(x), shift_msa, scale_msa)
    x = x + gate_msa.unsqueeze(1) * attention(modules_attn(modulated_attn))
    
    # Conditional MLP with gating
    modulated_mlp = modulate(normalize(x), shift_mlp, scale_mlp)
    x = x + gate_mlp.unsqueeze(1) * mlp(modulated_mlp)
    return x

The adaLN modification transforms standard layer normalization \\(\gamma \frac{x - \mu}{\sigma} + \beta\\) into a dynamic form \\(\text{scale} \cdot \frac{x - \mu}{\sigma} + \text{shift}\\), where both scale and shift are derived from the conditioning input. This enables the model to adjust its behavior based on timestep and class information.

UNet Architecture

UNet-based diffusion models employ encoder-decoder structure with skip connections at corresponding resolutions. The architecture consists of downsampling layers, middle processing, and upsampling stages.

from diffusers import UNet2DModel

unet_model = UNet2DModel(
    sample_size=128,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(128, 128, 256, 256, 512, 512),
    down_block_types=("DownBlock2D",) * 6,
    up_block_types=("UpBlock2D",) * 6
).to(device)

# Input: image (batch, 3, 128, 128), timestep (batch,)
# After initial convolution: (batch, 128, 128, 128)
# Timestep embedding: (batch, 512)

Processing flow for a 32×128×128 input batch:

• Downsampling: 128→128→64→32→16→8→4 dimensions progressively
• Middle Processing: Feature extraction at the bottleneck (batch, 512, 4, 4)
• Upsampling: 4→8→16→32→64→128 dimensions with skip connections

Attention module are inserted at specific resolutions (typically 16×16) to capture long-range dependencies. The time embedding uses sinusoidal positional encoding, similar to transformer attention mechanisms.

Training Considerations

Diffusion models exhibit a characteristic learning pattern: they initially capture low-frequency image components (overall structure and composition) before progressively learning high-frequency details (fine textures and edges). This sequential learning behavior explains why diffusion models require extensive training iterations—often thousands to hundreds of thousands of steps—to achieve high-quality results, particularly for fine details.

The training instability is typically lower than GANs, and the model naturally supports conditional generation through classifier-free guidance, where labels are randomly dropped during training to enable unconditional generation at inference time.