Normalizing flows and flow matching

Problem

A generative model should do two things: draw new samples, and say how likely a given sample is. VAEs give a bound on the likelihood, not the likelihood itself; diffusion models sample well but need hundreds of network calls, and their likelihood is a bound unless you solve an ODE.

A normalizing flow makes the likelihood exact by construction. Pick a base distribution you can sample and score, usually a standard normal, and an invertible, differentiable map f. Sample by pushing noise through f; score a data point by pulling it back through f⁻¹ and correcting for how much f stretched space. The hard part is building maps that are expressive, invertible and have a cheap Jacobian determinant. Continuous flows sidestep the architecture constraints by defining f as the solution of an ODE, and flow matching is the training objective that made them practical at scale.

Intuition

Picture probability as sand. The base distribution is a smooth pile; the data distribution is the shape you want. A flow is a set of instructions for moving every grain, and because grains are neither created nor destroyed, wherever the map squeezes space the density goes up by exactly the factor it squeezed. That bookkeeping is the change-of-variables formula.

A continuous flow gives each grain a velocity at every moment instead of one big jump. Flow matching learns that velocity field by regression: take a noise grain and a data grain, draw the straight line between them, and teach the network that a grain at a point on the line at time t should be moving along it. Many lines pass through the same point, so the network learns their average direction, and following those averaged directions turns out to carry the whole noise pile onto the whole data pile.

Mechanics

Change of variables

If z has density p_Z and y = f(z) with f invertible and differentiable, then

log p_Y(y) = log p_Z(f⁻¹(y)) − log |det ∂f/∂z|, evaluated at z = f⁻¹(y)

The determinant measures how much f expands a small volume around z; expanding spreads the same probability over more space, so the density drops. A composition f = f_K ∘ … ∘ f₁ adds one log-determinant per step, which is why flows are built from many simple layers. PyTorch implements this directly: a Transform provides the forward map, its inverse and log_abs_det_jacobian(x, y), and TransformedDistribution.log_prob walks the transforms in reverse, inverting each and subtracting its log-Jacobian before scoring the base. LogNormal is literally TransformedDistribution(Normal(loc, scale), ExpTransform()).

For a general D-dimensional map, the determinant costs O(D³). Normalizing-flow research is largely the search for layers whose Jacobians are triangular, so the determinant is the product of the diagonal.

Coupling layers

A coupling layer splits the input into halves x_a and x_b, leaves x_a alone and transforms x_b with parameters computed from x_a. The affine coupling of RealNVP is

y_a = x_a,   y_b = x_b ⊙ exp(s(x_a)) + t(x_a),   log |det J| = Σ s(x_a)

The Jacobian is block triangular with an identity block and a diagonal block exp(s), so its log-determinant is just the sum of s. Inverting needs no inverse of the networks s and t: x_b = (y_b − t(y_a)) ⊙ exp(−s(y_a)), because y_a = x_a. s and t can be any networks. Stacking layers that alternate which half is left alone lets every dimension be transformed. NICE used the additive special case (s = 0), which preserves volume; Glow added learned invertible 1×1 convolutions to mix channels between couplings. Autoregressive flows generalize the idea with one conditioner per dimension.

Continuous flows

Make the layers infinitely thin. A continuous normalizing flow defines the map as the solution of an ODE with a learned velocity field v_θ:

dx/dt = v_θ(x, t),   x(0) ~ p₀ (noise),   x(1) = the sample

Any reasonable v_θ gives an invertible map (run the ODE backwards), with no architectural constraint. The log density changes along a trajectory by the instantaneous change of variables from the Neural ODE paper:

d log p(x(t)) / dt = − tr(∂v_θ/∂x) = − div v_θ

Computing the trace exactly costs about D backward passes, so FFJORD estimated it with Hutchinson's trick, εᵀ (∂v/∂x) ε for random ε. Training these models by maximum likelihood still meant solving the ODE inside every training step, which was slow. Flow matching removes that.

The flow-matching objective

