6. Flow matching
Draft
Flows and the continuity equation
Flow matching starts with the flow we want to learn. Choose paths between data and noise, calculate their velocities, and train a network to predict the average velocity at each point and time. This reaches the probability-flow ODE of Section 4.3 without first constructing an SDE.
Lipman, Chen, Ben-Hamu, Nickel and Le developed flow matching alongside Liu, Gong and Liu's rectified flow and Albergo and Vanden-Eijnden's stochastic interpolants.Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel and M. Le, "Flow matching for generative modeling", ICLR 2023 (arXiv 2210.02747); X. Liu, C. Gong and Q. Liu, "Flow straight and fast: learning to generate and transfer data with rectified flow", ICLR 2023 (arXiv 2209.03003); M. S. Albergo and E. Vanden-Eijnden, "Building normalizing flows with stochastic interpolants", ICLR 2023 (arXiv 2209.15571). The two expositions the series follows are T. Fjelde, E. Mathieu and V. Dutordoir, "An introduction to flow matching" (Cambridge MLG blog, 20 January 2024), and Nakkiran et al. (2024), Section 4, whose time convention is the one used here. On independent Gaussian paths, the learned velocity is the one we already obtained from the score. The new viewpoint also makes the paths themselves a design choice. We will use that freedom to ask how straight, and how easy to sample, they can become.
The time convention needs care. Most flow-matching papers put noise at \(t = 0\) and data at \(t = 1\). This series keeps data at \(t = 0\) and noise at \(t = 1\), following the convention used in the earlier parts and in Nakkiran et al. To translate a formula from Lipman et al., replace \(t\) by \(1 - t\) and reverse the velocity's sign.
The object learned is a rule for moving whatever point currently arrives at \((x,t)\). A training pair supplies one possible direction there. At sampling time the network receives only the current point and time, so it must combine the directions of all compatible pairs. Understanding that combination is the central step in understanding flow matching.
Definition 6.1.1 (flow, push-forward). A time-dependent vector field \(u_t\) on \(\mathbb R^d\), locally Lipschitz in \(x\) and continuous in \(t\), defines a flow \(\psi_{s \to t}\) by \(\frac{d}{dt}\psi_{s \to t}(x) = u_t(\psi_{s \to t}(x))\), \(\psi_{s \to s}(x) = x\). The push-forward of a law \(p\) by a measurable map \(\psi\) is the law of \(\psi(X)\) for \(X \sim p\), written \(\psi_\# p\).
Theorem 6.1.2 (flows and the continuity equation). Let \(u_t\) be as in Definition 6.1.1 with its flow defined on \([0, 1]\). For a law \(p_0\), the family \(p_t = (\psi_{0 \to t})_\# p_0\) is a solution, in the sense of distributions, of the continuity equation \(\partial_t p_t + \nabla \cdot (p_t u_t) = 0\) with initial datum \(p_0\), and it is the only solution among families of laws with finite first moment that is continuous in \(t\).
Proof. Existence: for a test function \(\phi\), \(\frac{d}{dt}\int\phi\, dp_t = \frac{d}{dt}\mathbb E\phi(\psi_{0 \to t}(X)) = \mathbb E[\nabla\phi \cdot u_t](\psi_{0 \to t}(X)) = \int \nabla\phi \cdot u_t\, dp_t\), which is the weak form of the equation. Uniqueness, for a smooth \(u\): let \(q_t\) be any solution and pull it back along the flow; the function \(t \mapsto \int\phi(\psi_{t \to 1}(x))\, q_t(dx)\) has derivative zero for every \(\phi\), because the time derivative of \(\phi \circ \psi_{t \to 1}\) is \(-\nabla(\phi \circ \psi_{t \to 1}) \cdot u_t\), which cancels the transport term of the equation. Hence \((\psi_{t \to 1})_\# q_t = q_1\) for all \(t\), and in particular \(q_1 = (\psi_{0 \to 1})_\# q_0 = p_1\); the same argument on \([0, t]\) gives \(q_t = p_t\). The extension to Lipschitz fields is by regularization.C. Villani, Optimal Transport: Old and New (Springer, 2009), Theorem 5.34 and the surrounding discussion of the continuity equation, after R. J. DiPerna and P.-L. Lions, "Ordinary differential equations, transport theory and Sobolev spaces", Inventiones Mathematicae 98 (1989), 511–547, for fields of lower regularity; L. Ambrosio, N. Gigli and G. Savaré, Gradient Flows in Metric Spaces and in the Space of Probability Measures, 2nd ed. (Birkhäuser, 2008), Chapter 8. ∎
A continuous normalizing flow learns a field \(u_\theta\) that transports \(\mathcal N(0, I)\) to the data. Maximum-likelihood training uses Theorem 4.4.1: for every training example, solve an ODE, integrate its divergence, and differentiate through the solve. Flow matching replaces that expensive training calculation with regression. We still solve an ODE when sampling, but do not need to solve one to construct a training target.Chen, Rubanova, Bettencourt and Duvenaud (2018) and Grathwohl et al. (2019), cited in Section 4.4, for the maximum-likelihood training of continuous normalizing flows; the cost Lipman et al. identify in their introduction is exactly this solve.
Conditional paths and the marginalization theorem
It is easier to describe a path for one data point than for the whole data distribution. We can build those simple conditional paths first, then mix them. The key result tells us how to combine their velocities so that one field moves the mixture along the desired path.
Definition 6.2.1 (conditional path and field). A conditional probability path is a family of densities \(p_t(\cdot \mid x_0)\), \(t \in [0, 1]\), indexed by data points, with \(p_0(\cdot \mid x_0)\) concentrated at \(x_0\) and \(p_1(\cdot \mid x_0)\) a fixed law \(q\) that does not depend on \(x_0\), together with a field \(u_t(\cdot \mid x_0)\) that generates it: \(\partial_t p_t(\cdot \mid x_0) + \nabla \cdot (p_t(\cdot \mid x_0)\, u_t(\cdot \mid x_0)) = 0\). The marginal path is \(p_t(x) = \int p_t(x \mid x_0)\, p_0(dx_0)\), so that \(p_1 = q\).
The marginal path begins at the data law and ends at the fixed law \(q\), which is taken to be \(\mathcal N(0, I)\) below; what is needed is a field that generates it, and the next theorem builds one from the conditional fields.
Probability currents add when we combine populations. Their velocities must therefore be weighted by how much of each population is present at the current location. If three quarters of the local density moves right at unit speed and one quarter moves left at unit speed, the combined current corresponds to velocity \(1/2\) to the right. A plain average of the two velocities would incorrectly give zero. The posterior weights below perform exactly this local accounting.
top R1 under the linear path at t = 0.5: the densities of the two populations, 1/2 p_t(x | ±1), in
grey and their sum p_t in ink; at the probe x = 0.27 the two heights give the shares w− = 0.25
and w+ = 0.75, the weights of (6.1)
bottom the conditional velocities (x ∓ 1)/t, −1.45 and +2.55, their plain average x/t = +0.55 and the
net velocity u_t = −0.45 (green), whose sign the plain average gets wrong between −0.96 and
0.96; the currents add, −0.52 + 0.31 = −0.22 = p_t u_t; in the sampled view 1,024 pairs (seed 11):
94 land in the bin [0.2, 0.4), 68 from +1 and 26 from −1, mean target −0.28 against the exact
−0.47, and the fit's loss 1.755 sits at the floor 1.798
(the live figure needs JavaScript)
Theorem 6.2.2 (the marginal field). Define
\[u_t(x) \;=\; \mathbb E\bigl[u_t(x \mid x_0) \,\big|\, x_t = x\bigr] \;=\; \int u_t(x \mid x_0)\, \frac{p_t(x \mid x_0)\, p_0(dx_0)}{p_t(x)} . \tag{6.1}\]Then, under regularity permitting differentiation under the integral, \(u_t\) generates the marginal path: \(\partial_t p_t + \nabla \cdot (p_t u_t) = 0\). Consequently the flow of \(u_t\) from \(t = 1\) to \(t = 0\) carries \(q\) onto \(p_0\).
Proof. Differentiate the marginal under the integral and use the conditional continuity equations:
\[\begin{aligned} \partial_t p_t(x) &= \int \partial_t p_t(x \mid x_0)\, p_0(dx_0) \\ &= -\int \nabla \cdot \bigl(p_t(x \mid x_0)\, u_t(x \mid x_0)\bigr)\, p_0(dx_0) \\ &= -\nabla \cdot \Bigl(\int u_t(x \mid x_0)\, p_t(x \mid x_0)\, p_0(dx_0)\Bigr) = -\nabla \cdot \bigl(p_t(x)\, u_t(x)\bigr), \end{aligned}\]the last step being the definition (6.1). The consequence is Theorem 6.1.2 applied from \(t = 1\) backward. ∎Lipman et al. (2023), Theorem 1: "Given vector fields \(u_t(x \mid x_1)\) that generate conditional probability paths \(p_t(x \mid x_1)\), for any distribution \(q(x_1)\), the marginal vector field \(u_t\) in equation 8 generates the marginal probability path \(p_t\) in equation 6, i.e., \(u_t\) and \(p_t\) satisfy the continuity equation", with their \(x_1\) the data point. The proof is the one above; its only ingredient is the linearity of the continuity equation in the pair \((p, pu)\).
The structure is the same as in Theorem 2.1.4. There, the marginal score was the posterior average of the conditional scores. Here, the marginal velocity is the posterior average of the conditional velocities. Both identities come from equations that are linear in the law.
Evaluating the posterior average directly would require the unknown \(p_t(x)\). Regression avoids that calculation: sample a conditional path and use its velocity as the target. By Lemma 2.3.6, least squares learns the required average.
We never have to compute those posterior weights explicitly during training. Sampling a pair and a point on its path naturally visits locations in proportion to the conditional density. Repeated squared-error updates then average the target velocities with the right local frequencies. The theorem says that the disagreement among those targets adds a constant to the loss, without changing which marginal field minimizes it. In the sampled view of the figure below, the counts in a bin of \(x_t\) are those weights.
Theorem 6.2.3 (the conditional objective has the same gradient). Let
\[\begin{aligned} L_{\mathrm{FM}}(\theta) &= \mathbb E_{t,\, x \sim p_t}\bigl\|u_\theta(x, t) - u_t(x)\bigr\|^2, \\ L_{\mathrm{CFM}}(\theta) &= \mathbb E_{t,\, x_0 \sim p_0,\, x \sim p_t(\cdot \mid x_0)}\bigl\|u_\theta(x, t) - u_t(x \mid x_0)\bigr\|^2, \end{aligned}\]with \(t\) drawn from any law on \([0, 1]\). Then \(L_{\mathrm{FM}} - L_{\mathrm{CFM}}\) does not depend on \(\theta\), so \(\nabla_\theta L_{\mathrm{FM}} = \nabla_\theta L_{\mathrm{CFM}}\).
Proof. Expand both squares. The terms \(\|u_\theta\|^2\) have the same expectation, since \((x_0, x) \mapsto x\) has law \(p_t\). The cross terms agree by the tower property: \(\mathbb E[u_\theta(x_t) \cdot u_t(x_t \mid x_0)] = \mathbb E[u_\theta(x_t) \cdot \mathbb E[u_t(x_t \mid x_0) \mid x_t]] = \mathbb E[u_\theta(x_t) \cdot u_t(x_t)]\) by (6.1). The remaining terms do not involve \(\theta\). ∎Lipman et al. (2023), Theorem 2: "Assuming that \(p_t(x) > 0\) for all \(x \in \mathbb R^d\) and \(t \in [0, 1]\), then, up to a constant independent of \(\theta\), \(L_{\mathrm{CFM}}\) and \(L_{\mathrm{FM}}\) are equal. Hence, \(\nabla_\theta L_{\mathrm{FM}}(\theta) = \nabla_\theta L_{\mathrm{CFM}}(\theta)\)."
Algorithm 6.2.4 FLOW-MATCHING-TRAIN
(conditional flow matching; Lipman et al. 2023)
Picture: one training example. A pair (x_0, x_1) is
drawn, a point on the conditional path between them
is formed, and the network is pulled toward the
conditional velocity at that point.
x_1 (noise) o - - - - - x_t - - - - - o x_0 (data)
^ |
target u_t(x_t|x_0)| u_theta(x_t, t)
v
loss |u_theta - u_t(x_t|x_0)|^2
Input samples of p_0; the law q of x_1; a conditional
path with a sampler for x_t | x_0 and its field
u_t(x | x_0); a law for t; a network u_theta.
Output theta with u_theta ~ u_t, the marginal field.
1. repeat until converged:
2. draw x_0 from the data, t from its law
3. draw x_t ~ p_t(. | x_0)
4. L <- |u_theta(x_t, t) - u_t(x_t | x_0)|^2
5. take a gradient step on L in theta
Algorithm 6.2.5 FLOW-SAMPLE
(Euler on the learned field, from noise to data)
Picture: N steps against the flow, from t = 1 to 0.
x ~ q at t = 1 o---->o---->o---->o---->o t = 0
h h h h h = -1/N
Input a trained u_theta; a number of steps N.
Output a sample.
1. x ~ q; t <- 1
2. repeat N times:
3. x <- x - (1/N) u_theta(x, t); t <- t - 1/N
4. return x
Remark 6.2.6 (takeaway: the same regression in a new target). Algorithm 6.2.4 is Algorithm 2.3.7 with a conditional velocity in place of the clean point. Least squares learns its conditional mean, and the remaining conditional variance gives the loss floor. No density, divergence or ODE solve enters the construction of a training target. For independent Gaussian paths, the dictionary of Section 2.4 converts this objective into a weighted denoising loss. With other paths or couplings, the averaging principle survives, but the posterior, variance and conversion formulas must be worked out for that construction.
Gaussian paths, the linear path and the dictionary
Gaussian conditional paths connect flow matching to the diffusion models we already know. Their velocities can be written explicitly, and averaging those velocities recovers the probability-flow field. The linear path is especially convenient: each training pair travels at a constant velocity, even though the learned flow need not do so.
Proposition 6.3.1 (Gaussian conditional paths). Let \(p_t(\cdot \mid x_0) = \mathcal N(\alpha_t x_0, \sigma_t^2 I)\) with \(\alpha, \sigma\) as in Proposition 1.1.5 and \(\sigma_t > 0\) for \(t > 0\). The map \(\psi_t(z \mid x_0) = \alpha_t x_0 + \sigma_t z\) pushes \(\mathcal N(0, I)\) onto \(p_t(\cdot \mid x_0)\), and the field that generates the path along these maps is
\[u_t(x \mid x_0) \;=\; \frac{\dot\sigma_t}{\sigma_t}\,(x - \alpha_t x_0) + \dot\alpha_t\, x_0 .\]Proof. The push-forward is the definition of the Gaussian. Differentiating the map in \(t\) at fixed \(z\) gives \(\dot\alpha_t x_0 + \dot\sigma_t z\), and solving \(x = \alpha_t x_0 + \sigma_t z\) for \(z\) gives the stated field, which is the velocity of the point of the path currently at \(x\); a family of maps that moves every point of a law along a field generates that law's path, by Theorem 6.1.2 applied to the conditional law. ∎
The conditional field above is what one obtains by conditioning on \(x_0\) alone. Conditioning on the pair \((x_0, z)\) is simpler and more general.
Theorem 6.3.2 (two-sided conditioning, any coupling). Let \(\pi\) be any joint law of \((x_0, z)\) whose first marginal is \(p_0\) and whose second is \(\mathcal N(0, I)\), and let \(x_t = \alpha_t x_0 + \sigma_t z\). Then the marginals \(p_t\) of \(x_t\) are generated by
\[u_t^\pi(x) \;=\; \mathbb E_\pi\bigl[\dot\alpha_t x_0 + \dot\sigma_t z \,\big|\, x_t = x\bigr], \tag{6.2}\]and the regression of \(u_\theta(x_t, t)\) onto \(\dot\alpha_t x_0 + \dot\sigma_t z\) over pairs drawn from \(\pi\) has \(u_t^\pi\) as its minimizer. For the independent coupling, \(p_t\) is the noised law of Section 1 and \(u_t^\pi = u_t\) is the probability-flow velocity (4.3).
Proof. Condition on the pair: the path through \((x_0, z)\) is the deterministic curve \(t \mapsto \alpha_t x_0 + \sigma_t z\), whose law at time \(t\) is a point mass generated by the field equal to its velocity \(\dot\alpha_t x_0 + \dot\sigma_t z\) at that point; Theorem 6.2.2 with the pair as the conditioning variable, and \(\pi\) as the mixing law, gives (6.2) and the generation statement. The regression claim is Lemma 2.3.6. For the independent coupling the marginal of \(x_t\) is the forward kernel's, and (6.2) is Proposition 4.3.2. ∎
Corollary 6.3.3 (the dictionary). Flow matching on a Gaussian path with the independent coupling trains the probability-flow velocity of the diffusion with the same schedule, and its loss is the denoiser loss of Section 2 with the weight \((\alpha_t g_t^2/2\sigma_t^2)^2\) of Proposition 2.4.2; on the linear path \(\alpha_t = 1 - t\), \(\sigma_t = t\) the conditional velocity is \(z - x_0\), the marginal velocity is \(u_t = \hat\varepsilon_t - D_t\), and the weight is \(1/t^2\), heaviest at the data end. Conversely, the flow of a diffusion model trained by denoising is a flow-matching model with the diffusion's schedule, and Sections 4 and 5 apply to flow-matching models without change: the \(\lambda\)-family (4.2), DDIM, Heun and the exponential integrators, with \(f_t = \dot\alpha_t/\alpha_t\) and \(g_t^2 = -\sigma_t^2\dot\lambda_t\) read off the path.
The independence assumption is essential to this dictionary. A coupling can give the noise a Gaussian marginal while making its conditional law given the data non-Gaussian. Such a path still admits velocity regression, but the Gaussian denoising identities cannot be assumed without checking them. The table gives the dictionary for the independent Gaussian construction.
| object | in the diffusion vocabulary | in the flow-matching vocabulary |
|---|---|---|
| \(x_t = \alpha_t x_0 + \sigma_t z\) | the forward kernel, Proposition 1.1.5 | the Gaussian (affine) conditional path, Proposition 6.3.1 |
| \(\alpha_t = 1 - t\), \(\sigma_t = t\) | a schedule with \(f = -1/(1 - t)\), \(g^2 = 2t/(1 - t)\) | the linear path, the optimal-transport conditional path of Lipman et al. |
| \(\dot\alpha_t x_0 + \dot\sigma_t z\) | the derivative of the kernel path | the conditional velocity; \(z - x_0\) on the linear path |
| \(u_t(x) = f_t x - \tfrac12 g_t^2 s_t(x)\) | the probability-flow velocity, (4.3) | the marginal velocity, (6.1), (6.2) |
| \(u_t = \dot\alpha_t D_t + \dot\sigma_t \hat\varepsilon_t\) | Tweedie's formula applied twice, Proposition 4.3.2 | the marginalization theorem, Theorem 6.3.2 |
| the loss \(\lVert u_\theta - \dot x_t\rVert^2\) | the denoiser loss with weight \((\alpha_t g_t^2/2\sigma_t^2)^2\) | the conditional flow-matching loss |
| \(X_{1 \to 0}\) | the flow of the ODE; DDIM with \(\eta = 0\) is its Euler scheme | the flow map; Algorithm 6.2.5 |
| the family (4.2), \(\lambda > 0\) | the reverse SDE and its relatives | the stochastic interpolants' SDEs with tunable diffusion, Section 6.5 |
left 24 pairs (x_0 from R3, x_1 from N(0, I)) and their straight conditional paths, with the
point x_t of each at the chosen t as a cross; one pair is followed, its path and its x_t
green; the coupling is independent or a mini-batch OT assignment on 24, 64 or 256 pairs;
a probe shows the fan of conditional velocities through it
right the exact marginal velocity u_t of the independent coupling as streamlines, its flow lines,
and the contours of p_t; the flow line of the followed pair's own noise point in ink, with the
posterior over the 24 centres along it, which turns its direction by 175 degrees in all
strip the straightness integrand tr Cov(z - x_0 | x_t) against t, whose area is S = 1.94
(the live figure needs JavaScript)
A straight training segment has a fixed pair of endpoints. A learned trajectory has only its current position and time; as it moves, different endpoint pairs become plausible. Their average direction can therefore turn. Also, two segments crossing in a drawing creates conflicting targets only if they meet at the same time. The network can prescribe different velocities at the same location at different times. The figure of Section 6.3 follows one flow line as its aim point turns, and the figure below draws the segments in the \((t, x)\) plane, where a crossing is a conflict.
Example 6.3.4 (conditional and marginal paths on R3). In the figure's default view the twenty-four conditional paths of the independent coupling cross one another sixty times as curves in the plane, and their mean length is \(1.463\). The marginal velocity, which averages the conditional velocities of all pairs whose path passes through a point at a time, is not constant along its own flow lines: the thirty-six flow lines on the right have arc lengths \(2.338\) times their chords on average. Every target the regression ever sees is a straight segment, and the field it converges to is curved. The strip under the figure draws the straightness of Definition 6.4.1 for the independent coupling as a function of time: the conditional variance \(\mathrm{tr}\,\mathrm{Cov}(z - x_0 \mid x_t)\) peaks at \(2.469\) near \(t = 0.36\) and integrates to \(S = 1.944\), against \(\mathbb E\|z - x_0\|^2 = 2.986\); within the batch of twenty-four pairs the transport cost \(\mathbb E\|x_1 - x_0\|^2\) falls from \(2.751\) to \(0.652\) under the optimal assignment. Section 6.4 says what can be done about it. ◁
Remark 6.3.5 (confusion: the network does not learn the lines it is shown). The regression target at \((x_t, t)\) is \(z - x_0\) for the pair that produced \(x_t\), but other pairs produce the same \(x_t\) at the same \(t\) with other targets, and the minimizer is their average (Lemma 2.3.6). The straightness of the conditional paths is a property of the training data, not of the model; the model's paths are straight only when, at every \((x, t)\), all pairs through it agree, which is Definition 6.4.1 below. Click a point in the figure above to see the fan.
Remark 6.3.6 (trick: where to put the training time). By Theorem 6.3.2 the law of \(t\) in Algorithm 6.2.4 is free, and with it the weighting of noise levels, as in Section 3.4. Esser et al. compared schedules and samplers for their rectified-flow transformer and found that drawing \(t\) from a logit-normal law, \(\mathrm{logit}(t) \sim \mathcal N(0, 1)\), which concentrates the training time in the middle of the path, "consistently achieves a good rank" and beats uniform \(t\), with the caveat that the best choice depends on the number of sampling steps that will be used.P. Esser, S. Kulal, A. Blattmann et al., "Scaling rectified flow transformers for high-resolution image synthesis", ICML 2024 (arXiv 2403.03206), Section 3.1, Equation 19, and Section 5.1: "rf/lognorm(0.00, 1.00) consistently achieves a good rank. It outperforms a rectified flow formulation with uniform timestep sampling"; the caveat is theirs: a variant with location \(0.5\) "is the best-performing variant with 50 sampling steps but much worse … with 5 sampling steps". The same paper's architecture is the MM-DiT of Section 7.3.
Couplings, straightness and reflow
Straight training segments can produce a curved learned flow because different pairs prescribe different directions at the same place and time. Their conditional average changes as the flow moves. The quantity below measures that disagreement. Changing how data and noise are paired can reduce it and make larger sampling steps more accurate.
A coupling answers which noise point is paired with which data point. Rearranging those pairings leaves both endpoint distributions unchanged, but changes the interpolated clouds and the velocity targets between them. Pairings that send many trajectories through the same place in conflicting directions create a difficult regression problem. The quantity below measures that conflict rather than the geometric length of a single segment.
panel R2 and N(0, 1) under the linear path in the (t, x) plane, noise at the left: 16 pairs (seed 23),
whose straight segments cross 58 times (20 at t < 0.5), each crossing a conflict of two targets;
the 16 flow lines of the exact u_t from the same noise points cross nowhere; one is followed
(green); at the probed crossing, t = 0.56 and x = −0.03, the targets are +2.28 and −2.17 and
the marginal velocity is +0.35
strip the exact straightness integrand against t, peak 2.36 at t = 0.61, area S = 1.688 against
E|z − x_0|^2 = 2.462 and W_2^2 = 0.205; the sorted batch costs 0.520, the pairing (z, Φ(z)) 0.207
(the live figure needs JavaScript)
Definition 6.4.1 (straightness). For a coupling \(\pi\) of \((x_0, z)\) with the linear path, the straightness is
\[\begin{aligned} S(\pi) &= \int_0^1 \mathbb E_\pi\bigl\|(z - x_0) - u_t^\pi(x_t)\bigr\|^2\, dt \\ &= \int_0^1 \mathbb E\bigl[\mathrm{tr}\, \mathrm{Cov}_\pi(z - x_0 \mid x_t)\bigr]\, dt , \end{aligned}\]the integrated variance of the conditional velocity around the marginal one.
Proposition 6.4.2 (straight paths and one-step sampling). \(S(\pi) = 0\) if and only if, for almost every pair, \(u_t^\pi(x_t) = z - x_0\) along the whole path, in which case the flow lines of \(u^\pi\) are the conditional segments themselves and one Euler step of Algorithm 6.2.5 from \(z\) lands exactly on \(X_{1 \to 0}(z)\).
Proof. \(S = 0\) forces the integrand to vanish for almost every \((t, \text{pair})\), so the marginal field at \(x_t\) equals the pair's constant velocity; the path \(t \mapsto x_t\) then solves the ODE of \(u^\pi\), hence is its flow line by uniqueness, and the Euler step with a constant velocity is exact. The converse is immediate. ∎
Reflow uses the current flow to choose new pairs. Start from a noise point, follow the flow to its endpoint, and make those two points a training pair for the next round. The new training segments connect endpoints that the previous flow already matched. This creates a way to learn more direct routes while retaining the endpoint laws in the ideal construction; approximations in a trained teacher and its solver also enter the practical procedure.
Definition 6.4.3 (rectification). Given \(\pi\) with marginal field \(u^\pi\) and flow map \(X^\pi_{1 \to 0}\), the rectified coupling is the law of \((X^\pi_{1 \to 0}(z), z)\) for \(z \sim \mathcal N(0, I)\): the pairing of each noise point with the data point its flow line reaches. Repeating the construction gives the \(k\)-rectified flows of Liu, Gong and Liu; they call the passage from \(\pi\) to its rectification reflow.
Theorem 6.4.4 (Liu, Gong and Liu). (a) The rectified coupling has the marginals \(p_0\) and \(\mathcal N(0, I)\). (b) For every convex cost \(c\) on \(\mathbb R^d\), \(\mathbb E\, c(z - X^\pi_{1 \to 0}(z)) \le \mathbb E_\pi\, c(z - x_0)\): rectification does not increase any convex transport cost. (c) The straightness of the successive rectified couplings is summable, with \(\sum_{k \ge 0} S(\pi^{(k)}) \le \mathbb E_\pi\|z - x_0\|^2\), so it tends to zero.
Proof of (a) and (b). (a) is Theorem 6.3.2: the flow carries \(\mathcal N(0, I)\) onto \(p_0\). For (b), along the flow line from \(z\), \(z - X_{1 \to 0}(z) = \int_0^1 u_t^\pi(x_t)\, dt\) with \(x_t = X_{1 \to t}(z)\), so by Jensen's inequality for the convex \(c\), twice,
\[\begin{aligned} c\bigl(z - X_{1 \to 0}(z)\bigr) &\le \int_0^1 c\bigl(u_t^\pi(x_t)\bigr)\, dt = \int_0^1 c\bigl(\mathbb E_\pi[z - x_0 \mid x_t]\bigr) dt \\ &\le \int_0^1 \mathbb E_\pi\bigl[c(z - x_0) \mid x_t\bigr]\, dt . \end{aligned}\]Taking expectations, and using that \(x_t\) under the flow started from \(\mathcal N(0, I)\) has law \(p_t\), which is also its law under \(\pi\), the right side becomes \(\mathbb E_\pi c(z - x_0)\). (c) is Liu, Gong and Liu's straightening theorem. Its proof combines the variance decomposition of Lemma 2.3.6 integrated in \(t\), \(\mathbb E_\pi\|z - x_0\|^2 = \int_0^1 \mathbb E\|u_t^\pi(x_t)\|^2\, dt + S(\pi)\), with the first of the two Jensen inequalities of (b) for the quadratic cost, \(\mathbb E\|z - X_{1 \to 0}(z)\|^2 \le \int_0^1 \mathbb E\|u_t^\pi(x_t)\|^2\, dt\), which together give \(S(\pi) \le \mathbb E_\pi\|z - x_0\|^2 - \mathbb E\|z - X_{1 \to 0}(z)\|^2\), with equality exactly when \(u^\pi\) is constant along its flow lines; the right side telescopes over the rounds and the costs decrease by (b), so the straightnesses are summable. On R4 the second figure of this section prints the three terms: \(5.622 - 1.922 = 3.700\) against \(S = 2.906\), the gap being the bending of the flow lines. ∎Liu, Gong and Liu (2023), Section 2 and Theorems 3.3 to 3.7 of the arXiv version: rectification "turns an arbitrary coupling of \(\pi_0\) and \(\pi_1\) to a new deterministic coupling with provably non-increasing convex transport costs", and "recursively applying rectification allows us to obtain a sequence of flows with increasingly straight paths, which can be simulated accurately with coarse time discretization"; "perfectly straight paths can be simulated exactly with a single Euler step and is effectively a one-step model".
left round 0 on R4: 36 independent pairs (z, x_0), their straight segments, the curved flow lines
of the exact marginal velocity from each z to Phi(z), and 120 faint flow lines of the field
right round 1: the rectified pairs (z, Phi(z)), whose segments are the flow lines, one pair followed
(green, with its x_t in both rounds); rings mark the optimal map's images, which coincide
with Phi(z) here
strip the exact integrand of the straightness of round 0 against t, area S = 2.906; round 1's is 0
(the live figure needs JavaScript)
Example 6.4.5 (reflow on R4, exactly). For \(x_0 \sim \mathcal N(\mu, \Sigma)\) and \(z \sim \mathcal N(0, I)\) independent, the pair \((x_0, z, x_t)\) is jointly Gaussian, so the marginal velocity (6.2) is affine,
\[u_t(x) = -\mu + \bigl(tI - (1 - t)\Sigma\bigr)\, C_t^{-1}\bigl(x - (1 - t)\mu\bigr), \qquad C_t = (1 - t)^2\Sigma + t^2 I,\]the flow map is affine, \(\Phi(z) = Mz + c\), and the straightness is \(S = \int_0^1 \mathrm{tr}\bigl(\Sigma + I - B_t C_t^{-1} B_t\bigr) dt\) with \(B_t = tI - (1 - t)\Sigma\), the integrated conditional covariance of \(z - x_0\) given \(x_t\). In the figure's default view, \(\mu = (1.1, 0.5)\) and \(\Sigma\) has standard deviations \(1.4\) and \(0.45\) along axes rotated by \(35^\circ\); the independent coupling has transport cost \(\mathbb E\|z - x_0\|^2 = \|\mu\|^2 + \mathrm{tr}\,\Sigma + 2 = 5.622\) and straightness \(S = 2.906\), and its flow lines have arc lengths \(1.159\) times their chords. After one reflow the coupling is deterministic, \(x_t\) determines \(z\), \(S = 0\), the paths are the flow lines, one Euler step is exact, and the cost has fallen to \(1.922\). That number is also the squared Wasserstein-2 distance \(\|\mu\|^2 + \mathrm{tr}(\Sigma + I - 2\Sigma^{1/2})\), and the matrix of \(\Phi\) is \(\Sigma^{1/2}\) to three decimals: here one round of reflow produced the optimal transport map. The fall in cost, \(5.622 - 1.922 = 3.700\), exceeds \(S = 2.906\) by \(0.794\): the difference is the bending of the flow lines, by the inequality of Theorem 6.4.4(c), and the strip under the figure draws the integrand whose area is \(S\). The program below recomputes all of it. ◁
Remark 6.4.6 (warning: straight is not optimal, and one round is not the rule). Theorem 6.4.4 says that reflow lowers every convex cost and straightens; it does not say that the limit is the optimal transport map, and in general it is not. R4 reaches the optimum in one round because the source is isotropic and the target Gaussian, so the affine flow map commutes with everything in sight. For R3 the first rectified coupling is a deterministic map whose cost is lower than the independent coupling's and whose paths are curved; the second round straightens them further, and the sequence converges to a straight coupling that need not be the Wasserstein-2 optimal one.
Straightness is what a one-step sampler needs; optimality is a different property, and Pooladian et al. and Tong et al. obtain it approximately by a different route: couple noise and data within each training batch by the optimal assignment, which is legitimate by Theorem 6.3.2 since the marginals are untouched, and which lowers the variance of the regression target and the curvature of the learned flow at once. The figure of Section 6.3 shows the effect of such an assignment on twenty-four pairs.A.-A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman and R. T. Q. Chen, "Multisample flow matching: straightening flows with minibatch couplings", ICML 2023 (arXiv 2304.14772): non-trivial couplings "while satisfying the correct marginal constraints" let them "(i) reduce gradient variance during training, (ii) obtain straighter flows for the learned vector field … and (iii) obtain transport maps with lower cost in high dimensions"; A. Tong, K. Fatras, N. Malkin et al., "Improving and generalizing flow-based generative models with minibatch optimal transport", Transactions on Machine Learning Research (arXiv 2302.00482): "when the true OT plan is available, our OT-CFM method approximates dynamic OT". Reflow in practice: X. Liu, X. Zhang, J. Ma, J. Peng and Q. Liu, "InstaFlow: one step is enough for high-quality diffusion-based text-to-image generation", ICLR 2024 (arXiv 2309.06380), reflow then distillation of Stable Diffusion to one step at "an FID of 23.3 on MS COCO 2017-5k".
# R4: N(0, I) to N(mu, Sigma) in the plane under the linear path.
# independent coupling cost E|x_1 - x_0|^2 = |mu|^2 + tr Sigma + 2
# optimal cost W_2^2 = |mu|^2 + tr(Sigma + I - 2 Sigma^{1/2})
# straightness of round 0 S = int_0^1 tr Cov(x_1 - x_0 | x_t) dt
# the rectified map Phi(z) = M z + c, by integrating the affine flow
#
# Picture: round 0 pairs at random and its flow bends;
# round 1 pairs each z with Phi(z) and the paths are straight.
#
# z o----. o x_0 z o--------o Phi(z)
# '--/ crossing straight, S = 0
# round 0 round 1
import math
def mat(a, b, c, d): return [[a, b], [c, d]]
def mul(A, B): return [[sum(A[i][k] * B[k][j] for k in range(2)) for j in range(2)] for i in range(2)]
def add(A, B, c=1.0): return [[A[i][j] + c * B[i][j] for j in range(2)] for i in range(2)]
def scale(A, c): return [[A[i][j] * c for j in range(2)] for i in range(2)]
def inv(A):
det = A[0][0] * A[1][1] - A[0][1] * A[1][0]
return [[A[1][1] / det, -A[0][1] / det], [-A[1][0] / det, A[0][0] / det]]
def tr(A): return A[0][0] + A[1][1]
I = mat(1, 0, 0, 1)
mu = (1.1, 0.5)
s1, s2, th = 1.4, 0.45, math.radians(35)
R = mat(math.cos(th), -math.sin(th), math.sin(th), math.cos(th))
Sigma = mul(mul(R, mat(s1 * s1, 0, 0, s2 * s2)), mat(math.cos(th), math.sin(th), -math.sin(th), math.cos(th)))
det = Sigma[0][0] * Sigma[1][1] - Sigma[0][1] * Sigma[1][0]
root = scale(add(Sigma, I, math.sqrt(det)), 1 / math.sqrt(tr(Sigma) + 2 * math.sqrt(det))) # the symmetric square root
mu2 = mu[0] ** 2 + mu[1] ** 2
print(f"Sigma = {[[round(v, 4) for v in r] for r in Sigma]}, Sigma^1/2 = {[[round(v, 4) for v in r] for r in root]}")
print(f"independent cost {mu2 + tr(Sigma) + 2:.4f} optimal cost {mu2 + tr(Sigma) + 2 - 2 * tr(root):.4f}")
def C(t): return add(scale(Sigma, (1 - t) ** 2), I, t * t) # Var(x_t)
def A(t): return mul(add(scale(I, t), Sigma, -(1 - t)), inv(C(t))) # u_t(x) = A_t (x - (1-t) mu) - mu
S, n = 0.0, 2000 # straightness of round 0
for i in range(n):
t = (i + 0.5) / n
B = add(scale(I, t), Sigma, -(1 - t))
S += tr(add(add(Sigma, I), mul(mul(B, inv(C(t))), B), -1)) / n
print(f"straightness S of the independent coupling {S:.4f}")
M, c, h = I, [0.0, 0.0], -1 / 400 # the flow map from t = 1 to 0, RK4 on (M, c)
def F(Mx, cx, t):
a = A(t)
return mul(a, Mx), [a[0][0] * (cx[0] - (1 - t) * mu[0]) + a[0][1] * (cx[1] - (1 - t) * mu[1]) - mu[0],
a[1][0] * (cx[0] - (1 - t) * mu[0]) + a[1][1] * (cx[1] - (1 - t) * mu[1]) - mu[1]]
for k in range(400):
t = 1 + k * h
k1 = F(M, c, t)
k2 = F(add(M, k1[0], h / 2), [c[j] + h / 2 * k1[1][j] for j in range(2)], t + h / 2)
k3 = F(add(M, k2[0], h / 2), [c[j] + h / 2 * k2[1][j] for j in range(2)], t + h / 2)
k4 = F(add(M, k3[0], h), [c[j] + h * k3[1][j] for j in range(2)], max(0.0, t + h))
M = add(M, add(add(k1[0], k2[0], 2), add(k3[0], k4[0], 0.5), 2), h / 6)
c = [c[j] + h / 6 * (k1[1][j] + 2 * k2[1][j] + 2 * k3[1][j] + k4[1][j]) for j in range(2)]
print(f"rectified map: M = {[[round(v, 4) for v in r] for r in M]}, c = {[round(v, 4) for v in c]}")
D = add(M, I, -1)
print(f"cost of the rectified coupling {sum(D[i][j] ** 2 for i in range(2) for j in range(2)) + c[0] ** 2 + c[1] ** 2:.4f}")
Sigma = [[1.3818, 0.8258], [0.8258, 0.7807]], Sigma^1/2 = [[1.0875, 0.4464], [0.4464, 0.7625]]
independent cost 5.6225 optimal cost 1.9225
straightness S of the independent coupling 2.9060
rectified map: M = [[1.0875, 0.4464], [0.4464, 0.7625]], c = [1.1, 0.5]
cost of the rectified coupling 1.9225
Stochastic interpolants, bridges and generators
The conditional-average construction extends beyond independent Gaussian noise and deterministic flows. We can change the endpoint laws, pair the endpoints differently, or learn the generator of a more general Markov process. Each extension keeps the same basic strategy: construct a simple conditional problem, then learn the quantity that combines those problems into the desired marginal evolution.
Stochastic interpolants. Albergo, Boffi and Vanden-Eijnden join two arbitrary laws with \(x_t = I(t, x_0, x_1) + \gamma_t W\), using a deterministic interpolant \(I\) and optional latent noise. The resulting marginals satisfy a transport equation and a family of forward and backward Fokker–Planck equations with an adjustable diffusion coefficient. This recovers the freedom of the \(\lambda\)-family (4.2) from the interpolation side; neither endpoint has to be Gaussian.
Optimizing the interpolant connects this construction to the Schrödinger bridge between the endpoint laws, an entropic optimal transport problem. It gives a stochastic counterpart to the transport questions in Section 6.4.Albergo, Boffi and Vanden-Eijnden (2025, cited in Section 4.2): interpolants "bridge any two probability density functions exactly in finite time", their density satisfies "a transport equation as well as a family of forward and backward Fokker-Planck equations with tunable diffusion coefficient", and "stochastic interpolants recover the Schrödinger bridge between the two target densities when explicitly optimizing over the interpolant".
Couplings beyond independence. Theorem 6.3.2 allows any coupling with the required marginals. We can pair noise and data by an optimal assignment within each batch, use pairs produced by reflow, or start with paired observations such as an image and its corruption or a molecule and its conformer. The regression construction remains the same; the choice of pairs changes the path it learns.
Generator matching. The marginalization proof uses the continuity equation's linearity in the law. The forward equation of a general Markov process has the same property. Holderrieth et al. extend the construction by starting with a conditional generator, learning its conditional average, and sampling the resulting Markov process. Diffusions, flows, jump processes, and the discrete chains of Section 9 all fit this framework. The later chapters of Lipman et al.'s guide develop the theory in this form.P. Holderrieth, M. Havasi, J. Yim et al., "Generator matching: generative modeling with arbitrary Markov processes", ICLR 2025 (arXiv 2410.20587): "we construct conditional generators which generate single data points, then learn to approximate the marginal generator which generates the full data distribution … unifies various generative modeling methods, including diffusion models, flow matching and discrete diffusion models"; Lipman et al., "Flow matching guide and code" (2024), Chapters 6 (continuous-time Markov chain models), 7 (discrete flow matching), 8 (continuous-time Markov processes) and 9 (generator matching).
Remark 6.5.1 (takeaway: what flow matching added). With independent Gaussian paths, flow matching learns the probability-flow velocity of Section 4.3 and samples through the same ODE. Its broader contribution is to make the path a design choice: choose a coupling of the endpoints, then choose an interpolation between them.
The linear path, logit-normal time weighting, mini-batch couplings, and reflow use that freedom in different ways. Section 8 shows how it supports the few-step image models developed from 2024 to 2026.