core/nn/diffusion library

Denoising Diffusion Probabilistic Models — schedule + tiny U-Net.

Ho, Jain, Abbeel (2020). What's here:

  • NoiseSchedule.linear — linear beta schedule with the paper's defaults (beta_start = 1e-4, beta_end = 0.02, T = 1000). Precomputes alpha, alpha_bar, sqrt(alpha_bar), sqrt(1 - alpha_bar), and the posterior variance used in the reverse step.
  • NoiseSchedule.forwardDiffuse — closed-form q(x_t | x_0) = N(√α̅_t · x_0, (1 - α̅_t) I).
  • NoiseSchedule.reverseStep — single Langevin step of the Markov chain given a predicted ε̂.
  • TinyUNet — a minimal 2-down / 2-up U-Net whose upsampling path is our fresh ConvTranspose2d. Forward-only wiring showcase (no training loop here — that needs Conv2d input-grad support, currently missing). Loads pretrained tiny-DDPM checkpoints or serves as scaffolding for future ports.

See test/diffusion_test.dart for the schedule invariants, exact forward-diffusion means/variances, reverse-step algebra, and the U-Net shape checks.

Classes

NoiseSchedule
TinyUNet
Minimal U-Net for [N, 1, H, W] grayscale images. Two down blocks (stride-2 Conv2d), a mid block, then two up blocks (ConvTranspose2d, stride-2). A per-image scalar timestep is embedded via a linear projection and broadcast-added to the mid features. Predicts an ε-shape tensor [N, 1, H, W].