Medium End-to-End

DDPM Training Step End-to-End

Implement ddpm_train_step(x_clean, w_model, betas, alpha_bar, lr, seed). Train a small linear noise predictor for exactly one batch and one SGD update. Compute the gradient yourself using tensor operations, without automatic differentiation in either framework.

Inputs and output

x_clean has shape (N, d), with N and d positive.
w_model has shape (d + 1, d). Its last row weights the time feature.
betas and alpha_bar have shape (T,), with T >= 1. Read T from betas, but use the supplied alpha_bar values for noising; do not recompute them. Each cumulative value lies between 0 and 1 inclusive.
All tensors share a floating-point dtype and device.
lr is a finite, nonnegative Python float; seed is a nonnegative integer.

Return the updated weights flattened in row-major order, with shape ((d + 1) * d,), preserving the weights' dtype and device. Return neither the loss nor a prediction. Leave every input unchanged.

One training update

Draw one timestep t uniformly from 0 through T-1. Use this same t for every example in the batch.
Draw a standard normal tensor eps of shape (N, d), with one draw per element. The sampling recipes below fix the generators and draw order.
Construct the noisy input: x_t = sqrt(alpha_bar[t]) * x_clean + sqrt(1 - alpha_bar[t]) * eps.
Append a column containing t / T to x_t. This gives inp with shape (N, d + 1). Use T, not T - 1, as the denominator.
Compute predicted = inp @ w_model. There is no bias or activation.
Compute the gradient of the mean squared error between predicted and eps, then update the weights once.

The target is the sampled eps, not x_clean. Treat inp and eps as fixed when differentiating with respect to w_model:

loss = mean((predicted - eps)**2)
d_predicted = 2 * (predicted - eps) / (N * d)
grad_w = inp.T @ d_predicted
updated_weights = w_model - lr * grad_w

The mean includes every batch and feature element, so the divisor is N*d, not N alone. Compute the whole gradient from the original weights before updating. With lr=0, the returned values are the original weights flattened.

For T=1, the sampled timestep is always 0 and the time column is zero. The last row of weights then has zero gradient, while the other rows can still change.

Reproducible sampling

Use seed for the timestep and seed+1 for the noise, with separate fresh generators or keys. Draw the noise in float32, then convert it to the input dtype and device. Each framework has its own expected outputs, compared with floating-point tolerance.

In PyTorch

Create gen_t = torch.Generator().manual_seed(int(seed)) and sample torch.randint(0, T, (1,), generator=gen_t, dtype=torch.int64, device="cpu"). Convert its single value to a Python integer.

Create a second CPU generator seeded with int(seed)+1. Draw torch.randn(N, d, generator=gen_eps, dtype=torch.float32, device="cpu"), then move eps to x_clean's dtype and device. Leave the global RNG unchanged.

In JAX

Sample t from jax.random.randint(jax.random.PRNGKey(int(seed)), (1,), 0, T, dtype=jnp.int32), then convert its single value to a Python integer.

Draw eps with jax.random.normal(jax.random.PRNGKey(int(seed)+1), (N, d), dtype=jnp.float32) and convert to x_clean.dtype. Do not split keys. The runtime sets jax_threefry_partitionable=True for these fixtures.

The noising and noise-prediction objective follow Denoising Diffusion Probabilistic Models, using this exercise's 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…