A Benamou-Brenier Proximal Splitting Method for Constrained Unbalanced Optimal Transport

The dynamic formulation of optimal transport, also known as the Benamou-Brenier formulation, has been extended to the unbalanced case by introducing a source term in the continuity equation. When this source term is penalized based on the Fisher-Rao metric, the resulting model is referred to as the Wasserstein-Fisher-Rao (WFR) setting, and allows for the comparison between any two positive measures without the need for equalized total mass. In recent work, we introduced a constrained variant of this model, in which affine integral equality constraints are imposed along the measure path. In the present paper, we propose a further generalization of this framework, which allows for constraints that apply not just to the density path but also to the momentum and source terms, and incorporates affine inequalities in addition to equality constraints. We prove, under suitable assumptions on the constraints, the well-posedness of the resulting class of convex variational problems. The paper is then primarily devoted to developing an effective numerical pipeline that tackles the corresponding constrained optimization problem based on finite difference discretizations and parallel proximal schemes. Our proposed framework encompasses standard balanced and unbalanced optimal transport, as well as a multitude of natural and practically relevant constraints, and we highlight its versatility via several synthetic and real data examples.

💡 Research Summary

**
The paper introduces a highly general formulation of dynamic unbalanced optimal transport (UOT) that extends the Wasserstein‑Fisher‑Rao (WFR) model by allowing affine integral constraints not only on the density ρ(t,·) but also on the momentum ω(t,·) and the source term ζ(t,·). In the classical Benamou‑Brenier formulation, the transport problem is expressed as a minimization of the kinetic energy ∫₀¹∫Ω‖v‖² u dx dt under the continuity equation ∂ₜu+∇·v=0. The unbalanced version adds a source term w, leading to ∂ₜu+∇·v=w, and penalizes both transport and mass creation/destruction through the WFR Lagrangian fδ(ρ,ω,ζ) = (‖ω‖²+δ²ζ²)/(2ρ).

Previous work (ACUOT) imposed only linear equality constraints on the density path, of the form ∫Ω H(t,x) ρ(t,x)dx = F(t). This paper generalizes the constraint class in three major ways:

1. Constraint Targets – The affine functional may involve ρ, ω, and ζ simultaneously:
   ∫Ω Hρ_i ρ dx + ∫Ω Hω_i·ω dx + ∫Ω Hζ_i ζ dx ≤ F_i(t) (or = F_i(t)).
   This enables direct control of mass growth rates, directional flow, and local source intensity, which is essential in applications such as population dynamics (limiting birth/death rates), robotics (obstacle avoidance), and crowd management (congestion caps).

2. Inequality Constraints – By allowing “≤” constraints, the model can enforce upper or lower bounds on total mass, density, or flux, something not possible with pure equalities.

3. Time‑Dependent Linear Functionals – The functions Hρ_i, Hω_i, Hζ_i and the bounds F_i(t) may vary arbitrarily in time, offering great flexibility for modeling time‑varying resources or regulations.

The authors prove well‑posedness of the resulting convex variational problem (Theorem 3.4). The key ingredients are: (i) the WFR integrand fδ is 1‑homogeneous and jointly convex in (ρ,ω,ζ); (ii) the admissible set defined by the affine constraints is closed and convex under mild regularity (boundedness, measurability) assumptions on the H‑functions; (iii) the direct method of the calculus of variations yields existence of minimizers.

For computation, the continuous problem is discretized on a uniform space‑time grid using finite differences for the continuity equation (forward–backward scheme). The discretized energy and constraints remain convex, allowing the use of a Parallel Proximal Algorithm (PPA) rather than the Douglas‑Rachford splitting used previously. Each constraint contributes a separate proximal operator; because the operators are independent, they can be evaluated in parallel on GPUs or multi‑core CPUs, dramatically reducing runtime. The authors explicitly derive the proximal maps for the density, momentum, source, and each affine functional, and prove convergence of the algorithm by linking the discrete minimizers to the continuous existence result (Remark 4.2).

Numerical experiments cover both synthetic and real‑world scenarios:

• Synthetic tests compare unconstrained WFR, density‑only constrained, and fully constrained formulations, illustrating how the optimal transport path deforms to respect the added restrictions.
• Population dynamics on a real city map demonstrates simultaneous enforcement of a total‑population upper bound and a regional congestion limit, yielding realistic migration patterns that would be impossible under standard UOT.
• Additional examples include obstacle avoidance in robot path planning and color‑transfer tasks, showcasing the framework’s versatility.

The paper situates its contributions relative to prior work on static OT capacity constraints, martingale constraints, and deep‑learning based soft‑constraint methods. Unlike those approaches, the present method retains convexity, provides hard (exact) constraints, and offers provable convergence—features crucial for scientific and engineering applications where reliability outweighs sheer scalability.

In summary, the authors deliver a comprehensive extension of the WFR dynamic model that accommodates a broad class of linear equality and inequality constraints on all control variables, prove existence and convergence, and propose an efficient, parallelizable proximal‑splitting solver. The work opens avenues for further research on high‑dimensional scaling, non‑linear constraints, and integration with data‑driven parameter estimation.