Probability-flow ODE (spec layer)

This file defines a small continuous-time VP schedule (linear β(t)) and the corresponding probability-flow ODE drift field, using an ε_θ(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 ∈ [0,1]: β(t) = β0 + t(β1-β0).

• beta0 : α

  β(0).

• beta1 : α

  β(1).

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

Linear interpolation β(t) on t ∈ [0,1].

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

Closed-form ᾱ(t) for the VP SDE with linear β(t):

ᾱ(t) = exp(-∫₀ᵗ β(s) ds) = exp(-(β0 t + 0.5 (β1-β0) t^2)).

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

σ(t) = sqrt(1-ᾱ(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 ε_θ(x,t) model.

For VP SDE:

dx = -0.5 β(t) x dt + sqrt(β(t)) dW

The probability-flow ODE is:

dx = (-0.5 β(t) x - β(t) score(x,t)) dt.

Using the ε-parameterization, an approximate score is score ≈ -(1/σ(t)) ε̂, so:

dx = (-0.5 β(t) x + (β(t)/σ(t)) ε̂(x,t)) 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):

xNext = 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/steps, with dt = -1/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