Source-linked AI summary
Rectified Flow: A Marginal Preserving Approach to Optimal Transport
Qiang Liu
TL;DR
Continuous optimal transport requires minimizing a specified cost while satisfying difficult infinite-dimensional marginal constraints. The paper introduces c-rectified flow, which repeatedly solves unconstrained regression problems inside the valid-coupling set. Under mild conditions, recursive c-Rectify has c-optimal couplings as fixed points and drives the associated optimality loss to zero.
Problem
Continuous OT must optimize transport cost under infinite-dimensional marginal constraints, while minimax and regularization approaches have practical convergence or enforcement difficulties.
Method
The method iteratively learns neural ODE velocity fields through unconstrained regression, preserving marginals while applying a cost-dependent restriction for a specified convex c.
Results
Under mild conditions, c-Rectify fixed points are exactly c-optimal couplings, and recursive updates drive the c-optimality loss to zero.
Takeaways & Limitations
The approach provides an interior procedure for solving fixed-cost continuous OT without directly enforcing coupling constraints from outside.
Takeaways & Limitations
The recursive procedure accumulates approximation errors in practice, although one or two iterations were reported sufficient for practical applications.
Abstract
from arXiv · showhide
We present a flow-based approach to the optimal transport (OT) problem between two continuous distributions $π_0,π_1$ on $\mathbb{R}^d$, of minimizing a transport cost $\mathbb{E}[c(X_1-X_0)]$ in the set of couplings $(X_0,X_1)$ whose marginal distributions on $X_0,X_1$ equals $π_0,π_1$, respectively, where $c$ is a cost function. Our method iteratively constructs a sequence of neural ordinary differentiable equations (ODE), each learned by solving a simple unconstrained regression problem, which monotonically reduce the transport cost while automatically preserving the marginal constraints. This yields a monotonic interior approach that traverses inside the set of valid couplings to decrease the transport cost, which distinguishes itself from most existing approaches that enforce the coupling constraints from the outside. The main idea of the method draws from rectified flow, a recent approach that simultaneously decreases the whole family of transport costs induced by convex functions $c$ (and is hence multi-objective in nature), but is not tailored to minimize a specific transport cost. Our method is a single-object variant of rectified flow that guarantees to solve the OT problem for a fixed, user-specified convex cost function $c$.
1 Introduction
The paper frames continuous optimal transport as a difficult constrained problem and introduces an interior, flow-based alternative that preserves marginals while reducing transport cost. Its c-rectified variant adapts rectified flow to a specified convex cost and characterizes c-optimal couplings through recursive updates and fixed points.
- Problem: Continuous OT minimizes E[c(X1 − X0)] over couplings with prescribed marginal laws, but these constraints are infinite dimensional.The paper focuses on high-dimensional absolutely continuous distributions observed through empirical samples.
- Approach: The method starts from a valid coupling and repeatedly solves unconstrained nonlinear least-squares problems, automatically preserving marginals while monotonically reducing transport cost.This interior strategy traverses within the valid-coupling set rather than enforcing constraints externally.
- Rectified flow: Rectified flow maps a coupling to another coupling with the same marginals and no larger transport cost for every convex cost function.Its velocity field is learned by regressing the line direction X1 − X0 from points on linear interpolation paths, and the resulting ODE induces the rectified coupling.
- Limitation: Original rectification is cost-agnostic and need not converge to an optimum for a specified cost, especially when d ≥2.In one dimension, convex costs can share a common straight optimal coupling, but this property does not generally extend to higher dimensions.
- c-rectified flow: c-rectified flow restricts the velocity field according to the user-specified convex cost, yielding marginal preservation and non-increasing cost for that cost alone.For quadratic cost, the restriction is to gradient fields; more generally, the velocity has the form ∇c*(∇f).
- Guarantees: Under mild conditions, c-Rectify fixed points are exactly c-optimal couplings, while recursive updates drive the c-optimality loss toward zero.The minimum loss in the c-rectified-flow objective provides a criterion of c-optimality without directly solving the OT problem.
2 Background of Optimal Transport
Optimal transport seeks the lowest-cost coupling between fixed marginals, but continuous marginal constraints make direct optimization difficult. Dynamic formulations recast the problem through paths or ODE-induced flows, while retaining marginal constraints.
- Static formulations: The Monge–Kantorovich problem minimizes transport cost over all deterministic and stochastic couplings with prescribed marginals.The Monge problem restricts couplings to transport maps, whereas the Monge–Kantorovich relaxation allows all couplings.
- Dynamic formulations: Dynamic OT represents transport with continuous-time processes connecting π0 to π1 and uses path-wise costs derived from convex transport costs.For convex c, Jensen’s inequality yields an integral cost over the process velocity, minimized along linear interpolation paths.
- Dynamic formulations: Restricting dynamic optimization to deterministic ODE-induced processes yields the Benamou–Brenier formulation with a continuity-equation constraint.This restriction reduces the search space but remains computationally challenging in practice.
- Dynamic formulations: For differentiable stochastic processes, c-rectified flow yields no larger path-wise c-transport cost than the original process.The paper connects this reduction to Jensen’s inequality and later uses c-rectified flow as a coordinate-descent-like approach.
3 Rectified Flow: An Optimization-Based View
Rectified flow converts a stochastic interpolation process into an ODE while preserving its marginal laws, and it optimizes path-wise transport cost under a velocity-field constraint. Its convex-cost improvement is broad, but standard rectification is not tailored to a fixed cost and can select non-optimal couplings in higher dimensions.
- Flow construction: Rectified flow is defined as the ODE driven by a process’s expected velocity field, with marginal evolution characterized by a continuity equation.Under uniqueness assumptions, the initial law and expected velocity determine the marginal laws throughout time.
- Flow construction: Rectified flow preserves marginal laws because its ODE velocity field equals the original process’s expected velocity field.For rectifiable processes, the resulting flow and original process have identical marginal laws at every time.
- Optimization view: The rectified flow minimizes path-wise c-transport cost among time-differentiable processes sharing the original expected velocity field, for every convex c.The least-squares and Bregman-divergence views characterize the expected velocity as the optimizer.
- Optimization view: For a linear interpolation coupling, rectification produces a coupling with no larger transport cost for every convex cost function.Jensen’s inequality links the path-wise reduction to the endpoint transport-cost reduction.
- Limitations: A straight coupling is a fixed point of rectification, and its associated trajectories are straight lines that can be computed in closed form.In one dimension the straight coupling, when it exists, is unique and minimizes all convex costs for which an optimum exists; higher-dimensional straight couplings need not be cost-optimal.
- Limitations: Standard rectification is multi-objective: it improves all convex costs simultaneously but cannot optimize a preselected fixed cost.In dimensions d ≥ 2, straight couplings are not unique, and recursive rectification depends on the initial coupling.
4 Differentiable Processes with Equivalent Marginal Laws
Processes can share all marginal laws even when their velocity fields differ by a marginal-preserving rotational component. The paper characterizes this condition and uses it to motivate removing rotational components in dynamic OT.
- Marginal-preserving fields: Equal marginal laws do not require equal expected velocity fields because their difference can be a rotation-only vector field.Such a field changes dynamics without modifying the marginal distributions.
- Implication for OT: The resulting dynamic OT problem has a continuum of marginal constraints, and solving it removes rotational components that make standard rectified flow non-optimal.This characterization motivates the paper’s c-rectified flow construction.
- Marginal-preserving fields: A vector-field difference is marginal-preserving when its density-weighted field is divergence-free.Integration by parts gives ∇·(rtρt) = 0 as the classical characterization.
- Characterization: Two processes with the same initial distribution have identical marginal laws at all times if and only if their expected velocity difference is marginal-preserving.This equivalence holds under the stated rectifiability and uniqueness assumptions.
5 c-Rectified Flow
c-Rectified flow is a cost-dependent variant of rectified flow that preserves marginals while reducing a specified convex transport cost. Recursively applying it yields c-optimal couplings under the stated conditions.
- c-Rectified Flow: c-Rectified flow restricts velocities to ∇c∗∘∇f and replaces the quadratic objective with a matching loss.Together, these changes remove the rotation-only component of the velocity field.
- Marginal Preservation: The residual velocity field vX − gX,c is X-marginal-preserving, so c-Rectify preserves Law(Xt) for every t.This establishes automatic preservation of the process’s marginal laws.
- Optimality: c-Rectified flow attains the minimum of the infinite-marginal optimization problem and has strong duality with its regression formulation.The optimal solution fX,c satisfies LX,c(fX,c) = Fc(X) − Fc(Z).
- Coupling Updates: Applying c-Rectify produces a coupling with no larger c-transport cost than the input coupling.The monotonicity guarantee applies to the specific cost c used by c-Rectify, not to all convex costs.
- Fixed Points: A coupling is a fixed point of c-Rectify if and only if it is c-optimal.The paper also states that the surrogate optimum ℓ∗X,c = 0 exactly characterizes c-optimality.
- Convergence: The minimum surrogate ℓ∗X,c over the first k iterations decreases at an O(1/k) rate.This rate concerns the surrogate measure rather than a directly stated transport-cost optimality gap.
6 Discussion and Open Questions
The discussion identifies unresolved questions about convergence, practical error accumulation, statistical analysis, and when to prefer a cost-specific OT method over cost-agnostic rectified flow.
- Convergence: Corollary 5.7 bounds the surrogate measure ℓ∗Zk,c, but the discussion asks whether the optimality gap in c-transport cost can be bounded directly.It also asks whether a strong-convexity-like condition could yield exponential decay.
- Applications: The discussion asks when OT with a specified c should replace simpler cost-agnostic rectified flow in generative modeling and domain transfer.It also asks how to choose c optimally for those tasks.
- Practical Error: Recursive c-rectification accumulates training and ODE-simulation errors, potentially worsening the approximation of π1 as k increases.The proposed direction is to adjust each output sample toward samples from π1, ideally efficiently.
- Sample Adjustment: A proposed adjustment seeks a one-to-one matching between generated and target samples, with the challenging goal of near-linear-time computation.The discussion also calls for complete theoretical analysis of statistical error with or without this adjustment.
A Proofs
The proof analyzes a rotational coupling of identical Gaussian marginals and shows that its straightness does not imply c-optimality for sufficiently regular convex costs.
- Straightness: Equation (32) makes the interpolated velocity uniquely determined by the interpolated state, establishing that (X0, AX0) is straight.The argument invokes the straight-coupling result from rectified flow.
- Contradiction: If the rotational coupling were c-optimal, its Hessian relation would make A − I similar in structure to a diagonalizable matrix with real eigenvalues.This uses the positive definiteness of Hx and the symmetric-matrix lemma.
- Conclusion: Because a non-reflecting, non-identity rotation has complex eigenvalues, the required Hessian relation cannot hold.Thus the straight rotational coupling is not c-optimal for the specified class of convex costs.