Slicing Unbalanced Optimal Transport

TL;DR

This work addresses robustly comparing positive measures with unequal mass by merging unbalanced OT with sliced OT through two new losses, $SUOT$ and $USOT$. It develops GPU-friendly Frank-Wolfe algorithms that rely on translation-invariant duals to decompose into 1D sliced OT problems, achieving favorable convergence and dimension-free sample complexity. The authors prove existence, metric properties, duality, and relationships between the proposed losses and classical OT notions, and demonstrate substantial practical gains on document classification, color transfer, and large-scale geophysical barycenters. The approach offers a modular framework that extends prior work and enables efficient, scalable analysis of high-dimensional, unnormalized data with theoretical rigor and broad applicability.

Abstract

Optimal transport (OT) is a powerful framework to compare probability measures, a fundamental task in many statistical and machine learning problems. Substantial advances have been made in designing OT variants which are either computationally and statistically more efficient or robust. Among them, sliced OT distances have been extensively used to mitigate optimal transport's cubic algorithmic complexity and curse of dimensionality. In parallel, unbalanced OT was designed to allow comparisons of more general positive measures, while being more robust to outliers. In this paper, we bridge the gap between those two concepts and develop a general framework for efficiently comparing positive measures. We notably formulate two different versions of sliced unbalanced OT, and study the associated topology and statistical properties. We then develop a GPU-friendly Frank-Wolfe like algorithm to compute the corresponding loss functions, and show that the resulting methodology is modular as it encompasses and extends prior related work. We finally conduct an empirical analysis of our loss functions and methodology on both synthetic and real datasets, to illustrate their computational efficiency, relevance and applicability to real-world scenarios including geophysical data.

Figures (14)

• Figure 1: Toy illustration on the behaviors of $\text{\upshape{SUOT}}$ and $\text{\upshape{USOT}}$. (left) Original 2D samples and slices used for illustration. KDE density estimations of the projected samples: grey, original distributions, colored, distributions reweighed by $\text{\upshape{SUOT}}$(center), and reweighed by $\text{\upshape{USOT}}$(right).
• Figure 1: Accuracy on document classification
• Figure 2: Ablation on BBCSport of $\rho$.
• Figure 3: Color transfer between a source and a target image ( first and second columns). We compare $\text{\upshape{SOT}}$ gradient flows operated in the color space ( third column) and the same procedure with a reweighing of the distributions by $\text{\upshape{USOT}}$ ( fourth column). The last column shows a percentage of mass change given by $\text{\upshape{USOT}}$, i.e., $\frac{(\pi^*_2 - \beta)}{\beta}$, where red indicates mass creation and blue mass destruction.
• Figure 4: Barycenter of geophysical data. ( First row) Simulated output of 4 different climate models depicting different scenarios for the evolution of a tropical cyclone ( Second row) Results of different aggregation strategies.

Theorems & Definitions (54)

• Definition 2.1: $\varphi$-divergences
• Definition 2.2: Unbalanced OT liero2018optimal
• Proposition 2.3: Corollary 4.12 in liero2018optimal
• Definition 2.4: Sliced OT
• Definition 3.1: Sliced Unbalanced OT
• Proposition 3.2: $\text{\upshape{SUOT}}$: Existence of solutions
• Proposition 3.3: SUOT: Metric properties
• Theorem 3.4: SUOT: Sample complexity
• Theorem 3.5: $\text{\upshape{SUOT}}$: Strong duality
• Definition 3.6: Unbalanced Sliced OT