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

TL;DR

This work advances constrained unbalanced optimal transport by allowing affine equality and inequality constraints on density, momentum, and source terms within the Wasserstein-Fisher-Rao framework. It proves well-posedness under feasibility, derives a discrete, convex formulation, and develops a parallel proximal algorithm (PPXA) to solve the resulting problems efficiently. Through extensive synthetic and real-data experiments, the method demonstrates versatility across mass-preservation constraints, barriers, and control-constraint scenarios, while providing convergence guarantees. The approach offers a flexible, rigorous toolkit for constrained transport of measures in low-dimensional settings with potential extensions to stochastic Schrödinger-type models.

Abstract

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.

Figures (8)

• Figure 1: Spherical Hellinger-Kantorovich distance: Solution of the constrained WFR problem where the solution is restricted to the space of probability measures. Gray and blue indicate $\rho_0$ and $\rho_1$, respectively; the red curves show the solution at $t=0.25$ (left), $t=0.5$ (middle), and $t=0.75$ (right). The green dotted curve represents the solution of the unconstrained WFR geodesic projected onto the constraint according to the results of laschos2019geometric.
• Figure 2: Left panel: the inequality constraint $\int_\Omega d\rho_t \geq c$, with $c=0.8$ (purple), $c=1.0$ (green), and the unconstrained case (orange). Top figure: gray and blue again correspond to $\rho_0$ and $\rho_1$, while the remaining curves display the solution at $t=0.5$. Bottom figure: Total mass as a function of time for the inequality constraint, with colors matching those in the top figure. Right panel: In 2D, the top panel shows a constrained geodesic and the bottom panel an unconstrained geodesic (evaluated at $t=0$, $0.25$, $0.5$, $0.75$, and $1.0$) along with the evolution of total mass. Here, the total mass constraint is defined by $F(t)=3-8(t-0.5)^2$ for all $t\in[0,1]$.
• Figure 3: Barrier constraint: transporting a density through a domain with obstacles. The top row corresponds to a static barrier case, while the bottom one shows the results with a dynamic (i.e. time changing) barrier. The solution is shown at $t = 0, 0.1, 0.3, 0.5, 0.7, 0.9$ and $1.0$.
• Figure 4: Effect of a convex-curve constraint on the geodesic evolution. From left to right we show $t=0, 0.25, 0.5, 0.75,$ and $1$. Each panel displays the density on $\mathbb{S}^1$ (orange) together with the corresponding convex curve (blue). The top row enforces a convex-curve constraint; the bottom row shows the unconstrained geodesic.
• Figure 5: Illustration of a flow constraint: the brighter region represents where the momentum is constrained to align with the given vector field. The bottom row shows the constrained geodesic, while the top row shows the unconstrained geodesic. Each panel displays the density at $t=0, 0.25, 0.5, 0.75,$ and $1$.

Theorems & Definitions (18)

• Definition 2.1: Distributional Continuity Equation
• Definition 2.2: WFR distance
• Remark 2.3
• Definition 2.4: WFR problem with time-varying constraints
• Remark 2.5
• Definition 3.1
• Definition 3.2
• Remark 3.3
• Theorem 3.4
• Remark 3.5: Feasibility of the constrained problem