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

TL;DR

This paper introduces a flexible numerical method for constrained unbalanced optimal transport that handles various constraints on measure paths, including density, momentum, and source terms, extending existing models.

Contribution

It generalizes the Benamou-Brenier formulation to include diverse constraints on multiple components, with proven well-posedness and an efficient parallel proximal splitting numerical scheme.

Findings

The framework encompasses standard and unbalanced optimal transport.

The method effectively handles affine inequality and equality constraints.

Numerical experiments demonstrate versatility on synthetic and real data.

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.