Flow matching, rectified flow and optimal transport

Before this lesson: score matching and the "conditional target" idea (lesson 13); flow matching for images, if you want the ML side first (generative models 10); matrix square roots (linear algebra 07); Lagrange multipliers (optimization 08); conditional expectation (probability 04). Time: about 120 minutes, plus the notebook.

---

Why this matters #

• Fast generative ensembles. A generative weather model samples many possible future states. The cost of one sample is (solver steps) × (one network call). Straighter paths mean fewer steps. ArchesWeatherGen (Couairon et al., 2024) is a weather ensemble model trained with flow matching. Stable Diffusion 3 (Esser et al., 2024) trains rectified-flow transformers for images.
• Comparing distributions. The Wasserstein distance $W_2$ compares two distributions by how far mass must move. It compares a generated ensemble with observations, and its Gaussian version (the Bures formula) is the formula behind FID, the standard score for image generators.
• Seismic imaging. A Wasserstein misfit between observed and simulated seismograms is less prone to cycle skipping than the least-squares misfit in full-waveform inversion (see lesson 04).
• Data assimilation. The ensemble transform particle filter (Reich, 2013) turns a weighted ensemble into an equally weighted one with an optimal transport plan (lesson 06).

1. Predict first #

2. The picture #

One pair: a straight line #

Fix a noise point $x_0$ and a data point $x_1$. The straight line $x_t = (1 - t)x_0 + t x_1$ has the constant velocity $\dot x_t = x_1 - x_0$. One Euler step of size 1 from $x_0$ lands exactly on $x_1$.

Many pairs: the lines cross, the flow bends #

Now draw many pairs. Their lines cross. An ODE $\dot x = u_t(x)$ cannot have crossing paths: through each point $(x, t)$ passes exactly one solution. So the best single velocity at $(x, t)$ is the average direction of all lines that pass through $x$ at time $t$:

Averaging over crossing lines bends the paths. In the figure, grey lines are the pairs and orange curves are the paths of the flow with this average velocity, which is what a perfect model learns.

Interactive figure: loading…

The 1D example, in full #

Take noise $\mathcal N(0, 1)$ and data $\pm 2$. At $t = 0$ the line through $x_0$ goes to $-2$ or to $+2$ with equal chance. So $u_0(x) = \mathbb E[x_1] - x = -x$: every path starts toward the mean $0$, and only later turns to $-2$ or $+2$. Section 4(d) derives the whole field: $u_t(x) = \big(2\tanh(2tx/(1 - t)^2) - x\big)/(1 - t)$.

Interactive figure: loading…

A better pairing gives straight paths #

Send negative noise to $-2$ and positive noise to $+2$. Now no lines cross. The average over the lines through a point is the one line itself, and the paths are straight. This sorted pairing is the optimal transport coupling. There are two ways to straighten the paths:

1. choose a better coupling before training (exact OT, or minibatch OT);
2. train once, then re-pair each noise point with the end point of the model's own path, and train again (reflow).

3. Definitions and main results #

Flows and the continuity equation #

A velocity field $u_t(x)$ defines a flow: the solution of $\dot x = u_t(x)$. If $x(0) \sim p_0$, the density $p_t$ of $x(t)$ solves the continuity equation

In words: the density at a point changes only because mass flows in or out; no mass is created. We say "$u_t$ generates $p_t$". (Lesson 12 derived the Fokker–Planck equation; with zero diffusion it is this equation.)

Flow matching needs a path of densities $p_t$ from noise $p_0$ to data $p_1$ and a field $u_t$ that generates it. The flow matching (FM) loss is

We do not know $u_t$, so we cannot compute this loss. The way out uses pairs.

A coupling $\pi$ of $p_0$ and $p_1$ is a joint distribution of $(x_0, x_1)$ whose marginals are $p_0$ and $p_1$. The independent coupling is $p_0 \otimes p_1$: draw $x_0$ and $x_1$ separately. Given a coupling, set $x_t = (1 - t)x_0 + t x_1$, let $p_t$ be the law of $x_t$, and define the marginal velocity

Theorem 1 (the marginalization trick) #

The marginal velocity $u_t$ generates $p_t$. When $u_t$ is regular enough (for example, Lipschitz in $x$), the ODE $\dot x = u_t(x)$ started from $p_0$ has the marginals $p_t$ and ends at the data distribution $p_1$.

Why it is true, in one line. Each pair moves mass along its own line. At a point $x$, the total flow of mass is the sum over the lines through $x$; dividing by the density $p_t(x)$ gives the average velocity there.

Proof of Theorem 1 (weak form)

Take any smooth test function $\varphi$ with compact support.

