Domain decomposition for entropic unbalanced optimal transport

TL;DR

This work extends domain decomposition (DD) methods to entropic unbalanced optimal transport (UOT), enabling scalable solutions on large grids where standard Sinkhorn becomes costly. The authors introduce sequential and parallel DD algorithms that handle soft marginal constraints by incorporating a Y-marginal background term ν_{-J} and using adaptive convex-combination updates GetWeights to guarantee global descent. They establish existence, regularity, and convergence results for the cell subproblems and the overall DD schemes, including scenarios with finitely supported measures and parallel updates. Practical implementation details—cost/divergence choices, Newton-based Y-steps, log-domain stability, and a balancing step—are paired with extensive numerical experiments that compare DD against global Sinkhorn, reveal conditioning on the regularization param λ, and demonstrate strong performance on large-scale problems. The findings indicate that unbalanced DD, especially with staggered partitions and adaptive weighting, delivers substantial speedups and memory efficiency while preserving solution quality, making it viable for large-scale unbalanced OT applications.

Abstract

Solving large scale entropic optimal transport problems with the Sinkhorn algorithm remains challenging, and domain decomposition has been shown to be an efficient strategy for problems on large grids. Unbalanced optimal transport is a versatile variant of the balanced transport problem and its entropic regularization can be solved with an adapted Sinkhorn algorithm. However, it is a priori unclear how to apply domain decomposition to unbalanced problems since the independence of the cell problems is lost. In this article we show how this difficulty can be overcome at a theoretical and practical level and demonstrate with experiments that domain decomposition is also viable and efficient on large unbalanced entropic transport problems.

Figures (11)

• Figure 1: Illustration of $\nu^n_{-J}$ and $\pi^{n}_J$ in the counterexample of Remark \ref{['remark:continuousY']}.
• Figure 2: The left and center panels show the evolution of two notions of the marginal error for an unbalanced transport problem between two images of size $256 \times 256$ and regularization strength $\varepsilon = (8\Delta x)^2$ (where $\Delta x=1/256$ is the image grid resolution). The $\mathop{\mathrm{KL}}\nolimits$ criterion on the left is the primal-dual gap, normalized by $\lambda$ (see \ref{['eq:PD-gap-KL']}), while the $\mathrm{L}^1$ criterion in the center is a heuristic adaptation of the $\mathrm{L}^1$ error that is often used for balanced transport. For finite $\lambda$, after an initial stage of fast convergence the rate decreases sharply. On the right, we show the impact that the plateaus in the left and center plots have on the number of iterations needed to reach a given tolerance for different ranges of $\lambda$.
• Figure 3: Behavior of the different parallelization approaches for $N = 32$. The column titles denote the iteration number $k$.
• Figure 4: Evolution of the sub-optimality gap for the different parallelization approaches.
• Figure 5: Evolution of the coupling sparsity, which we define as the number of plan entries above the truncation threshold $\textsc{Thr} = 10^{-14}$, normalized to the number of entries (above the truncation threshold) of the optimal plan.

Theorems & Definitions (22)

• Proposition 2.2: ChizatUnbalanced2016, Proposition 2.3
• Example 2.3
• Definition 2.4: Unbalanced entropic optimal transport with background measure $\nu_{-}$
• Theorem 2.5: ChizatUnbalanced2016 Solutions to entropic UOT
• Definition 2.6: Sinkhorn algorithm for unbalanced entropic optimal transport
• Definition 2.7: Basic and composite partitions
• Definition 3.1: Unbalanced cell subproblem
• Remark 3.2
• Lemma 3.5: Properties of $\varphi^*$
• Lemma 3.6: Properties of cell optimizers on finite spaces