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.

References #

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).
source
@[extern lean_ddpm_flow_step_inputs]
opaque Ddpm.flowStepInputs (x0 epsIn : ByteArray) (B n seed Tmax mode : USize) :
IO (ByteArray × ByteArray × ByteArray)

Flow matching's per-step inputs on the linear interpolant. The one thing that changes from stepInputs: per row t ~ U(0, 1), x_t = (1-t)·x0 + t·ε and the target is the velocity v = ε - x0 instead of ε. Returns (x_t, v, t_idx) with t_idx = round(t · Tmax) as the int32 index prependSinCosT already takes, so the time channel is untouched. mode 0 draws ε; 1 takes it from epsIn (reflow's fixed pairs); 2 draws it and then pairs rows by a minimum-cost assignment on |x0_i - ε_j|² (minibatch OT). The MSE block does not know whether the target is ε or v, so the train step is trainStepAdamF32Ddpm unchanged.

source
def Ddpm.flowSamplerNfe :
List (String × Nat)

The flow-matching samplers, beside samplerNfe: Euler and Heun on dx/dt = v(x, t) from t = 1 to 0 on a uniform grid. Kept as a separate list so the image-side probes that parse samplerNfe do not start accepting names their VP-schedule math has no meaning for.

Equations
Instances For

    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.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
                source
                def Ddpm.ddimCoefs (abT abP eta : Float) :
                Float × Float × Float

                One generalized-DDIM step from ᾱ_t to ᾱ_prev (Song, Meng & Ermon 2020, eq. 12), as the coefficients (a, b, σ) of x ← a·x + b·ε̂ + σ·z:

                σ = η·√((1-ᾱ_prev)/(1-ᾱ_t))·√(1 - ᾱ_t/ᾱ_prev), b = √(1 - ᾱ_prev - σ²) − a·√(1-ᾱ_t)

                with a = √ᾱ_prev/√ᾱ_t. η = 0 is deterministic DDIM (σ = 0); η = 1 is ancestral DDPM sampling, the MNIST demo's default — 0.0067 energy distance at 50 evaluations against deterministic DDIM's 0.0188 at 200. The inner root is clamped because η → 1 can drive 1 − ᾱ_prev − σ² marginally negative at the ends of the schedule. Every image and 2-D sampler takes its step from here, so the drivers cannot disagree on it.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For