📘

One DDPM Ancestral Sampling Step

MediumDiffusion Models 📚 Diffusion Models 1

Problem Statement

Take one step of the DDPM reverse chain: given $x_t$ and the network's predicted noise, produce $x_{t-1}$ .

Background

Substituting $\hat{x}_0 = (x_t - \sqrt{1-\bar{\alpha}_t}\,\hat{\varepsilon}) / \sqrt{\bar{\alpha}_t}$ into the posterior mean and simplifying gives the form Algorithm 2 of DDPM actually uses:

$x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\hat{\varepsilon}_\theta(x_t, t)\right) + \sigma_t z, \qquad z \sim \mathcal{N}(0, I)$

The mean subtracts a small fraction of the predicted noise -- $\beta_t / \sqrt{1-\bar{\alpha}_t}$ , not all of it -- and then rescales. Removing all the noise at once would be the DDIM jump straight to $\hat{x}_0$ , not an ancestral step.

For the variance DDPM reports two choices that work about equally well:

"posterior": $\sigma_t^2 = \tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$ (optimal for $x_0 \sim \mathcal{N}(0,I)$ )
"beta": $\sigma_t^2 = \beta_t$ (optimal for deterministic $x_0$ )

At $t = 0$ no noise is added at all -- the final step must output a clean sample, and $\tilde{\beta}_0$ is zero anyway.

Your Task

Implement:

def ddpm_step(x_t, eps, t, betas, z, variance_type="posterior"):

Derive alphas and alpha_bars from betas inside the function. Return $x_{t-1}$ as an array shaped like x_t. Raise ValueError for an unknown variance_type.

Input Format

x_t, eps, z: NumPy arrays of matching shape (z is the pre-drawn standard normal sample).
t (int): 0-based timestep index.
betas: 1-D NumPy array.
variance_type (str): "posterior" or "beta".

Output Format

A NumPy array shaped like x_t.

Sample

betas = np.array([0.1, 0.2, 0.3])
x_t = np.array([0.5, -0.5])
eps = np.array([0.2, 0.1])
z = np.array([1.0, -1.0])
print(np.round(ddpm_step(x_t, eps, 1, betas, z), 4).tolist())

alpha_1 = 0.8, alpha_bar_1 = 0.72, so the mean is (x_t - 0.2/sqrt(0.28) * eps) / sqrt(0.8), and *sigma = sqrt((0.1/0.28)0.2) is added times z.

Example:

Input:

betas = np.array([0.1, 0.2, 0.3])
x_t = np.array([0.5, -0.5])
eps = np.array([0.2, 0.1])
z = np.array([1.0, -1.0])
print(np.round(ddpm_step(x_t, eps, 1, betas, z), 4).tolist())

Output:

[0.7418, -0.8685]

Reasoning:

At t = 1: beta = 0.2, alpha = 0.8, alpha_bar = 0.72, so 1 - alpha_bar = 0.28. The bracket is x_t - (0.2/sqrt(0.28))*eps, divided by sqrt(0.8). Since t > 0, noise is added with sigma = sqrt((1-0.9)/0.28 * 0.2) = 0.2673, scaled by z = [1, -1].

Constraints:

0 <= t < len(betas).
The noise coefficient inside the bracket is $\beta_t/\sqrt{1-\bar{\alpha}_t}$ , and the whole bracket is divided by $\sqrt{\alpha_t}$ (note: $\alpha_t$ , not $\bar{\alpha}_t$ ).
When t == 0, return the mean with no noise term.
Use $\bar{\alpha}_{-1} = 1$ in $\tilde{\beta}_t$ .
An unrecognised variance_type must raise ValueError.
Do not round inside the function.

More from Diffusion Models 1

Classifier-Free Guidance: Uncond/Cond MixingEasy DDPM Posterior Mean and Variance from x0Medium Adaptive CFG Scheduling over the TrajectoryHard Classifier Gradient Guidance StepMedium Negative Prompt as the Unconditional BranchMedium

1. Background Knowledge

Diffusion Models operate by gradually adding noise to data (forward process) and then learning to reverse this process (reverse process). The DDPM (Denoising Diffusion Probabilistic Model) reverse step is a probabilistic operation that estimates the previous state $x_{t-1}$ given the current noisy state $x_t$ and a neural network's prediction of the noise $\hat{\varepsilon}$ . Unlike deterministic samplers like DDIM, DDPM is an ancestral sampler, meaning it explicitly adds stochastic noise at each step to approximate the true posterior distribution $p(x_{t-1} | x_t)$ .

The core update rule for DDPM is derived from the Gaussian posterior of the forward process. It consists of two main components: a mean term and a variance term. The mean term effectively denoises the input by subtracting a scaled version of the predicted noise and rescaling by the square root of the current alpha coefficient. The variance term controls the amount of randomness injected, ensuring the sample remains diverse and follows the learned distribution.

Crucially, the coefficients $\alpha_t$ and $\bar{\alpha}_t$ are derived from the beta schedule ( $\beta_t$ ). $\alpha_t = 1 - \beta_t$ represents the retention factor at step $t$ , while $\bar{\alpha}_t$ is the cumulative product of alphas up to step $t$ , representing the total signal retention from $x_0$ to $x_t$ . Understanding how to compute these cumulative products efficiently is key to implementing the sampler correctly.

2. Algorithm Approach

The approach involves implementing the exact mathematical formula provided in the problem description using NumPy for vectorized operations. The algorithm follows a direct computational graph:

