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.