1. Along each line, $\frac{d}{dt}\varphi(x_t) = \nabla\varphi(x_t) \cdot (x_1 - x_0)$. Take the expectation: $\frac{d}{dt}\mathbb E[\varphi(x_t)] = \mathbb E[\nabla\varphi(x_t) \cdot (x_1 - x_0)]$.
2. Condition on $x_t$ (the tower property $\mathbb E[Z] = \mathbb E[\mathbb E[Z \mid x_t]]$): this equals $\mathbb E[\nabla\varphi(x_t) \cdot u_t(x_t)] = \int \nabla\varphi \cdot u_t\, p_t\, dx$.
3. The left side is $\frac{d}{dt}\int \varphi\, p_t\, dx$. Integrate the right side by parts: $\int \varphi\, \partial_t p_t\, dx = -\int \varphi\, \nabla\cdot(p_t u_t)\, dx$ for every $\varphi$. This is the weak form of $\partial_t p_t + \nabla\cdot(p_t u_t) = 0$. $\blacksquare$

Theorem 2 (conditional flow matching) #

Define the conditional flow matching (CFM) loss

Then $L_{\mathrm{CFM}}(\theta) = L_{\mathrm{FM}}(\theta) + V$, where $V$ does not depend on $\theta$. So the two losses have the same gradients and the same minimizer $v^\star = u_t$.

1. Expand: $\lVert v - (x_1 - x_0)\rVert^2 = \lVert v\rVert^2 - 2\langle v, x_1 - x_0\rangle + \lVert x_1 - x_0\rVert^2$.
2. The term $\mathbb E\lVert v_\theta(x_t, t)\rVert^2$ is the same in both losses, because $x_t \sim p_t$.
3. Cross term, by the tower property: $\mathbb E\langle v_\theta(x_t, t), x_1 - x_0\rangle = \mathbb E\langle v_\theta(x_t, t), \mathbb E[x_1 - x_0 \mid x_t]\rangle = \mathbb E\langle v_\theta(x_t, t), u_t(x_t)\rangle$.
4. What remains does not contain $\theta$: $V = \int_0^1 \mathbb E\lVert x_1 - x_0 - u_t(x_t)\rVert^2\, dt \ge 0$. $\blacksquare$

This is the argument of denoising score matching in lesson 13: a conditional target whose conditional mean is the marginal target. The floor $V$ is the variance of the direction given the position. It is zero exactly when no two lines meet at the same time.

Link to diffusion. With source $p_0 = \mathcal N(0, I)$ and the independent coupling, $u_t(x) = \big(x + (1 - t)\nabla\log p_t(x)\big)/t$ (check yourself 6). A velocity model and a score model carry the same information.

Theorem 3 (rectified flow lowers convex transport costs) #

Let $u_t$ be the marginal velocity of a coupling $(X_0, X_1)$, and let $Z$ solve $\dot Z_t = u_t(Z_t)$ with $Z_0 = X_0$. Assume this ODE has a unique solution (Liu, Gong and Liu, 2022, call the coupling rectifiable). Then $(Z_0, Z_1)$ is again a coupling of $p_0$ and $p_1$ (Theorem 1), called the rectified flow, and for every convex cost $c$:

Proof of Theorem 3

1. $Z_1 - Z_0 = \int_0^1 u_t(Z_t)\, dt$, so by Jensen's inequality over $t$: $c(Z_1 - Z_0) \le \int_0^1 c(u_t(Z_t))\, dt$.
2. $Z_t$ and $X_t$ have the same law (Theorem 1), so $\mathbb E\, c(u_t(Z_t)) = \mathbb E\, c(u_t(X_t))$.
3. $u_t(X_t) = \mathbb E[X_1 - X_0 \mid X_t]$, so by conditional Jensen: $c(u_t(X_t)) \le \mathbb E[c(X_1 - X_0) \mid X_t]$.
4. Take expectations and integrate over $t$: $\mathbb E\, c(Z_1 - Z_0) \le \mathbb E\, c(X_1 - X_0)$. $\blacksquare$

The new coupling is deterministic: $Z_1$ is a function of $Z_0$.

Straightness. Liu et al. define $S(Z) = \int_0^1 \mathbb E\lVert (Z_1 - Z_0) - u_t(Z_t)\rVert^2\, dt$. $S = 0$ means straight paths at constant speed, so one Euler step is exact. For $c = \lVert\cdot\rVert^2$ the two Jensen steps become equalities plus variance terms, and

Reflow. Train a new flow on the pairs $(Z_0, Z_1)$, then repeat. The transport cost falls in each round by $S + V$ and cannot fall below $0$. Summing over $K$ rounds gives $\min_{1 \le k \le K} S(Z^k) \le \mathbb E\lVert X_1 - X_0\rVert^2/K$ for the $k$-th rectified flows $Z^k$. The price: each round must simulate the current flow to make new pairs, and each round fits a model to model output, so errors add up. In practice one or two reflows are used, often followed by distillation into a one-step model (generative models 12).

Optimal transport: Monge and Kantorovich #

Let $\mu$ and $\nu$ be probability distributions and $c(x, y)$ a cost of moving a unit of mass from $x$ to $y$.

