Flow and Diffusion Models

Chapter 1: Vector Fields

A vector field is a function that assigns a velocity vector to every point in space and time:

At position x and time t, the vector field u_t(x) tells you the velocity — which direction and how fast to move. Think of it as a wind map: at every location and every moment, there is an arrow showing which way the wind blows and how strongly.

Worked example — linear vector field. Consider u_t(x) = −θx for θ = 2. This field always points toward the origin. At x = 3, the velocity is −6 (move left). At x = −1, the velocity is +2 (move right). Every particle is pulled toward zero, and the further away you are, the stronger the pull.

Worked example — circular vector field. In 2D, u_t(x₁, x₂) = (−x₂, x₁) creates circular trajectories. At (1, 0) the velocity is (0, 1) — pointing upward. At (0, 1) the velocity is (−1, 0) — pointing left. Particles orbit the origin.

Worked example — saddle point vector field. In 2D, u_t(x₁, x₂) = (θ x₁, −θ x₂) creates a saddle point. At (1, 1): velocity is (2, −2) — pushed right and up. The x₁ axis is unstable (grows), the x₂ axis is stable (shrinks). Particles spread along x₁ and compress along x₂.

Mathematical properties of vector fields. A vector field u_t(x) maps every point in R^d to a velocity vector in R^d. For our purposes, the vector field depends on both position x and time t. This time-dependence is crucial — the neural network needs to know "what time it is" to provide the correct velocity. At t = 0 (noise), the field should push particles apart. At t = 1 (data), it should guide them to the right locations.

python
# Defining vector fields in code
import torch

def linear_vf(x, theta=2.0):
    """u(x) = -theta * x (pulls toward origin)"""
    return -theta * x

def circular_vf(x, theta=1.0):
    """u(x1,x2) = (-theta*x2, theta*x1) (rotation)"""
    return theta * torch.stack([-x[...,1], x[...,0]], dim=-1)

# Neural network vector field (what we actually use)
class NeuralVF(torch.nn.Module):
    def forward(self, x, t):
        """x: [batch, d], t: [batch, 1] -> velocity [batch, d]"""
        return self.net(torch.cat([x, t], dim=-1))

Chapter 2: ODEs and Flows

An ordinary differential equation (ODE) says: "follow the vector field." Given a starting point x₀, the ODE asks for a trajectory X_t whose velocity at every moment equals the vector field:

The first equation says the trajectory's velocity must match the vector field. The second equation says we start at x₀. Together, they fully determine the trajectory — there is exactly one path through each starting point.

The solution to an ODE is called a flow, denoted ψ_t:

The flow ψ_t(x₀) tells you: "if you start at x₀ and follow the vector field for time t, where do you end up?" The trajectory is X_t = ψ_t(X₀).

Worked example — linear ODE. Let u_t(x) = −θx for θ = 2. The ODE is dX/dt = −2X. The solution is:

Verification: At t = 0, ψ₀(x₀) = e⁰ · x₀ = x₀. Check. The derivative is dψ/dt = −2e^−2tx₀ = −2ψ_t(x₀) = u_t(ψ_t(x₀)). Check. So the flow satisfies the ODE.

Numerical example. Start at x₀ = 4, θ = 2:

The particle decays exponentially toward zero. The velocity decreases as the particle gets closer to the origin.

Moreover, the flow ψ_t is a diffeomorphism — a smooth, invertible transformation with a smooth inverse. This means the flow "warps" space without tearing or folding it. Think of it as stretching a rubber sheet: every point maps to exactly one other point, and you can always undo the deformation.

Why diffeomorphisms matter: The flow maps the noise distribution N(0, I) to the data distribution p_data. Because the map is invertible, we can also go backwards: given a data point z, we can find the noise point ψ₁⁻¹(z) that generated it. This invertibility is useful for computing likelihoods and for encoding data into a latent space.

Worked example — 2D rotation flow. For the circular vector field u(x₁, x₂) = (−x₂, x₁), the flow is a rotation:

At t = π/2: the flow rotates everything 90 degrees counterclockwise. At t = π: 180 degrees. The inverse is simply rotating in the opposite direction.

Chapter 3: The Euler Method

We can analytically solve simple ODEs like dX/dt = −θX. But when the vector field u_t(x) is a neural network with millions of parameters, there is no formula for the flow. We need to simulate the ODE numerically.

The Euler method is the simplest simulation technique. The idea: take small discrete steps along the vector field. Starting at X₀ = x₀, update iteratively:

where h = 1/n is the step size and n is the number of steps. We evaluate the velocity at the current position, multiply by the step size, and add it to get the new position. After n steps, we reach t = 1.

