Generative Modelling with Flow Matching, Optimal Transport, and Schrödinger bridge

Flow matching, a generative modelling technique, learns a velocity field that transports noise directly onto data in one continuous motion, avoiding the stochastic reversal of diffusion models. The method, detailed in a technical tour, regresses a neural network onto conditional velocities of endpoint pairs, with a linear interpolant reducing the loss to teaching the network the average displacement through each point. The approach unifies flow matching, optimal transport, and Schrödinger bridges as configurations of the same object, enabling efficient ODE-based sampling.

Generative modelling has quietly converged on a beautifully simple question. How do you move one probability distribution onto another? Diffusion models answer it by slowly corrupting data with noise and learning to reverse the process. Flow matching answers it more directly. It learns a velocity field that transports a cloud of noise straight onto the data in one continuous motion with no stochastic reversal to unwind. This article is a guided tour of that idea and the family of algorithms it generates. We start from a single displacement vector, build up flow matching, sharpen it with optimal-transport couplings, bend the path into a Schrödinger bridge, and finally reuse the trained field to solve inverse problems. The running theme is that these are not five unrelated methods, but they are five configurations of the same object. We begin by fixing two distributions. The first one is a source p0 that we can sample trivially, usually a standard Gaussian N 0, I , and the other is a target distribution p1, which we only ever see through a finite sample. Our goal is a map that carries the source density onto the target density. Suppose we pair a source sample x0 ∼ p0 with a data sample x1 ∼ p1. The most literal way to get from one to the other is to walk in a straight line. We call this displacement the delta carried from source to target. Interpolating along that line gives a time-indexed state whose velocity is constant and equal to the displacement itself. So each pair of endpoints defines a tiny trivial dynamical system. Start at x0, move with constant velocity ∆, and arrive at x1 at time t = 1. The catch is that at generation time we will not know x1. We only know where we currently are. The whole concept is to learn a field vθ x,t that, given only the current position x and time t, predicts the velocity we should be moving with. Once we have such a field, sampling is integration over a small portion of the known velocities. We start from drawing noise and then solve the ordinary differential equation forward from t = 0 to t = 1. The endpoint x 1 is a sample from the learned data distribution. That is the entire inference procedure: an ODE solver . The two objects matter in this concept. A probability path determines how x0 and x1 are connected over time and what the target velocity vt along that path is. A velocity field vθ x,t is the neural network we train to match it. Everything else is a choice of how to wire these two together. The conceptual leap of flow matching is that we never need to know about the intractable global dynamics. We only regress the network onto the conditional velocity of individual endpoint pairs. Concretely, we sample a data point x1 , sample noise x0 ∼N 0, I , and sample a random time t ∼U 0,1 , and then form the interpolated state xt before asking the network vθ x,t to predict the known conditional velocity ut . The objective is a plain regression The remarkable result behind flow matching is that the field minimising this conditional loss also generates the correct marginal probability path. Averaging over all endpoint pairs that route through a given point x,t , the network learns the expected velocity there, and that expected field is exactly the one whose ODE transports p0 onto p1. We get a tractable per-sample loss, whose optimum solves an intractable transport problem for free. For the straight-line rectified path , the target is simply the delta, so the loss reduces to teaching the network the average displacement passing through each point in space and time. With the library’s primitives, that training step is the following two-line loop. python from deltaflow.interpolants import LinearInterpolantfrom deltaflow.losses import FlowMatchingLossloss fn = FlowMatchingLoss interpolant=LinearInterpolant loss = loss fn model, x1 samples x0, t, builds x t, regresses v x t, t onto u tloss.backward The interpolant owns the path geometry . The loss owns the regression. The model owns the field. Keeping them separate is what lets us maintain modularity by changing one without touching the others, which is the thread we pull on for the rest of the article. The straight line is only the simplest way to connect noise to data. An interpolant is any rule that, given endpoints x0 , x1 and a time t , returns the interpolated state xt and its conditional target velocity ut . Swap the interpolant, yet maintain the same loss and solver, and train a different generative model. This is the baseline path we introduced above. It has constant velocity, straight trajectories in the ideal case, and the cheapest possible target to compute. This is the rectified-flow geometry. Diffusion models do not walk in straight lines. They follow a curved schedule that keeps the marginal variance controlled. We can reproduce that geometry as an interpolant with a trigonometric schedule. so that α2t + σ2t = 1. At every time. At t = 0, the state is pure noise, while at t = 1 , it is pure data. The path bends between them like a diffusion trajectory. Crucially, only the interpolant changes, whereas the loss and the ODE solver stay untouched. python from deltaflow.interpolants import VariancePreservingInterpolantfrom deltaflow.losses import FlowMatchingLossloss fn = FlowMatchingLoss interpolant=VariancePreservingInterpolant loss = loss fn model, x1 same loss, a curved diffusion-style pat The straight line is deterministic. Given x0 and x1 , the trajectory is fixed. The Schrödinger bridge asks a softer question. Among all stochastic processes that start at x0 and end at x1 , which path is closest in a relative-entropy sense to pure Brownian motion? Conditioned on a fixed pair of endpoints, the answer is a Brownian bridge. The straight line with a controlled bulge of noise that swells in the middle and vanishes at both ends, as defined by the equation below, where σ sets the diffusivity. Noticeably, if σ = 0, it recovers the linear path. Regressing the field onto the conditional velocity of this bridge trains a stochastic interpolant that approaches the entropic optimal-transport bridge between noise and data. python from deltaflow.interpolants import SchrodingerBridgeInterpolantfrom deltaflow.losses import ConditionalFlowMatchingLossfrom deltaflow.trainer import OTCouplingloss fn = ConditionalFlowMatchingLoss interpolant=SchrodingerBridgeInterpolant sigma=1.0 ,coupling=OTCoupling , loss = loss fn model, x1 The above snippet introduces a second knob, the coupling , and it has the biggest practical payoff. Return to the very first step. We paired a source sample x0 with a data sample x1 and drew a line between them. But which x0 pairs with which x1 was never specified. The default is to pair them independently, drawing x0 ∼ N 0, I with no regard for the x1 it will be matched against. This is the classical setup. It works, but it is wasteful. The problem is crossing paths. If noise is assigned to data arbitrarily, the straight lines between pairs cut across one another. The true marginal field at a point where many trajectories cross needed average conflicting velocities, so the learned field is curved even though every individual conditional path is straight. Curved fields need many small integration steps to follow accurately. Optimal-Transport Coupling is a method used to pair each source to destination points. Within each mini-batch, it requires to solve the discrete assignment problem that matches the n noise samples to the n data samples at minimum total squared distance where π ranges over permutations. This is the mini-batch optimal-transport coupling solved by the Hungarian algorithm and a greedy fallback for speed. Pairing the nearest available endpoints keeps the displacements short and, more importantly, keeps the bundle of trajectories from crossing. It indicates straighter marginal paths determining the sampling ODE to be integrated in far fewer steps at the same quality. python from deltaflow.losses import ConditionalFlowMatchingLossfrom deltaflow.trainer import OTCouplingloss fn = ConditionalFlowMatchingLoss coupling=OTCoupling same objective, OT pairsloss = loss fn model, x1 There is a deeper connection worth stating. This hard zero-entropy assignment is the limiting case of the static Schrödinger bridge as the diffusivity goes to zero. So the coupling knob and the interpolant knob are two views of the same underlying transport geometry. Pairing endpoints by optimal transport is what makes the Schrödinger-bridge approximation tight in practice. Everything above is the same training loop. Only two arguments change, which are the interpolant and the coupling . Train an identical network on each configuration for the same number of steps and the differences become concrete. The independent pairing sweeps long crossing arcs while the OT pairing stays an orderly near-straight bundle at lower transport cost. Training produces a field, while a solver turns it into data by integrating the ODE dx/dt = vθ x,t . The simplest choice is explicit Euler , which is stepping forward with a fixed increment h = 1/N . It is cheap, but its error accumulates, so it needs many steps when the field is curved. A second-order Heun step corrects the Euler prediction with a trapezoidal average of the velocity at both ends of the interval. Heun costs two field evaluations per step but is markedly more accurate at low step counts, which is exactly the regime that straight OT-coupled paths unlock. python from deltaflow.samplers import FlowSamplerimport torchsamples = FlowSampler model .sample torch.randn 1000, 2 , n steps=50 This is where the earlier design choices pay off. A straighter field from OT coupling integrated with a second-order solver reaches the data manifold in a handful of steps rather than the hundreds a diffusion model typically demands. Here is the contribution of a modular design contribute. Suppose you have a field trained once, unconditionally, on clean data. Now you are handed a degraded measurement, a masked image, a blurred scan, or a downsampled signal, and asked to reconstruct the clean source that produced it. Formally, given a measurement operator A and an observation, we want to sample from the posterior p x1 | y rather than the prior p x1 . The trick is to steer the sampling ODE. At each solver step, we already have the current state xt and the predicted velocity vt . A Tweedie decomposition converts that pair into an estimate of the clean endpoint ˆx1 xt, vt, t , which is the field’s best guess of where this trajectory is heading. We then measure how well that clean estimate explains the observation through the Gaussian likelihood and nudge the step along the gradient that improves the fit. where γ is a guidance scale. The prior flow proposes plausible clean data. The likelihood gradient keeps that proposal consistent with what was actually measured. Critically, this reuses the very same Euler solver and the very same pretrained field, which are wrapped rather than re-implemented with no retraining for each new degradation. python from deltaflow.inverse import GaussianLikelihood, LinearTweedie, MaskOperatorfrom deltaflow.solvers import EulerSolver, PosteriorSolverimport torchlikelihood = GaussianLikelihood y=y, operator=MaskOperator mask , sigma=1.0 solver = PosteriorSolver base solver=EulerSolver model , the same integrator, reused likelihood=likelihood, tweedie=LinearTweedie , guidance scale=0.5, x = solver.sample torch.randn 16, 1, 16, 16 , n steps=60 posterior samples DeltaFlow’s generative machinery is not limited to producing samples. The same conditional field trained for synthesis can be repurposed for detection and representation learning without architectural changes. Framing landmark detection as a conditional flow turns a discriminative task into a sampling problem. For instances, the field transports noise onto landmark coordinates conditioned on the image, and detection becomes inference-time sampling with sample spread giving uncertainty for free. The delta alignment idea extends this further. It runs a conditionally trained field twice, with and without conditioning, isolates what the conditioning changed and cancels shared structure. In a medical-imaging setting, this yields an anatomy-cancelling feature, and aligning it across noise levels trains a backbone whose representation holds up under a linear probe. Both cases show the same underlying principle, which is a distributional difference as the learning signal. It is applied outward for detection and inward for representation learning. The same machinery reaches previous reconstruction into discriminative tasks. We can frame anatomical landmark detection as a conditional flow p landmarks | image . Here the X-ray image is the condition and the stacked landmark coordinates are the generative target, so the field transports noise onto landmark positions while conditioning on the image. Sampling at inference is detection, and the spread of the samples is a built-in uncertainty estimate. The displacement idea reaches prior sample synthesis. We take a conditionally trained field and run it twice on the same input, once with the conditioning signal and once without the signal. The difference between the two internal feature sets that isolates precisely what the conditioning changed and largely cancels the structure both passes share. In a medical-imaging setting, that shared structure is the anatomy, so ∆h becomes a feature that suppresses shared anatomical content leaving a signal that isolates what the conditioning specifically changed. A small conditional field learns p image | y for a dataset label y, and aligning ∆h across two noise-level views of the same image with the Delta Alignment Loss the summation of flow matching loss and alignment loss . It trains a backbone whose representation is consistent regardless of the anatomy underneath. Once again, the learning signal is a difference between two distributions of features and the same delta principle turned inward. To validate whether that pretrained representation is discriminative. We freeze the backbone and fit a linear probe on its features. On a synthetic four-dataset task, the probe reaches near-perfect held-out accuracy while a random-init baseline sits far lower, and a PCA of the frozen features separates cleanly by dataset. From the algorithms above — interpolant, coupling, solver, and likelihood — the whole landscape collapses onto one diagram. Pick a path for how noise connects to data. Pick a coupling for which noise connects to which datum. Regress a field onto the path’s velocity with a single objective. Integrate that field with a solver. Optionally inject a measurement likelihood to solve an inverse problem with the velocity field already trained. Flow matching, OT-coupling rectified flow, variance-preserving interpolant, the Schrödinger bridge, and posterior sampling are not five libraries worth of code, instead they are five ways of wiring these six components together. Implement each component once, and the method becomes a module configuration rather than rewrites. DeltaFlow is a small PyTorch library built around the delta between two distributions. Every component in the table above is a composable primitive. Swap pip install torchdeltaflow Code, examples, and the animations behind these figures are available at github.com/phrugsa-limbunlom/deltaflow https://github.com/phrugsa-limbunlom/deltaflow , with documentation at phrugsa-limbunlom.github.io/deltaflow https://phrugsa-limbunlom.github.io/deltaflow/ and PyPi torchdeltaflow https://pypi.org/project/torchdeltaflow/ . 1 Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow Matching for Generative Modeling. ICLR, 2023. arXiv:2210.02747. 2 A. Tong, K. Fatras, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. TMLR, 2024. arXiv:2302.00482. 3 N. Ma, M. Goldstein, M. S. Albergo, N. M. Boffi, E. Vanden-Eijnden, and S. Xie. SiT: Exploring Flow and Diffusion-Based Generative Models with Scalable Interpolant Transformers. 2024. arXiv:2401.08740. 4 M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden. Stochastic Interpolants: A Unifying Framework for Flows and Diffusions. 2023. arXiv:2303.08797. 5 V. De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling. NeurIPS, 2021. arXiv:2106.01357. 6 A. Tong, N. Malkin, K. Fatras, L. Atanackovic, Y. Zhang, G. Huguet, G. Wolf, and Y. Bengio. Simulation-Free Schrödinger Bridges via Score and Flow Matching. AISTATS, 2024. arXiv:2307.03672. 7 J. Kim et al. FlowDPS: Flow-Driven Posterior Sampling for Inverse Problems. 2025. arXiv:2503.08136. 8 M. Pourya et al. Flower: A Flow-Matching Solver for Inverse Problems. 2025. arXiv:2509.26287. 9 J. Ho and T. Salimans. Classifier-Free Diffusion Guidance. 2022. arXiv:2207.12598. Generative Modelling with Flow Matching, Optimal Transport, and Schrödinger bridge https://pub.towardsai.net/generative-modelling-with-flow-matching-optimal-transport-and-schr%C3%B6dinger-bridge-3bbfe986b4de was originally published in Towards AI https://pub.towardsai.net on Medium, where people are continuing the conversation by highlighting and responding to this story.