• Monge problem: find a map $T$ with $T_\#\mu = \nu$ (if $x \sim \mu$ then $T(x) \sim \nu$) that minimizes $\int c(x, T(x))\, d\mu(x)$. A map may not exist: $\mu = \delta_0$ cannot be split into $\tfrac12\delta_{-1} + \tfrac12\delta_1$.
• Kantorovich problem: minimize $\int c\, d\pi$ over all couplings $\pi \in \Pi(\mu, \nu)$. Mass may now be split. This is a linear program, and under mild conditions (for example, a continuous cost bounded below) a minimizer exists. For discrete weights $a \in \R^n$, $b \in \R^m$ and a cost matrix $C$: minimize $\langle C, P\rangle = \sum_{ij} C_{ij}P_{ij}$ over transport plans $P \ge 0$ with $P\mathbf 1 = a$ and $P^\top\mathbf 1 = b$.
• Birkhoff–von Neumann: with $n = m$ and weights $1/n$, some optimal $P$ is a permutation matrix divided by $n$. Then Monge and Kantorovich agree, and the Hungarian algorithm (scipy.optimize.linear_sum_assignment) solves the problem in $O(n^3)$.
• Wasserstein distance: $W_p(\mu, \nu) = \big(\min_\pi \mathbb E_\pi\lVert x - y\rVert^p\big)^{1/p}$. It is a metric on distributions with finite $p$-th moments.
• Brenier's theorem: for the cost $\lVert x - y\rVert^2$, $\mu$ with a density, and finite second moments, the optimal map exists, is unique ($\mu$-almost everywhere), and is the gradient of a convex function, $T = \nabla\phi$. Conversely, a map $\nabla\phi$ with convex $\phi$ that pushes $\mu$ to $\nu$ is optimal.

1D OT is sorting. On a line, for every convex cost $c(x - y)$, the monotone coupling is optimal: the $i$-th smallest source point goes to the $i$-th smallest target point. In terms of the quantile functions $F^{-1}$ and $G^{-1}$, $W_p^p = \int_0^1 \lvert F^{-1}(s) - G^{-1}(s)\rvert^p\, ds$. The reason: if two pairs cross, $x < x'$ but $y > y'$, then swapping partners does not raise the cost, and for a strictly convex cost it lowers it.

Interactive figure: loading…

Theorem 4 (Gaussian $W_2$, the Bures formula) #

For $\mu = \mathcal N(m_1, S_1)$ and $\nu = \mathcal N(m_2, S_2)$,

The first term moves the centre; the second term (the Bures term) reshapes the covariance. The optimal map is affine: $T(x) = m_2 + A(x - m_1)$ with the symmetric positive definite matrix

Proof sketch of Theorem 4 (for $S_1 \succ 0$; the singular case follows by continuity)

1. For any coupling, $\mathbb E\lVert X - Y\rVert^2 = \lVert m_1 - m_2\rVert^2 + \operatorname{tr}S_1 + \operatorname{tr}S_2 - 2\operatorname{tr}K$ with $K = \operatorname{Cov}(X, Y)$. So OT means: make $\operatorname{tr}K$ as large as possible.
2. Try $Y = m_2 + A(X - m_1)$ with $A$ symmetric positive semidefinite. Then $Y \sim \mathcal N(m_2, AS_1A)$, so we need $AS_1A = S_2$. The matrix $A$ above solves it: substitute and use $S_1^{-1/2}S_1S_1^{-1/2} = I$.
3. Then $K = S_1A$, and by the cyclic property of the trace, $\operatorname{tr}(S_1A) = \operatorname{tr}\big(S_1^{1/2}S_2S_1^{1/2}\big)^{1/2}$.
4. The map is $\nabla\phi$ for the convex $\phi(x) = \tfrac12(x - m_1)^\top A(x - m_1) + m_2^\top x$, so it is optimal by Brenier's theorem. Put $\operatorname{tr}K$ into step 1. $\blacksquare$

