Beckmann Transport Models: From Autonomous Flows to One-Step Maps
Lee Cheuk-Kit, Florentin Coeurdoux, Peter Potaptchik, Yilun Du, Michael Samuel Albergo, Eric Vanden-Eijnden
cs.LG
2026-08-03
Proves a time-independent flow-matching drift still transports, unifying Poisson Flow and Equilibrium Matching while fixing the latter's weight bias; one-step ImageNet FID 17.58.
Flow matching and diffusion drive modern image generation by learning a velocity field and integrating an ODE at inference to turn noise into an image. That integration typically takes tens to hundreds of neural-network calls (NFE, number of function evaluations), which is slow.
Equilibrium Matching (Wang and Du, 2025) offered a tempting simplification: make the velocity field time-independent, a fixed "autonomous flow", which simplifies both the network and the dynamics. The problem is that the EM paper never proved this time-independent drift actually transports the base distribution μ₀ to the target μ₁, and its training loss is itself biased. This paper closes that gap and generalizes the idea into a full framework.
The core is a divergence equation. The authors show that if a positive weight ν and a drift b satisfy
∇·(νb) = μ₀ − μ₁
then the autonomous flow Ẋₜ = b(Xₜ) sends almost every trajectory to the target's lower-dimensional manifold M₁, defining a map T that pushes μ₀ forward to μ₁. μ₀ is a source, μ₁ a sink, and b the flow lines joining them.
The key assumption is that the target μ₁ must be "singular", supported on a manifold of lower dimension than the ambient space. This is rarely restrictive: any distribution embeds into a higher-dimensional space by zero-padding, landing on an affine subspace. Discrete atomic distributions are the k=0 special case.
The framework unifies two ideas that looked separate. The weight ν is a genuine design knob: set ν to the time-averaged interpolant distribution and you recover standard flow matching; set ν≡1 and b becomes the Coulomb potential of the signed charge μ₀−μ₁, which is exactly the training-free Poisson Flow generative model (Xu et al., 2022). Poisson Flow and Equilibrium Matching are two instances of one principle.
The "Beckmann" name comes from optimal transport. The divergence equation is the flux constraint of Beckmann's 1952 transportation problem, which until now had no dynamic interpretation of the kind Benamou and Brenier gave to Monge-Kantorovich. BTM supplies one, with a quantitative bonus: the autonomous action ∫|b|²ν upper-bounds the W₂² transport cost and is itself bounded by the standard flow-matching action.
The other half of the framework is a one-step map. The endpoint T is constant along flow lines, so it satisfies the conservation equation b·∇T = 0 off M₁ with T(x)=x on M₁. Discretizing this into a regression loss lets you learn T directly from samples, so inference is a single forward pass with no ODE integration. If T is under-trained, the semigroup property lets you iterate the network a few times (T³ ≈ three forward passes) to approach the full map, trading extra forward passes for less training.
A 2D atomic experiment exposes EM's bias directly. The authors use five points with deliberately unequal weights p=(0.30,0.30,0.15,0.15,0.10) so weight errors are visible: BTM's basin areas match the target weights with MAE=0.005, while EM gives 0.102, the largest basins swelling and the smallest nearly vanishing, a 20× error. The mechanism is clear: EM's loss forces the regression target to zero at t=1, which violates the divergence equation, so the flow arrives with wrong weights. BTM ties the target to the interpolant derivative İₜ, making the source-sink balance hold by construction.
On ImageNet 256×256 with the XL/2 architecture and the same budget as EM:
| Model | Setup | FID ↓ |
| SiT-XL/2 (flow matching) | 250 NFE | 2.06 |
| EqM-XL/2 (original, biased) | autonomous | 1.90 |
| EqM-XL/2 (BTM corrected) | autonomous | 1.87 |
The correction moves FID from 1.90 to 1.87 essentially for free, with no change to architecture, training, or inference. The small gain is expected: when class weights are nearly uniform, the weight-allocation error is small to begin with.
The one-step table is more telling:
| Model | NFE | Guidance | FID ↓ |
| iCT | 1 | none | 30.10 |
| BTM (this work) | 1 | none | 17.58 |
| Shortcut Models | 1 | CFG | 10.60 |
| MeanFlow | 1 | CFG | 3.43 |
| SiT (with CFG) | 2×250 | CFG | 2.06 |
Among single-step methods without classifier-free guidance (CFG), BTM's 17.58 cuts iCT's 30.10 nearly in half. But once competitors turn on CFG, Shortcut and MeanFlow pull ahead. The authors are candid: bringing CFG into BTM is still open.
For generative-model practitioners the value is threefold. First, the theory establishes that dropping the time axis from the velocity field is sound, filling the proof gap in Equilibrium Matching and flagging that the original loss is biased; anyone using EM or a similar autonomous-flow method should check whether they share the same flaw. Second, it unifies flow matching, Poisson Flow, and Equilibrium Matching under one divergence equation, turning ν into a tunable knob for future design. Third, the conservation equation offers a route to learn a one-step map directly from samples, collapsing inference to one forward pass, useful for latency-sensitive uses.
The cold water: this is not a benchmark-topping paper. The ImageNet correction over EM is 0.03 FID, and the one-step 17.58 sits well above CFG-equipped state of the art (single-digit FID). Its real contribution is theoretical cleanup plus a framework that might grow faster generators, not a drop-in model.
The authors name their own: the framework needs a singular target (bypassable by zero-padding but still structural); CFG is not yet integrated, so the one-step numbers are not strong; the optimal split between training and inference compute (how many iterations) is uncharacterized; larger-scale direct map learning is left to future work.
A critical read: the 1.90 to 1.87 gap is small enough to sit inside evaluation noise, so "the correction works" rests more on the 2D 20× weight-error comparison than on the 0.03 FID. The one-step comparison is also uneven: BTM uses no CFG while most rivals do, so 17.58 against 3.43 measures who can use guidance more than the map itself. The framework's real competitiveness will show once CFG lands.