Choose a probability path p_t from noise p₀ to data p₁, and regress v_θ onto the velocity u_t that generates it. That velocity is unknown, but it is easy for one data point at a time. With the straight-line path

x_t = (1 − t) x₀ + t x₁,   x₀ ~ N(0, I), x₁ ~ data, t ~ U[0, 1]

each pair moves at constant velocity x₁ − x₀. The conditional flow-matching loss regresses onto that per-pair target:

L_CFM(θ) = E_{t, x₀, x₁} ‖ v_θ(x_t, t) − (x₁ − x₀) ‖²

The minimizer of a squared loss is a conditional mean, so the best v_θ is v*(x, t) = E[x₁ − x₀ | x_t = x], the average direction of all lines through x at time t. Lipman et al. show that this marginal velocity generates the marginal path p_t, and that the conditional loss has the same gradients as regressing on the intractable marginal velocity directly. The loss never reaches zero: what is left is the variance of x₁ − x₀ among the lines through each point. Lipman et al. write the path with a small floor, x_t = (1 − (1 − σ_min) t) x₀ + t x₁, calling it the optimal-transport conditional path; σ_min = 0 gives the line above. Training needs no ODE solve, no simulation and no noise schedule: one network call and one squared error per example, like a denoising loss.

Rectified flow and reflow