Special cases. In 1D, $W_2^2 = (m_1 - m_2)^2 + (s_1 - s_2)^2$ with standard deviations $s_1, s_2$. For commuting covariances, the Bures term is $\lVert S_1^{1/2} - S_2^{1/2}\rVert_F^2$. For non-Gaussian distributions with the same means and covariances, the formula is a lower bound on $W_2^2$ (Gelbrich's bound).

Interactive figure: loading…

Entropic OT and the Sinkhorn algorithm #

Exact OT between $n$ points costs about $O(n^3)$. A popular alternative adds an entropy term. With $H(P) = -\sum_{ij} P_{ij}\log P_{ij}$ and $\varepsilon > 0$, entropic OT is

The problem is strictly convex, so the solution $P_\varepsilon$ is unique. The entropy term rewards spread-out plans, so $\varepsilon$ sets how blurred the plan is.

Form of the solution. Use multipliers $f_i$ for the row constraints and $g_j$ for the column constraints (optimization 08). Setting the derivative of the Lagrangian with respect to $P_{ij}$ to zero gives $C_{ij} + \varepsilon(\log P_{ij} + 1) - f_i - g_j = 0$. Absorb the constant into $f$:

Equivalently $P = \operatorname{diag}(u)\, K \operatorname{diag}(v)$ with the Gibbs kernel $K = e^{-C/\varepsilon}$, $u = e^{f/\varepsilon}$, $v = e^{g/\varepsilon}$.

Sinkhorn alternates two scalings. Each one makes one marginal exact: $u \leftarrow a \oslash (Kv)$, then $v \leftarrow b \oslash (K^\top u)$ (elementwise division). In the log domain, with $\operatorname{LSE}(z) = \log\sum e^{z}$:

Each iteration costs $O(nm)$ and is all matrix–vector work, so it runs well on a GPU. The number of iterations grows as $\varepsilon$ shrinks. Why the log domain: $e^{-746} = 0$ in float64, so kernel entries with $C_{ij}/\varepsilon$ above about 745 underflow. When a whole row of $K$ is zero, $a \oslash (Kv)$ divides by zero. The LSE with its maximum subtracted never underflows.

Interactive figure: loading…

Theorem 5: the bias of entropic OT, and the dual problem

Theorem 5. $0 \le \langle C, P_\varepsilon\rangle - \mathrm{OT} \le \varepsilon\min\big(H(a), H(b)\big)$, where OT is the exact optimal cost and $H(a) = -\sum_i a_i\log a_i$.

1. $P_\varepsilon$ is a coupling and OT is the minimum over couplings, so the gap is $\ge 0$.
2. $P_\varepsilon$ beats an exact OT plan $P^\star$ on the entropic objective: $\langle C, P_\varepsilon\rangle - \varepsilon H(P_\varepsilon) \le \mathrm{OT} - \varepsilon H(P^\star)$. So the gap is $\le \varepsilon\big(H(P_\varepsilon) - H(P^\star)\big)$.
3. For every coupling, $\max(H(a), H(b)) \le H(P) \le H(a) + H(b)$: a joint entropy is at least each marginal entropy and at most their sum.
4. So the gap is $\le \varepsilon\big(H(a) + H(b) - \max(H(a), H(b))\big) = \varepsilon\min(H(a), H(b))$. $\blacksquare$

As $\varepsilon \to 0$, $P_\varepsilon$ tends to an exact OT plan (the one of largest entropy if there are several); as $\varepsilon \to \infty$, it tends to $ab^\top$.

The dual. Entropic OT has the unconstrained concave dual problem $\max_{f, g}\ \langle f, a\rangle + \langle g, b\rangle - \varepsilon\sum_{ij} e^{(f_i + g_j - C_{ij})/\varepsilon}$ (up to a constant). Setting the gradient in $f$ to zero gives exactly the $f$-update above, and the same for $g$. So Sinkhorn is block coordinate ascent on the dual: each half-step maximizes over one block exactly. See Peyré and Cuturi, Computational Optimal Transport, ch. 4.

Theorem 6 (optimal transport paths do not cross) #

Let $\pi$ be optimal for the cost $\lVert x - y\rVert^2$. If two pairs $(x_0, x_1)$ and $(x_0', x_1')$ in the support of $\pi$ are at the same point at some time $t \in (0, 1)$, then they are the same pair.

1. Optimality means that swapping partners cannot lower the cost (c-cyclical monotonicity): $\lVert x_0 - x_1\rVert^2 + \lVert x_0' - x_1'\rVert^2 \le \lVert x_0 - x_1'\rVert^2 + \lVert x_0' - x_1\rVert^2$.
2. Expand both sides. The squared norms cancel, and what is left is $\langle a, b\rangle \ge 0$ with $a = x_0 - x_0'$ and $b = x_1 - x_1'$.
3. Meeting at time $t$ means $(1 - t)a + tb = 0$, so $b = -\frac{1 - t}{t}a$.
4. Then $\langle a, b\rangle = -\frac{1 - t}{t}\lVert a\rVert^2 \ge 0$ forces $a = 0$, and so $b = 0$. $\blacksquare$

Consequences. Under the OT coupling, the position $x_t$ determines its pair, so $u_t(x_t) = x_1 - x_0$: the floor $V$ is $0$, the flow moves along the straight lines, $S = 0$, and one Euler step is exact. An optimal coupling is a fixed point of reflow. Its kinetic energy $\int_0^1 \mathbb E\lVert u_t\rVert^2\, dt$ equals $W_2^2$, the smallest possible value: by the Benamou–Brenier formula, $W_2^2(p_0, p_1)$ is the minimum of $\int_0^1 \mathbb E_{p_t}\lVert v_t\rVert^2\, dt$ over all pairs $(p_t, v_t)$ that solve the continuity equation from $p_0$ to $p_1$.

Minibatch OT (Tong et al., 2023; Pooladian et al., 2023). In each training batch, draw $B$ noise points and $B$ data points and re-pair them by exact OT (an assignment problem). This coupling lies between the independent one and the OT one: its transport cost is lower than the independent coupling's, and it approaches the OT coupling as $B \to \infty$. Limits: an assignment costs about $O(B^3)$, and in high dimension a batch is a very coarse sample of the two distributions (empirical OT converges slowly, at a rate like $B^{-1/d}$), so the paths become only partly straight.

Answers to section 1. (1) All samples land at $0$, the data mean, since $x_0 + u_0(x_0) = x_0 + (\mathbb E[x_1] - x_0) = \mathbb E[x_1]$: a complete collapse, at distance 2 from every data point. (2) No, by Theorem 6. Their lines may still cross in space at different times, but two particles never occupy the same place at the same time.

4. Worked example #

(a) 1D OT by sorting. Take $x = (0, 1, 3)$ and $y = (5, 2, 4)$, weights $\tfrac13$, cost $\lvert x - y\rvert^2$.

1. Sort both: $x = (0, 1, 3)$, $y = (2, 4, 5)$.
2. Pair in order: $0 \to 2$, $1 \to 4$, $3 \to 5$.
3. Cost: $\big((0 - 2)^2 + (1 - 4)^2 + (3 - 5)^2\big)/3 = (4 + 9 + 4)/3 = 17/3 \approx 5.667 = W_2^2$.
4. The pairing in the given order, $0 \to 5$, $1 \to 2$, $3 \to 4$, costs $(25 + 1 + 1)/3 = 9$.

(b) The Bures formula. $m_1 = 0$, $S_1 = I$; $m_2 = (1, 0)$, $S_2 = \begin{pmatrix} 2 & 1 \\ 1 & 2 \end{pmatrix}$.

1. Mean term: $\lVert m_1 - m_2\rVert^2 = 1$.
2. Since $S_1 = I$: $\big(S_1^{1/2}S_2S_1^{1/2}\big)^{1/2} = S_2^{1/2}$.
3. $S_2$ has eigenvalues $3$ (eigenvector $(1, 1)/\sqrt2$) and $1$ (eigenvector $(1, -1)/\sqrt2$). So $S_2^{1/2}$ has eigenvalues $\sqrt3$ and $1$, and $\operatorname{tr}S_2^{1/2} = \sqrt3 + 1$.
4. Bures term: $\operatorname{tr}(I) + \operatorname{tr}(S_2) - 2\operatorname{tr}(S_2^{1/2}) = 2 + 4 - 2(\sqrt3 + 1) = 4 - 2\sqrt3 \approx 0.536$.
5. $W_2^2 = 1 + 4 - 2\sqrt3 = 5 - 2\sqrt3 \approx 1.536$. The OT map is $T(x) = m_2 + S_2^{1/2}x$, with $S_2^{1/2} = \tfrac12\begin{pmatrix} \sqrt3 + 1 & \sqrt3 - 1 \\ \sqrt3 - 1 & \sqrt3 + 1\end{pmatrix} \approx \begin{pmatrix} 1.37 & 0.37 \\ 0.37 & 1.37\end{pmatrix}$ (the "worked example" preset of the Bures figure).

(c) Sinkhorn on a 2×2 problem. $a = b = (\tfrac12, \tfrac12)$, $C = \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}$.

