📘

DDPM Posterior Mean and Variance

MediumDiffusion Models 📚 Diffusion Models 1

Problem Statement

Compute the exact reverse-process posterior $q(x_{t-1} \mid x_t, x_0)$ -- the distribution the denoising network is trained to imitate.

Background

The reverse of the forward chain is intractable in general, but if you are also told $x_0$ it is available in closed form and it is Gaussian:

$q(x_{t-1} \mid x_t, x_0) = \mathcal{N}\!\left(x_{t-1};\ \tilde{\mu}_t(x_t, x_0),\ \tilde{\beta}_t I\right)$

$\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t}\,x_0 \;+\; \frac{\sqrt{\alpha_t}\,(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\,x_t$

$\tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\,\beta_t$

with the convention $\bar{\alpha}_{-1} = 1$ for the very first step.

Both coefficients are positive, so the mean sits between $x_0$ and $x_t$ : early in the reverse trajectory (large $t$ ) almost all the weight is on $x_t$ and the posterior barely moves, while near $t = 0$ it swings onto $x_0$ . (They sum to slightly less than 1, not exactly 1 -- the forward process shrinks the signal a little at every step.) At $t = 0$ the convention $\bar{\alpha}_{-1} = 1$ makes the coefficients exactly $(1, 0)$ and $\tilde{\beta}_0 = 0$ : knowing $x_0$ pins $x_{-1} = x_0$ with no uncertainty left. Handling that boundary is the part implementations get wrong.

Your Task

Implement:

def posterior(x0, x_t, t, betas):

Derive alphas and alpha_bars from betas inside the function. Return (mean, var) where mean is an array shaped like x0 and var is a Python float.

Input Format

x0, x_t: NumPy arrays of the same shape.
t (int): 0-based timestep index, 0 <= t < len(betas).
betas: 1-D NumPy array of the noise schedule.

Output Format

A tuple (mean, var): a NumPy array shaped like x0, and a float.

Sample

betas = np.array([0.1, 0.2, 0.3])
x0 = np.array([1.0, -1.0])
x_t = np.array([0.4, 0.2])
m, v = posterior(x0, x_t, 1, betas)
print(np.round(m, 4).tolist())
print(round(v, 4))

Here alpha_bar_1 = 0.72, alpha_bar_0 = 0.9, so the coefficients are *sqrt(0.9)0.2/0.28 on x0 and *sqrt(0.8)0.1/0.28 on x_t, and the variance is *(0.1/0.28)0.2.

Example:

Input:

betas = np.array([0.1, 0.2, 0.3])
x0 = np.array([1.0, -1.0])
x_t = np.array([0.4, 0.2])
m, v = posterior(x0, x_t, 1, betas)
print(np.round(m, 4).tolist())
print(round(v, 4))

Output:

[0.8054, -0.6137]
0.0714

Reasoning:

alphas = [0.9, 0.8, 0.7], so alpha_bar_0 = 0.9 and alpha_bar_1 = 0.72, giving 1 - alpha_bar_1 = 0.28. The x0 coefficient is sqrt(0.9)*0.2/0.28 = 0.6776 and the x_t coefficient is sqrt(0.8)*0.1/0.28 = 0.3194, so the mean is 0.6776*1.0 + 0.3194*0.4 = 0.8054 and 0.6776*(-1.0) + 0.3194*0.2 = -0.6137. The variance is (0.1/0.28)*0.2 = 0.0714.

Constraints:

0 <= t < len(betas).
Use $\bar{\alpha}_{-1} = 1$ when t == 0; the mean is then exactly x0 and the variance exactly 0.0.
The $x_0$ coefficient uses the cumulative $\bar{\alpha}_{t-1}$ , while the $x_t$ coefficient uses the per-step $\alpha_t$ together with $1-\bar{\alpha}_{t-1}$ ; mixing the two up is the usual bug.
Return var as a plain Python float.
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 Gaussian noise to data until it becomes pure noise (the forward process) and then learning to reverse this process (the reverse process). The forward process is defined by a variance schedule $\beta_t$ . At each step $t$ , the transition $q(x_t \mid x_{t-1})$ is a Gaussian distribution. A key property of this chain is that we can sample $x_t$ directly from $x_0$ using the cumulative product of the signal retention terms, denoted as $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$ , where $\alpha_t = 1 - \beta_t$ .

The core challenge in training diffusion models is that the true reverse posterior $q(x_{t-1} \mid x_t)$ is intractable because it depends on the entire dataset. However, if we condition on the original clean data $x_0$ , the posterior $q(x_{t-1} \mid x_t, x_0)$ becomes analytically tractable and remains Gaussian. This is derived using Bayes' theorem and the Markov property of the forward chain. The resulting distribution has a mean $\tilde{\mu}_t$ that is a weighted average of $x_t$ and $x_0$ , and a variance $\tilde{\beta}_t$ that represents the residual uncertainty.

Understanding the reparameterization trick is crucial here. In practice, we don't just compute the mean; we use the fact that $x_t$ can be written as $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$ , where $\epsilon \sim \mathcal{N}(0, I)$ . This allows the neural network to predict noise $\epsilon$ instead of $x_0$ directly, which stabilizes training. For this specific problem, you are implementing the exact mathematical ground truth for the posterior mean and variance, which serves as the target for the model's predictions during training.

2. Algorithm Approach

The approach is a direct numerical implementation of closed-form Gaussian posterior equations. You are not optimizing or iterating; you are evaluating static formulas given specific inputs.

Precompute Schedule Parameters: Convert the input betas array into alphas ( $\alpha_t = 1 - \beta_t$ ) and then into alpha_bars ( $\bar{\alpha}_t$ ). This requires a cumulative product operation.
Index Selection: Extract the specific values $\alpha_t$ , $\bar{\alpha}_t$ , $\alpha_{t-1}$ , and $\bar{\alpha}_{t-1}$ corresponding to the current timestep $t$ .
Coefficient Calculation: Compute the two weights for the mean equation and the single weight for the variance equation using the provided formulas.
Vectorized Arithmetic: Apply these scalar coefficients to the input arrays x0 and x_t using NumPy broadcasting to compute the final mean array.

3. Step-by-Step Strategy

Derive Alphas and Alpha Bars:

Calculate alphas as 1.0 - betas.
Calculate alpha_bars using np.cumprod(alphas). This gives you $\bar{\alpha}_1, \bar{\alpha}_2, \dots$ .
Crucial Step: The formula for $\bar{\alpha}_{t-1}$ when $t=0$ requires $\bar{\alpha}_{-1} = 1$ . You should prepend 1.0 to your alpha_bars array so that index $t$ corresponds to $\bar{\alpha}_t$ and index $t-1$ corresponds to $\bar{\alpha}_{t-1}$ naturally. Let's call this padded array alpha_bars_padded.

Extract Timestep Values:

Retrieve $\beta_t = \text{betas}[t]$ .
Retrieve $\alpha_t = \text{alphas}[t]$ .
Retrieve $\bar{\alpha}_t = \text{alpha\_bars\_padded}[t]$ .
Retrieve $\bar{\alpha}_{t-1} = \text{alpha\_bars\_padded}[t-1]$ .

Compute Mean Coefficients:

Calculate the coefficient for $x_0$ : $c_0 = \frac{\sqrt{\bar{\alpha}_{t-1}} \cdot \beta_t}{1 - \bar{\alpha}_t}$
Calculate the coefficient for $x_t$ : $c_1 = \frac{\sqrt{\alpha_t} \cdot (1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}$
Compute the mean: mean = c0 * x0 + c1 * x_t.

Compute Variance:

Calculate the posterior variance $\tilde{\beta}_t$ : $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t$
Return this as a Python float.

Handle Edge Cases:

Ensure that when $t=0$ , the logic holds. With the padded array, $\bar{\alpha}_{-1}$ is accessed at index 0 (value 1.0) and $\bar{\alpha}_0$ at index 1. The denominator $1 - \bar{\alpha}_0$ will be non-zero (unless $\beta_0=0$ , which is rare in standard schedules). The formula naturally collapses to $x_0$ with variance 0 if implemented correctly.

4. Common Pitfalls

Off-by-One Errors in Indexing: The most common mistake is misaligning $\bar{\alpha}_t$ and $\bar{\alpha}_{t-1}$ . Remember that np.cumprod on alphas gives $[\bar{\alpha}_1, \bar{\alpha}_2, \dots]$ . It does not include $\bar{\alpha}_0 = 1$ at the start. You must explicitly prepend 1.0 to handle the $t=0$ case and the $t-1$ index correctly.
Division by Zero: The term $1 - \bar{\alpha}_t$ appears in the denominator. Ensure that $\bar{\alpha}_t$ is not exactly 1.0 for $t > 0$ . In standard schedules, $\beta_t > 0$ , so $\bar{\alpha}_t < 1$ , but numerical precision issues can sometimes cause problems if betas are extremely small.
Shape Mismatch: Ensure that mean retains the exact shape of x0. NumPy broadcasting usually handles this, but if c0 or c1 are accidentally arrays instead of scalars, it might cause unexpected behavior.
Variance Type: The problem specifies var should be a Python float. Do not return a NumPy scalar or array for the variance. Use float(var_value).
Square Root Domain: Ensure arguments to np.sqrt are non-negative. Since $\bar{\alpha}$ values are products of numbers between 0 and 1, they are always positive, but good practice to check.

5. Time & Space Complexity

Time Complexity: $O(T + N)$ , where $T$ is the length of the betas array (number of timesteps) and $N$ is the number of elements in x0 (or x_t).
Computing alpha_bars via cumprod takes $O(T)$ .
Computing the mean involves element-wise operations on arrays of size $N$ , taking $O(N)$ .
Since $T$ is typically much smaller than $N$ (e.g., $T=1000$ , $N=10^6$ for an image), the dominant factor is often $N$ .
Space Complexity: $O(T + N)$ .
$O(T)$ to store the alphas and alpha_bars arrays.
$O(N)$ to store the output mean array.
Auxiliary space is minimal as we reuse input arrays for broadcasting.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.