📊 statistics

Optimal and Diffusion Transports in Machine Learning

This survey explores the mathematical connections between diffusion methods and optimal transport in machine learning, demonstrating how their shared Lagrangian framework for modeling time-evolving probability distributions unifies applications ranging from generative AI sampling and neural network optimization to the analysis of large language model dynamics.

Original authors: Gabriel Peyré

Published 2026-05-28

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Gabriel Peyré

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to move a massive pile of sand from one shape to another. Maybe you want to turn a pile of sand shaped like a mountain into a pile shaped like a castle. In the world of Machine Learning, this "sand" isn't just dirt; it's data, weights in a computer brain, or even the words (tokens) in a sentence.

This paper, written by Gabriel Peyré, acts as a map for understanding how these piles of data move and change over time. It argues that instead of looking at data as static pictures, we should view it as a flowing river. The paper focuses on two main ways to steer this river: Diffusion (like letting ink spread in water) and Optimal Transport (like finding the most efficient truck route to move furniture).

Here is a breakdown of the paper's core ideas using simple analogies:

1. The Two Ways to Watch the River Flow

The paper explains that we can watch our data river in two different ways:

The Eulerian View (The Satellite): You stand on a bridge and watch the water flow past you. You see the density of the water at specific spots. This is good for seeing the "big picture" of where the data is concentrated.
The Lagrangian View (The Raft): You hop on a raft and float with the water. You track individual particles (or data points) as they move. This is better for understanding how a specific piece of data gets from point A to point B.

The paper's main trick is switching between these two views. It suggests that if we can figure out the "wind" (a vector field) that pushes the raft, we can control the whole river.

2. The Two Main Methods

Method A: Diffusion and Flow Matching (The "Blender" Approach)

This is the engine behind modern AI that creates images, music, and text (Generative AI).

The Analogy: Imagine you have a clear glass of water (simple data) and a glass of muddy water (complex data).
- Diffusion is like slowly adding mud to the clear water until it's a uniform brown soup. Then, you try to reverse the process: you slowly filter the mud out until you get clear water again.
- Flow Matching is a smarter version. Instead of just guessing the reverse path, it draws a straight line between a drop of clear water and a drop of muddy water. It calculates the exact speed and direction needed to push the clear drop to become the muddy drop.
The Catch: This method is very popular and works great, but the path it takes isn't always the most efficient. It's like taking a winding scenic route instead of a straight highway. The paper notes that while it works, we still don't fully understand the geometry of these winding paths.

Method B: Optimal Transport (The "Moving Company" Approach)

This method is rooted in math from the 18th century.

The Analogy: Imagine you are a moving company. You have a pile of boxes (data) in one room and need to move them to a new room. You want to move them using the least amount of energy possible.
The Rule: You don't just mix the boxes; you find the perfect partner for every single box. Box A in the old room goes to Spot A in the new room. This creates a "straight line" path for every particle.
The Benefit: This is the most efficient, "shortest distance" way to transform data. The paper shows that this method provides a very strict, geometric structure that helps us understand how to move data without wasting energy.

3. Where This Applies in Machine Learning

The paper shows that this "river flow" idea explains three different things in AI:

Creating New Things (Generative Models): As mentioned above, this is how AI draws pictures or writes songs. It learns the "flow" to turn random noise into a masterpiece.
Training Neural Networks (The "Brain" Learning): Imagine a neural network as a crowd of people (neurons) trying to solve a puzzle.
- The paper suggests that as the network learns, the crowd moves together like a fluid.
- If the network is "shallow" (not very deep), we can mathematically prove that this fluid flow will eventually find the best solution (the global minimum). It's like a ball rolling down a hill until it hits the bottom.
- However, for very deep networks, the math gets messy, and we aren't sure yet if the "ball" will always find the bottom or get stuck.
Transformers (The "Language" Models): Transformers (like the ones powering chatbots) process words (tokens) as a group.
- The paper models the layers of a Transformer as a continuous flow. As a word passes through layer 1, then layer 2, then layer 3, it changes.
- This is modeled as a "Vlasov equation" (a type of physics equation for particles that interact). The words interact with each other (like a crowd at a concert) to decide what the next word should be.
- The paper shows that if you have enough words, their distribution follows a predictable mathematical curve, almost like gas molecules in a box.

