Diffusion Forward Process - Generative Models

Diffusion Forward Process (Adding Noise)

What you'll build

The closed-form forward process q(x_t | x_0): jump straight from clean data to its noisy version at step t, and return the noise you used (it is the training target of the denoiser).

Background

Adding N(0, β_t) noise T times in a row is equivalent to adding a single Gaussian whose size depends only on ᾱ_t. That is what makes diffusion training cheap: pick a random t, noise the batch in one shot, and train the network to predict ε.

Formula

ε    ~ N(0, I)                                      # same shape as x_0
x_t  = sqrt(ᾱ_t) · x_0 + sqrt(1 − ᾱ_t) · ε

• ᾱ_t = alpha_bars[t] (0-indexed, as in Diffusion Noise Schedule)
• ᾱ_t = 1 returns x_0 unchanged; ᾱ_t → 0 gives (almost) pure noise
• The variance is preserved: if x_0 has unit variance, so does x_t

Inputs

• x_0: float array of any shape (batch, ...)
• t: int, an index into alpha_bars
• alpha_bars: float array (T,) with values in [0, 1]

Expected return format

A tuple (x_t, noise), both float arrays with the shape of x_0.

Conventions (the tests rely on these)

The starter code calls np.random.seed(42). Keep it and draw the noise once with np.random.randn(*x_0.shape).

Constraints

NumPy only, vectorised.

Follow-up

• In training, t is sampled per example. How would you vectorise with a (batch,) array of timesteps? (Index alpha_bars[t] and reshape to (batch, 1, ...).)
• Why predict ε instead of x_0 directly?

Input: x_0 = np.array([[1.0, -1.0]]), t = 1, alpha_bars = np.array([0.9, 0.5])
Output: x_t = [[1.0583, -0.8049]], noise = [[0.4967, -0.1383]]

Input: x_0 = np.array([[1.0, 2.0, 3.0]]), t = 0, alpha_bars = np.array([1.0, 0.5, 0.1])
Output: x_t = [[1.0, 2.0, 3.0]]

import numpy as np

def diffusion_forward(x_0, t, alpha_bars):
    """
    Sample x_t ~ q(x_t | x_0) in closed form.

    Args:
        x_0: Clean data, shape (batch, ...)
        t: Timestep (int, 0-indexed into alpha_bars)
        alpha_bars: Cumulative alpha products, shape (T,)

    Returns:
        (x_t, noise): both with the shape of x_0
    """
    np.random.seed(42)  # keep: draw the noise once with np.random.randn(*x_0.shape)
    # Your code here
    pass