Worked example — Euler for dX/dt = −2X. Let X₀ = 4, n = 4 steps, so h = 0.25:

Euler gives X₁ ≈ 0.250. The exact answer is 4e⁻² ≈ 0.541. The error is significant because h = 0.25 is coarse. With n = 100 steps (h = 0.01), Euler gives X₁ ≈ 0.536 — much closer.

Error analysis. The Euler method has a local truncation error of O(h²) per step and a global error of O(h). This means:

Heun's method has O(h²) global error but costs 2 NFE per step. So 25 Heun steps = 50 NFE, comparable to 50 Euler steps. In practice, Heun often gives better results at the same compute budget.

Heun's method (a refinement): Instead of blindly trusting the velocity at the current point, Heun's method takes a trial step, evaluates the velocity at the trial point, and averages the two velocities:

This costs 2 function evaluations per step but is much more accurate.

python
def euler_sample(u_theta, n_steps=50, d=2):
    """Sample from a flow model using Euler method."""
    h = 1.0 / n_steps
    x = torch.randn(d)  # X_0 ~ N(0, I)
    t = 0.0
    for _ in range(n_steps):
        v = u_theta(x, t)  # evaluate neural network
        x = x + h * v       # Euler step
        t += h
    return x  # X_1 ~ p_data

Chapter 4: Flow Models

We now have everything needed to define a flow model — a generative model based on an ODE. The recipe:

The ODE is deterministic — given the same starting noise X₀, you always get the same output X₁. The randomness comes entirely from the random initialization. Different noise samples produce different outputs.

What neural architecture is used? The neural network u_t^θ must take a noisy input x and a time t and output a velocity of the same shape. Common architectures:

The time t is typically embedded using sinusoidal or learned embeddings, similar to positional encodings in Transformers. This time embedding is added to intermediate activations so the network "knows" what stage of the denoising process it is in. At t = 0, the input is pure noise, so the network should output a large velocity. At t = 1, the input is nearly clean, and the velocity should be small (just fine-tuning details).

Why the initial distribution matters. We usually choose p_init = N(0, I_d) because:

1. We can sample from it easily (just generate random numbers).

2. It has full support on R^d — every point has positive probability.

3. It is isotropic — no preferred direction, so the model does not need to "undo" any structure in the noise.

4. It is well-studied mathematically, making theoretical analysis tractable.

Could we use a different p_init? Yes! Some models use a uniform distribution on a hypersphere, or a truncated Gaussian. But N(0, I) is the simplest and most common choice. The choice of p_init affects the difficulty of the learning problem: if p_init is very different from p_data, the flow must perform a more dramatic transformation, which requires more neural network capacity and more ODE steps.

Computational cost of generation. Each Euler step requires one forward pass through the neural network. The total cost of generating one sample is n × (cost of one forward pass). For reference:

Research in "distillation" aims to reduce the number of steps needed. Consistency models and progressive distillation can achieve good quality in 1-4 steps, at the cost of additional training. This is an active area of research as of 2025-2026.

Latent vs. pixel space. In practice, flow/diffusion models rarely operate directly in pixel space (too high-dimensional). Instead, a VAE first encodes images into a lower-dimensional latent space. The dimensions given above are for the latent representation. The actual generation pipeline is: noise → ODE/SDE in latent space → VAE decoder → pixels.

Classifier-free guidance. One of the most impactful practical techniques is not part of the basic theory but uses the score function. By training the model to sometimes drop the conditioning (text prompt), it learns both conditional and unconditional generation. At inference, the two predictions are combined to amplify the conditioning signal. This dramatically improves text-to-image alignment. We will see the mathematical details in Chapter 5 of the book, but it relies on the score function concepts from Chapter 4.

The Euler method vs. higher-order solvers. Beyond Euler and Heun, there are more advanced ODE solvers like DPM-Solver, DPM-Solver++, and UniPC that take advantage of the specific structure of the flow matching ODE. These can achieve the same quality as 50 Euler steps in just 10-20 steps by using higher-order approximations and specialized scheduling. In production systems, these advanced solvers are standard.

Step schedule matters. Instead of using uniform step sizes h = 1/n, production models use non-uniform schedules. The key insight: more steps should be allocated where the vector field changes rapidly. For CondOT paths, the field is nearly constant, so uniform steps work well. For VP (DDPM-style) paths, more steps are needed near t = 1 where fine details emerge.

python
# Non-uniform step schedules
import torch

# Uniform schedule (basic)
t_uniform = torch.linspace(0, 1, n_steps + 1)

