ML

Continuous Normalizing Flows & Flow Matching

From discrete layers to continuous dynamics: neural ODEs for generation, the instantaneous change of variables, flow matching training, rectified flows, and optimal transport connections.

Advanced

Prerequisites

Table of Contents

  1. Learning Objectives
  2. Notation
  3. Core Intuition
  4. Neural ODEs for Generation
  5. Instantaneous Change of Variables
  6. Flow Matching Objective
  7. Conditional Flow Matching
  8. Rectified Flows
  9. Optimal Transport Connection
  10. Comparison with Diffusion
  11. Common Pitfalls
  12. Summary
  13. Exercises

Learning Objectives

  1. Derive the continuous-time change of variables (instantaneous form).
  2. Explain why flow matching avoids simulation during training.
  3. Derive the conditional flow matching objective.
  4. Connect rectified flows to optimal transport.
  5. Compare flow matching vs diffusion in terms of sample quality and speed.

Notation

  • vθ(x,t)\mathbf{v}_\theta(\mathbf{x}, t) — velocity field
  • ψt\psi_t — flow (diffeomorphism)
  • p0=N(0,I)p_0 = \mathcal{N}(0, \mathbf{I}), p1=pdatap_1 = p_{\text{data}}
  • OT — optimal transport

Core Intuition

Instead of stacking many discrete invertible layers, define a continuous transformation via an ODE: dx/dt=vθ(x,t)d\mathbf{x}/dt = \mathbf{v}_\theta(\mathbf{x}, t). Integrating from t=0t=0 to t=1t=1 maps noise to data. The velocity field vθ\mathbf{v}_\theta is learned to transport the noise distribution to the data distribution along efficient (ideally straight) paths.

Continuous Normalizing Flow

t=0.50 · log|det|≈-0.40
Time t
0.50
Base N(0,I)ODE trajectory
Explore: CNFs integrate dx/dt = f(x,t) from t=0 to 1. The instantaneous change of variables d(log p)/dt = −Tr(∂f/∂x) tracks density along the path.

Neural ODEs for Generation

Define the generative process as an ODE:

dψt(x)dt=vθ(ψt(x),t),ψ0(x)=x.(1)\frac{d\psi_t(\mathbf{x})}{dt} = \mathbf{v}_\theta(\psi_t(\mathbf{x}), t), \quad \psi_0(\mathbf{x}) = \mathbf{x}. \tag{1}

Sampling: Start from z∼N(0,I)\mathbf{z} \sim \mathcal{N}(0, \mathbf{I}), solve ODE forward to get x=ψ1(z)\mathbf{x} = \psi_1(\mathbf{z}).

Density: The log-probability evolves according to the instantaneous change of variables.

Instantaneous Change of Variables

dlog⁡pt(x)dt=−∇⋅vθ(x,t)=−tr(∂vθ∂x).(2)\frac{d\log p_t(\mathbf{x})}{dt} = -\nabla \cdot \mathbf{v}_\theta(\mathbf{x}, t) = -\text{tr}\left(\frac{\partial\mathbf{v}_\theta}{\partial\mathbf{x}}\right). \tag{2}

Log-likelihood:

log⁡p1(x1)=log⁡p0(x0)−∫01tr(∂vθ∂x)dt.(3)\log p_1(\mathbf{x}_1) = \log p_0(\mathbf{x}_0) - \int_0^1 \text{tr}\left(\frac{\partial\mathbf{v}_\theta}{\partial\mathbf{x}}\right)dt. \tag{3}

Problem: Computing the trace requires O(d)O(d) evaluations or Hutchinson's estimator — expensive during training.

Flow Matching Objective

Key insight: We don't need to simulate the ODE during training. Instead, directly regress the velocity field:

LFM=Et,pt(x)[∥vθ(x,t)−ut(x)∥2],(4)\mathcal{L}_{\text{FM}} = \mathbb{E}_{t, p_t(\mathbf{x})}\left[\|\mathbf{v}_\theta(\mathbf{x}, t) - \mathbf{u}_t(\mathbf{x})\|^2\right], \tag{4}

where ut\mathbf{u}_t is the target velocity field that generates the desired flow.

Problem: We don't know ptp_t or ut\mathbf{u}_t for the marginal flow.

Conditional Flow Matching

Solution: Condition on endpoints. For each pair (x0,x1)(\mathbf{x}_0, \mathbf{x}_1), define a simple conditional path:

