Lesson 4: Diffusion Models — Mathematics & Intuition

Generative AI: Create Images & Videos with AI

Part 2: Diffusion Models — Revolutionary Image Creation

xdev.asia

Introduction

Diffusion Models is the foundation of Stable Diffusion, DALL-E 2/3, Midjourney — achieving superior image generation quality than GANs. Genius idea: add noise gradually until the image is completely noisy, then learn how to reverse that process.

---

1. Forward Process — Add Noise

$$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t} x_{t-1}, \beta_t I)$$

import torch

def forward_diffusion(x_0, t, betas):
    """Thêm noise vào image tại timestep t"""
    alphas = 1 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)

    sqrt_alpha_bar = torch.sqrt(alpha_bar[t])
    sqrt_one_minus = torch.sqrt(1 - alpha_bar[t])

    noise = torch.randn_like(x_0)
    x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus * noise
    return x_t, noise

t=0 (clean)  →  t=250  →  t=500  →  t=750  →  t=1000 (noise)
  [clear]       [hazy]    [noisy]   [very noisy]  [pure noise]

---

2. Reverse Process — Denoise (Study)

$$p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 I)$$

class UNet(nn.Module):
    """Noise predictor: dự đoán noise đã thêm vào"""
    def __init__(self):
        super().__init__()
        # Simplified UNet
        self.time_embed = nn.Sequential(
            nn.Linear(1, 128),
            nn.SiLU(),
            nn.Linear(128, 128),
        )
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.SiLU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128 + 128, 64, 4, stride=2, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, 1, 3, padding=1),
        )

    def forward(self, x, t):
        t_emb = self.time_embed(t.float().unsqueeze(-1))
        h = self.encoder(x)
        t_emb = t_emb.view(-1, 128, 1, 1).expand_as(h)
        h = torch.cat([h, t_emb], dim=1)
        return self.decoder(h)

Training

def train_step(model, x_0, optimizer, betas, T=1000):
    """DDPM training step"""
    t = torch.randint(0, T, (x_0.size(0),))

    # Forward: add noise
    x_t, noise = forward_diffusion(x_0, t, betas)

    # Predict noise
    noise_pred = model(x_t, t)

    # Simple MSE loss
    loss = nn.functional.mse_loss(noise_pred, noise)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

---

3. Sampling — Generate Image

@torch.no_grad()
def sample(model, shape, betas, T=1000):
    """DDPM sampling: noise → image"""
    alphas = 1 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)

    # Start from pure noise
    x = torch.randn(shape)

    for t in reversed(range(T)):
        noise_pred = model(x, torch.tensor([t]))

        # Compute x_{t-1}
        alpha = alphas[t]
        alpha_bar_t = alpha_bar[t]

        mean = (1 / torch.sqrt(alpha)) * (
            x - (betas[t] / torch.sqrt(1 - alpha_bar_t)) * noise_pred
        )

        if t > 0:
            noise = torch.randn_like(x)
            x = mean + torch.sqrt(betas[t]) * noise
        else:
            x = mean

    return x  # Generated image!

---

4. Noise Scheduling

# Linear schedule (DDPM original)
betas = torch.linspace(1e-4, 0.02, T)

# Cosine schedule (improved, OpenAI)
def cosine_schedule(T, s=0.008):
    steps = torch.arange(T + 1)
    f = torch.cos((steps / T + s) / (1 + s) * torch.pi / 2) ** 2
    alpha_bar = f / f[0]
    betas = 1 - alpha_bar[1:] / alpha_bar[:-1]
    return torch.clamp(betas, 0.0001, 0.999)

# Cosine schedule → better quality ở low-noise timesteps

---

5. DDIM — Faster Sampling

# DDPM: 1000 steps → slow
# DDIM: 50 steps → fast, deterministic

@torch.no_grad()
def ddim_sample(model, shape, alpha_bar, steps=50):
    """DDIM: deterministic, fewer steps"""
    # Sub-sample timesteps
    timesteps = torch.linspace(999, 0, steps).long()
    x = torch.randn(shape)

    for i in range(len(timesteps) - 1):
        t = timesteps[i]
        t_prev = timesteps[i + 1]

        noise_pred = model(x, torch.tensor([t]))

        # DDIM update (deterministic, η=0)
        x0_pred = (x - torch.sqrt(1 - alpha_bar[t]) * noise_pred) / torch.sqrt(alpha_bar[t])
        x = torch.sqrt(alpha_bar[t_prev]) * x0_pred + \
            torch.sqrt(1 - alpha_bar[t_prev]) * noise_pred

    return x

Method	Steps	Speed ​​	Quality	Deterministic
DDPM	1000	Slow	Best	No
DDIM	50	20x faster	Good	Yes
DPM++	20-30	35x faster	Great	Yes
Euler	20-30	35x faster	Great	Configurable

---

6. Score-based Models

$$\nabla_x \log p(x) \approx s_\theta(x)$$

Score function: gradient of log-probability
- Không cần biết p(x) chính xác
- Chỉ cần biết "hướng nào density tăng"
- Langevin dynamics: follow score → sample

Liên hệ: score matching ≡ denoising score matching ≡ diffusion
→ Các framework thực ra equivalent!

---

Summary

Concept	Description
Forward process	Add Gaussian noise progressively
Reverse process	Learn denoise step-by-step
DDPM	Denoising Diffusion Probabilistic Models
UNet	Neural network predicts noise
DDIM	Deterministic sampling, fewer steps
Noise schedule	Linear vs cosine — quality influence
Score-based	Learn gradients of log-density

📌 Next article: Stable Diffusion Deep Dive — latent diffusion, UNet, CLIP conditioning.