How Stochastic Differential Equations Power Modern Generative AI Models

Overview

This note summarizes the mathematical framework of the two most widely used generative‑AI algorithms: denoising diffusion probabilistic models (DDPM) and continuous flow‑matching models. Both approaches view generation as the evolution of a probability density under a differential equation that transforms a simple prior (usually a Gaussian) into the data distribution.

Diffusion and Flow Models

Let x_0 denote a data sample and z\sim\mathcal N(0, I) a latent variable drawn from a standard Gaussian. A flow model defines an ordinary differential equation (ODE) dx_t = v_θ(x_t, t)\,dt,\qquad t\in[0,1] where the time‑dependent vector field v_θ is parameterized by a neural network. Solving the ODE from t=0 (prior) to t=1 yields a sample x_1 approximating the data distribution.

A diffusion model defines a stochastic differential equation (SDE) dx_t = f(x_t, t)\,dt + g(t)\,dW_t,\qquad t\in[0,1] with drift f, diffusion coefficient g, and Brownian motion W_t. The forward SDE gradually adds Gaussian noise; the reverse‑time SDE, obtained by applying the Fokker‑Planck equation, removes noise and generates data.

Training Objectives and Sampling Algorithms

Both models are trained by minimizing the mean‑squared error between the neural network output and the true vector (or score) field that drives the desired probability path. For diffusion models the target is the score function ∇_x log p_t(x); for flow models it is the velocity field v(x,t).

Sampling from a flow model (Algorithm 1) consists of numerically integrating the ODE with a chosen solver (e.g., Euler or Runge‑Kutta) starting from a Gaussian sample. Sampling from an SDE (Algorithm 2) uses a discretized reverse‑time SDE, typically with the Euler‑Maruyama scheme.

Algorithm 1 (Flow Sampling)
1. Sample z∼𝒩(0,I)
2. Set x←z
3. For t from 0 to 1 (step size Δt):
   x←x+v_θ(x,t)·Δt
4. Return x

Algorithm 2 (Diffusion Sampling)
1. Sample z∼𝒩(0,I)
2. Set x←z
3. For t from 1 down to 0 (step size Δt):
   x←x+\big[f_θ(x,t)‑g(t)^2∇_x log p_t(x)\big]·Δt + g(t)·√Δt·ε,
   where ε∼𝒩(0,I)
4. Return x

Conditional Generation and Classifier‑Free Guidance

Conditional generation introduces a conditioning variable y (e.g., a text prompt). The joint model learns a conditional vector field v_θ(x,t|y) or conditional score ∇_x log p_t(x|y). During sampling, the conditioning is supplied, allowing the model to generate samples that satisfy the desired attribute.

Classifier‑free guidance replaces an explicit classifier with a weighted combination of conditional and unconditional networks:

v_guided(x,t|y)=v_uncond(x,t)+γ·(v_cond(x,t|y)‑v_uncond(x,t))

Increasing the guidance scale γ improves perceptual fidelity at the cost of diversity.

Network Architectures

Most diffusion models employ a U‑Net backbone. The encoder‑decoder structure preserves spatial resolution through skip connections, enabling the network to output a full‑resolution velocity or score field.

Recent work replaces convolutions with attention‑only blocks, yielding the Diffusion Transformer (DiT). DiT treats an image as a sequence of patches, processes them with a Vision‑Transformer, and has become the core of models such as Stable Diffusion 3.

Memory‑Efficient Latent Diffusion

High‑resolution generation (e.g., 1024×1024) creates millions of pixel dimensions, which quickly exhaust GPU memory. The standard remedy is to operate in a compressed latent space:

Train an auto‑encoder (E, D) that maps images x to low‑dimensional latents z=E(x) and reconstructs them via \hat{x}=D(z).

Train the diffusion or flow model on the latent distribution p(z).

At inference, sample a latent ẑ with the trained generative model and decode it with D.

Applications

Robotics

Robotic control can be cast as a diffusion process over command sequences. Each timestep predicts a short‑horizon command vector (e.g., 10 Hz joint torques). A conditional diffusion model receives sensory observations as conditioning and generates coherent trajectories. A safety layer buffers low‑level commands before execution, ensuring stability.

Protein Design

Diffusion on the SE(3) manifold enables generative modeling of protein structures. Three parameterizations are common:

Direct 3‑D atomic coordinates.

Torsion angles (fixed bond lengths, variable angles).

Backbone‑centered frames (rigid‑body frames for each residue).

Training proceeds by adding isotropic Gaussian noise to the chosen representation, learning a score network that predicts the denoising direction, and then sampling from pure noise. Combining latent diffusion with a pretrained structure predictor (e.g., AlphaFold‑like) yields long, high‑quality protein sequences.

Future Outlook

The authors anticipate that diffusion‑based generative models will soon dominate molecular design, extending beyond proteins to small molecules and other chemical entities. Iterative loops that incorporate experimental feedback and ever larger pretrained backbones are expected to accelerate progress in both AI‑driven biology and robotics.