1. Exact OT keeps each mass in place: cost 0.
2. By symmetry $u = v = (c, c)$, so $P_{11} = c^2$ and $P_{12} = c^2e^{-1/\varepsilon}$.
3. The row sum must be $\tfrac12$: $c^2(1 + e^{-1/\varepsilon}) = \tfrac12$. So $P_{12} = \dfrac{1}{2(e^{1/\varepsilon} + 1)}$, and the cost is $\langle C, P\rangle = 2P_{12} = \dfrac{1}{e^{1/\varepsilon} + 1}$.
4. For $\varepsilon = 1$: $P_{11} = 0.3655$, $P_{12} = 0.1345$, cost $0.2689$, below the bound $\varepsilon\log 2 = 0.693$ of Theorem 5. For $\varepsilon = 0.1$: cost $4.54 \times 10^{-5}$.

(d) The collapse, in numbers. Noise $\mathcal N(0, 1)$, data $\pm 2$, independent coupling.

1. Given $x_1$, $x_t = (1 - t)x_0 + tx_1 \sim \mathcal N(tx_1, (1 - t)^2)$.
2. Bayes' rule with equal prior weights: the posterior odds of $x_1 = +2$ against $-2$ are $\exp\!\big(-\frac{(x - 2t)^2 - (x + 2t)^2}{2(1 - t)^2}\big) = \exp\!\big(\frac{4tx}{(1 - t)^2}\big)$. So $\mathbb E[x_1 \mid x_t = x] = 2\tanh\big(2tx/(1 - t)^2\big)$.
3. Since $x_1 - x_0 = (x_1 - x_t)/(1 - t)$, the marginal velocity is $u_t(x) = \big(2\tanh(2tx/(1 - t)^2) - x\big)/(1 - t)$.
4. One Euler step: $x_0 + u_0(x_0) = x_0 + (0 - x_0) = 0$ for every $x_0$.
5. The exact flow from $x_0 = 0.5$ first moves toward $0$ (down to about $0.451$ at $t \approx 0.17$), then turns and reaches $+2$. Since 1D flows keep the order of points, the rectified coupling is $x_1 = 2\operatorname{sign}(x_0)$: the OT coupling.
6. Transport costs. Independent: $\mathbb E\lvert x_1 - x_0\rvert^2 = \mathbb E x_1^2 + \mathbb E x_0^2 = 4 + 1 = 5$. OT: $\mathbb E\lvert 2\operatorname{sign}(x_0) - x_0\rvert^2 = 4 - 4\,\mathbb E\lvert x_0\rvert + 1 = 5 - 4\sqrt{2/\pi} \approx 1.808 = W_2^2$. The difference $5 - 1.808 = 3.192$ is $S + V$ of the first rectification.

