DDPM Posterior Mean and Variance

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.