# Quadratic schedule (more steps near t=1)
t_quad = torch.linspace(0, 1, n_steps + 1) ** 2

# Karras schedule (empirically optimized)
sigma_min, sigma_max = 0.002, 80
rho = 7
ramp = torch.linspace(0, 1, n_steps + 1)
t_karras = (sigma_max ** (1/rho) + ramp * (sigma_min ** (1/rho) - sigma_max ** (1/rho))) ** rho

The choice of step schedule can improve FID by 10-30% at the same number of steps. This is one of the most impactful "free" improvements in practice.

Adaptive step sizes. Some ODE solvers (like Dopri5) automatically choose the step size based on a local error estimate. They take larger steps where the vector field is smooth and smaller steps where it changes rapidly. This can achieve the same accuracy as 50 Euler steps in just 15-20 adaptive steps, with no manual tuning of the schedule.

Implementation in practice. The torchdiffeq library provides ODE solvers that work with PyTorch:

python
# Using torchdiffeq for adaptive ODE solving
from torchdiffeq import odeint

class FlowODE(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, t, x):
        """Required signature for torchdiffeq: (t, x) -> dx/dt"""
        t_batch = t.expand(x.shape[0], 1)
        return self.model(x, t_batch)

# Sample with adaptive solver
x0 = torch.randn(16, d)  # batch of 16 noise samples
t_span = torch.tensor([0.0, 1.0])
x1 = odeint(FlowODE(model), x0, t_span,
    method='dopri5', atol=1e-5, rtol=1e-5)[-1]

The adaptive solver automatically determines how many function evaluations (NFE) are needed, typically 15-40 depending on the vector field complexity. This often beats fixed-step Euler at the same compute budget.

Numerical stability considerations. When implementing ODE/SDE solvers for generative models:

1. Avoid t = 0 and t = 1 exactly. Some vector fields diverge at the endpoints. Use t ∈ [ε, 1−ε] with ε = 10⁻⁵.

2. Use float32 or higher. Float16 can cause numerical issues in the early steps (t near 0) where the velocity is large.

3. Gradient clipping during training. Clip gradients to norm 1.0 to prevent training instabilities.

4. EMA of model weights. Use exponential moving average of parameters for sampling (decay 0.9999).

A note on continuous vs. discrete time. In this course, we treat time as continuous (t ∈ [0, 1]). The original DDPM paper used discrete time steps (t ∈ {0, 1, ..., T}). The continuous formulation is more elegant and general — the discrete version is recovered as a special case when you discretize. Most modern implementations use continuous time internally and only discretize for the Euler/Heun solver.

Forward process vs. reverse process. In the DDPM literature, the "forward process" adds noise to data (going from t = 0 to t = 1) and the "reverse process" removes noise (going from t = 1 to t = 0). In the flow matching literature, the convention is reversed: t = 0 is noise and t = 1 is data. Both conventions are valid; just be careful when reading papers to check which convention is used.

In this course, we follow the flow matching convention: t = 0 is noise, t = 1 is data. The ODE flows from t = 0 to t = 1 during sampling. During training, we sample random t values and compute the loss at that timestep.

The complete lifecycle of a generated image. To tie everything together, here is what happens end-to-end when you type "a sunset over the ocean" into Stable Diffusion 3:

Total time: ~3 seconds on a modern GPU. Total neural network forward passes: 2 × 28 = 56 (two per step for classifier-free guidance). The model has ~2 billion parameters and was trained on ~1 billion images for ~500,000 GPU-hours. This is the machinery that Ch 3 and Ch 4 will teach you to build.

The mathematical elegance. Step back and appreciate the framework we have built. The entire generative modeling pipeline reduces to: (1) choose a vector field parametrized by a neural network, (2) simulate a differential equation. That is it. The complexity of generating images, videos, proteins, and music is encoded in the neural network weights, which are learned from data using the simple regression losses of Ch 3-4. The differential equation framework provides the mathematical guarantee that this process is well-defined (existence and uniqueness of solutions) and produces valid probability distributions (via the continuity/Fokker-Planck equations). This marriage of deep learning and classical mathematics is what makes flow and diffusion models both practical and beautiful.

Key equations to remember from this chapter:

Summary of ODE/SDE solvers for generative models:

Worked example — Algorithm 1 by hand. Suppose d = 1, n = 4 steps, h = 0.25, and X₀ = 1.5 (sampled from N(0,1)). Suppose the neural network outputs these velocities:

The output X₁ = 0.600 is our generated sample. The noise X₀ = 1.5 was transformed into the data sample 0.600 by the learned vector field.

