DDPM Sampler Loop

Implement the full DDPM reverse sampling loop — starting from pure Gaussian noise and iteratively denoising down to a clean sample.

This is Algorithm 2 from Ho et al. 2020 (“Denoising Diffusion Probabilistic Models”), using the same tiny linear noise-prediction model as ddpm-train-step-end-to-end: concat (x_t, t/T) → linear → predicted noise.

Model architecture

input = concat([x_t, (t/T) * ones(N, 1)])   # shape (N, d+1)
predicted_noise = input @ w_model             # shape (N, d)

w_model has shape (d+1, d).

Reverse sampling pipeline

1. Initialize x ~ N(0, I) with shape [N, d] using the provided seed:

   gen_init = torch.Generator().manual_seed(int(seed))
   x = torch.randn(N, d, generator=gen_init)

2. Set T = betas.shape[0].
3. Loop t from T-1 down to 0:
   • Predict noise:

     t_col = (t / T) * ones(N, 1)
     inp   = concat([x, t_col], dim=-1)   # (N, d+1)
     predicted_noise = inp @ w_model

   • Compute mean (reverse step):

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

   • Add noise (except at the final step):

     if t > 0:
         gen_step = torch.Generator().manual_seed(int(seed) + 100 + t)
         z = torch.randn(N, d, generator=gen_step)
         x = mean + sqrt(betas[t]) * z
     else:
         x = mean   # no noise on the last step

4. Return x with shape (N, d).

Inputs

• w_model: shape (d+1, d) — trained noise-prediction model weights.
• shape: 1-D tensor [N, d] (float values; cast to int before use).
• betas: shape (T,) — noise schedule.
• alpha_bar: shape (T,) — cumulative product prod(1 - betas[:t+1]).
• seed: int — initial noise seed; per-step noise uses seed + 100 + t.

Output

Shape (N, d) — the final denoised sample.

Notes

• In production (e.g., DALL·E, Stable Diffusion), T=1000 and the model is a U-Net. Here T=4 and a tiny linear model keep the loop tractable for testing.
• Each step’s noise uses a fresh seeded Generator (seed + 100 + t) so results are fully reproducible step by step.
• PRNG note: expected outputs are generated using PyTorch PRNG. Reference: Ho et al., “Denoising Diffusion Probabilistic Models” (2020).