Diffusion and Generative Models

Generative Models Overview

A generative model learns to approximate the data distribution p(x) and can draw new samples from it. Different families make different tradeoffs between sample quality, diversity, training stability, and inference speed.

Generative Adversarial Networks (GANs)

A GAN pits two networks against each other: - Generator G: maps noise z ~ p(z) → fake sample G(z) ≈ real data - Discriminator D: binary classifier — real or fake

Minimax objective:

min_G max_D E[log D(x)] + E[log(1 − D(G(z)))]

Equilibrium: G produces samples indistinguishable from real data (D(x) = 0.5 everywhere).

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, latent_dim=128, img_dim=784):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.LeakyReLU(0.2),
            nn.Linear(256, 512),        nn.LeakyReLU(0.2),
            nn.Linear(512, img_dim),    nn.Tanh()
        )
    def forward(self, z): return self.net(z)

class Discriminator(nn.Module):
    def __init__(self, img_dim=784):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(img_dim, 512), nn.LeakyReLU(0.2),
            nn.Linear(512, 256),     nn.LeakyReLU(0.2),
            nn.Linear(256, 1),       nn.Sigmoid()
        )
    def forward(self, x): return self.net(x)

# Training step
def train_step(real_x, G, D, g_opt, d_opt, criterion):
    batch = real_x.size(0)
    z = torch.randn(batch, 128)
    fake = G(z).detach()

    # Train D
    d_opt.zero_grad()
    d_loss = criterion(D(real_x), torch.ones(batch,1)) + \
             criterion(D(fake),   torch.zeros(batch,1))
    d_loss.backward(); d_opt.step()

    # Train G
    g_opt.zero_grad()
    g_loss = criterion(D(G(z)), torch.ones(batch,1))
    g_loss.backward(); g_opt.step()

GAN Training Challenges

• Mode collapse: G maps all z to few modes, ignoring D's full distribution
• Instability: G and D can cycle without converging
• Training imbalance: if D is too strong, G gradient vanishes

GAN Variants

Frechet Inception Distance (FID): measures distance between distributions of real and generated images in Inception-v3 feature space. Lower = better. SOTA is ~2 FID on FFHQ.

Variational Autoencoders (VAEs)

VAEs learn a latent space Z from which samples can be decoded into data. Unlike standard autoencoders, the encoder outputs a distribution q(z|x) = N(μ, σ²).

Training Objective (ELBO)

ELBO = E_{q(z|x)}[log p(x|z)] − KL(q(z|x) ‖ p(z))

• Reconstruction term: decoder should regenerate x from z
• KL term: force posterior q(z|x) toward standard normal prior p(z) = N(0,I)

This enables sampling: z ~ N(0,I), then decode x ~ p(x|z).

Reparameterization trick: to backpropagate through sampling, write z = μ + σ⊙ε where ε ~ N(0,I). The randomness is separated from the parameters.

