Diffusion Models

Config-style diffusion model constructors plus reusable, dataset-independent DDPM/DDIM helpers.

The runnable examples decide where data comes from (CIFAR-10, ImageNet-style folders, synthetic artifacts). The definitions here are shape-parametric and can be reused by tests, examples, and future proof layer specifications.

structure TorchLean.nn.models.Diffusion.NoisePredictor.Config (d : ℕ) : Type

Configuration for a convolutional diffusion-noise predictor.

• dataChannels : ℕ

  Number of channels in the denoised sample.

• spatial : Tensor ℕ [d]

  Size of each sample axis. Values such as [32, 32] work directly.

• hiddenChannels : ℕ

  Hidden channel width.

• kernelRadius : Tensor ℕ [d]

  Radius of the same-padding convolution kernel on each axis.

  A radius of 1 gives the usual kernel size 3; every residual branch therefore preserves the sample grid by construction.

def TorchLean.nn.models.Diffusion.NoisePredictor.Config.validate {d : ℕ} (config : Config d) : Except String Unit

Validate the complete epsilon-predictor geometry before allocating convolution parameters.

@[reducible, inline]

abbrev TorchLean.nn.models.Diffusion.NoisePredictor.Config.inputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Input shape, with one extra channel carrying diffusion time.

@[reducible, inline]

abbrev TorchLean.nn.models.Diffusion.NoisePredictor.Config.outputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Output shape matching the denoised data channels.

def TorchLean.nn.models.Diffusion.NoisePredictor.Internal.conv {d : ℕ} (config : Config d) (batchShape : Shape) (inputChannels outputChannels : ℕ) : Builder (Sequential (batchShape.concat ((config.spatial.to Shape).prependDim inputChannels)) (batchShape.concat ((config.spatial.to Shape).prependDim outputChannels)))

Implementation helper for the shape-preserving convolutions in the epsilon predictors.

def TorchLean.nn.models.Diffusion.NoisePredictor.Internal.residualBlock {d : ℕ} (config : Config d) (batchShape : Shape) : Builder (Sequential (batchShape.concat ((config.spatial.to Shape).prependDim config.hiddenChannels)) (batchShape.concat ((config.spatial.to Shape).prependDim config.hiddenChannels)))

Build one independently seeded two-convolution residual branch.

