DDPM Reverse Sampling Step - Generative Models

One Reverse Step of DDPM

What you'll build

A single denoising step of DDPM sampling (Ho et al., 2020, Algorithm 2): given the current noisy sample x_t and the network's noise prediction, produce x_{t-1}.

Background

Sampling starts from pure noise x_T ~ N(0, I) and applies this step T times. A network ε_θ(x_t, t) predicts the noise contained in x_t; the step removes the right fraction of it, rescales, and (except at the very end) injects a little fresh noise so the chain stays stochastic.

Formula

Schedule, with timesteps 0-indexed into betas (t = 0 is the last, least noisy step):

α_t = 1 − β_t
ᾱ_t = α_0 · α_1 · ... · α_t                    # cumulative product

Reverse step:

mean    = ( x_t − β_t / sqrt(1 − ᾱ_t) · ε_θ ) / sqrt(α_t)
x_{t-1} = mean + σ_t · z,     σ_t = sqrt(β_t),  z ~ N(0, I)

• Use σ_t² = β_t (the paper's first choice; the alternative is the posterior variance β̃_t = (1 − ᾱ_{t−1}) / (1 − ᾱ_t) · β_t).
• At t = 0 add no noise: return mean, and ignore z.
• z is passed in, so the step is deterministic and testable.

Inputs

• x_t, eps_pred, z: float arrays of the same (any) shape
• t: int in [0, T−1]
• betas: float array (T,)

Expected return format

x_{t-1}: float array with the shape of x_t.

Sanity check

At t = 0, ᾱ_0 = α_0. If x = sqrt(α_0)·x_0 + sqrt(β_0)·ε and the model predicts ε exactly, the step returns x_0 exactly.

Constraints

NumPy only, vectorised over the sample.

Follow-up

• Why does DDIM allow 50 steps instead of 1000? (A deterministic, non-Markovian sampler on the same trained model.)
• Where would classifier-free guidance plug in? (Replace eps_pred; see the next problem.)

Input: x_t = [1.0, -0.5], eps_pred = [0.2, 0.4], t = 2, betas = [0.1, 0.2, 0.3], z = [0.5, -1.0]
Output: [1.3673, -1.349]

Input: Same inputs with t = 0
Output: [0.9874, -0.6604]

import numpy as np

def ddpm_step(x_t, eps_pred, t, betas, z):
    """
    One reverse DDPM step with sigma_t^2 = beta_t.

    Args:
        x_t: current noisy sample
        eps_pred: predicted noise ε_θ(x_t, t), same shape as x_t
        t: timestep (int, 0-indexed into betas)
        betas: (T,) noise schedule
        z: standard normal noise, same shape as x_t (ignored at t = 0)

    Returns:
        x_{t-1}, same shape as x_t
    """
    # Your code here
    pass