DDPM Sampling
Denoising Diffusion Probabilistic Models
Medium
Problem
Implement the complete deterministic-input DDPM sampling loop. Begin from the supplied noisy sample x_T and process timesteps T through 1. At loop index i, epsilon_preds[i] is the model prediction for timestep T-i. For steps T through 2, z_values[i] is the supplied stochastic noise.
At each timestep, compute
\mu_t = \frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\hat{\epsilon}_t\right)
Then update the sample with
x_{t-1} = \mu_t + \sqrt{\beta_t}\,z_t
For t=1, omit the noise term and set x_0=mu_1. Here, T is the number of beta values, alpha_t=1-beta_t, bar-alpha_t is the cumulative product through t, hat-epsilon_t is the supplied model prediction, and z_t is the supplied stochastic noise. Use NumPy float64 calculations. Return the final x_0 as a nested list with the same shape as x_T, rounded to four decimals.
Theory
Sampling is the generative act of a diffusion model: starting from pure Gaussian noise and progressively removing it to produce a clean data sample. In the Denoising Diffusion Probabilistic Model (DDPM) framework introduced by Ho, Jain, and Abbeel (2020), sampling runs the learned reverse Markov chain from timestep T all the way down to timestep 1, inverting the forward diffusion process that gradually destroyed the data's structure.
Unlike GANs, which generate samples in a single forward pass through a generator network, DDPM sampling is an iterative procedure. Each step asks the model to predict and subtract a small amount of noise, gradually revealing the underlying data distribution one denoising step at a time.
What It Is / What It Does
The DDPM sampling pipeline is a complete generation procedure that transforms a random noise vector into a realistic data sample. It consists of three stages: (1) sample initial noise \mathbf{x}_T from a standard Gaussian, (2) iteratively apply the learned reverse denoising transition from $ = T$ down to t = 1, and (3) output the final \mathbf{x}_0 as the generated sample.
At each step, a neural network \epsilon_\theta(\mathbf{x}_t, t) predicts the noise component present in \mathbf{x}_t. This prediction is used to compute the mean of the reverse transition distribution, from which \mathbf{x}_{t-1} is sampled. The noise network was trained to minimize the difference between its prediction and the actual noise that was added during the forward process.
The procedure is stochastic: at every step except the last, fresh Gaussian noise \mathbf{z} is added to the denoised estimate. This injects randomness that allows the model to explore diverse modes of the data distribution. At t = 1, no noise is added because we want the final output to be a clean sample, not a noisy one.
Key Equations
The noise schedule defines a sequence of values \beta_1, \beta_2, \dots, \beta_T (small positive constants, typically \beta_1 = 10^{-4} and \beta_T = 0.02 with linear interpolation). From these, we derive:
- \alpha_t = 1 - \beta_t: The per-step signal retention factor.
- \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s: The cumulative signal retention from step 0 to step t.
The initial noise sample is drawn from a standard Gaussian:
\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})
For each timestep t from T down to 1, the reverse step computes:
\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \, \epsilon_\theta(\mathbf{x}_t, t) \right) + \sigma_t \, \mathbf{z}
where \mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) for t > 1 and \mathbf{z} = \mathbf{0} for t = 1.
- \frac{1}{\sqrt{\alpha_t}}: Rescales the signal back up to compensate for the contraction applied during the forward process at step t.
- \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}: The coefficient that converts the noise prediction \epsilon_\theta into the correct noise magnitude to subtract from \mathbf{x}_t.
- \epsilon_\theta(\mathbf{x}_t, t): The neural network's estimate of the noise component in \mathbf{x}_t. It takes both the noisy input and the timestep as arguments.
- \sigma_t: The noise scale for the reverse step. Ho et al. set \sigma_t^2 = \beta_t, matching the forward process variance. An alternative is \sigma_t^2 = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t, the posterior variance.
The term inside the parentheses computes the estimated mean \mu_\theta(\mathbf{x}_t, t) of the reverse distribution p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t). Adding \sigma_t \mathbf{z} samples from that distribution rather than taking the point estimate.
The Sampling Algorithm
The full procedure, following Algorithm 2 from Ho et al. (2020):
Step 1. Sample \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}). This is a random tensor with the same shape as the desired output (e.g., 3 \times 64 \times 64 for a 64 \times 64 RGB image).
Step 2. For t = T, T-1, \dots, 1:
- If t > 1, sample \mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}). If t = 1, set \mathbf{z} = \mathbf{0}.
- Run the neural network forward pass: \hat{\epsilon} = \epsilon_\theta(\mathbf{x}_t, t).
- Compute the denoised estimate: \mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \hat{\epsilon} \right) + \sigma_t \mathbf{z}.
Step 3. Return \mathbf{x}_0 as the generated sample.
Each iteration of the loop requires exactly one forward pass through the denoising network. The schedule values \alpha_t, \bar{\alpha}_t, and \sigma_t are precomputed once and stored as lookup arrays indexed by t. No gradient computation is needed during sampling, so inference mode is used.
Why T Steps is Slow
The default DDPM uses T = 1000 timesteps, meaning 1000 sequential neural network forward passes per sample. On a modern GPU, each U-Net pass for a $ \times 256$ image takes 20-50 ms, putting total generation time at 20-50 seconds per image.
The steps are inherently sequential: \mathbf{x}_{t-1} depends on \mathbf{x}_t, which depends on \mathbf{x}_{t+1}. You cannot parallelize across timesteps for a single sample. Batch parallelism across multiple samples is possible, but per-sample latency remains T forward passes.
For comparison, a GAN generates a sample in a single forward pass (~5-20 ms). This 1000x slowdown is DDPM's primary practical limitation and the main focus of subsequent research.
Paper Context
Ho, Jain, and Abbeel (2020) presented DDPM's sampling procedure in Algorithm 2 of "Denoising Diffusion Probabilistic Models." They described it as "sampling from the model is performed by running the reverse Markov chain, starting from \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) and iteratively sampling \mathbf{x}_{t-1} \sim p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t)."
The paper achieved an FID of 3.17 and Inception Score of 9.46 on CIFAR-10, competitive with GANs at the time. On LSUN bedrooms (256 \times 256), DDPM produced sharp, coherent room structures.
Compared to GANs, DDPM offered several advantages despite slower sampling. Training is stable with no adversarial min-max game, no mode collapse, and a simple MSE loss. DDPM also provides better mode coverage, generating more diverse samples rather than concentrating on a few high-quality modes.
Ho et al. chose \sigma_t^2 = \beta_t for the reverse process variance and a linear schedule from \beta_1 = 10^{-4} to \beta_T = 0.02 with T = 1000. Subsequent work (Nichol and Dhariwal, 2021) found that a cosine schedule produces better results, especially for higher-resolution images.
Numerical Example
Consider a toy example with T = 3 and a single scalar value (1D) to trace through the full sampling loop. We use the following simplified schedule:
- \beta_1 = 0.1, \quad \beta_2 = 0.2, \quad \beta_3 = 0.3
- \alpha_1 = 0.9, \quad \alpha_2 = 0.8, \quad \alpha_3 = 0.7
- \bar{\alpha}_1 = 0.9, \quad \bar{\alpha}_2 = 0.72, \quad \bar{\alpha}_3 = 0.504
- \sigma_t = \sqrt{\beta_t}, so \sigma_1 \approx 0.3162, \quad \sigma_2 \approx 0.4472, \quad \sigma_3 \approx 0.5477
Step 1: Initialize \mathbf{x}_3
Sample \mathbf{x}_3 \sim \mathcal{N}(0, 1). Suppose we draw \mathbf{x}_3 = 1.5.
Step 2: Reverse from t = 3 to t = 2
The model predicts \epsilon_\theta(\mathbf{x}_3, 3) = 0.8 (its estimate of the noise in \mathbf{x}_3). Sample \mathbf{z} \sim \mathcal{N}(0, 1); suppose \mathbf{z} = -0.3.
\mathbf{x}_2 = \frac{1}{\sqrt{0.7}} \left( 1.5 - \frac{1 - 0.7}{\sqrt{1 - 0.504}} \cdot 0.8 \right) + 0.5477 \cdot (-0.3)
= \frac{1}{0.8367} \left( 1.5 - \frac{0.3}{0.7043} \cdot 0.8 \right) - 0.1643
= 1.1952 \times (1.5 - 0.3408) - 0.1643
= 1.1952 \times 1.1592 - 0.1643 = 1.3855 - 0.1643 = 1.2212
Step 3: Reverse from t = 2 to t = 1
The model predicts \epsilon_\theta(\mathbf{x}_2, 2) = 0.5. Sample \mathbf{z} \sim \mathcal{N}(0, 1); suppose \mathbf{z} = 0.1.
\mathbf{x}_1 = \frac{1}{\sqrt{0.8}} \left( 1.2212 - \frac{1 - 0.8}{\sqrt{1 - 0.72}} \cdot 0.5 \right) + 0.4472 \cdot 0.1
= \frac{1}{0.8944} \left( 1.2212 - \frac{0.2}{0.5292} \cdot 0.5 \right) + 0.0447
= 1.1180 \times (1.2212 - 0.1890) + 0.0447
= 1.1180 \times 1.0322 + 0.0447 = 1.1540 + 0.0447 = 1.1987
Step 4: Reverse from t = 1 to t = 0
The model predicts \epsilon_\theta(\mathbf{x}_1, 1) = 0.2. At t = 1, we set \mathbf{z} = \mathbf{0} (no noise added).
\mathbf{x}_0 = \frac{1}{\sqrt{0.9}} \left( 1.1987 - \frac{1 - 0.9}{\sqrt{1 - 0.9}} \cdot 0.2 \right) + 0
= \frac{1}{0.9487} \left( 1.1987 - \frac{0.1}{0.3162} \cdot 0.2 \right)
= 1.0541 \times (1.1987 - 0.0632)
= 1.0541 \times 1.1355 = 1.1968
The final generated sample is \mathbf{x}_0 \approx 1.197. Starting from pure noise (\mathbf{x}_3 = 1.5), the model progressively refined it through 3 denoising steps. With a well-trained model on real data, \mathbf{x}_0 would land in a high-density region of the training distribution.
The Role of Stochasticity
The noise term \sigma_t \mathbf{z} added at each reverse step is not an implementation detail but a fundamental part of the generative process. It serves as a source of diversity: different random draws of \mathbf{z} at each step lead to different final samples, even from the same initial \mathbf{x}_T. Without this noise, the same starting point would always produce the same output.
The magnitude \sigma_t controls how much exploration happens at each step. Ho et al. used \sigma_t^2 = \beta_t, which matches the forward process variance and corresponds to the upper bound of the reverse process entropy. The alternative \sigma_t^2 = \tilde{\beta}_t (the posterior variance) gives a lower bound. Both produce valid samples, but the choice affects sample diversity and quality.
DDIM (Song, Meng, and Ermon, 2020) showed that the noise term can be removed entirely, yielding a deterministic sampling process. The mapping from \mathbf{x}_T to \mathbf{x}_0 becomes a fixed bijection, enabling meaningful interpolation in the latent space.
Temperature scaling provides a continuous control lever. Multiplying \sigma_t by a factor \eta (where \eta = 1 recovers DDPM and \eta = 0 gives DDIM) interpolates between stochastic and deterministic sampling. Lower \eta produces sharper but less diverse samples.
Modern Speedups
The 1000-step sampling requirement of DDPM motivated extensive research into faster alternatives. These methods reduce the number of neural network evaluations while maintaining sample quality.
DDIM (Song et al., 2020) reinterprets the diffusion process as a non-Markovian chain, enabling sampling with a subset of timesteps. Instead of all 1000 steps, DDIM uses a subsequence like \{1, 51, 101, \dots, 951\} (20 steps) with minimal quality loss. The noise prediction network generalizes across timesteps, so skipping intermediate steps works.
DPM-Solver (Lu et al., 2022) treats the reverse diffusion as an ODE and applies high-order numerical solvers. While DDPM uses first-order Euler-like updates, DPM-Solver uses second and third-order methods that take larger, more accurate steps, achieving strong results in 10-20 steps.
Consistency Models (Song et al., 2023) learn to map any point on the diffusion trajectory directly to \mathbf{x}_0 in a single step, trained either by distilling a pre-trained diffusion model or from scratch.
Progressive Distillation (Salimans and Ho, 2022) trains a student to combine two teacher steps into one, then repeats. Each round halves the step count: 1024, 512, 256, ..., down to 4 steps.
Latent Diffusion (Rombach et al., 2022) runs diffusion in a compressed latent space rather than pixel space. An encoder compresses images (e.g., 512 \times 512 to 64 \times 64 latents), diffusion operates there, and a decoder maps back to pixels. Stable Diffusion is the most prominent example.
Pitfalls
Wrong Loop Direction
The reverse process must iterate from t = T down to t = 1. A common implementation error is looping from t = 1 to t = T, which runs the forward (noising) direction instead of the reverse (denoising) direction. The result is progressively noisier outputs rather than cleaner ones.
Forgetting \mathbf{z} = \mathbf{0} at t = 1
At the final step (t = 1), no noise should be added. If \mathbf{z} is sampled normally at t = 1, the final output \mathbf{x}_0 will have unnecessary Gaussian noise overlaid on it, producing a visibly grainy image. This is a one-line bug with significant visual impact.
Using Wrong Schedule Values
Confusing \alpha_t with \bar{\alpha}_t is a frequent source of bugs. The reverse step formula uses both: \alpha_t (the per-step value) appears as \frac{1}{\sqrt{\alpha_t}} and in the numerator 1 - \alpha_t, while \bar{\alpha}_t (the cumulative product) appears in the denominator \sqrt{1 - \bar{\alpha}_t}. Swapping them produces incorrect denoising magnitudes and corrupted samples.
Not Starting from Pure Gaussian Noise
The derivation assumes \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}). If \mathbf{x}_T is initialized from a different distribution (e.g., uniform noise, or Gaussian with wrong variance), the reverse process produces samples outside the learned distribution. The variance must be exactly \mathbf{I}, not scaled.
Clipping Output Values
Intermediate \mathbf{x}_t values can drift outside the expected range (e.g., beyond [-1, 1]). Aggressive clipping at every step introduces artifacts. The recommended practice is to clip only the final \mathbf{x}_0, or use dynamic thresholding (Saharia et al., 2022) which rescales outlier values rather than hard clipping.
Examples
Example 1
- Input
x_T=[[0.5,-0.3]], betas=[0.001,0.0105,0.02], epsilon_preds=[[[0,0]],[[0,0]],[[0,0]]], z_values=[[[0.1,-0.2]],[[0.3,0.1]]]- Output
[[0.553,-0.323]]- Explanation
- Three reverse steps are applied; supplied noise is added during the first two and omitted during the final step.
Example 2
- Input
x_T=[[1,-1]], betas=[0.001,0.0105,0.02], epsilon_preds=[[[0.5,0.5]],[[0.5,0.5]],[[0.5,0.5]]], z_values=[[[0,0]],[[0,0]]]- Output
[[0.8935,-1.1386]]
Example 3
- Input
x_T has shape (2,2), betas=[0.001,0.0105,0.02], epsilon_preds has shape (3,2,2), z_values has shape (2,2,2)- Output
[[0.1949,-0.4234],[0.632,0.0624]]
Hints
- Use enumerate(range(T, 0, -1)) to align each prediction with its timestep.
- Precompute alphas and their cumulative product before the reverse loop.
- Read z_values[i] only when t is greater than 1.
Requirements
- Start from the supplied x_T; do not generate a replacement sample.
- Process beta values in reverse timestep order from T through 1.
- Pair epsilon_preds in order with timesteps T through 1.
- Pair z_values in order with timesteps T through 2 and omit noise at t=1.
- Return the final nested list rounded to four decimals.
Constraints
- x_T is a nonempty nested list of real numbers.
- betas is a nonempty list of T float values strictly between 0 and 1.
- epsilon_preds contains T arrays, each with the same shape as x_T.
- z_values contains T-1 arrays, each with the same shape as x_T.
- Numerical calculations use NumPy float64.
Starter Code
import numpy as np
def ddpm_sample(x_T: list, betas: list[float], epsilon_preds: list, z_values: list) -> list:
"""
Returns the final denoised sample rounded to four decimals.
"""
passTest Cases
| Case | Matches | |
|---|---|---|
| Zero model, T=3 | — | public |
| Constant model, T=3 | — | public |
| 2x2 shape, T=3 | — | public |