Precompute Coefficients: Calculate $\alpha_t$ and $\bar{\alpha}_t$ from the input betas. Since $\bar{\alpha}_t$ is a cumulative product, use np.cumprod.
Select Variance: Determine $\sigma_t^2$ based on the variance_type argument. This involves a conditional check to select between the "posterior" formula (using $\tilde{\beta}_t$ ) or the "beta" formula (using $\beta_t$ ).
Compute Mean: Apply the mean update formula: $\frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\hat{\varepsilon}\right)$ .
Compute Noise Term: Calculate $\sigma_t z$ . Note that if $t=0$ , this term must be zero.
Combine: Sum the mean and noise terms to produce $x_{t-1}$ .

This is a direct evaluation pattern. There are no loops over data dimensions; all operations should be element-wise or broadcasted across the array shape.

3. Step-by-Step Strategy

Derive Alphas:

Compute alphas as 1.0 - betas.
Compute alpha_bars using np.cumprod(alphas). Ensure alpha_bars has the same length as betas. You may need to prepend a 1.0 if the indexing logic requires $\bar{\alpha}_0 = 1$ , but typically for step $t$ , you access index $t$ . Check the problem's indexing convention: usually $\bar{\alpha}_t$ corresponds to the cumulative product up to index $t$ .

Handle Variance Type:

If variance_type is "posterior":
Calculate $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$ .
Be careful with indexing: $\bar{\alpha}_{t-1}$ is the previous cumulative alpha. For $t=0$ , this term is undefined or zero, but the noise is zero anyway.
If variance_type is "beta":
Set $\sigma_t^2 = \beta_t$ .
Raise ValueError if the type is unknown.

Compute the Mean Term:

Extract alpha_t and alpha_bar_t at index t.
Calculate the noise coefficient: noise_coeff = betas[t] / np.sqrt(1 - alpha_bar_t).
Calculate the mean: mean = (x_t - noise_coeff * eps) / np.sqrt(alpha_t).

Compute the Noise Term:

If t == 0: Set sigma = 0.0.
Else: Calculate sigma = np.sqrt(variance).
Compute noise_term = sigma * z.

Final Output:

Return mean + noise_term.

4. Common Pitfalls

Indexing Errors: Confusing $t$ with $t-1$ when accessing alpha_bars. Remember that $\bar{\alpha}_t$ is the cumulative product including $\beta_t$ . If your alpha_bars array starts with $\bar{\alpha}_0$ , ensure you access the correct index.
Division by Zero: At $t=0$ , $1 - \bar{\alpha}_0$ might be zero if $\bar{\alpha}_0=1$ . However, the noise term is zero at $t=0$ , so the division in the mean term uses $\sqrt{1-\bar{\alpha}_0}$ ? No, the mean term uses $\sqrt{1-\bar{\alpha}_t}$ . At $t=0$ , $\bar{\alpha}_0=1$ , so $1-\bar{\alpha}_0=0$ . This causes a division by zero in the noise coefficient $\frac{\beta_0}{\sqrt{1-\bar{\alpha}_0}}$ . Wait, look at the formula: the mean term is $\frac{1}{\sqrt{\alpha_t}}(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\hat{\varepsilon})$ . At $t=0$ , $\bar{\alpha}_0=1$ , so the denominator is 0. However, in practice, the reverse process usually starts from $T$ down to $1$ , and the step from $1 \to 0$ is handled carefully. But the problem asks for a general step. If $t=0$ , the formula might be ill-defined numerically. However, note that at $t=0$ , we are predicting $x_{-1}$ ? No, $x_{t-1}$ where $t=0$ implies $x_{-1}$ which doesn't exist. Usually, the loop runs for $t$ from $T$ down to $1$ . If $t=0$ is passed, it might be an edge case. The problem states "At $t=0$ no noise is added". It does not explicitly say the mean term is skipped. But mathematically, if $t=0$ , $\bar{\alpha}_0=1$ , so $\sqrt{1-\bar{\alpha}_0}=0$ . This suggests $t=0$ might not be a valid input for the noise subtraction part, or the formula simplifies. Actually, standard DDPM implementations often handle $t=0$ by just returning the mean prediction of $x_0$ without noise. But the formula provided has $\sqrt{1-\bar{\alpha}_t}$ in the denominator. If $t=0$ , this is 0. This implies the problem likely assumes $t \ge 1$ for the noise subtraction, or you must handle the $t=0$ case specifically to avoid division by zero. Given the sample uses $t=1$ , focus on $t \ge 1$ . If $t=0$ is tested, check if eps is ignored or if the term becomes 0.
Variance Calculation: For "posterior", ensure you use $\bar{\alpha}_{t-1}$ in the numerator. If $t=0$ , $\bar{\alpha}_{-1}$ is undefined. Again, $t=0$ is a special case.
Data Types: Ensure all calculations are done in float64 or float32 consistently to avoid precision issues.

5. Time & Space Complexity

Time Complexity: $O(N)$ , where $N$ is the number of timesteps (length of betas), due to the cumulative product calculation. The actual step computation is $O(D)$ , where $D$ is the number of elements in x_t, as it involves element-wise operations. Since $N$ is typically small (e.g., 1000) and $D$ can be large (image pixels), the dominant factor is the array size $D$ .
Space Complexity: $O(N)$ to store the alphas and alpha_bars arrays. The output array takes $O(D)$ space. Auxiliary space is minimal.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.