Medium Primitives

DDPM Forward Noising

Create a noisy version of a clean input at one chosen diffusion step. Instead of simulating every earlier step, use the supplied cumulative alpha value to combine the clean input with one Gaussian noise draw.

Implement ddpm_forward(x0, t, alpha_bar, seed).

Inputs and output

x0 has shape (N, d): a batch of clean inputs.
alpha_bar has shape (T,), with values between 0 and 1 inclusive. It contains the cumulative values, not the beta row of the noise schedule.
x0 and alpha_bar share a floating-point dtype and device.
t is an integer index with 0 <= t < T, shared by the whole batch.
seed is a nonnegative integer used to create a fresh random generator.

Return the noisy tensor with the shape, dtype and device of x0. Leave both input tensors unchanged.

Combination rule

a = alpha_bar[t]
output = sqrt(a) * x0 + sqrt(1 - a) * noise

Draw one standard normal noise tensor with the full shape of x0. Each element gets its own draw; do not share a single noise value across a row or batch. Do not recalculate the cumulative product.

The square roots are essential: 1 - a is the noise variance, so sqrt(1 - a) is its standard deviation. At a=1, the result is x0; at a=0, it is the sampled noise. Small a reduces the clean input's coefficient, but the actual contribution also depends on the input values.

The schedule uses zero-based indices. Its index 0 already includes the first noising step, so t=0 does not generally return a clean input. Use alpha_bar[t] directly, without shifting the index.

Reproducible noise

Use the recipes below so the grader can compare a seeded sample. Generate float32 noise, then convert it to the dtype and device of x0 before combining. The two frameworks produce different samples for the same seed; each has its own expected values, compared with floating-point tolerance.

In PyTorch

Create a fresh CPU generator with torch.Generator().manual_seed(int(seed)). Draw torch.randn(x0.shape, generator=gen, dtype=torch.float32, device="cpu"), then move the result to x0's dtype and device. Do not reseed the global generator or draw on the GPU, since that changes the sampling recipe.

In JAX

Create jax.random.PRNGKey(int(seed)) and draw jax.random.normal(key, x0.shape, dtype=jnp.float32). Convert the result to x0.dtype. The runtime sets jax_threefry_partitionable=True, matching these fixtures. Draw the whole tensor in one call without splitting the key.

See Denoising Diffusion Probabilistic Models 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…