Technical Summary: Beckmann Transport Models

Problem Statement

Standard generative modeling via flow matching and diffusion models relies on time-dependent velocity fields (drifts) to transport samples from a base distribution $\mu_0$ to a target distribution $\mu_1$. While effective, this approach requires integrating a time-dependent ordinary differential equation (ODE) over a fixed interval $[0, 1]$, which is computationally expensive at inference.

An alternative approach, proposed by Wang and Du [2025] as "Equilibrium Matching" (EM), constrains the drift to be time-independent (autonomous). This simplifies the architecture and the dynamical system. However, the original EM formulation lacked a rigorous proof that an autonomous flow with a time-independent drift could exactly transport $\mu_0$ to $\mu_1$. Furthermore, the loss function proposed in the original EM paper was shown to introduce biases in the transport, failing to enforce the correct mass allocation, particularly when the target distribution is supported on a lower-dimensional manifold.

Methodology

The paper introduces Beckmann Transport Models (BTM), a framework that instantiates flow matching using a time-independent velocity field $b: \mathbb{R}^d \to \mathbb{R}^d$. The core theoretical insight is that if the target distribution $\mu_1$ is singular (supported on a lower-dimensional manifold $M_1$), a time-independent drift can define a valid transport map.

1. Autonomous Flow and Divergence Condition

The authors establish that a time-independent drift $b$ defines a valid transport if it satisfies a specific divergence equation involving a positive weight function $\nu$:
$$\nabla \cdot (\nu b) = \mu_0 - \mu_1$$
Here, $\mu_0$ acts as the source and $\mu_1$ as the sink. The weight $\nu$ represents the time-averaged interpolant distribution in the standard flow matching setup.

• Convergence: Under this condition, trajectories of the autonomous ODE $\dot{X}_t = b(X_t)$ starting from $\mu_0$ converge to the support $M_1$ of $\mu_1$.
• Transport Map: The hitting time $\tau(x_0)$ (which may be finite or infinite depending on the drift behavior at $M_1$) defines a map $T(x_0) = \lim_{t \to \tau} X_t(x_0)$ that pushes $\mu_0$ forward to $\mu_1$ ($T_\sharp \mu_0 = \mu_1$).

2. Correction of Equilibrium Matching

The paper identifies that the original Equilibrium Matching loss minimizes a regression objective that does not strictly satisfy the divergence condition required for exact transport. By treating the schedule of the interpolant as part of the interpolant itself (rather than a separate scaling factor on the target), the authors derive a consistent loss (Equation 5) that ensures the learned drift satisfies the divergence condition. This correction eliminates the mass-allocation bias observed in the original EM method.

3. One-Step Map via Conservation Equation

A key contribution is the characterization of the transport map $T$ not as the result of ODE integration, but as the unique solution to a conservation equation along the flow trajectories:
$$b(x) \cdot \nabla T(x) = 0 \quad \text{for } x \notin M_1$$
with the boundary condition $T(x) = x$ for $x \in M_1$.
This equation implies that the value of the map $T$ is constant along the flow lines of $b$. This structural insight allows for the direct learning of $T$ without solving the ODE during inference.

4. Direct Learning Objective

To learn $T$ directly from samples, the authors discretize the Eulerian equation governing the flow map. Using a forward-Euler step, they derive a regression objective (Equation 15) that trains a network $T_\theta$ to predict the next state in the flow sequence:
$$\mathcal{L}(T) = \mathbb{E}_{t, x_0, x_1} \left[ \left| T(I_t) - \text{sg}\left( T(I_t) + \dot{I}_t \cdot \nabla T(I_t) \right) \right|^2 \right] + \lambda \mathbb{E}_{x_1} \left[ |T(x_1) - x_1|^2 \right]$$
where $\text{sg}$ denotes the stop-gradient operator. This enables one-step generation ($x_1 = T(x_0)$) or iterative refinement ($x_1 = T^k(x_0)$) without numerical ODE integration.

Key Contributions

1. Theoretical Justification: Proves that time-independent flow matching drifts define valid transports to singular targets, recovering Equilibrium Matching up to a loss correction.
2. General Framework: Shows that any pair $(\nu, b)$ satisfying the divergence equation $\nabla \cdot (\nu b) = \mu_0 - \mu_1$ yields a valid autonomous transport. This unifies Flow Matching (where $\nu$ is the time-averaged interpolant) and the training-free Poisson Flow Generative Model (where $\nu \equiv 1$).
3. Conservation Equation: Derives a unique conservation equation ($b \cdot \nabla T = 0$) characterizing the transport map, providing a dynamic interpretation of the flux constraint in Beckmann's transportation problem.
4. Direct Map Learning: Proposes a method to learn the transport map $T$ directly from samples using a residual loss derived from the flow's Eulerian equation, bypassing ODE integration at inference.
5. Bias Correction: Demonstrates that the original Equilibrium Matching loss introduces bias in mass allocation and provides a corrected, consistent interpolant that resolves this issue at no additional cost.

Experimental Results

The authors validate their theory on 2D atomic targets and image benchmarks:

• 2D Atomic Targets: The corrected BTM loss achieves near-perfect mass allocation (MAE $\approx$ 0.005) across modes, whereas the original Equilibrium Matching loss showed significant bias (MAE $\approx$ 0.102) due to incorrect basin structures.
• Iterative Convergence: On a spiral target embedded in high dimensions, a partially trained map $T$ produces diffuse samples, but iterating the map (e.g., $T^3$) sharpens the generation, converging toward the autonomous flow trajectory. A fully trained $T$ matches the flow in a single step.
• ImageNet 256x256:
  • Correction: Applying the consistent loss to the Equilibrium Matching architecture (EqM-XL/2) improved the FID from 1.90 to 1.87 without changing the architecture or training budget.
  • One-Step Generation: Training a direct map on the SD-VAE latent space achieved an FID of 17.58 with one forward pass (no classifier-free guidance). This outperforms other one-step methods like iCT (FID 30.10) and is competitive with methods requiring guidance, though it remains below multi-step diffusion models.

Significance and Claims

The paper claims to provide a unifying framework that recovers closed-form models (Poisson Flow) and recent equilibrium-based methods under a single theoretical umbrella. By linking autonomous flows to Beckmann's transportation problem, the authors provide a dynamic interpretation for a flux constraint that previously lacked one, distinct from the Benamou–Brenier formulation of Monge–Kantorovich optimal transport.

The significance lies in:

• Theoretical Rigor: Providing the missing proof that autonomous flows can exactly transport to singular manifolds.
• Practical Efficiency: Enabling one-step generation by learning the map directly, removing the need for ODE integration at inference.
• Correcting Existing Methods: Identifying and fixing a fundamental bias in the Equilibrium Matching loss function.
• Design Flexibility: Introducing the weight $\nu$ as a design knob to trade off between different generative models (e.g., FM vs. Poisson Flow) and control the behavior of the drift near the target manifold.

The authors remain modest regarding the current performance of one-step BTM on ImageNet compared to state-of-the-art multi-step diffusion models, noting that the FID of 17.58 is "modest" but demonstrates the feasibility of the approach. They identify extending classifier-free guidance to BTM and optimizing the trade-off between training time and inference iterations as open directions for future work.