python
# Algorithm 1: Sampling from a Flow Model
def sample_flow_model(u_theta, n=50, d=784):
    t = 0.0
    h = 1.0 / n
    x = torch.randn(d)          # X_0 ~ N(0, I)
    for i in range(n):
        x = x + h * u_theta(x, t)  # Euler step
        t += h
    return x                      # X_1 ~ p_data

Chapter 5: Brownian Motion

Flow models are deterministic — once you fix the noise X₀, the trajectory is fully determined. Diffusion models add stochasticity during the trajectory itself, using a fundamental mathematical object called Brownian motion.

A Brownian motion W = (W_t)_{0 ≤ t ≤ 1} is a continuous random walk. Think of it as a drunk person stumbling around — at every instant, they take a tiny random step. It has three defining properties:

Worked example — simulating Brownian motion. In 1D, with step size h = 0.01, we simulate by setting W₀ = 0 and updating:

With h = 0.01: √h = 0.1. At each step, we add a Gaussian with standard deviation 0.1. After 100 steps, we have a Brownian motion path from t = 0 to t = 1.

Numerical trace. Let ε₁ = 0.83, ε₂ = −1.21, ε₃ = 0.47:

The path zigzags randomly. Run the simulation again and you get a completely different path. This randomness is the ingredient that makes diffusion models stochastic.

Statistics of Brownian motion. Key properties that follow directly from the definition:

Numerical check. Simulate 10,000 Brownian paths to t = 1. The empirical distribution of W₁ should be approximately N(0, 1). Mean ≈ 0, Std ≈ 1. This is because W₁ − W₀ ~ N(0, 1) by the normal increments property.

python
# Simulating Brownian motion
import torch

def simulate_brownian(n_paths=5, n_steps=200, d=1):
    h = 1.0 / n_steps
    W = torch.zeros(n_paths, d)
    paths = [W.clone()]
    for _ in range(n_steps):
        W = W + (h ** 0.5) * torch.randn(n_paths, d)
        paths.append(W.clone())
    return torch.stack(paths)  # [n_steps+1, n_paths, d]

Chapter 6: Stochastic Differential Equations

An SDE extends an ODE by adding a Brownian motion term. At each step, the particle follows the vector field and gets a random kick:

The first term u_t(X_t) dt is the drift — the deterministic part, same as an ODE. The second term σ_t dW_t is the diffusion — the random part, scaled by the diffusion coefficient σ_t. When σ_t = 0, the SDE reduces to an ODE.

We simulate SDEs using the Euler-Maruyama method, the stochastic analog of the Euler method:

Compare to Euler for ODEs: X_t+h = X_t + h · u_t(X_t). The only difference is the added noise σ_t√h · ε_t.

Worked example — Ornstein-Uhlenbeck process. The OU process uses u_t(x) = −θx and constant σ:

The drift −θx pushes the particle back toward zero (a "spring"). The diffusion σ adds noise. The two forces balance: the particle bounces around zero, eventually settling into a Gaussian distribution N(0, σ²/(2θ)).

Numerical example. Let θ = 2, σ = 1, X₀ = 3, h = 0.1:

Notice how the trajectory is jagged (unlike the smooth ODE trajectory) because of the random noise at each step.

The √h scaling is crucial. Why does the noise scale as σ√h rather than σ h? Because Brownian increments have variance proportional to the time gap. If we used σ h, the total variance after n steps would be n · (σh)² = σ²h, which vanishes as h → 0. With σ√h, the total variance is n · (σ√h)² = n · σ²h = σ², which is finite. This is the mathematically correct scaling for continuous-time stochastic processes.

Convergence of the OU process. The Ornstein-Uhlenbeck process converges to the stationary distribution N(0, σ²/(2θ)). For our example with θ = 2, σ = 1: the stationary distribution is N(0, 0.25). Starting at X₀ = 3, the process decays toward zero (due to the drift −2X) and fluctuates around zero (due to the noise), eventually settling into N(0, 0.25).

python
# Euler-Maruyama method for SDE simulation
def euler_maruyama(u_theta, sigma_t, n=50, d=2):
    h = 1.0 / n
    x = torch.randn(d)
    t = 0.0
    for _ in range(n):
        eps = torch.randn_like(x)
        x = x + h * u_theta(x, t) + sigma_t(t) * (h**0.5) * eps
        t += h
    return x

Chapter 7: Diffusion Models

A diffusion model is a generative model based on an SDE, just as a flow model is based on an ODE. The recipe is identical except we add noise during sampling:

The neural network u_t^θ parameterizes the vector field (exactly as in a flow model). The diffusion coefficient σ_t is a fixed schedule — not learned. The goal is still X₁ ~ p_data.