5. In code #

NumPy and SciPy check examples (a)–(c). The Hungarian algorithm in linear_sum_assignment solves the assignment problem exactly; logsumexp gives a stable LSE.

import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.linalg import sqrtm
from scipy.special import logsumexp

# (a) 1D OT: sorting agrees with the exact assignment problem.
x, y = np.array([0.0, 1.0, 3.0]), np.array([5.0, 2.0, 4.0])
C = (x[:, None] - y[None, :]) ** 2                 # cost matrix |x_i - y_j|^2
rows, cols = linear_sum_assignment(C)              # Hungarian algorithm: exact OT for equal weights
print(C[rows, cols].mean(), cols)                  # 5.6667 [1 2 0]: pairs 0->2, 1->4, 3->5
print(np.mean((np.sort(x) - np.sort(y)) ** 2))     # 5.6667: sort both and pair in order
print(np.mean((x - y) ** 2))                       # 9.0: the pairing in the given order

# (b) Gaussian W2 by the Bures formula.
def bures_w2sq(m1, S1, m2, S2):
    R = sqrtm(S1)
    return np.sum((m1 - m2) ** 2) + np.trace(S1 + S2 - 2 * sqrtm(R @ S2 @ R)).real
m1, S1 = np.zeros(2), np.eye(2)
m2, S2 = np.array([1.0, 0.0]), np.array([[2.0, 1.0], [1.0, 2.0]])
print(bures_w2sq(m1, S1, m2, S2), 5 - 2 * np.sqrt(3))   # 1.5359 1.5359

# (c) Log-domain Sinkhorn on the 2 x 2 problem; closed form P12 = 1 / (2 (e^{1/eps} + 1)).
def sinkhorn_log(a, b, C, eps, iters=2000):
    f, g = np.zeros(len(a)), np.zeros(len(b))
    for _ in range(iters):
        f = eps * np.log(a) - eps * logsumexp((g[None, :] - C) / eps, axis=1)
        g = eps * np.log(b) - eps * logsumexp((f[:, None] - C) / eps, axis=0)
    return np.exp((f[:, None] + g[None, :] - C) / eps)
a = b = np.array([0.5, 0.5])
C = np.array([[0.0, 1.0], [1.0, 0.0]])
for eps in (1.0, 0.1):
    P = sinkhorn_log(a, b, C, eps)
    print(eps, P[0, 1], 1 / (2 * (np.exp(1 / eps) + 1)), (C * P).sum())
# 1.0 0.13447 0.13447 0.26894    (exact OT cost is 0; the gap is below eps*log 2 = 0.693)
# 0.1 2.2699e-05 2.2699e-05 4.5398e-05

And example (d): the closed-form marginal velocity, one Euler step against an accurate ODE solve, Euler with a few steps, and the two transport costs by Monte Carlo:

from scipy.integrate import solve_ivp

def u(t, x):                                       # exact marginal velocity of example (d)
    return (2 * np.tanh(2 * t * x / (1 - t) ** 2) - x) / (1 - t)

for x0 in (0.5, -0.1, 1.5):
    one_step = x0 + u(0.0, x0)                     # one Euler step from t = 0 to t = 1
    sol = solve_ivp(u, (0.0, 1 - 1e-6), [x0], method="Radau", rtol=1e-9, atol=1e-12)
    print(x0, one_step, round(sol.y[0, -1], 4))
# 0.5 0.0 2.0 / -0.1 0.0 -2.0 / 1.5 0.0 2.0: one step gives 0; the exact flow ends at ±2

z = np.random.default_rng(0).standard_normal(200_000)
for N in (1, 2, 4, 8):                             # Euler with N steps from many noise samples
    xs = z.copy()
    for k in range(N):
        xs = xs + u(k / N, xs) / N
    print(N, np.mean(np.minimum(np.abs(xs - 2), np.abs(xs + 2))))
