Optimal Transport for Generative Models

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Monge Problem
5. Kantorovich Relaxation
6. Wasserstein Distance
7. Brenier's Theorem
8. Entropic Regularization & Sinkhorn
9. Mini-Batch OT for Flow Matching
10. Dynamic Optimal Transport
11. Wasserstein Gradient Flows
12. Common Pitfalls
13. Summary
14. Exercises

Learning Objectives

1. State and compare the Monge and Kantorovich formulations of OT.
2. Define the Wasserstein-$p$ distance and its metric properties.
3. State Brenier's theorem and its connection to convex functions.
4. Derive the Sinkhorn algorithm for entropic OT.
5. Apply mini-batch OT to improve flow matching training.

Notation

• $\mu, \nu$ — source and target distributions
• $T: \mathbb{R}^d \to \mathbb{R}^d$ — transport map
• $\gamma \in \Pi(\mu, \nu)$ — coupling (joint distribution)
• $W_p$ — Wasserstein-$p$ distance
• $\varepsilon$ — entropic regularization strength

Core Intuition

Optimal transport asks: what's the cheapest way to move mass from one distribution to another? For generative models, this means finding the most efficient mapping from noise to data. OT maps follow straight, non-crossing paths — exactly what flow matching needs for fast generation. The Wasserstein distance gives a geometrically meaningful metric between distributions, unlike KL divergence.

The Monge Problem

Monge (1781): Find a transport map $T$ pushing $\mu$ to $\nu$ that minimizes total cost:

Constraint: $T_\#\mu = \nu$ means for any set $A$: $\mu(T^{-1}(A)) = \nu(A)$.

Problem: The Monge map may not exist (e.g., if $\mu$ is a Dirac mass and $\nu$ is not).

Kantorovich Relaxation

Kantorovich (1942): Relax to probabilistic couplings:

where $\Pi(\mu, \nu) = \{\gamma : \pi_{1\#}\gamma = \mu, \pi_{2\#}\gamma = \nu\}$.

Advantages over Monge:

• Always has a solution (linear program over convex set).
• Allows mass splitting (one source point can go to multiple targets).
• Duality theory available.

Kantorovich dual:

where $\Phi_c = \{(\phi, \psi): \phi(x) + \psi(y) \leq c(x,y)\}$.

Wasserstein Distance

Definition:

Properties:

• $W_p$ is a metric on probability distributions.
• Metrizes weak convergence (unlike KL which can be infinite).
• For $p=2$: connected to displacement convexity and gradient flows.

$W_2$ for Gaussians (closed form):

Brenier's Theorem

Theorem (Brenier, 1991): If $\mu$ is absolutely continuous and $c(x,y) = \|x-y\|^2$, then the optimal transport map exists, is unique, and is the gradient of a convex function:

for some convex $\phi: \mathbb{R}^d \to \mathbb{R}$.

Implication: The OT map is curl-free (irrotational). Paths don't "twist" — they follow the gradient of a potential. This is why OT maps give straight, non-crossing paths.

For generative models: If we could find the Brenier map from $\mathcal{N}(0, I)$ to $p_{\text{data}}$, we'd have the optimal one-step generator. But computing it exactly is intractable for complex distributions.

Entropic Regularization & Sinkhorn

Entropic OT: Add entropy to make the problem smooth:

where $H(\gamma) = -\int \gamma\log\gamma$.

Solution form: $\gamma^*(x, y) = a(x) \cdot e^{-c(x,y)/\varepsilon} \cdot b(y)$.

Sinkhorn algorithm (discrete case): Given cost matrix $\mathbf{C} \in \mathbb{R}^{n \times n}$, kernel $\mathbf{K} = e^{-\mathbf{C}/\varepsilon}$:

Iterate until convergence. Transport plan: $\gamma_{ij} = u_i K_{ij} v_j$.

Cost: $O(n^2)$ per iteration (matrix-vector products). Converges in $O(1/\varepsilon)$ iterations.

Mini-Batch OT for Flow Matching

Problem: Full OT between $p_0$ and $p_1$ is intractable. But we can approximate within each training batch.

Mini-batch OT (Tong et al., 2023):

1. Sample batch: $\{x_0^i\}_{i=1}^B \sim p_0$, $\{x_1^j\}_{j=1}^B \sim p_1$.
2. Compute cost matrix: $C_{ij} = \|x_0^i - x_1^j\|^2$.
3. Solve assignment: $\pi^* = \arg\min_\pi \sum_i C_{i,\pi(i)}$ (Hungarian, $O(B^3)$).
4. Train on paired batch: $(x_0^i, x_1^{\pi^*(i)})$.

Effect on training:

• Without OT: paths cross, velocity is multi-valued, need many sampling steps.
• With mini-batch OT: paths approximately non-crossing within each batch, straighter learned flow.

Quality improvement: 2-5 FID points better, especially at few sampling steps.

Dynamic Optimal Transport

Benamou-Brenier formulation: OT as a fluid dynamics problem:

subject to the continuity equation $\partial_t\rho + \nabla\cdot(\rho v) = 0$ with $\rho_0 = \mu, \rho_1 = \nu$.

Connection to flow matching: The OT velocity field $v_t^*$ gives the kinetic-energy-minimizing flow. This is exactly what we want: minimum-energy paths = straight paths = few-step generation.

Wasserstein Gradient Flows

Many processes in ML can be viewed as gradient flows in Wasserstein space:

where $\mathcal{F}$ is a functional (energy).

Examples:

• $\mathcal{F} = \text{KL}(\rho\|p^*)$: converges to target $p^*$ (Langevin dynamics).
• $\mathcal{F} = H(\rho)$ (entropy): heat equation.
• $\mathcal{F} = W_2^2(\rho, \nu)$: displacement interpolation.

Common Pitfalls

Pitfall 1. Applying full OT (Hungarian algorithm) to large batches. $O(B^3)$ is expensive; use Sinkhorn with $\varepsilon > 0$ for $B > 256$.

Pitfall 2. Setting entropic regularization $\varepsilon$ too large. The coupling becomes too blurry (close to independent); OT benefit vanishes. Use $\varepsilon \approx 0.01$–$0.1$ relative to cost scale.

Pitfall 3. Assuming mini-batch OT equals global OT. It's an approximation; pairs within a batch don't guarantee global optimality. Use larger batches for better approximation.

Summary

• Monge problem: Find cheapest map $T$ from $\mu$ to $\nu$.
• Kantorovich relaxation: Probabilistic couplings; always solvable.
• Wasserstein distance: Geometry-aware metric on distributions.
• Brenier's theorem: OT map = gradient of convex function (no crossing).
• Sinkhorn: Efficient entropic OT via alternating scaling.
• Mini-batch OT improves flow matching by reducing path crossing.
• Dynamic OT connects to minimum-energy flows.

Exercises

Exercise 1. Compute $W_2$ between $\mathcal{N}(0, 1)$ and $\mathcal{N}(3, 4)$ using the closed-form Gaussian formula.

Exercise 2. Implement Sinkhorn for two empirical distributions of 100 points each in 2D. Visualize the transport plan.

Exercise 3. Prove that the Brenier map between two 1D distributions is the quantile function composition: $T = F_\nu^{-1} \circ F_\mu$.

Exercise 4. For mini-batch OT with $B=64$: compare the computational cost of Hungarian vs Sinkhorn (50 iterations) algorithms.

Exercise 5. Show that the displacement interpolation $\rho_t = [(1-t)\text{Id} + tT]_\# \mu$ satisfies the continuity equation with constant velocity field $v(x) = T(x) - x$.