Flow Matching: The Training Objective Behind SD3 and FLUX
Diffusion training is a tower of Jenga blocks: pick a noise schedule, derive the ELBO across noise levels, switch to score parameterization, re-derive the loss as a noise-prediction MSE, then patch in classifier-free guidance, EDM preconditioning, v-prediction, and learned variance heads. Most of those blocks exist because the underlying object — the score of a noisy distribution — is hard to access directly. Flow matching (Lipman 2023) throws the tower out. There is one regression target — a velocity field along an interpolation path — and one loss — MSE between the model's prediction and that target. The objective fits on a single line. The framework was originally a niche alternative; in 2024 Stability and Black Forest Labs put it behind SD3 and FLUX, and it is now the default training objective for state-of-the-art image and video generators.
Learning Objectives
After this lesson, you will be able to:
Understand why flow matching exists — a simpler training objective for continuous normalizing flows than ELBO
Derive the Conditional Flow Matching (CFM) objective from the marginal-vs-conditional equivalence
Explain why Optimal Transport CFM uses straight-line paths and what that buys you
Connect flow matching to diffusion via the score-based view, and to Rectified Flow (used in SD3/FLUX)
Implement and run OT-CFM training in numpy on a toy 2D dataset
The generative-AI track has a lesson on flow matching that emphasizes the inference picture — straight paths, fewer ODE steps, faster sampling, the taxi-vs-helicopter intuition. This lesson is its complement. We stay on the training objective and treat flow matching as a different way to set up the supervised regression problem that a U-Net or DiT actually solves. If you have ever stared at the diffusion ELBO and asked "is there a simpler thing the network could regress against?", flow matching is the answer that 2025 settled on.
The math here is rigorous but the conceptual load is light: flow matching is just MSE against a target vector field. The hard part is choosing the path that defines the target. Everything in the next 6,000 words is decoration around that one sentence.
#The Premise: Diffusion Training Is Procedurally Heavy
Recall what diffusion training requires. You pick a noise schedule β_t (linear, cosine, sigmoid, EDM-style — there is no canonical choice). You derive a forward process q(x_t | x_0) whose marginals interpolate from data to Gaussian noise. You write down the ELBO across T noise levels, simplify it by training a noise-prediction network ε_θ(x_t, t), and discover that the loss reduces to a denoising MSE. You add v-prediction or x-prediction parameterizations because they have better SNR behavior at extreme noise. You add classifier-free guidance dropout. You add EDM preconditioning so the network sees inputs and targets with normalized scales. You sample from the trained model by running an SDE or a probability-flow ODE in reverse over 50–250 steps, with DPM-Solver, DPM-Solver++, EDM samplers, LCM distillation, and so on.
Every one of those choices is a research paper. Each piece was invented to compensate for a difficulty: the score ∇ log q(x_t) is hard to access, so you train against ε; the per-timestep loss has wildly different magnitudes at low vs high noise, so you reweight; the reverse-time SDE is hard to discretize, so you build dedicated solvers. The complexity is real, and most of it is architecture-independent. The same U-Net works under any of these recipes.
#CNF Refresher: ODEs, Vector Fields, and the Old Problem
Flow matching lives inside the framework of continuous normalizing flows (Chen 2018, Neural Ordinary Differential Equations). The setup is a time-dependent vector field
v_t(x): R^d × [0, 1] → R^d
that defines an ordinary differential equation dx/dt = v_t(x). Pick an initial point x_0 ~ p_0 (a base distribution, typically a standard Gaussian) and integrate forward in time:
x(t) = x_0 + integral_0^t v_s(x(s)) ds
The map ϕ_t: x_0 ↦ x(t) is a diffeomorphism — smooth, invertible, with a smooth inverse. Pushing the base density p_0 through ϕ_t gives a time-dependent density p_t, and by construction p_t is the law of x(t). If you train v_t such that the terminal density p_1 matches the data distribution, you have a generative model: sample x_0 ~ p_0, integrate the ODE forward, get x_1 ~ p_1 ≈ data.
This was the dream of Neural ODEs in 2018. The problem was training. The natural objective is maximum likelihood, but evaluating the likelihood under a CNF requires the change of variables formula:
The trace of the Jacobian inside the integral is the killer. For a d-dimensional flow with d = 49,152 (a 128×128×3 image flattened, or its VAE latent in SD3), computing that trace exactly costs O(d²) — infeasible. People tried Hutchinson estimators, divergence-tricks, augmented losses; the methods worked for small problems and were unreliable at image scale. CNFs trained by maximum likelihood never scaled. Six years later, flow matching is the workaround.
What Do You Think?
Why does training a CNF by maximum likelihood require the Jacobian trace ∂v/∂x?
#Lipman's Reformulation: Regression Instead of Likelihood
The trick in Lipman 2023, Flow Matching for Generative Modeling, is to side-step likelihood entirely. Instead of asking "which v_t makes the model density match the data density?" — a question that drags in the Jacobian trace — Lipman asks "if I specify the target density path p_t directly, what is the vector field u_t that generates it?" Then he trains the neural network v_θ to regress against u_t.
Specify a smooth probability path p_t(x) that interpolates between p_0 (base) and p_1 (data). The vector field u_t(x) that pushes p_0 to p_1 along this path is uniquely determined (up to a divergence-free perturbation, which we ignore). The flow-matching objective is
Compare this to the diffusion ELBO. There is no schedule β_t (the path p_t plays that role implicitly, and we will pick it explicitly in a moment). There is no ε vs v vs x parameterization debate (the target is one specific vector u_t(x)). There is no per-timestep reweighting (the expectation is uniform, though we will revisit that). It is just MSE.
The remaining question is: what is u_t(x)? Specifying a probability path p_t in closed form is easy (e.g., a Gaussian whose mean and variance vary smoothly with t). Specifying the marginal velocity field u_t(x) that pushes p_0 to p_1 along that path is, in general, intractable. To compute u_t(x), you would have to marginalize over all (x_0, x_1) pairs that could have produced x at time t, weight them by their joint density, and average their conditional velocities. For an image distribution this is hopeless.
This is where the second half of Lipman 2023 — and the actual practical algorithm — comes in.
#Conditional Flow Matching: The Trick That Made It Tractable
Suppose instead of trying to specify the marginal path p_t(x) directly, you specify a conditional path p_t(x | x_1) — the path that interpolates from the base noise distribution to a specific data point x_1. This is much easier. A conditional path can be a simple parameterized curve. The corresponding conditional vector field u_t(x | x_1) is then a closed-form expression you can write down on a napkin.
The marginal probability path is recovered by mixing over data:
p_t(x) = ∫ p_t(x | x_1) p_data(x_1) dx_1
The marginal velocity field is the conditional-expectation:
Then ∇_θ L_CFM(θ) = ∇_θ L_FM(θ). The two objectives differ by a constant in θ, so training against the conditional loss optimizes the same parameters as training against the (intractable) marginal loss.
The proof is a one-page computation using the tower property of conditional expectation: under squared loss, regressing against a target equals regressing against the conditional expectation of that target. Since u_t(x) is by definition E[u_t(x | x_1) | x, t], the two MSE problems have the same minimizer and the same gradient.
This is the algorithmic breakthrough. Sampling from the conditional path p_t(· | x_1) is trivial — draw x_0 from the base noise, draw x_1 from the data, then x_t = some closed-form interpolation. The conditional velocity u_t(x_t | x_1) is also a closed-form expression. Training reduces to: sample a noise–data pair, interpolate, ask the network to predict the velocity along that interpolation, compute MSE.
Quick check
Why does the conditional flow matching objective L_CFM have the same gradient with respect to θ as the marginal objective L_FM?
#Optimal Transport CFM: The Path That Picks Itself
The conditional flow matching recipe leaves one decision: which conditional path p_t(x | x_1) do you use? The natural answer — straight lines — has remarkable properties.
Optimal Transport Conditional Flow Matching (OT-CFM). Pick a noise sample x_0 ~ N(0, I) and a data sample x_1 ~ p_data. Define the conditional path as the straight line between them:
xt=(1−t),x0+t,x1,qquadtin[0,1]
The conditional velocity field along this straight line is obtained by differentiating x_t with respect to t:
ut(xtmidx1)=fracdxtdt=x1−x0
This is striking. The training target is the end-to-end displacement between noise and data — a quantity that does not depend on the time t at all. Time enters only through the network's input: at small t, the network sees x_t close to pure noise; at large t, close to clean data. But the desired output is always the same vector x_1 − x_0.
That property has a name: the conditional velocity field is constant along each straight-line path. This is exactly the optimal-transport map from x_0 to x_1 under squared cost. OT-CFM is so named because, conditional on the pair (x_0, x_1), the straight-line path realizes the L²-optimal transport between the point masses δ_ and δ_.
The marginal vector field u_t(x) is not constant along straight lines in general — it is an average of straight-line velocities over all (x_0, x_1) pairs that could have produced x at time t. But because the conditional field is constant, the learned marginal field is empirically much straighter than what you would get from a diffusion-style curved Gaussian path. Straighter paths can be integrated by an ODE solver with fewer steps. That is the inference speedup.
What Do You Think?
Why is OT-CFM's conditional velocity field u_t(x_t | x_1) constant along each conditional path (independent of t)?
Liu, Gong, and Liu 2023 — Flow Straight and Fast — arrived at the same algorithm independently and from a different angle. Rectified Flow starts from the question "what is the simplest possible coupling between two distributions?" and answers "the straight-line interpolation." Mathematically, RF and OT-CFM converge on the same training objective: regress the network output against x_1 − x_0 on the linear path x_t = (1 − t) x_0 + t x_1.
Two contributions distinguished Rectified Flow as a paper.
First, the reflow procedure discussed in the DeepDive above. Lipman's paper showed that OT-CFM gives reasonably straight paths after one training round; Liu et al. showed that you can iterate the procedure to drive curvature toward zero. After two reflows, FID matches diffusion at 4–8 steps with a 1-step generator. This is now standard practice in production image models.
Second, the empirical demonstration on real-world data. Lipman's paper was largely theoretical and showed results on small toy distributions and CIFAR-10. Liu's paper showed competitive ImageNet-scale results, including image-to-image translation (FID 7.93 on cat→dog at 1 step), which forced the community to take flow matching seriously as a production-grade alternative.
In 2024, Esser et al. — Scaling Rectified Flow Transformers for High-Resolution Image Synthesis — combined Rectified Flow training, the MMDiT (Multimodal Diffusion Transformer) architecture, and several modifications to the time-sampling distribution to produce Stable Diffusion 3. Two months later, Black Forest Labs released FLUX, also based on rectified-flow-style training. The 2024 generation of state-of-the-art open-weights image models is built on this objective.
Euler is the simplest ODE solver. With N = 4 to N = 20 Euler steps, an OT-CFM model trained on ImageNet-scale data produces samples competitive with a diffusion model that took 50–250 score-network evaluations to sample. The 5×–25× speedup is the entire reason the field switched to flow matching.
You can do better than Euler. Heun's second-order method — a midpoint correction — uses two network evaluations per step but halves the discretization error:
v_1 = v_θ(x_t, t) # velocity at the start of the step
x_pred = x_t + Δt · v_1 # Euler-predict the endpoint
v_2 = v_θ(x_pred, t + Δt) # velocity at the predicted endpoint
x_{t+Δt} = x_t + (Δt/2)(v_1 + v_2) # average the two velocities
Heun-4-step is roughly as accurate as Euler-8-step for the same wall-clock cost (each Heun step costs two network evaluations vs Euler's one). For very straight paths (post-reflow), Euler wins on wall clock because the linear extrapolation is already accurate; for curved paths (vanilla OT-CFM on hard distributions), Heun's midpoint correction pays off.
The big picture: flow matching has the same number of decisions to make at inference as diffusion (which solver, how many steps), but the underlying ODE is much friendlier to coarse discretization because the velocity field is nearly straight by construction.
#Connection to Diffusion: One Framework, Many Paths
Lipman's framework subsumes diffusion. Pick the conditional path
p_t(x | x_1) = N(α_t · x_1, σ_t^2 · I)
with α_t and σ_t chosen to match the variance-preserving SDE schedule, and the corresponding conditional velocity field is exactly what a diffusion model's probability-flow ODE follows. Train against that velocity and you get a model that, at inference, behaves identically to a VP-SDE diffusion model sampled by the probability-flow ODE. The math worked out by Song et al. 2021 for the score-based view and the math worked out by Lipman 2023 for the flow-matching view are the same math, expressed in two different parameterizations.
The conceptual unification is:
Diffusion corresponds to flow matching with a Gaussian conditional path whose mean and variance follow a fixed SDE schedule. The conditional velocity field along this path is not constant — it curves with t.
OT-CFM and Rectified Flow correspond to flow matching with a straight-line conditional path. The conditional velocity field is constant along each path.
Stochastic interpolants (Albergo & Vanden-Eijnden 2023) are a generalization that mixes Gaussian and OT paths and adds a stochastic component during sampling.
You can mix and match. SD3 uses a straight-line OT path during training but injects a non-uniform time-sampling distribution to upweight the difficult parts of the trajectory; we cover that next.
Quick check
Diffusion's probability-flow ODE uses a Gaussian conditional path with a curved velocity field; OT-CFM uses a straight-line conditional path with a constant velocity field. What is the practical consequence for sampling?
#2024–2025 Refinements: Time Sampling, Timestep Shift, EDM Preconditioning
The basic OT-CFM recipe — uniform t, straight-line path, MSE loss — works. State-of-the-art flow-matching models add three refinements that move FID by 5–15%.
#Logit-Normal Time Sampling (SD3, Esser et al. 2024)
Uniform t ~ U(0, 1) gives every time step the same weight in the loss. Empirically, the network has different difficulty at different t. Near t = 0 (close to pure noise) the conditional velocity x_1 − x_0 is dominated by x_0, which is random — the network can do little better than predict zero, and the loss is bounded below by the variance of x_0. Near t = 1 (close to clean data) the conditional velocity is dominated by x_1 − x_0 where x_1 is structured — the network can use detailed image content to predict it. Both extremes are easy: at low t the network gives up, at high t the network has nearly full information.
The hard regime is the middle — t ≈ 0.5, where x_t is a roughly equal mixture of noise and data and the network must use partial structure to predict the displacement. SD3 upweights this region by sampling t from a logit-normal distribution:
t=frac11+e−z,qquadzsimmathcalN(m,s2)
The marginal density of t under this sampling is a sigmoid-shaped bump. At t = 0.5 the density is 4 / (e^0 · 2π)^ ≈ 1.6× higher than uniform; at t = 0.05 or t = 0.95 the density is ≈ 0.3× uniform. The net effect is that the network's optimizer focuses 2–3× more of its gradient updates on the hard middle range. SD3 reports a 5–10% FID improvement from this single change.
What Do You Think?
Logit-normal time sampling (SD3) vs uniform time sampling — which gives the network more gradient signal at t = 0.5?
When you train at higher resolutions (1024×1024 vs 256×256), the per-pixel signal-to-noise ratio at a given t changes — there is more total signal in a 1024×1024 image, so the same noise level is relatively less destructive. This means the high-noise regime (small t) effectively gets easier as resolution grows, and the network underinvests in it. SD3's fix is a resolution-dependent timestep shift that pushes the time distribution toward smaller t when training at high resolution:
t_shifted = shift_fn(t, resolution)
= s · t / (1 + (s − 1) · t), where s = sqrt(resolution / 256)
For 1024×1024 training, s = 2, and the shift function pushes a uniform t toward smaller values, concentrating training on the high-noise regime where the network needs more help. This is a separate axis from logit-normal sampling and the two compose: first apply logit-normal to t, then apply the resolution-dependent shift.
#EDM-Style Preconditioning (Karras et al. 2024 for FM)
The EDM preconditioning trick (Karras 2022 for diffusion) applies cleanly to flow matching. Instead of asking the network to directly output the velocity v_θ(x, t), parameterize:
v_θ(x, t) = c_skip(t) · x + c_out(t) · F_θ( c_in(t) · x , c_noise(t) )
where c_skip, c_out, c_in, c_noise are hand-chosen time-dependent scalars that normalize the input scale, output scale, and effective noise to the inner network F_θ. The result is that F_θ always sees inputs with O(1) scale and produces outputs with O(1) scale, regardless of t — which makes optimization much better behaved than letting the raw scale vary 1000× across the trajectory.
For flow matching, c_skip and c_out are chosen so that v_θ at t = 0 equals x_0 (pass-through) and at t = 1 equals x_1 − x_0 (the full displacement). The intermediate values interpolate smoothly. Karras et al.'s 2024 paper on EDM2 reports that this preconditioning closes ~50% of the remaining FID gap between OT-CFM and the best diffusion models at the same compute.
Time to write the actual loop. The following playground trains an OT-CFM model on a 2D toy distribution — a four-mode Gaussian mixture — using a tiny MLP. The dataset is small enough to fit in memory and the training takes ~10 seconds in Pyodide, which lets you watch the velocity field straighten as training progresses.
Loading visualization...
What to look for when you run the cell: the loss drops to ~0.4 within 500 steps (it bottoms at the variance of x_0, which is 1.0 per dimension × 2 dimensions × an averaging factor — the irreducible noise floor), and the Euler-8 samples cluster around the four mode centers with the correct standard deviation. The preset "Compare Euler 4 vs Heun 4 steps" uses a closed-form velocity to remove the network-approximation error and shows what pure solver error looks like — Heun lands more accurately at the cost of 2× the network evaluations per step.
A practical comparison, holding architecture and data fixed:
Axis
Diffusion (DDPM / VP-SDE)
Flow Matching (OT-CFM)
Training target
Predict ε (noise) or v (Karras-v)
Predict velocity x_1 − x_0
Loss
Weighted MSE across noise levels
Plain MSE, optionally with logit-normal t
Noise schedule
Required (linear / cosine / EDM / sigmoid)
None (path defines the geometry directly)
Hyperparameters
β_t schedule, prediction parameterization, loss weighting, EDM σ_data
Conditional path family, time-sampling distribution
Inference
Reverse SDE or probability-flow ODE
Forward ODE
Steps for high quality
50–250 (DPM-Solver: 20–30)
4–20 (post-reflow: 1)
Single-step possible?
Only via distillation (LCM, ADD)
Built-in via reflow (Liu 2023)
Quality at convergence
The two columns are not so much "diffusion vs flow matching" as "two parameterizations of the same generative-modeling problem with different design surfaces." Diffusion exposes more knobs at training time (the schedule and its many variants), flow matching exposes fewer; flow matching exposes more knobs at the path-design level (you can pick any smooth path, not just Gaussian).
You want fewer inference steps at the same quality. The straighter ODE geometry pays off most when sampling is your bottleneck (real-time video, interactive generation, on-device inference).
You want a cleaner training pipeline. No schedule tuning, no parameterization debates, one MSE loss.
You want to do 1-step generation. Reflow gives you that without an extra distillation stage.
Diffusion still wins when
You need rigorous likelihood estimates. The diffusion ELBO gives a tight likelihood bound; the OT-CFM objective does not directly correspond to a likelihood (you can recover one via the change-of-variables formula post-training, but you needed to learn the velocity, not the score).
You want to use mature samplers (DPM-Solver++, EDM2). The flow-matching solver ecosystem is younger; there is no flow-matching analog of EDM2's level of solver engineering yet.
You want classifier-free guidance to behave in a specific characterized way. CFG works for flow matching but the literature on its strength scaling is thinner than for diffusion.
For 2025-vintage image and video generation, the field has substantially converged on flow matching (SD3, SD3.5, FLUX, Hunyuan-DiT, Stable Video Diffusion 1.5, Sora-like models). For 3D, protein, and discrete-data generation, the picture is mixed — diffusion variants still dominate molecular and protein modeling, where the score-based formulation maps cleanly onto physical noise processes.
Continuous normalizing flows parameterize a generative model as the solution of an ODE dx/dt = v_t(x). Training them by maximum likelihood requires the Jacobian trace ∂v/∂x, which is intractable at image scale.
Lipman's flow matching sidesteps maximum likelihood: instead, specify a target vector field u_t(x) by choosing a probability path p_t, and train v_θ by MSE regression against u_t. The marginal u_t(x) is intractable, but the conditional field u_t(x | x_1) is closed-form.
Conditional Flow Matching theorem. Training against the conditional objective gives the same gradient as training against the (intractable) marginal objective. Sample (x_0, x_1) pairs, evaluate u_t along the conditional path, do MSE.
OT-CFM chooses the conditional path to be the straight line x_t = (1−t) x_0 + t x_1. Its conditional velocity is the constant vector x_1 − x_0. This produces marginal flows that are nearly straight, which makes the ODE integrable in few steps.
Rectified Flow is the parallel discovery (Liu 2023) and adds the reflow procedure that iteratively straightens the marginal flow toward 1-step generation.
Inference integrates the forward ODE from t=0 (noise) to t=1 (data) using Euler or Heun. 4–20 steps suffice for high quality (vs 50–250 for diffusion's reverse SDE).
SD3 / FLUX 2024 refinements. Logit-normal time sampling to upweight the difficult middle of the trajectory; resolution-dependent timestep shift; EDM-style preconditioning. Each is worth 5–10% FID.
Conceptual unification. Flow matching with a Gaussian path recovers diffusion's probability-flow ODE; with a straight-line path it is OT-CFM. The same network can be trained as either, with one line of code different in the loss.
Flow matching closes out the Deep Learning track — twenty-seven lessons from the tensor and the autograd graph to SD3-grade generative training. The architectures and tricks you just learned (attention, normalization, mixed precision, FSDP, MoE, SSMs, flow matching) are the substrate; the next track shows what happens when one of them — the transformer — gets scaled to a trillion parameters and trained to predict the next token. Track 5 (NLP & Transformers) covers tokenization, scaling laws, RLHF, and the modern LLM stack; revisit Diffusion ModelsDiffusion ModelsDiffusion models generate data by learning to reverse a gradual noise-addition process, iteratively denoising random noise into coherent samples.Learn more → in Track 6 when you want the generative-image analogue of what comes next.
State-of-the-art (Imagen, DALL-E 3)
State-of-the-art (SD3, FLUX)
Training cost
Similar
Similar
Esser, Kulal, Blattmann, Entezari, Müller, Saini, Levi, Lorenz, Sauer, Boesel, Podell, Dockhorn, English, Lacey, Goodwin, Marek & Rombach 2024 — Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. The Stable Diffusion 3 paper. Combined rectified-flow training, the MMDiT architecture, logit-normal time sampling, and resolution-dependent timestep shift. Set the bar for the 2024 open-weights image-generation era.
Black Forest Labs 2024 — FLUX.1 release (technical report pending). Built on similar rectified-flow foundations as SD3 with architecture and data refinements; FLUX Schnell ships as a 4-step generator trained via flow-matching distillation.
The 6 years between Chen 2018 and SD3 2024 are a case study in how a beautiful but impractical mathematical idea — continuous normalizing flows — eventually finds the training trick (Lipman's conditional reformulation) that makes it scale.