# mean distance to the nearest data point: 2.0, 0.51, 0.011, 1.2e-05

x1 = np.random.default_rng(1).choice([-2.0, 2.0], size=z.size)
print(np.mean((x1 - z) ** 2), np.mean((2 * np.sign(z) - z) ** 2), 5 - 4 * np.sqrt(2 / np.pi))
# about 5.0, about 1.807, 1.8085 (= W2^2)

6. Where it shows up #

• Generative models. Stable Diffusion 3 trains rectified-flow transformers (Esser et al., 2024). Meta's Movie Gen is trained with flow matching (2024).
• Weather ensembles. ArchesWeatherGen (Couairon et al., 2024) uses flow matching to sample weather states around a deterministic forecast. Ensemble models pay per network call and per member, so the number of solver steps matters (lesson 16 compares model families).
• Seismology. Wasserstein misfits for full-waveform inversion (Engquist and Froese, 2014; Métivier et al., 2016) compare seismograms as distributions of mass along time; they are more convex with respect to time shifts than the least-squares misfit, which reduces cycle skipping (lesson 04).
• Evaluation. FID is the Bures $W_2^2$ between Gaussian fits of two sets of image features. The same formula compares an ensemble's mean and covariance with observations; by Gelbrich's bound it never overstates the true $W_2$.
• Data assimilation. The ensemble transform particle filter (Reich, 2013) turns a weighted ensemble into an equally weighted one with an OT plan.

7. Common mistakes #

• Thinking the model learns $x_1 - x_0$. It learns the conditional mean $u_t(x)$. When paths cross, the best CFM loss is $V > 0$, and the learned paths are curved even though every conditional path is straight.
• Mixing time conventions. Flow matching uses $t = 0$ noise, $t = 1$ data; score-based diffusion often uses the opposite. Converting flips the sign of the velocity.
• $W_2$ versus $W_2^2$. The optimal cost for the squared distance is $W_2^2$, and so is the Bures formula. Take the square root before you call it a distance.
• sqrtm on singular or nearly singular covariances. It can return NaN or complex values. Use eigh and clip tiny negative eigenvalues to 0.
• Sinkhorn with $e^{-C/\varepsilon}$ for small $\varepsilon$. It underflows; use the log domain. And remember that $\langle C, P_\varepsilon\rangle$ is above the exact OT cost (Theorem 5).
• Minibatch OT done once, or reflow done many times. The pairing must be recomputed in every batch. Each reflow round fits a model to model output, so errors pile up.

8. Check yourself #

1. Compute $W_2$ between $\mathcal N(1, 4)$ and $\mathcal N(-2, 9)$ in 1D (the second numbers are variances).
   Answer

   $W_2^2 = (1 - (-2))^2 + (2 - 3)^2 = 9 + 1 = 10$, so $W_2 = \sqrt{10} \approx 3.162$.
2. For the independent coupling, why does the CFM loss of the best possible model stay above zero, and what is its value?
   Answer

   Lines meet, so the direction $x_1 - x_0$ is random given the position $x_t$. The best model predicts its conditional mean $u_t(x_t)$; the loss left over is $V = \int_0^1 \mathbb E\lVert x_1 - x_0 - u_t(x_t)\rVert^2\,dt$, the conditional variance.
3. Costs $C_{ij}$ up to 10 and $\varepsilon = 0.01$. What goes wrong in the scaling form of Sinkhorn, and what fixes it?
   Answer

   $e^{-C_{ij}/\varepsilon}$ is as small as $e^{-1000}$, which is $0$ in float64; $Kv$ can be $0$ and $a \oslash (Kv)$ divides by zero. The log-domain updates with a stable LSE (subtract the maximum) fix it.
4. Where does the proof of Theorem 3 use convexity, and why are there two uses?
   Answer

   Step 1: Jensen over time, because $Z_1 - Z_0$ is the time average of $u_t(Z_t)$. Step 3: conditional Jensen, because $u_t(X_t)$ is the conditional mean of $X_1 - X_0$ given $X_t$.
5. Source $\mathcal N(0, 1)$, target $\mathcal N(3, 4)$ in 1D. Give the OT map, the velocity field $u_t(x)$ of the straight OT flow, and say whether one Euler step is exact.
   Answer

   $T(x) = 3 + 2x$. On the path, $x_t = x_0 + t(3 + x_0) = 3t + (1 + t)x_0$, so $x_0 = (x - 3t)/(1 + t)$ and $u_t(x) = 3 + x_0 = (3 + x)/(1 + t)$. One step: $x_0 + u_0(x_0) = 3 + 2x_0 = T(x_0)$, so it is exact.
