DDPM MoG Example

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

# Reproducibility: identical seed across runs so this MoG matches the one
# used in the EM tutorial.
torch.manual_seed(42)
np.random.seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Build the target distribution as 300 i.i.d. samples from a 3-component MoG.
# DDPM never sees the cluster labels, means, or covariances - only the
# 300x2 tensor of points. Recovering the structure from raw points alone
# is the whole point.
n_samples = 300
cluster1 = np.random.multivariate_normal([0, 0],   [[1, 0.5],   [0.5, 1]],   n_samples // 3)
cluster2 = np.random.multivariate_normal([6, 6],   [[1.5, -0.7], [-0.7, 1.5]], n_samples // 3)
cluster3 = np.random.multivariate_normal([-4, 5],  [[0.8, 0],   [0, 2.0]],   n_samples // 3)
data_np = np.vstack([cluster1, cluster2, cluster3]).astype(np.float32)
data = torch.from_numpy(data_np).to(device)

# Forward-process hyperparameters. T is the number of noising steps;
# `betas` is the variance schedule beta_1..beta_T (small at t=0, larger at t=T-1).
# Linear schedule from 1e-4 to 0.02 is the original DDPM choice (Ho et al. 2020).
T = 500
betas = torch.linspace(1e-4, 0.02, T, device=device)
alphas = 1.0 - betas                          # alpha_t = 1 - beta_t
alpha_bars = torch.cumprod(alphas, dim=0)     # bar_alpha_t = prod_{s=1}^t alpha_s


def q_sample(x0, t, eps=None):
    """Closed-form jump from x_0 directly to x_t.

    Because every forward step is linear-Gaussian, the marginal q(x_t | x_0)
    has a closed form and we never have to simulate the chain step by step:
        x_t = sqrt(bar_alpha_t) * x_0 + sqrt(1 - bar_alpha_t) * eps,   eps ~ N(0, I).
    """
    if eps is None:
        eps = torch.randn_like(x0)
    # Index the schedule at each example's timestep, then unsqueeze so the
    # scalar coefficient broadcasts across the 2-D point dimension.
    sqrt_ab   = alpha_bars[t].sqrt().unsqueeze(-1)
    sqrt_omab = (1 - alpha_bars[t]).sqrt().unsqueeze(-1)
    return sqrt_ab * x0 + sqrt_omab * eps

class SinusoidalTimeEmbedding(nn.Module):
    """Transformer-style positional encoding of the diffusion timestep.

    A single network handles every t in {0, ..., T-1}, so it needs to know
    *which* noise level it is denoising. The sin/cos embedding (used in
    transformers and the original DDPM paper) gives a smooth, continuous
    representation of t at many frequencies.
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        # Geometric series of frequencies spanning short to long timescales.
        freqs = torch.exp(-np.log(10000) * torch.arange(half, device=t.device) / half)
        emb = t.float().unsqueeze(-1) * freqs
        return torch.cat([emb.sin(), emb.cos()], dim=-1)


class NoisePredictor(nn.Module):
    """eps_theta(x_t, t): predicts the noise that produced x_t.

    The network output is NOT the reverse-step mean mu_theta; it is the
    estimated 2-D noise vector eps_hat. The mean is computed analytically
    from eps_hat in the sampling loop via the Ho et al. reparameterisation
    (see the formula in the sampling cell). Equivalent parameterisations
    exist (predict mu directly, or predict x_0), but eps-prediction is
    better-conditioned because eps ~ N(0, I) at every timestep.

    A small MLP suffices for 2-D data; for images this becomes a U-Net.
    """
    def __init__(self, time_dim=64, hidden=128):
        super().__init__()
        # Project the raw sin/cos embedding through a tiny MLP (standard DDPM trick).
        self.time_emb = nn.Sequential(
            SinusoidalTimeEmbedding(time_dim),
            nn.Linear(time_dim, time_dim), nn.SiLU(),
            nn.Linear(time_dim, time_dim),
        )
        # Concatenate (x_t, time-embedding) -> 4-layer MLP. Output is the
        # 2-D noise estimate eps_hat (same shape as x_t), NOT the reverse-step mean.
        self.net = nn.Sequential(
            nn.Linear(2 + time_dim, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden),       nn.SiLU(),
            nn.Linear(hidden, hidden),       nn.SiLU(),
            nn.Linear(hidden, 2),            # eps_hat in R^2; reverse mean derived from it
        )

    def forward(self, x, t):
        # Concatenate the noisy point with its timestep embedding.
        h = torch.cat([x, self.time_emb(t)], dim=-1)
        return self.net(h)


model = NoisePredictor().to(device)

optim = torch.optim.Adam(model.parameters(), lr=1e-3)
batch_size = 256
n_steps = 4000
losses = []

# DDPM training is stochastic over (x_0, t, eps) triples.
# At each step we sample a clean point, a random diffusion timestep, and a
# noise vector, then ask the network to recover that noise.
for step in range(n_steps):
    # 1. Random minibatch of clean data points (with replacement; tiny dataset).
    idx = torch.randint(0, data.shape[0], (batch_size,), device=device)
    x0 = data[idx]

    # 2. Independent random timestep per example. Mixing all noise levels
    #    in one batch is what teaches the network ALL t in one model.
    t = torch.randint(0, T, (batch_size,), device=device)

    # 3. Form x_t = sqrt(bar_alpha_t) x_0 + sqrt(1 - bar_alpha_t) eps in closed form.
    eps = torch.randn_like(x0)
    xt = q_sample(x0, t, eps)

    # 4. Predict the noise and minimise MSE against the true eps.
    #    This is L_simple from Ho et al.: a reweighted ELBO that works
    #    dramatically better in practice than the raw variational bound.
    eps_hat = model(xt, t)
    loss = F.mse_loss(eps_hat, eps)

    optim.zero_grad()
    loss.backward()
    optim.step()
    losses.append(loss.item())
    if (step + 1) % 1000 == 0:
        print(f"step {step + 1:>5d}  loss {float(np.mean(losses[-500:])):.4f}")

step  1000  loss 0.5713

step  2000  loss 0.5401

step  3000  loss 0.5375

step  4000  loss 0.5341

@torch.no_grad()
def sample(n=1000):
    """Run the reverse chain from pure noise back to data.

    Start at x_T ~ N(0, I) and iteratively denoise down to x_0 by applying
    the learned reverse transition p_theta(x_{t-1} | x_t).
    """
    # 1. Initialise n samples as pure standard Gaussian noise. This is x_T.
    x = torch.randn(n, 2, device=device)
    trajectory = [x.cpu().numpy().copy()]

    # 2. Loop t = T-1, T-2, ..., 0  (reversed range produces this order).
    for t_step in reversed(range(T)):
        t = torch.full((n,), t_step, device=device, dtype=torch.long)

        # Predicted noise at the current step.
        eps = model(x, t)
        beta_t  = betas[t_step]
        alpha_t = alphas[t_step]
        ab_t    = alpha_bars[t_step]

        # 3. Reverse-transition mean under the Ho et al. parameterisation:
        #    mu_theta(x_t, t) = (1/sqrt(alpha_t)) * (x_t - (1-alpha_t)/sqrt(1-bar_alpha_t) * eps_theta).
        coef = (1 - alpha_t) / (1 - ab_t).sqrt()
        mean = (x - coef * eps) / alpha_t.sqrt()

        # 4. Add Langevin noise scaled by sqrt(beta_t), except at the final
        #    step where we want a deterministic x_0 rather than a noisy draw.
        if t_step > 0:
            x = mean + beta_t.sqrt() * torch.randn_like(x)
        else:
            x = mean

        trajectory.append(x.cpu().numpy().copy())

    return x.cpu().numpy(), trajectory


generated, traj = sample(n=1000)