class VAE(nn.Module):
    def __init__(self, input_dim=784, latent_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(input_dim, 512), nn.ReLU())
        self.fc_mu    = nn.Linear(512, latent_dim)
        self.fc_logvar= nn.Linear(512, latent_dim)
        self.decoder  = nn.Sequential(
            nn.Linear(latent_dim, 512), nn.ReLU(),
            nn.Linear(512, input_dim),  nn.Sigmoid()
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decoder(z)
        return recon, mu, logvar

def vae_loss(recon_x, x, mu, logvar, beta=1.0):
    recon_loss = F.binary_cross_entropy(recon_x, x, reduction="sum")
    kl_loss    = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + beta * kl_loss   # beta-VAE: beta > 1 → more disentangled

Normalizing Flows

Flows transform a simple base distribution (e.g., N(0,I)) through a sequence of invertible, differentiable transformations f₁, f₂, …, fₖ:

z₀ ~ p(z₀), z₁ = f₁(z₀), …, x = fₖ(…f₁(z₀)…)

The exact likelihood is tracked via the change-of-variables formula:

log p(x) = log p(z₀) − Σₖ log|det(∂fₖ/∂zₖ₋₁)|

Advantages: exact log-likelihood, fast sampling, invertible for encoding. Challenge: Jacobian determinant must be tractable → architecture constraints.

Diffusion Models

Diffusion models are the current state of the art for image, audio, and video generation. They outperform GANs on sample diversity and stability.

Forward Process (Adding Noise)

Gradually add Gaussian noise over T steps until the image becomes pure noise:

q(xₜ | xₜ₋₁) = N(xₜ; √(1−βₜ) xₜ₋₁, βₜ I)

Using cumulative variance ᾱₜ = Πₛ₌₁ᵗ (1−βₛ), we can directly sample xₜ from x₀:

q(xₜ | x₀) = N(xₜ; √ᾱₜ x₀, (1−ᾱₜ) I)

So: xₜ = √ᾱₜ x₀ + √(1−ᾱₜ) ε, where ε ~ N(0,I)

Reverse Process (Denoising)

A neural network ε_θ(xₜ, t) learns to predict the noise ε that was added:

L = E_{t, x₀, ε} ‖ε − ε_θ(xₜ, t)‖²

Sampling: start from xₜ ~ N(0,I), then iteratively denoise:

x_{t-1} = (1/√αₜ)(xₜ − (βₜ/√(1−ᾱₜ)) ε_θ(xₜ,t)) + σₜ z

where z ~ N(0,I) and αₜ = 1 − βₜ.

DDPM (Denoising Diffusion Probabilistic Models)

The original formulation (Ho et al. 2020). 1000 steps → slow sampling (~1 min per image on GPU).

DDIM (Denoising Diffusion Implicit Models)

Non-Markovian formulation allows sampling in 20–50 steps with nearly identical quality.

Score-Based Diffusion

Equivalent formulation via score function s_θ(x) = ∇_x log p(x). Score matching + Langevin dynamics.

Latent Diffusion Models (LDM) / Stable Diffusion

Running diffusion in pixel space is expensive. Latent diffusion runs it in a compressed latent space z = E(x) from a pre-trained VAE:

1. Train a VAE encoder-decoder (8× compression)
2. Run DDPM in the latent space (64×64 instead of 512×512)
3. Decode final latent to image

from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",  # maintained SD 1.5 mirror
    torch_dtype=torch.float16
).to("cuda")

image = pipe(
    prompt="a photorealistic cat wearing sunglasses on a beach",
    negative_prompt="blurry, ugly, watermark",
    num_inference_steps=30,
    guidance_scale=7.5,
    height=512, width=512
).images[0]
image.save("output.png")

Classifier-Free Guidance (CFG)

Guides the diffusion process toward a condition (text prompt) without a separate classifier:

ε̃_θ(xₜ, c) = (1+w) ε_θ(xₜ, c) − w ε_θ(xₜ, ∅)

where c is the condition (text embedding), ∅ is the unconditioned prediction, and w is the guidance scale. - High w (7.5–12): more faithful to prompt, less diverse - Low w (1–3): more diverse, may stray from prompt

Modern Diffusion Variants

Flow Matching

A cleaner, more general framework that subsumes diffusion:

• Define a probability path from data to noise (can be straight lines, not curved)
• Train a velocity field v_θ to match this path
• Integrate v_θ with an ODE solver (fewer steps than DDPM SDE)

Straight-flow paths (Rectified Flow, Flux) allow 1-8 step sampling.

Autoregressive Image Generation

Quantize images to discrete tokens (VQ-VAE), then apply next-token prediction: - DALL-E (2021): transformer on VQVAE tokens - ImageGPT: pixel-level autoregression - VQGAN + transformer: higher quality via learned codebooks - MAR, RandAR: masked/random-order autoregressive (2024)

Evaluation of Generative Models