Documentation

LeanMlir.Ddpm

DDPM noise schedule + per-step input plumbing.

A DDPM trainer:

  1. Precomputes the cumulative-α table once via cosineSchedule, stored as a [T] f32 ByteArray.
  2. Per training step, calls stepInputs to:
    • sample a timestep t_b ∈ [0, T) per image
    • sample Gaussian noise ε ~ N(0, I)
    • compute x_t = √ᾱ_t · x_0 + √(1-ᾱ_t) · ε and returns (x_t, ε, t) as three ByteArrays.
  3. Trains the model to predict ε from x_t (per-pixel MSE).

The loss + backward live in the useDdpm codegen branch of MlirCodegen.generateTrainStep.

The schedule is the cosine variant from Nichol & Dhariwal 2021, which trains more stably than Ho et al.'s original linear schedule.

source
@[extern lean_ddpm_cosine_schedule]
opaque Ddpm.cosineSchedule (T : USize) :
IO ByteArray

Build the cosine noise schedule (Nichol & Dhariwal 2021) as a [T] f32 LE ByteArray of ᾱ_t values. With s = 0.008 the schedule keeps log-SNR roughly linear in t.

source
@[extern lean_ddpm_sample_noise]
opaque Ddpm.sampleNoise (n seed : USize) :
IO ByteArray

Sample n f32 from N(0, 1) via Box–Muller. Used for sampling-time noise (training noise comes from stepInputs).

source
@[extern lean_ddim_step]
opaque Ddpm.ddimStep (xt eps : ByteArray) (a b : Float) (n : USize) :
IO ByteArray

DDIM (η = 0) deterministic update: x_{t-1} = a · x_t + b · ε, where a = √ᾱ_{t-1} / √ᾱ_t and b = √(1-ᾱ_{t-1}) - a·√(1-ᾱ_t). Caller precomputes a, b from the schedule.

source
@[extern lean_ddpm_prepend_t_channel]
opaque Ddpm.prependTChannel (xt t : ByteArray) (B C H W Tmax : USize) :
IO ByteArray

Prepend a timestep-encoding channel to each image. Output is [B, C+1, H, W] (flat) where channels 0..C-1 are the input image and channel C is filled with t[i] / T_max. Lets the UNet condition on the diffusion timestep without a new codegen primitive — it just sees a (C+1)-channel input. t is a [B] int32 array (one timestep per image).

source
@[extern lean_ddpm_prepend_t_channel_scalar]
opaque Ddpm.prependTChannelScalar (xt : ByteArray) (B C H W t Tmax : USize) :
IO ByteArray

Scalar variant of prependTChannel for the sampler — broadcasts a single timestep to all images in the batch.

source
@[extern lean_ddpm_prepend_sincos_t]
opaque Ddpm.prependSinCosT (xt t : ByteArray) (B C H W nFreq Tmax : USize) :
IO ByteArray

Sinusoidal time embedding: prepend 2 * nFreq channels of [sin(t · ω_k), cos(t · ω_k)] at log-spaced frequencies (Vaswani / NeRF convention). Replaces the cruder single-channel t/T_max tile with multi-frequency information. Output: [B, C + 2·nFreq, H, W] flat.

source
@[extern lean_ddpm_prepend_sincos_t_scalar]
opaque Ddpm.prependSinCosTScalar (xt : ByteArray) (B C H W t nFreq Tmax : USize) :
IO ByteArray

Scalar variant of prependSinCosT for the sampler.

source
@[extern lean_ddpm_step_inputs]
opaque Ddpm.stepInputs (x0 alphaBar : ByteArray) (B npixels seed : USize) :
IO (ByteArray × ByteArray × ByteArray)

Per training step: sample t_b ∈ [0, T) per image, sample ε, compute x_t. Returns (x_t, ε, t) where:

  • x_t is [B, npixels] f32
  • ε is [B, npixels] f32 (the loss target — what the model should learn to predict)
  • t is [B] int32 LE (the per-image timesteps; useful for future time-conditioning, currently unused by codegen).

The continuous-time (VP-SDE) view of the same schedule #

`cosineSchedule` tabulates ᾱ at the integers `t = 0 … T-1`. Score-SDE
(Song et al. 2021) treats the same schedule as a function of continuous
`t ∈ [0,1]`, which is what lets an ODE or SDE solver choose its own steps.
⭐ **No retraining is needed to get there.** Under the VP SDE the score is
`∇ log p_t(x) = -ε̂(x,t)/σ_t`, so an ε-predicting network *is* a score model
up to that factor. Everything here is sampler-side arithmetic, shared by the
2-D and MNIST drivers so the two cannot drift.
source
def Ddpm.piF :
Float

π. Lean core has no Float.pi; the C side spells the same digits inline.

Equations
Instances For
    source
    def Ddpm.sBias :
    Float

    Nichol & Dhariwal's s, the same 0.008 cosineSchedule uses.

    Equations
    Instances For
      source
      def Ddpm.abarC (t : Float) :
      Float

      ᾱ(t) = cos²θ(t) / cos²θ(0) — the closed form of the tabulated schedule.

      Equations
      Instances For
        source
        def Ddpm.sigC (t : Float) :
        Float

        σ(t) = √(1-ᾱ(t)), the noise scale of the marginal at time t.

        Equations
        Instances For
          source
          def Ddpm.betaC (t : Float) :
          Float

          β(t) = -d/dt log ᾱ(t) = π·tan θ(t)/(1+s), differentiated in closed form rather than differenced — the schedule is analytic, so approximating its derivative would be inventing error. ⚠ β diverges as t → 1. That stiffness is why an explicit solver in x-space struggles here and why DDIM, which integrates the linear part exactly, does not.

          Equations
          Instances For
            source
            def Ddpm.tOfAbar (ab : Float) :
            Float

            Inverse of abarC: t = (2(1+s)/π)·arccos(cos θ₀·√ᾱ) - s. Lets a solver lay its grid out in ᾱ (or σ, or log-SNR) and land on exact times.

            Equations
            Instances For
              source
              def Ddpm.samplerNfe :
              List (String × Nat)

              The samplers the drivers can run, and what one step of each costs in network evaluations. ⚠ Comparisons are made at matched NFE, not matched steps: Heun is second-order and pays two evaluations per step, so on a fixed budget it takes half as many. Reporting steps instead would flatter it.

              Equations
              Instances For