E2E: Diffusion Model - End-to-End Implementations

Implement E2E: Diffusion Model in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

End-to-End Diffusion Model (DDPM)

What you'll build

The full DDPM machinery as one class: the noise schedule, closed-form forward noising, the ε-prediction training loss, the reverse step and the sampling loop. A real model would plug a neural network in as noise_predictor; here any callable works, so you can test the plumbing. It combines Diffusion Noise Schedule, Diffusion Forward Process and DDPM Reverse Sampling Step from Generative Models.

Pipeline

Phase	Step	Formula
setup	schedule (0-indexed `t = 0..T−1`)	`β = linspace(β_start, β_end, T)`, `α = 1 − β`, `ᾱ_t = Π_{s≤t} α_s`
training	noise data at a random t	`x_t = sqrt(ᾱ_t)·x_0 + sqrt(1 − ᾱ_t)·ε`
training	loss	`mean((ε − ε_θ(x_t, t))²)`
sampling	reverse step	`mean = (x_t − β_t / sqrt(1 − ᾱ_t)·ε_θ) / sqrt(α_t)`; `x_{t−1} = mean + sqrt(β_t)·z` for t > 0, `mean` at t = 0
sampling	loop	start from `x ~ N(0, I)` and apply the reverse step for `t = T−1, T−2, ..., 0`

Method contract

Method	Returns
`__init__(T=1000, beta_start=1e-4, beta_end=0.02)`	stores `self.T` and the schedule arrays `self.betas`, `self.alphas`, `self.alpha_bars`
`get_noise_schedule()`	dict `{'betas', 'alphas', 'alpha_bars'}`, each `(T,)`
`forward_diffusion(x_0, t)`	tuple `(x_t, noise)`, both shaped like `x_0`. Keep the starter's `np.random.seed(42)` and draw `noise = np.random.randn(*x_0.shape)`
`reverse_step(x_t, t, predicted_noise)`	`x_{t−1}` shaped like `x_t`. For t > 0 draw `z = np.random.randn(*x_t.shape)` (no re-seeding); at t = 0 draw nothing
`compute_loss(true_noise, predicted_noise)`	float, the MSE over all elements
`sample(shape, noise_predictor)`	final `x` of the given shape; calls `noise_predictor(x, t)` once per step with `t` from `T−1` down to `0`

Inputs

x_0, x_t, predicted_noise, true_noise: float arrays of any matching shape
t: int in [0, T−1]; shape: tuple; noise_predictor(x, t): callable returning an array shaped like x

Constraints

NumPy only; vectorised inside each step (the only loop is over timesteps in sample).

Follow-up

Sampling needs T network calls. How do DDIM or a shorter flow-matching ODE cut this to 20–50?
Why do training batches sample a random t per example instead of looping over all t?

Examples

Input: DiffusionModel(T=3, beta_start=0.1, beta_end=0.3).get_noise_schedule()
Output: {'betas': [0.1, 0.2, 0.3], 'alphas': [0.9, 0.8, 0.7], 'alpha_bars': [0.9, 0.72, 0.504]}

ᾱ is the running product of α

Input: dm.forward_diffusion(np.array([[1.0, -1.0]]), t=1) for the same model
Output: (array([[1.1114, -0.9217]]), array([[0.4967, -0.1383]]))

x_1 = sqrt(0.72)·x_0 + sqrt(0.28)·ε with ε the first draws after seed 42

Input: dm.reverse_step(np.array([1.0, -0.5]), t=0, predicted_noise=np.array([0.2, 0.4]))
Output: array([0.9874, -0.6604])

At t = 0 there is no added noise: (x − 0.1 / sqrt(0.1)·ε) / sqrt(0.9)

Practice in Python

import numpy as np

class DiffusionModel:
    def __init__(self, T: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02):
        """
        Initialize diffusion model.

        Args:
            T: Number of diffusion timesteps
            beta_start: Starting noise level
            beta_end: Ending noise level
        """
        self.T = T

        # Compute the linear schedule: self.betas, self.alphas, self.alpha_bars (each (T,))
        # Your code here
        pass

    def get_noise_schedule(self) -> dict:
        """Return {'betas', 'alphas', 'alpha_bars'}."""
        # Your code here
        pass

    def forward_diffusion(self, x_0: np.ndarray, t: int) -> tuple:
        """
        Add noise to data at timestep t.

        Args:
            x_0: Original data (batch, ...)
            t: Timestep (0-indexed)

        Returns:
            (x_t, noise): x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise
        """
        np.random.seed(42)  # keep: draw noise once with np.random.randn(*x_0.shape)
        # Your code here
        pass

    def reverse_step(self, x_t: np.ndarray, t: int, predicted_noise: np.ndarray) -> np.ndarray:
        """
        Single reverse diffusion step.

        Args:
            x_t: Noisy data at timestep t
            t: Current timestep (0-indexed)
            predicted_noise: Noise predicted by the model

        Returns:
            x_{t-1}: mean + sqrt(beta_t) * z for t > 0 (z = np.random.randn(*x_t.shape)),
                     just the mean at t = 0
        """
        # Your code here
        pass

    def compute_loss(self, true_noise: np.ndarray, predicted_noise: np.ndarray) -> float:
        """MSE between true and predicted noise (Python float)."""
        # Your code here
        pass

    def sample(self, shape: tuple, noise_predictor) -> np.ndarray:
        """
        Generate samples using reverse diffusion.

        Args:
            shape: Shape of samples to generate
            noise_predictor: Function that predicts noise given (x_t, t)

        Returns:
            Generated samples: start from np.random.randn(*shape), then
            reverse_step for t = T-1, ..., 0
        """
        # Your code here
        pass

The interactive editor loads with Python and NumPy. Run your implementation against 8 tests, or use the progressive hints.

More exercises in this section →