4. The Big Picture: What's Still Missing?

The paper concludes with a few open questions:

Efficiency vs. Reality: Optimal Transport gives us the mathematically perfect, shortest path, but Diffusion models (which are currently more popular) take a slightly longer, "wobbly" path. We don't fully understand the cost of taking that wobbly path.
Deep Networks: We have good math for shallow networks, but for the massive, deep networks used today, we still don't have a complete mathematical proof of why they work so well.
The "Flow" of Words: We are just starting to understand the complex physics of how words interact in Transformers. It's a new frontier where math meets language.

In Summary:
This paper unifies different parts of Machine Learning under one umbrella: moving probability distributions. Whether you are generating an image, training a brain, or processing a sentence, you are essentially pushing a cloud of data from one shape to another. The paper provides the mathematical tools to understand the speed, direction, and efficiency of that push.

Technical Summary: Optimal and Diffusion Transports in Machine Learning

Problem Statement

The paper addresses the challenge of mathematically unifying diverse machine learning problems that can be formulated as the time-evolution of probability distributions. These problems span generative modeling (sampling), neural network optimization (training dynamics), and the analysis of large language models (token evolution). While the target objects differ—samples, weights, or tokens—their underlying dynamics share a common structure: the advection of particles or measures under time-dependent vector fields.

The core difficulty lies in the transition between the Eulerian representation (tracking the density $\rho_t$ of the measure $\alpha_t$ ) and the Lagrangian representation (tracking individual particles $x(t)$ ). A fundamental issue is the non-uniqueness of the vector field $v_t$ driving a specific measure evolution $\alpha_t$ ; infinitely many fields can satisfy the continuity equation. This ambiguity presents both a challenge for well-posedness and an opportunity to design flows with favorable properties regarding regularity, stability, and computational tractability.

Methodology

The survey synthesizes two complementary mathematical frameworks to analyze these evolutions: Diffusion/Flow Matching and Optimal Transport (OT).

1. Flow Matching and Diffusion Models

This approach constructs a transport map by defining an explicit interpolation $(\alpha_t)_{t \in [0,1]}$ between a reference distribution $\alpha_0$ (e.g., Gaussian) and a target $\alpha_1$ .

Stochastic Interpolation: The method often employs linear interpolation of independent random vectors, $X_t = (1-t)X_0 + tX_1$ , where $X_0 \sim \alpha_0$ and $X_1 \sim \alpha_1$ .
Conditional Expectation Field: Instead of solving the inverse problem of finding a vector field from a density evolution (which requires inverting a Laplacian), the paper highlights a result where the vector field $v_t$ is expressed as a conditional expectation: $v_t(z) = \mathbb{E}[X_1 - X_0 \mid (1-t)X_0 + tX_1 = z]$ .
Least-Squares Formulation: This conditional expectation is estimated by solving an unconstrained least-squares problem, parameterized by neural networks (e.g., U-Nets, Transformers). This avoids explicit density estimation and allows for training on finite samples.
Connection to Diffusion: When $\alpha_0$ is isotropic Gaussian, this framework is shown to be equivalent to diffusion models via Tweedie's identity, linking the velocity field to the score function $\nabla \log \alpha_t$ .

2. Optimal Transport (OT)

In contrast to prescribing an interpolation, OT seeks an "optimal" interpolation that minimizes displacement cost.

Static and Dynamic Formulations: The paper reviews the Monge problem (finding a map $T$ ) and the Kantorovich relaxation (finding a coupling $\pi$ ). It emphasizes the Benamou-Brenier dynamic formulation, which views OT as minimizing kinetic energy along a flow: $W_2^2(\alpha_0, \alpha_1) = \inf \int \int |v_t|^2 d\alpha_t dt$ .
Gradient Structure: A key theoretical insight is that the optimal velocity field in OT is a gradient field ( $v_t = \nabla \psi_t$ ), derived from a convex potential. This contrasts with flow matching, where the vector field is not necessarily a gradient unless specific conditions (like commuting covariances) are met.
Gaussian Case: The paper details closed-form solutions for Gaussian measures, showing that OT maps are linear and coincide with flow matching maps only when covariance matrices commute.

