Hard End-to-End

DDPM Sampler Loop

Implement ddpm_sample_loop(w_model, shape, betas, alpha_bar, seed). Start with Gaussian noise, then apply every reverse diffusion step in descending order. Use the supplied linear model to predict noise from the current sample at each step.

This is a small sampling exercise. The weights are already supplied, and the function does not train them. Arbitrary weights need not produce a realistic or clean sample.

Inputs and output

shape is a one-dimensional tensor containing [N, d]. Its two values are positive whole numbers stored as floats; convert them to integers.
w_model has shape (d + 1, d).
betas and alpha_bar have shape (T,), with T >= 1. Each beta lies strictly between 0 and 1; each cumulative alpha lies between 0 inclusive and 1 exclusive. Use these supplied values directly.
The weights and schedules share a floating-point dtype and device.
seed is a nonnegative integer.

Return only the final tensor of shape (N, d), with the dtype and device of w_model. Leave all inputs unchanged.

Predict noise at each step

For each t from T-1 down through 0, append a time column to the current x:

time_column = full((N, 1), t / T)
model_input = concatenate([x, time_column], axis=1)
predicted_noise = model_input @ w_model

Use t / T, not t / (T - 1). The time feature is last, so the last row of w_model contains its weights. There is no separate bias or activation. Recompute the prediction using the updated sample on every iteration.

Reverse update

Initialize x with one standard normal draw of shape (N, d). At each step:

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

If t > 0, draw a new noise tensor z and set x = mean + sqrt(beta) * z. This exercise uses beta as the reverse variance. At t=0, set x=mean without drawing noise. The result of one step becomes the input to the next; return x after index 0 has been processed.

For T=3, process indices 2, 1, 0 and use time features 2/3, 1/3, 0. For T=1, still draw the initial sample, predict noise with time feature zero, and apply the mean update once. There is no per-step noise draw.

Reproducible draws

Use seed for the initial draw and seed + 100 + t for the draw at each t > 0. These are fresh generators or keys, not a single advancing stream. For T=3, the draw seeds are seed, seed+102 and seed+101, in that order.

Generate each complete noise tensor in float32, then convert it to the weights' dtype and device. The frameworks have different expected samples, compared with floating-point tolerance.

In PyTorch

For each draw, create torch.Generator().manual_seed(int(draw_seed)). Use torch.randn(N, d, generator=gen, dtype=torch.float32, device="cpu"), then move the result to w_model's dtype and device. Keep the global random generator state unchanged.

In JAX

For each draw, create jax.random.PRNGKey(int(draw_seed)). Use jax.random.normal(key, (N, d), dtype=jnp.float32), then convert to w_model.dtype. Do not split keys. The runtime sets jax_threefry_partitionable=True for these fixtures.

The reverse update follows Algorithm 2 in Denoising Diffusion Probabilistic Models, using this exercise's small linear model.

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…