ψt(x0∣x1)=(1−t)x0+tx1.(5)\psi_t(\mathbf{x}_0 | \mathbf{x}_1) = (1-t)\mathbf{x}_0 + t\mathbf{x}_1. \tag{5}

Conditional velocity: ut(x∣x1)=x1−x0\mathbf{u}_t(\mathbf{x}|\mathbf{x}_1) = \mathbf{x}_1 - \mathbf{x}_0.

Training:

LCFM=Et,x0∼p0,x1∼p1[∥vθ((1−t)x0+tx1,t)−(x1−x0)∥2].(6)\boxed{\mathcal{L}_{\text{CFM}} = \mathbb{E}_{t, \mathbf{x}_0 \sim p_0, \mathbf{x}_1 \sim p_1}\left[\|\mathbf{v}_\theta((1-t)\mathbf{x}_0 + t\mathbf{x}_1, t) - (\mathbf{x}_1 - \mathbf{x}_0)\|^2\right].} \tag{6}

Simulation-free: No ODE solving during training. Just sample tt, interpolate, predict velocity.

Rectified Flows

Straight-line paths: The interpolation (1−t)x0+tx1(1-t)\mathbf{x}_0 + t\mathbf{x}_1 defines straight lines in data space.

Reflow: Iteratively straighten trajectories:

  1. Train flow matching model vθ\mathbf{v}_\theta.
  2. Generate pairs: sample x0∼p0\mathbf{x}_0 \sim p_0, integrate ODE to get x1=ψ1(x0)\mathbf{x}_1 = \psi_1(\mathbf{x}_0).
  3. Retrain on the generated pairs (better coupling than random).
  4. After 2-3 iterations: nearly straight trajectories.

One-step generation: With perfectly straight flows, a single Euler step suffices: x1≈x0+vθ(x0,0)\mathbf{x}_1 \approx \mathbf{x}_0 + \mathbf{v}_\theta(\mathbf{x}_0, 0).

Optimal Transport Connection

Random coupling (x0,x1\mathbf{x}_0, \mathbf{x}_1 independent): Paths may cross, requiring curved trajectories.

OT coupling (minimize transport cost): Paths don't cross, trajectories are straighter.

Mini-batch OT: Approximate OT within each training batch by solving a linear assignment problem. Significantly improves training convergence and sample quality.

Comparison with Diffusion

  • Diffusion: Curved paths (VP-SDE), needs 20-50 ODE steps.
  • Flow matching: Straight-line paths, needs 5-20 Euler steps.
  • Quality: Comparable FID scores; flow matching slightly better at few steps.
  • Training: Both simulation-free; flow matching objective is simpler (no noise schedule).
  • Used in: Stable Diffusion 3, Flux, DALL-E 3.

Common Pitfalls

Pitfall 1. Random coupling without OT leads to crossing paths, which are hard to learn and require more steps.

Pitfall 2. Using too few Euler steps without reflow. Initial flow matching models still have curved trajectories; single-step quality is poor without reflow.

Pitfall 3. Assuming flow matching replaces diffusion. For highest quality, multi-step diffusion/flow methods are still superior; flow matching excels at the speed-quality tradeoff.

Summary

  • CNF: Continuous transformation via neural ODE; exact likelihood via trace formula.
  • Flow matching: Simulation-free training; regress velocity on interpolated points.
  • Conditional FM: Straight-line interpolation between noise and data.
  • Rectified flows: Iterative straightening for few-step generation.
  • OT coupling: Reduces path crossing, improves straightness.
  • Dominant in modern image/video generation (SD3, Flux).

Exercises

Exercise 1. Derive the instantaneous change of variables (equation 2) from the discrete formula.

Exercise 2. Show that conditional flow matching (equation 6) is equivalent to marginal flow matching (equation 4) in expectation.

Exercise 3. Compute the number of Euler steps needed for FID less than 5 on CIFAR-10 with and without reflow.

Exercise 4. Implement mini-batch OT: given 64 noise samples and 64 data samples, find the optimal pairing.

Exercise 5. Derive the velocity field for the VP-SDE interpolation xt=αˉtx1+1−αˉtx0\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_1 + \sqrt{1-\bar{\alpha}_t}\mathbf{x}_0 and compare to the straight-line velocity.

Next

Coupling Flows: NICE, RealNVP & Glow