Rectified Flow Math

The setup: two distributions and a line between them

Fix a source distribution π_0 (usually N(0, I)) and a target π_1 (the data). Draw an independent pair X_0 ∼ π_0, X_1 ∼ π_1 and define the linear interpolation for t ∈ [0, 1]:

X_t = (1 − t) X_0 + t X_1            dX_t/dt = X_1 − X_0   (constant in t)

Shapes are trivial: if an image latent is X: [C, H, W], so is every X_t and so is the velocity. This is not a generative process: computing X_t requires knowing X_1 — the very thing we want to produce. The interpolation is only a training scaffold. What we extract from it is the family of marginal laws p_t of X_t, interpolating smoothly from π_0 at t = 0 to π_1 at t = 1. The trick is to find a causal ODE whose marginals match p_t exactly — one integrable forward knowing only where it currently is.

The velocity field is a conditional expectation — so training is regression

The interpolation paths cross. Many (X_0, X_1) pairs pass through the same point z at the same time t, each carrying a different velocity X_1 − X_0. A deterministic ODE cannot have two velocities at one point, so rectified flow averages, and regresses a network v_θ onto that average:

v(z, t) = E[ X_1 − X_0 | X_t = z ]        dZ_t = v(Z_t, t) dt,   Z_0 ∼ π_0
L(θ) = E_t E_{X_0, X_1} || v_θ( (1−t)X_0 + tX_1 , t ) − (X_1 − X_0) ||^2

Since the minimiser of a squared loss over all measurable functions is the conditional mean, the optimum of L is exactly v(z, t). One training step: sample data, sample noise, sample t, blend, subtract, MSE. No noise schedule, no SNR reweighting, no α-bar table, no ε-versus-x_0-versus-v parameterisation debate. Averaging also makes the field single-valued, so trajectories are deterministic and, by uniqueness of ODE solutions, cannot cross: a bundle of crossing straight segments becomes a bundle of non-crossing curves. That bending is the price of causality, and it is what reflow removes.

Why the ODE preserves the marginals

The claim that makes it all work: Z_t has the same marginal law as X_t for every t. The argument is the continuity equation. Any flow with velocity field u transports density according to

∂p_t/∂t + ∇·( p_t u ) = 0

The interpolation’s marginals p_t satisfy this with the (multi-valued) velocity of the segments, and taking the conditional expectation over the segments through z collapses it to a single field while leaving the divergence term p_t · E[·|X_t = z] unchanged. So v drives the same density evolution, and at t = 1 the endpoint Z_1 is distributed as π_1 — even though the pairing (Z_0, Z_1) generally differs from the independent pairing (X_0, X_1) we trained on. That rewiring of the coupling is the mechanism reflow exploits.

The coupling improves: transport cost never increases

Rectified flow does more than preserve marginals. The induced coupling (Z_0, Z_1) is no worse than the one you started with, simultaneously for every convex cost:

E[ c(Z_1 − Z_0) ] ≤ E[ c(X_1 − X_0) ]   for every convex c

Jensen’s inequality supplies the intuition: Z_1 − Z_0 is an integral of averaged displacements, and averaging inside a convex function can only lower its value. That is stronger than reducing one particular cost: it covers squared distance, absolute distance and every other convex functional at once, which is why the coupling is described as closer to optimal transport without ever solving an OT problem. It also explains the geometry: crossing segments cost more than the rewired, non-crossing bundle replacing them.

Reflow: apply the operator to its own output

Having trained a first model v^1, generate pairs from it: draw Z_0 ∼ π_0, integrate the ODE accurately to get Z_1, keep the pair. Now retrain from scratch on the same objective with this synthetic coupling instead of the independent one:

v^{k+1} = argmin_v E || v( (1−t)Z_0 + tZ_1 , t ) − (Z_1 − Z_0) ||^2 ,   (Z_0, Z_1) ∼ ODE_{v^k}

Call the result the k-rectified flow. Because (Z_0, Z_1) is already a deterministic, non-crossing pairing, far fewer segments collide, so the conditional expectation averages much less and the learned field is closer to constant along each path. Crucially, π_1 is preserved at every round — reflow is not distillation and does not target a fixed teacher output per input; it re-solves the same problem on a better-behaved coupling.

Measuring straightness

Straightness has a definition, not just a vibe:

S(Z) = ∫_0^1 E || (Z_1 − Z_0) − v(Z_t, t) ||^2 dt