Rectified flow (Liu et al.) uses the same straight-line interpolation and loss. Its focus is what happens to the paths. Because the pairs (x₀, x₁) are drawn independently, their lines cross, and at a crossing the learned velocity averages the two directions. So the ODE trajectories of the trained model do not cross (an ODE's solutions never do) and are curved, and a coarse solver following them has to take many steps. Reflow fixes this: sample noise x₀, run the trained ODE to get its endpoint T(x₀), and train a new flow on the pairs (x₀, T(x₀)). Those pairs already come from a flow whose paths do not cross, so the new lines cross less, and the new model's paths are straighter. A perfectly straight flow is solved exactly by a single Euler step, which is why reflow, followed by distillation, gives few-step and one-step samplers. Stable Diffusion 3 trains its image model with the rectified-flow objective.

Relation to diffusion

Diffusion corrupts data with x_t = √ᾱ x₀ + √(1 − ᾱ) ε, with data at the start of time; flow matching writes x_t = t x₁ + (1 − t) x₀ with data at the end. Both are Gaussian paths x = α · data + σ · noise, and only α and σ differ: dividing the flow's x_t by √(t² + (1 − t)²) gives a diffusion point with ᾱ = t² / (t² + (1 − t)²). The optimal predictions are linked too. With x̂₁ = E[x₁ | x_t] the predicted data and x̂₀ = E[x₀ | x_t] the predicted noise,

v*(x_t, t) = x̂₁ − x̂₀ = (x̂₁ − x_t) / (1 − t)

so a velocity network, a denoiser and a noise predictor are three parametrizations of one function. Every diffusion model also has a deterministic probability-flow ODE (Song et al.) with the same marginals as its stochastic sampler, and DDIM's update is a first-order solver for it written in other variables. Flow matching picks the straight path and the velocity target directly, which in practice gives straighter trajectories and a simpler training recipe.

Sampling with an ODE solver

Sampling integrates dx/dt = v_θ(x, t) from t = 0 to 1. Euler takes N steps of x ← x + h · v(x, t) with h = 1/N, one network call each. Heun (a second-order Runge-Kutta method) takes an Euler step, evaluates the velocity at its end, and moves by the average of the two slopes, two calls per step and error falling as h² instead of h. Adaptive solvers such as Dormand-Prince pick their own step sizes from an error estimate. PyTorch has no ODE solver in the core library; the loop is a few lines, as in the implementation below, and libraries such as torchdiffeq provide adaptive ones. Straighter paths need fewer steps, which is the practical case for the straight-line path and for reflow. Likelihoods need the trace integral above solved alongside x, which is why flow-matching models are mostly used for sampling.

Worked example

Implementation

import torch
from torch import nn
from torch.distributions import Normal, TransformedDistribution
from torch.distributions.transforms import AffineTransform, ExpTransform

# A discrete flow: change of variables through two transforms.
flow = TransformedDistribution(Normal(0.0, 1.0), [AffineTransform(0.2, 0.5), ExpTransform()])
flow.log_prob(torch.tensor(1.5))   # −0.7157, the same as LogNormal(0.2, 0.5)

# A continuous flow trained with flow matching.
D = 2
net = nn.Sequential(nn.Linear(D + 1, 256), nn.SiLU(),
                    nn.Linear(256, 256), nn.SiLU(),
                    nn.Linear(256, D))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

def velocity(x, t):                       # x: (B, D), t: (B, 1)
    return net(torch.cat([x, t], dim=-1))

def train_step(x1):                       # x1: a batch of data, (B, D)
    x0 = torch.randn_like(x1)             # noise
    t = torch.rand(x1.shape[0], 1)        # one time per example
    xt = (1 - t) * x0 + t * x1            # point on the straight line
    loss = ((velocity(xt, t) - (x1 - x0)) ** 2).mean()
    opt.zero_grad()
    loss.backward()
    opt.step()
    return loss.item()

@torch.no_grad()
def sample(n, steps=32):                  # Euler from noise (t = 0) to data (t = 1)
    x = torch.randn(n, D)
    for k in range(steps):
        t = torch.full((n, 1), k / steps)
        x = x + velocity(x, t) / steps
    return x

The training step has no ODE in it: the only randomness is the noise, the time and the batch, and the target is a subtraction. Swapping the Euler loop for a Heun loop, or for an adaptive solver from a library, changes only sample. For reflow, run sample on a fixed batch of noise, keep each (noise, output) pair, and train a fresh net with x0 and x1 taken from those pairs instead of drawn independently.

Tradeoffs

• Straight paths buy cheap sampling. The number of solver steps is set by how curved the paths are, not by the network. Reflow straightens them at the cost of generating a dataset of pairs and training again, and each round adds some error of its own.
• The path is a design choice. The straight line and diffusion's variance-preserving path describe the same family with different α and σ; which works better also depends on how t is sampled during training, for which Stable Diffusion 3 reports gains from weighting the middle of the path.
• Exact likelihood is rarely the point. Discrete flows are used where an exact density matters, and as components of other models; for sample quality, large image models have moved to diffusion and flow matching.

Connections

Diffusion models are the closest relative: the same Gaussian paths in other coordinates, and DDIM is an ODE sampler like the ones here. VAEs also map between noise and data but approximate the likelihood, and flows have been used inside them as richer approximate posteriors. The training loss is ordinary regression, so everything on optimizers and backpropagation applies unchanged.

References

• Dinh, Krueger and Bengio. NICE: Non-linear Independent Components Estimation. 2014. arXiv:1410.8516
• Dinh, Sohl-Dickstein and Bengio. Density Estimation Using Real NVP. ICLR 2017. arXiv:1605.08803
• Kingma and Dhariwal. Glow: Generative Flow with Invertible 1x1 Convolutions. NeurIPS 2018. arXiv:1807.03039
• Chen, Rubanova, Bettencourt and Duvenaud. Neural Ordinary Differential Equations. NeurIPS 2018. arXiv:1806.07366
• Grathwohl et al. FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models. ICLR 2019. arXiv:1810.01367
• Song et al. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021. arXiv:2011.13456
• Song, Meng and Ermon. Denoising Diffusion Implicit Models. ICLR 2021. arXiv:2010.02502
• Lipman et al. Flow Matching for Generative Modeling. ICLR 2023. arXiv:2210.02747
• Liu, Gong and Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR 2023. arXiv:2209.03003
• Albergo and Vanden-Eijnden. Building Normalizing Flows with Stochastic Interpolants. ICLR 2023. arXiv:2209.15571
• Esser et al. Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. 2024. arXiv:2403.03206
• PyTorch v2.15.0: Transform, AffineTransform, ExpTransform, TransformedDistribution, LogNormal.