6. With source $\mathcal N(0, I)$ and linear paths, express the marginal velocity through the score of $p_t$.
   Answer

   $x_t = tx_1 + (1 - t)x_0$, and Tweedie's formula ([lesson 13](/ml-for-science/13-score-matching-diffusion/)) for the Gaussian part $(1 - t)x_0$ gives $\mathbb E[x_0 \mid x_t] = -(1 - t)\nabla\log p_t(x_t)$. Then $u_t(x) = \mathbb E[x_1 - x_0 \mid x_t = x] = (x - \mathbb E[x_0 \mid x_t = x])/t = \big(x + (1 - t)\nabla\log p_t(x)\big)/t$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the continuity equation for a density $p_t$ moved by a velocity field $u_t$.

   Show answer

2. What is the conditional flow matching (CFM) loss for linear paths?

   Show answer

3. State the marginalization trick.

   Show answer

4. Why are the learned paths curved for an independent coupling?

   Show answer

5. What does reflow do, and why does it help?

   Show answer

6. What is the difference between the Monge and the Kantorovich problem?

   Show answer

7. How do you solve OT between two equal-size samples on a line?

   Show answer

8. State the Bures formula for $W_2^2$ between two Gaussians.

   Show answer

9. Write the log-domain Sinkhorn update for $f$.

   Show answer

10. What does $\varepsilon$ control in entropic OT?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Flow matching or diffusion: what is the difference?" With Gaussian paths they learn the same object up to a linear change (velocity ↔ score). They differ in the path, the loss weighting and the sampler. Flow matching regresses a velocity and needs no simulation during training.
• "Why can rectified flow sample in few steps?" Its paths are straighter (Theorem 3, reflow); for $S = 0$ one Euler step is exact. In practice reflow is followed by distillation.
• "What is $W_2$ between two Gaussians? What is FID?" The Bures formula; FID is that formula applied to Gaussian fits of Inception features of real and generated images.
• "Monge versus Kantorovich?" A map versus a coupling. A map may not exist; the coupling problem is a linear program that has a solution under mild conditions.
• "How do you compute OT between $10^5$ points?" Not with the exact linear program (about cubic cost). Use Sinkhorn on a GPU (with its $\varepsilon$ bias), minibatch OT, or sliced OT (1D sorting along random directions).

10. Practice #

The notebook for this lesson runs in Colab or in the browser. It has five exercises, each with a check cell and a folded solution:

1. 1D OT by sorting, checked against the exact assignment (linear_sum_assignment) and brute force.
2. The Bures formula and the OT map between two Gaussians, checked by sampling.
3. Log-domain Sinkhorn on a small problem; its cost approaches the exact OT cost as $\varepsilon \to 0$, within the bound of Theorem 5.
4. The exact marginal velocity for a 1D mixture, and Euler sampling with few and many steps.
5. Minibatch OT pairing: the transport cost and the swap test for batch sizes from 4 to 256.

11. Watch and read #

MIT 6.S184: Flow Matching and Diffusion Models - Lecture 02: Flow Matching (2026)Peter Holderrieth · 83 min

The marginalization trick and the CFM loss of Theorems 1 and 2, with the same time convention.

Flow Matching: Simplifying and Generalizing Diffusion Models | Yaron LipmanProxima · 60 min

Flow matching explained by its first author; compare with Theorems 1 and 2.

Marco Cuturi - A primer on Optimal Transport Theory and Algorithms | MLSS Kraków 2023ML in PL · 192 min

A long tutorial on OT theory and algorithms; use it as a reference for the OT part of section 3.

Optimal transport for machine learning - Gabriel Peyre, Ecole Normale SuperieureThe Alan Turing Institute · 42 min

Optimal transport as a tool in machine learning, by a co-author of Computational Optimal Transport.

• Lipman et al., Flow Matching Guide and Code (2024): §3.5 (the continuity equation), §4.4–4.5 (the marginalization trick; FM and CFM have the same gradient), §4.7 (optimal transport and the linear conditional flow), §4.9.2 (multisample couplings).
• Lipman et al., Flow Matching for Generative Modeling (ICLR 2023): the original paper.
• Liu, Gong and Liu, Flow Straight and Fast (ICLR 2023): rectified flow, Theorems 3.3–3.7 (marginals, convex costs, straightness, reflow).
• Tong et al., Improving and generalizing flow-based generative models with minibatch optimal transport (2023) and Pooladian et al., Multisample Flow Matching (2023): minibatch OT couplings.
• Peyré, Optimal Transport for Machine Learners (2025): §2.6 (1D OT by sorting), §2.3–2.4 (Monge and Brenier), §2.7 (the Bures formula), §3 (Kantorovich), §8 (entropic OT and Sinkhorn), §16 (flow matching).
• Peyré and Cuturi, Computational Optimal Transport (2019): ch. 2 (Kantorovich), ch. 4 (entropic regularization and Sinkhorn).
• Cuturi, Sinkhorn Distances (NeurIPS 2013): entropic OT for machine learning.
• Couairon et al., ArchesWeather & ArchesWeatherGen (2024): flow matching for weather ensembles.