Reverse DDPM sampler (spec layer) #
This file defines a standard ε-prediction reverse sampler step for discrete VP/DDPM schedules.
We expose:
ddpmStep: one reverse step $x_t\to x_{t-1}$ with explicit noise input $z$,ddpmSample: run all $T$ reverse steps, given a noise stream $z_0,\ldots,z_{T-1}$.
We keep everything scalar-polymorphic (Context α). The intended use is:
- execute with
Float/IEEE32Exec/NeuralFloatfor concrete runs, and - reuse the same definitions with
ℝin proofs.
References (informal pointers):
- Ho, Jain, Abbeel (2020), DDPM, Algorithm 2 (reverse process).
def
Generative.Diffusion.x0Pred
{α : Type}
[Context α]
{T : ℕ}
{s : Spec.Shape}
(sched : VPSchedule α T)
(model : EpsModel α s)
(x_t : Spec.Tensor α s)
(t : Fin (T + 1))
:
Spec.Tensor α s
Predict $x_0$ from $x_t$ and an $\varepsilon$-prediction model.
Formula (ε-pred parameterization):
$$ x_0=\frac{x_t-\sqrt{1-\bar\alpha_t}\,\hat\varepsilon} {\sqrt{\bar\alpha_t}}. $$
Instances For
def
Generative.Diffusion.ddpmStep
{α : Type}
[Context α]
{T : ℕ}
{s : Spec.Shape}
(sched : VPSchedule α T)
(model : EpsModel α s)
(k : Fin T)
(x_t z : Spec.Tensor α s)
:
Spec.Tensor α s
One reverse DDPM step $x_t\to x_{t-1}$ with explicit noise $z$ (intended as $\mathcal{N}(0,I)$).
We index reverse steps by k : Fin T corresponding to the transition $t=k+1\to k$.
Implementation details:
- time embedding passed to the model is $t/T$ (see
VPSchedule.timeOfIndex). - we use epsilon-protected scalar division in the coefficient formulas to stay total.
Instances For
def
Generative.Diffusion.ddpmSample
{α : Type}
[Context α]
{T : ℕ}
{s : Spec.Shape}
(sched : VPSchedule α T)
(model : EpsModel α s)
(x_T : Spec.Tensor α s)
(noise : Fin T → Spec.Tensor α s)
:
Spec.Tensor α s
Run the full reverse DDPM sampler for T steps.
Inputs:
- $x_T$: starting state (typically standard normal noise),
noise: per-step noise streamz_kfork = 0..T-1.
Output:
- the terminal sample $x_0$.
Order note:
noise (T-1)is used first (for the step $T\to T-1$),noise 0is used last (for the step $1\to0$).