3. Wasserstein Gradient Flows

The survey extends OT to optimization by interpreting the evolution of measures as gradient flows in the Wasserstein space.

JKO Scheme: The evolution is discretized via the Jordan-Kinderlehrer-Otto (JKO) scheme, minimizing a functional $f(\alpha)$ plus a Wasserstein distance term.
Applications: This framework interprets nonlinear PDEs (e.g., porous medium equation) as gradient flows of entropy functionals. In machine learning, it models the training of shallow neural networks as the evolution of a probability measure over neuron weights.

Key Contributions and Results

1. Unification of Generative Models

The paper establishes that modern generative models (Diffusion and Flow Matching) rely on learning a vector field via least-squares regression on a stochastic interpolation. It clarifies that while these methods are computationally efficient and scalable, their geometric properties differ from Optimal Transport. Specifically, flow matching interpolations are generally not gradient flows and do not minimize displacement cost in the OT sense.

2. Neural Network Training as Measure Evolution

Mean-Field Limits: The training of two-layer MLPs is analyzed as a Wasserstein gradient flow over the distribution of weights. The paper cites the Chizat-Bach theorem, which guarantees that under mean-field assumptions (infinite width) and specific initialization, the flow converges to a global minimizer.
Linear Networks: For linear networks with Gaussian neuron distributions, the training dynamics are shown to preserve the Gaussian structure, reducing the infinite-dimensional flow to a Riccati-type ODE on the covariance matrix.
Deep Networks: The paper notes that extending these convergence guarantees to deep networks remains an open challenge, though ResNets can be modeled as Neural ODEs with a generalized Wasserstein metric over depth.

3. Transformer Dynamics

The paper proposes a measure-theoretic view of Transformers, treating tokens as particles in a probability distribution evolving through depth.

Transformer PDE: In the infinite-depth limit, the discrete layer updates converge to a continuity equation (a "Transformer PDE").
Non-Variational Nature: Unlike standard gradient flows, the Transformer dynamics generally do not admit a Wasserstein gradient flow structure because the attention mechanism is not the first variation of a functional. The dynamics are more akin to Vlasov or aggregation equations.
Gaussian Closure: For Gaussian token distributions, the evolution of the covariance matrix is governed by a nonlinear quadratic ODE, which can exhibit finite-time blow-up, mirroring clustering phenomena.
Universality: The paper cites a universality theorem stating that Transformer mappings can approximate any continuous mapping from a distribution and a token to an output, formalizing the concept of in-context learning.

4. Applications in Genomics

The survey highlights the application of Wasserstein flows to single-cell genomics, where OT-based methods are used to infer developmental trajectories and cell-fate transitions from snapshot data, effectively reconstructing the velocity field of cell population evolution.

Significance and Claims

The paper claims that viewing machine learning systems as evolutions of probability measures provides a unified mathematical understanding of computation, optimization, and generalization.

Theoretical Insight: It bridges the gap between stochastic generative models (diffusion) and geometric optimization (OT), clarifying their respective strengths: diffusion offers algorithmic scalability and empirical success, while OT provides principled geometric frameworks and convergence guarantees.
Open Challenges: The paper modestly acknowledges significant gaps. It notes that while Wasserstein flow theory offers convergence guarantees for shallow mean-field models, extending these to deep, expressive architectures remains an outstanding challenge. Similarly, the sample complexity of score estimation and the propagation of discretization errors in diffusion models are identified as open research questions.
Future Directions: The survey suggests that bridging the theoretical gaps between these paradigms is crucial for designing future architectures that combine the flexibility of diffusion with the structural rigor of optimal transport.

The paper does not propose new experimental algorithms or claim to solve the general convergence problem for deep networks; rather, it serves as a survey synthesizing existing theoretical frameworks to highlight common structures and remaining open problems in the field.