def TorchLean.nn.models.Diffusion.NoisePredictor.basic {d : ℕ} (config : Config d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a minimal epsilon-predictor conv net: conv -> relu -> conv -> relu -> conv -> relu -> conv.

The first three convolutions use hiddenChannels; the last returns dataChannels noise values at every spatial position.

def TorchLean.nn.models.Diffusion.NoisePredictor.residual {d : ℕ} (config : Config d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a same-resolution residual epsilon predictor.

Architecture:

stem conv -> relu -> residual block -> relu -> residual block -> relu -> output conv

Each residual block preserves hiddenChannels :: spatial and computes $x+\operatorname{conv}(\operatorname{relu}(\operatorname{conv}(x)))$. The skip adds the block's input unchanged, while its branch learns a correction at the same resolution. This constructor omits U-Net downsampling, upsampling, and multi-scale skip concatenation.

def TorchLean.diffusion.unitToSignedUnit {shape : Shape} (x01 : Tensor Float shape) : Tensor Float shape

Map a tensor from the unit interval to the signed unit interval.

def TorchLean.diffusion.normalNoise {shape : Shape} (seed step : ℕ) : Tensor Float shape

Deterministic Gaussian epsilon tensor for an arbitrary diffusion shape.

The (seed, step) pair is turned into the runtime RNG key, so examples and artifact generation can reproduce the same noising path without ambient randomness.

@[reducible, inline]

abbrev TorchLean.diffusion.sampleShape (batchShape : Shape) (channels : ℕ) {d : ℕ} (spatial : Tensor ℕ [d]) : Shape

Diffusion sample layout: arbitrary batch axes followed by channels and spatial axes.

def TorchLean.diffusion.Internal.linearBetaAt (T : ℕ) (betaStart betaEnd : Float) (t : ℕ) : Float

Linear beta-schedule worker for a natural-number timestep.

def TorchLean.diffusion.Internal.linearAlphaBarAt (T : ℕ) (betaStart betaEnd : Float) (t : ℕ) : Float

Cumulative linear-schedule worker for a natural-number timestep.

def TorchLean.diffusion.Internal.validAlphaBar (value : Float) : Bool

Whether a loaded cumulative diffusion coefficient is finite and probabilistically valid.

def TorchLean.diffusion.Internal.alphaBarsNonincreasing {steps : ℕ} (values : Tensor Float [steps]) : Bool

Cumulative diffusion coefficients cannot increase as noise is added.

def TorchLean.diffusion.linearBeta {T : ℕ} (betaStart betaEnd : Float) (t : Fin T) : Float

Linear beta-schedule value at a timestep that belongs to the schedule.

def TorchLean.diffusion.linearAlphaBar {T : ℕ} (betaStart betaEnd : Float) (t : Fin T) : Float

The cumulative coefficient $\bar\alpha_t=\prod_{s=0}^{t}(1-\beta_s)$.

def TorchLean.diffusion.linearAlphaBars (T : ℕ) (betaStart betaEnd : Float) : Tensor Float [T]

The T cumulative coefficients of a linear beta schedule.

The length belongs to the return type, so a consumer cannot pair the coefficients with a different timestep count.

structure TorchLean.diffusion.Schedule (steps : ℕ) : Type

A validated, nonempty diffusion schedule.

The constructor is private so runnable code never carries a separate proof that the coefficient tensor can be indexed. Use Schedule.from for loaded coefficients or Schedule.linear for the standard linear beta schedule.

• alphaBars : Tensor Float [steps]

  Cumulative coefficients indexed by diffusion timestep.

• nonempty : steps ≠ 0

  Internal invariant used to cycle natural-number training steps safely.

opaque TorchLean.diffusion.Schedule.from {steps : ℕ} (alphaBars : Tensor Float [steps]) : Except String (Schedule steps)

Validate a coefficient tensor as a runnable diffusion schedule.

def TorchLean.diffusion.Schedule.linear (steps : ℕ) (betaStart betaEnd : Float) : Except String (Schedule steps)

Build and validate the standard linear beta schedule.

def TorchLean.diffusion.Schedule.index {steps : ℕ} (schedule : Schedule steps) (step : ℕ) : Fin steps

Cycle an arbitrary logical step through this schedule.

def TorchLean.diffusion.Schedule.alphaBar {steps : ℕ} (schedule : Schedule steps) (step : ℕ) : Float

Cumulative coefficient selected by a logical training step.

def TorchLean.diffusion.Schedule.normalizedTime {steps : ℕ} (schedule : Schedule steps) (step : ℕ) : Float

Normalize a logical training step to the unit interval used by the time channel.

def TorchLean.diffusion.appendTimeChannel {α : Type} [Storage α] (batchShape : Shape) {d c : ℕ} (spatial : Tensor ℕ [d]) (x : Tensor α (sampleShape batchShape c spatial)) (tNorm : α) : Tensor α (sampleShape batchShape (c + 1) spatial)

Append a constant time channel to every sample in batchShape.

The input layout is batchShape × channels × spatial. The result preserves the batch and spatial axes and changes only the channel count from c to c + 1.

def TorchLean.diffusion.sample (batchShape : Shape) {d c T : ℕ} (spatial : Tensor ℕ [d]) (schedule : Schedule T) (x0 : Tensor Float (sampleShape batchShape c spatial)) (step : ℕ) (seed : ℕ := 0) (noise : Tensor Float (sampleShape batchShape c spatial) := normalNoise seed step) : Sample.Supervised Float (sampleShape batchShape (c + 1) spatial) (sampleShape batchShape c spatial)

Build an epsilon-prediction training sample.

By default, (seed, step) determines reproducible Gaussian noise. Pass noise := ... to use a supplied tensor instead, without changing the noising formula:

$x_t=\sqrt{\bar{\alpha}_t}\,x_0+\sqrt{1-\bar{\alpha}_t}\,\varepsilon$, with target $\varepsilon$.

def TorchLean.diffusion.ddimPrev {shape : Shape} (abPrev ab : Float) (x_t epsHat : Tensor Float shape) (postprocess : Tensor Float shape → Tensor Float shape := fun (reconstruction : Tensor Float shape) => reconstruction.clamp (-1) 1) (denominatorFloor : Float := 1e-12) : Tensor Float shape

One deterministic DDIM reverse update ($\eta=0$).

Given $x_t$, predicted epsilon, and adjacent schedule values, this estimates $x_0$ and remixes it to the previous timestep.

The default postprocessor clamps the intermediate $x_0$ estimate to the training image range $[-1,1]$. This stabilizes compact image models whose reconstruction can drift outside the data range. Supply postprocess := id for an unclipped reconstruction, or a shape-preserving transform for another data domain. Postprocessing happens before remixing and does not recompute epsilon.

The reconstruction denominator uses the explicit sqrtAb > denominatorFloor branch, with the floor in the else branch. The default is 1e-12. For finite schedule coefficients, a positive finite floor keeps the denominator finite and nonzero.