Hard Primitives

DDPM Denoising Step

Implement ddpm_reverse_step(x_t, predicted_noise, t, betas, alpha_bar, seed) to take one reverse diffusion step. The model's noise prediction is already supplied; use it to compute the reverse mean, then add fresh noise unless this is the final step.

Inputs and output

x_t and predicted_noise have shape (N, d).
betas and alpha_bar have shape (T,). alpha_bar contains cumulative alphas, not the per-step values 1 - betas.
At the selected index, 0 < betas[t] < 1 and 0 <= alpha_bar[t] < 1.
All tensors share a floating-point dtype and device.
t is an integer with 0 <= t < T, shared by all batch rows.
seed is a nonnegative integer, used only when t > 0.

Return one tensor with the shape, dtype and device of x_t. Leave inputs unchanged. Use the supplied schedule values directly.

Reverse update

beta = betas[t]
alpha = 1 - beta
mean = (x_t - beta / sqrt(1 - alpha_bar[t]) * predicted_noise) / sqrt(alpha)

if t > 0:
    output = mean + sqrt(beta) * z
else:
    output = mean

Here z is a standard normal tensor with the full shape of x_t. Each element gets its own noise draw.

This exercise chooses beta as the reverse-step variance. Use sqrt(beta) as the noise multiplier; do not substitute the alternative posterior variance that appears in some DDPM samplers. The reverse mean depends on both the current beta and the cumulative alpha.

The final step

The array indices run backward from T-1 to 0. Index 0 is the last reverse step, undoing the first forward noising step. Return the mean at t=0 without drawing noise. Do not return x_t unchanged, and do not look up a schedule entry at index -1. The result estimates clean data but need not recover it exactly.

For example, at t=0 with beta=0.25, alpha_bar=0.75, x_t=[[1, -1]] and predicted_noise=[[2, 0]], the numerator is [[0, -1]]. Divide by sqrt(0.75) to get approximately [[0, -1.154701]], regardless of the seed.

Reproducible noise for t > 0

Draw one float32 noise tensor, then convert it to x_t's dtype and device. The frameworks produce different seeded samples and have separate expected values, compared with floating-point tolerance.

In PyTorch

Use a fresh CPU generator, torch.Generator().manual_seed(int(seed)). Draw torch.randn(x_t.shape, generator=gen, dtype=torch.float32, device="cpu"), then move the result to x_t's dtype and device. Leave the global generator state unchanged.

In JAX

Use jax.random.PRNGKey(int(seed)) and jax.random.normal(key, x_t.shape, dtype=jnp.float32), then convert to x_t.dtype. Draw the full tensor in one call without splitting the key. The runtime sets jax_threefry_partitionable=True for these fixtures.

See Denoising Diffusion Probabilistic Models, Algorithm 2, by Ho, Jain and Abbeel.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Train a model to recognize added noise

A diffusion model learns from data with different amounts of added noise. Let beta_t specify the noise schedule, alpha_t = 1 - beta_t, and alpha_bar_t be the product of the alpha values up to step t. A noisy sample can be drawn from clean data x_0 directly:

$$x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon, \qquad\epsilon\sim\mathcal N(0,I).$$

If alpha_bar_t = 0.64, the coefficients are 0.8 for the clean input and 0.6 for the noise. These are amplitude coefficients; their squares determine the corresponding variance factors. When alpha_bar_t is near one, much of the clean signal remains. Near zero, noise dominates. The chosen schedule determines whether the final training step reaches nearly pure noise.

For noise-prediction training, compare the model's output with the actual epsilon used to construct x_t. Also supply the time step so the model knows the noise level. A model that predicts the clean sample instead needs different targets and sampling formulas.

Generate by taking reverse steps

Generation usually starts from noise. At each step, the model helps compute a less noisy sample. A stochastic sampler can also add fresh noise, using the variance specified by its schedule. Using beta_t as that variance is different from using the forward process's posterior variance. The final step commonly omits added noise.

Keep array indices separate from mathematical time-step labels. At array index zero, a cumulative product already includes the first factor. Follow the task's random-seed or key sequence too: changing the order of draws can change expected outputs. A small linear model can demonstrate these mechanics without being a realistic image generator.

Detecting runtime…