Why would we want noise during sampling? Two reasons:

1. Error correction. Neural networks are imperfect — u_t^θ only approximates the true vector field. The added noise can help "shake" trajectories out of error states, similar to how simulated annealing helps optimization escape local minima.

2. Diversity. Different amounts of noise can produce more diverse samples. The same initial noise X₀ can lead to different outputs depending on the Brownian motion realization.

3. Theoretical guarantees. Under certain conditions, SDE sampling can converge faster than ODE sampling, especially when the target distribution has many well-separated modes.

Worked example — comparing ODE and SDE sampling. Suppose our target has two sharp modes at −5 and +5. With ODE sampling, a particle starting at X₀ = 0.1 always ends up at the same mode (say +5). With SDE sampling, the Brownian noise can "push" the particle across the boundary, allowing it to reach either mode. This can improve mode coverage.

python
# Algorithm 2: Sampling from a Diffusion Model
def sample_diffusion_model(u_theta, sigma, n=50, d=784):
    t = 0.0
    h = 1.0 / n
    x = torch.randn(d)                   # X_0 ~ N(0, I)
    for i in range(n):
        eps = torch.randn_like(x)          # fresh noise
        x = x + h * u_theta(x, t) \
            + sigma(t) * (h**0.5) * eps  # Euler-Maruyama
        t += h
    return x                              # X_1 ~ p_data

Chapter 8: ODE vs SDE Showcase

Let's see everything in action. The canvas below lets you watch flow models (ODE) and diffusion models (SDE) generate samples from a 2D distribution in real time. Multiple particles start as Gaussian noise and evolve toward the target data distribution.

Practical considerations for real models:

In practice, modern systems like Stable Diffusion 3 and FLUX use ODE sampling (flow matching) by default because it is faster. SDE sampling is used when higher quality is needed or when guidance techniques require it.

The key takeaway from this showcase: A flow model (ODE) and a diffusion model (SDE) are two ways to use the same trained neural network. The training is identical — only the sampling differs. This flexibility is one of the most powerful aspects of the flow/diffusion framework: train once, sample with either method depending on your quality/speed requirements.

Interpolation between ODE and SDE. You can even mix the two approaches. Start with SDE sampling (large σ) for the early steps (to get good mode coverage) and switch to ODE sampling (σ = 0) for the later steps (for precise, deterministic refinement). This "stochastic early, deterministic late" strategy is used by some state-of-the-art samplers.

python
# Complete comparison: ODE vs SDE sampling
def sample_ode(model, x0, n_steps=50):
    """Deterministic ODE sampling."""
    x, h = x0, 1.0 / n_steps
    for i in range(n_steps):
        t = i * h
        x = x + h * model(x, t)
    return x

def sample_sde(model, x0, sigma=0.5, n_steps=100):
    """Stochastic SDE sampling."""
    x, h = x0, 1.0 / n_steps
    for i in range(n_steps):
        t = i * h
        eps = torch.randn_like(x)
        x = x + h * model(x, t) + sigma * (h**0.5) * eps
    return x

# Same model, two sampling methods:
x0 = torch.randn(1, 784)
img_ode = sample_ode(model, x0)   # always same output for this x0
img_sde = sample_sde(model, x0)   # different each time

Chapter 9: Connections

We have built the complete mathematical machinery for generative modeling with differential equations. Let's recap:

What we know vs. what we need. Here is the gap between what we have and what we need:

Chapter 3 (flow matching) closes this gap with a beautifully simple idea: define a probability path from noise to data, compute the vector field that follows this path analytically, and train the neural network to match it via MSE regression. No ODE simulation during training. No adversarial training. Just regression.

The key equations of this entire course fit on one card:

python
# Summary: the full generative modeling stack

# 1. REPRESENTATION (Ch 1)
z = data_sample   # z ∈ R^d (image, video, protein, ...)

# 2. MACHINERY (Ch 2 — this chapter)
u_theta = NeuralNet(d, hidden)  # u: R^d × [0,1] → R^d
# ODE: dX/dt = u_theta(X, t)  → flow model
# SDE: dX = u_theta(X,t)dt + σ dW  → diffusion model
# Simulate via Euler or Euler-Maruyama

# 3. TRAINING (Ch 3-4 — next)
# Flow matching: L = ||u_theta(x,t) - (z - eps)||^2
# Score matching: L = ||s_theta(x,t) + eps/beta||^2

# 4. GENERATION
# X_0 ~ N(0, I) → simulate ODE/SDE → X_1 ~ p_data