One DDPM Ancestral Sampling Step

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.