DDPM Sampling Step

The problem

To generate an image, a diffusion model runs the noising process backwards: it starts from pure noise $x_T$ and, at each step $t$ from $T$ down to 1, asks the network for a prediction of the noise in $x_t$ and removes a bit of it, getting $x_{t-1}$. You will write one reverse step, the clean-image estimate, and the whole loop (Ho et al., 2020, Algorithm 2).

One reverse step

Given the current sample $x_t$ and the network's noise prediction $\hat\epsilon$:

where:

• $\alpha_t = 1 - \beta_t$ and $\bar\alpha_t$ come from the schedule of the previous exercise;
• the first part removes the predicted noise from $x_t$;
• $\sigma_t = \sqrt{\beta_t}$, and $z$ is fresh standard Gaussian noise for $t > 1$, but $z = 0$ at the last step, $t = 1$: the final image has no noise added.

Example. Schedule $\beta = (0.1, 0.2, 0.3, 0.4)$, so $\bar\alpha = (0.9, 0.72, 0.504, 0.3024)$. Step $t = 2$ with $x_2 = 1.0$, prediction $\hat\epsilon = 0.5$ and fresh noise $z = 0.7$:

• first part: $\frac{1}{\sqrt{0.8}}\left(1.0 - \frac{0.2}{\sqrt{0.28}} \cdot 0.5\right) = 0.907$;
• noise: $\sqrt{0.2} \cdot 0.7 = 0.313$;
• $x_1 = 1.220$.

Step $t = 1$ with the same $x$ and $\hat\epsilon$: no noise term, and $x_0 = 0.887$ (with $\beta_1 = 0.1$).

The clean-image estimate

From the noise prediction one can estimate the clean image directly, inverting the forward formula:

Example. The forward example: $x_0 = 2$ and noise 1 gave $x_2 = 2.226$. If the network predicts the noise exactly ($\hat\epsilon = 1$): $\hat x_0 = (2.226 - 0.529 \cdot 1) / \sqrt{0.72} = 2$. A wrong prediction gives a wrong estimate.

What to implement

With betas and abar indexed from 0 as in the previous exercise (step $t$ is index $t - 1$):

• predict_x0(x_t, eps_hat, t, abar): the clean-image estimate.
• reverse_step(x_t, eps_hat, t, betas, abar, z): $x_{t-1}$. z is given (an array shaped like x_t), so the step is deterministic; ignore it when $t = 1$.
• sample(eps_model, x_T, betas, abar, rng): the full loop from $t = T$ down to 1, calling eps_model(x, t) for the prediction and drawing z = rng.standard_normal(x.shape) at every step with $t > 1$ (and none at $t = 1$).

Requirements

• $\alpha_t = 1 - \beta_t$.
• Draw the random numbers exactly as described so that the tests can reproduce the samples.