DDPM Training Step End-to-End

Implement one complete DDPM training step — forward noising, noise prediction, and a hand-computed gradient update — for a tiny linear noise-prediction model.

No torch.autograd. No loss.backward(). Every partial derivative is derived by hand.

Model

A minimal linear noise predictor (no activations, for gradient-checking clarity):

input = concat([x_t, t_normalized])    # shape (N, d+1)
predicted_noise = input @ w_model      # shape (N, d)

where w_model has shape (d+1, d).

Pipeline

1. Sample timestep t uniformly from [0, T):

   gen_t = torch.Generator().manual_seed(int(seed))
   t = int(torch.randint(0, T, (1,), generator=gen_t).item())

2. Sample noise eps ~ N(0, I) with shape (N, d):

   gen_eps = torch.Generator().manual_seed(int(seed) + 1)
   eps = torch.randn(N, d, generator=gen_eps)

3. Forward noising (DDPM q(x_t | x_0)):

   x_t = sqrt(alpha_bar[t]) * x_clean + sqrt(1 - alpha_bar[t]) * eps

4. Build model input:

   t_norm = t / T
   input = concat([x_t, t_norm * ones(N, 1)], dim=-1)   # (N, d+1)

5. Predict: predicted_noise = input @ w_model
6. MSE loss: mean((predicted_noise - eps)**2)
7. Gradient by hand:

   d_predicted = (2 / (N * d)) * (predicted_noise - eps)   # (N, d)
   grad_w_model = input.T @ d_predicted                     # (d+1, d)

8. SGD update: w_model -= lr * grad_w_model
9. Return w_model.flatten().

Inputs

• x_clean: shape (N, d) — clean training samples.
• w_model: shape (d+1, d) — noise-prediction model weights.
• betas: shape (T,) — diffusion noise schedule.
• alpha_bar: shape (T,) — cumulative product of (1 - betas).
• lr: float — SGD learning rate.
• seed: int — used for both timestep and noise sampling.

In JAX

JAX’s PRNG does not match PyTorch’s, so the expected outputs for the JAX cases come from this exact derivation. Two independent keys, not a split. Follow it literally or the values will not match:

key_t = jax.random.PRNGKey(int(seed))                 # the timestep
t = int(jax.random.randint(key_t, (1,), 0, T)[0])
key_eps = jax.random.PRNGKey(int(seed) + 1)           # the noise
eps = jax.random.normal(key_eps, (N, d))

Output

Returns shape ((d+1)*d,) — w_model.flatten() after the SGD update.

PRNG note

PyTorch and JAX use different pseudo-random number generators, so each has its own expected outputs here and both are checked exactly. The JAX values are pinned under jax_threefry_partitionable: true, which the local runtime sets for you. Write the idiomatic solution for the framework you are in. Reference: Ho et al., “Denoising Diffusion Probabilistic Models” (2020).