It measures how far the instantaneous velocity deviates from the straight-line displacement; S(Z) = 0 exactly when every trajectory is a straight line traversed at constant speed. The central theoretical result is that the reflow sequence drives it down at rate S(Z^k) = O(1/k). Practically, the first reflow buys most of the improvement and later rounds have diminishing returns, which is why published pipelines run one or two rounds, not ten. Straightness connects directly to solver error — for a perfectly straight flow the Euler update z + v·Δt is exact, so the discretisation error of one-step sampling is controlled by S.

A closed-form example, and Euler error before and after reflow

Take π_0 = N(0, 1), π_1 = N(μ, s^2) with μ = 3, s = 0.5, independent coupling. Then X_t and X_1 − X_0 are jointly Gaussian, so the conditional expectation is exact. With σ^2(t) = (1−t)^2 + t^2 s^2:

Cov(X_t, X_1−X_0) = t s^2 − (1−t)
v(z, t) = μ + [ ( t s^2 − (1−t) ) / σ^2(t) ] · (z − tμ)
solution:  z(t) = tμ + z_0 · σ(t)     ⇒   z(1) = μ + s·z_0

The flow lands precisely on the optimal transport map z → μ + s z, yet the path is curved, because σ(t) is not linear: σ(0.5) = 0.559, so a trajectory from z_0 = 1.0 passes through z(0.5) = 2.059 on its way to z(1) = 3.5. Integrate it numerically and that curvature costs accuracy — at t = 0 the field is v(z, 0) = μ − z:

1 Euler step  (Δt = 1.0):  1.0 + 2.0 = 3.00      error 0.50
2 Euler steps (Δt = 0.5):  1.0 → 2.00 → 3.20   error 0.30
exact:                          3.50

Error shrinks roughly like Δt, the usual first-order rate. Now reflow: the rectified coupling is (z_0, μ + s z_0), so the new interpolation Z_t = z_0 + t(μ + (s−1) z_0) is already a straight line with constant velocity 2.5, and one Euler step gives 1.0 + 2.5 = 3.5, exact. The 1-D Gaussian is the clean special case where one reflow straightens completely. In high-dimensional multimodal problems straightening is gradual and O(1/k) is what you get.

Rectified flow, flow matching, and diffusion

Flow matching is the general framework: pick any conditional probability path, regress the conditional velocity. Rectified flow is the specific instance with a linear interpolant and independent coupling — plus the reflow operator, which flow matching does not have. Diffusion trains a score model whose probability-flow ODE has an equivalent velocity field, so a diffusion model can be rewritten as a flow; the substantive difference is the schedule. Variance-preserving diffusion paths curve strongly near t = 0 because the noise level changes non-linearly, which is why they need many solver steps. The linear interpolant is the flat schedule, and reflow then attacks the residual curvature caused by the coupling rather than the schedule. Both sources of curvature must be handled to reach one-step sampling.

Inference cost, and where the idea shows up

Sampling cost is dominated by the number of function evaluations, not solver arithmetic: each NFE is a full forward pass of the velocity network, while the Euler update is a fused multiply-add over the latent. Total cost is roughly NFE × (forward-pass FLOPs), and straightening cuts the first factor. On CPU that decides usability. If a small latent generator takes 1.2 s per forward pass, 50 NFEs is a 60-second batch job, 4 NFEs is 4.8 s, and 1 NFE is interactive — same weights, same memory footprint. Straightening moves compute from inference to training.

The linear-interpolant objective is now mainstream: Stable Diffusion 3 adopts it, and FLUX builds on a rectified-flow transformer. Adopting the objective is not the same as running reflow — large text-to-image models usually reach few-step variants through additional, often adversarial, distillation. And reflow composes with distillation: InstaFlow reflows first, then distils the straightened flow into one step, the straightening being what makes distillation tractable.

Pitfalls

Reflow trains on synthetic data. The pairs come from your own model, so its biases and artefacts are baked in and compound across rounds. Expect a small quality drift per round.

Generate reflow pairs with an accurate solver. If you produce Z_1 with a sloppy 4-step integration, you are teaching the next model to reproduce discretisation error as if it were the target distribution.

Straightness is not quality. S(Z) = 0 only says the paths are lines; a straight flow with a poor coupling still generates poor samples.

Guidance breaks straightness. Classifier-free guidance evaluates a different velocity field than the one that was rectified, so heavily guided sampling re-introduces curvature and needs more steps.

Watch the convention. This article uses t = 0 for noise and t = 1 for data; several codebases reverse it, which flips the sign of every velocity.