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.Internal.noisePredictorConvolution {d : ℕ} (config : Diffusion.NoisePredictor.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.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.

This stays compact enough for the eager CUDA example while giving the CIFAR trainer more denoising capacity than a bare two-layer network.

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

Build a stronger 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)))$. This compact residual denoiser omits U-Net downsampling, upsampling, and multi-scale skip concatenation. It is still a useful compact architecture because residual paths make the denoising problem much easier than a plain conv chain while staying within the eager CUDA memory envelope used by examples.

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.Internal.appendTimeChannel (batchShape : Shape) {d c : ℕ} (spatial : Tensor ℕ [d]) (x : Tensor Float (sampleShape batchShape c spatial)) (tNorm : Float) : Tensor Float (sampleShape batchShape (c + 1) spatial)

Recursive worker for diffusion.appendTimeChannel.

def TorchLean.diffusion.appendTimeChannel (batchShape : Shape) {d c : ℕ} (spatial : Tensor ℕ [d]) (x : Tensor Float (sampleShape batchShape c spatial)) (tNorm : Float) : Tensor Float (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.noisedSampleFromNoise (batchShape : Shape) {d c T : ℕ} (spatial : Tensor ℕ [d]) (schedule : Schedule T) (x0 eps : Tensor Float (sampleShape batchShape c spatial)) (step : ℕ) : Sample.Supervised Float (sampleShape batchShape (c + 1) spatial) (sampleShape batchShape c spatial)

Build an epsilon-prediction training sample from explicit noise.

The caller supplies eps, usually from the runtime RNG. Keeping randomness outside this helper makes the transformation reusable:

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

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

Build a deterministic epsilon-prediction training sample.

This is the common DDPM training step used by examples: draw reproducible Gaussian noise from (seed, step), corrupt $x_0$, and use that same noise as the target.

def TorchLean.diffusion.ddimPrev {shape : Shape} (abPrev ab : Float) (x_t epsHat : Tensor Float shape) : 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.

We clamp the intermediate $x_0$ estimate to the training image range $[-1,1]$. This is the standard "clipped denoised" stabilizer used by many DDPM/DDIM samplers: without it, a compact model can drive one color channel far outside the data range and the final PPM exporter merely clips the damage into saturated color blobs.