Probability-flow ODE (spec layer)

This file defines a small continuous-time VP schedule (linear $\beta(t)$) and the corresponding probability-flow ODE drift field, using an $\varepsilon_\theta(x,t)$ model.

Why include this in the spec layer:

• the ODE is a deterministic dynamical system derived from the same diffusion model, and
• it is the natural interface for inference-time verification tooling (corridor certificates, IBP/CROWN bounds on the RHS, etc.).

We keep the implementation scalar-polymorphic (Context α) so it can be:

• executed with Float / IEEE32Exec / NeuralFloat, and
• reasoned about with ℝ.

References (informal pointers):

• Song et al. (2021), "Score-Based Generative Modeling through Stochastic Differential Equations". The VP SDE and its probability-flow ODE share the same marginals.

structure Generative.Diffusion.VPLinearSchedule (α : Type) [Context α] : Type

Continuous-time linear VP schedule on $t\in[0,1]$: $\beta(t)=\beta_0+t(\beta_1-\beta_0)$.

• beta0 : α

  $\beta(0)$.

• beta1 : α

  $\beta(1)$.

def Generative.Diffusion.VPLinearSchedule.beta {α : Type} [Context α] (sch : VPLinearSchedule α) (t : α) : α

Linear interpolation $\beta(t)$ on $t\in[0,1]$.

def Generative.Diffusion.VPLinearSchedule.alphaBar {α : Type} [Context α] (sch : VPLinearSchedule α) (t : α) : α

Closed-form $\bar\alpha(t)$ for the VP SDE with linear $\beta(t)$:

$$ \bar\alpha(t) =\exp\!\left(-\int_0^t\beta(s)\,ds\right) =\exp\!\left[-\left(\beta_0t+\tfrac12(\beta_1-\beta_0)t^2\right)\right]. $$

def Generative.Diffusion.VPLinearSchedule.sigma {α : Type} [Context α] (sch : VPLinearSchedule α) (t : α) : α

$\sigma(t)=\sqrt{1-\bar\alpha(t)}$ (clamped to stay total).

def Generative.Diffusion.pfOdeRhs {α : Type} [Context α] {s : Spec.Shape} (sch : VPLinearSchedule α) (model : EpsModel α s) (x : Spec.Tensor α s) (t : α) : Spec.Tensor α s

Probability-flow ODE drift for a VP schedule, expressed via an $\varepsilon_\theta(x,t)$ model.

For VP SDE:

$$ dx=-\tfrac12\beta(t)x\,dt+\sqrt{\beta(t)}\,dW. $$

The probability-flow ODE is:

$$ dx=\left[-\tfrac12\beta(t)x-\tfrac12\beta(t)\operatorname{score}(x,t)\right]dt. $$

Using the $\varepsilon$-parameterization, an approximate score is $\operatorname{score}\approx-\hat\varepsilon/\sigma(t)$, so:

$$ dx=\left[ -\tfrac12\beta(t)x +\tfrac12\frac{\beta(t)}{\sigma(t)}\hat\varepsilon(x,t) \right]dt. $$

def Generative.Diffusion.eulerStep {α : Type} [Context α] {s : Spec.Shape} (f : Spec.Tensor α s → α → Spec.Tensor α s) (x : Spec.Tensor α s) (t dt : α) : Spec.Tensor α s

One explicit Euler step for an ODE $x'=f(x,t)$:

$x_{\mathrm{next}}=x+dt\,f(x,t)$.

To integrate the probability-flow ODE backwards from $t=1$ to $t=0$, use a negative $dt$.

def Generative.Diffusion.pfOdeSampleEuler {α : Type} [Context α] {s : Spec.Shape} (sch : VPLinearSchedule α) (model : EpsModel α s) (steps : ℕ) (x1 : Spec.Tensor α s) : Spec.Tensor α s

Deterministic probability-flow sampler using Euler integration on a uniform grid.

Inputs:

• steps: number of Euler steps (typically large, e.g. 1000),
• x1: initial state at $t=1$ (typically standard normal noise).

We integrate backwards in time on the grid: $t_i=1-i/\mathtt{steps}$, with $dt=-1/\mathtt{steps}$.

def Generative.Diffusion.pfOdeSampleEuler.loop {α : Type} [Context α] {s : Spec.Shape} (sch : VPLinearSchedule α) (model : EpsModel α s) (n : ℕ) (dt : α) : ℕ → Spec.Tensor α s → Spec.Tensor α s

noncomputable def Generative.Diffusion.pfOdeEulerSystem {s : Spec.Shape} (sch : VPLinearSchedule Spec.SpecScalar) (model : EpsModel Spec.SpecScalar s) (t dt : Spec.SpecScalar) : NN.Spec.Dynamics.DynamicalSystem s

Real-valued probability-flow Euler step as a DynamicalSystem.

This is the formal hook used by trajectory/fixed-point/contraction lemmas in NN.Spec.Dynamics.System: at a fixed time and step size, Euler integration is an autonomous discrete update on the current sample.

@[simp]

theorem Generative.Diffusion.pfOdeEulerSystem_step {s : Spec.Shape} (sch : VPLinearSchedule Spec.SpecScalar) (model : EpsModel Spec.SpecScalar s) (t dt : Spec.SpecScalar) (x : Spec.SpecTensor s) : (pfOdeEulerSystem sch model t dt).step x = eulerStep (pfOdeRhs sch model) x t dt