Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport

## Abstract

Abstract A common challenge in aggregating data from multiple sources can be formalized as an Optimal Transport (OT) barycenter problem, which seeks to compute the average of probability distributions with respect to OT discrepancies. However, the presence of outliers and noise in the data measures can significantly hinder the performance of traditional statistical methods for estimating OT barycenters. To address this issue, we propose a novel, scalable approach for estimating the robust continuous barycenter, leveraging the dual formulation of the (semi-)unbalanced OT problem. To the best of our knowledge, this paper is the first attempt to develop an algorithm for robust barycenters under the continuous distribution setup.
Our method is framed as a min \min roman_min - max \max roman_max optimization problem and is adaptable to general cost function.
We rigorously establish the theoretical underpinnings of the proposed method and demonstrate its robustness to outliers and class imbalance through a number of illustrative experiments.

## 1 Introduction

The world is full of data, and accurate analysis of this data allows for solving a number of problems in the world related to science, medicine, engineering, etc.
A common challenge in data analysis arises from two factors: (i) a huge quantity of data, and (ii) the data originating from multiple sources, such as different scientific experiments, hospitals, or engineering trials.
One promising way to address these difficulties is to perform data aggregation. This means reducing the particular source-specific characteristics and leaving only source-agnostic information for solving a practical case on hand.

Recently, a fruitful branch of research (Li et al., ; Fan et al., ; Korotin et al., ; Kolesov et al., ; ) has emerged that focuses on addressing the challenge of data aggregation through the framework of the Optimal Transport (OT) barycenter problem.
Given a number of reference distributions representing data from different sources, the problem is to discover a geometrically meaningful average of these distributions by minimizing the average OT costs from the reference distributions to the desired one.

Originally introduced for a rather specific setup (quadratic Euclidean OT cost functions) (Agueh & Carlier, ), the subfield of OT barycenter has evolved tremendously over the last decade. At present, it is the subject of deep learning and can be adapted to high-dimensional data setups, e.g., images. However, current barycenter research tends to ignore a typical and important property of data appearing in real-world applications: the presence of undesirable noise and outliers (Le et al., ). This motivates the development of a robust OT barycenter framework which (i) inherits all the advances of current barycenter researches, e.g., permits deep learning and could be adapted to high dims; (ii) wisely deals with aforementioned inconveniences of real-world data. To satisfy both aims, we propose to opt for recent continuous unbalanced OT techniques (Gazdieva et al., ; Choi et al., ; ) which are naturally aligned with data imperfection. Our contributions are as follows:

- 1.
We propose a novel semi-unbalanced OT (SUOT) barycenter framework that is robust to outliers and class imbalance in the reference distributions (§4).
- 2.
We develop a solid theoretical foundation of our proposed approach (§4.1).
- 3.
We conduct a number of experiments on toy and image data setups to demonstrate the performance of our method and showcase its robustness to outliers and imbalance of classes. (§5)

To the best of our knowledge, our work is the first attempt to construct robust OT barycenters under continuous computational setup (§2.3). Existing works with similar problematic consider exclusively discrete setting (Le et al., ; Wang et al., ).

Notations. Throughout the paper, we denote $\overline{K}=\{1,...,K\}$ for $K\in\mathbb{N}$ and use $s_{[1:K]}$ to denote the tuplet of objects $(s_{1},...,s_{K})$. We define $\mathcal{X}\subset\mathbb{R}^{D^{\prime}}$, $\mathcal{X}_{k}\subset\mathbb{R}^{D_{k}}$ and $\mathcal{Y}\subset\mathbb{R}^{D}$ as the compact subsets of Euclidian spaces. We use $\mathcal{C}(\mathcal{Y})$ to denote the set of continuous functions on $\mathcal{Y}$. The sets of absolutely continuous Borel probability (or non-negative) measures on $\mathcal{X}$ are denoted as $\mathcal{P}(\mathcal{X})$ (or $\mathcal{M}_{+}(\mathcal{X})$).
The joint probability distributions on $\mathcal{X}\times\mathcal{Y}$ with the marginals $\mathbb{P}$ and $\mathbb{Q}$ are denoted as $\Pi(\mathbb{P},\mathbb{Q})$ and usually called the transport plans. All joint positive measures on $\mathcal{X}\times\mathcal{Y}$ are referred to as $\mathcal{M}_{+}(\mathcal{X}\times\mathcal{Y})$. For a given measure $\gamma(x,y)\in\mathcal{P}(\mathcal{X}\times\mathcal{Y})$ (or $\mathcal{M}_{+}(\mathcal{X}\times\mathcal{Y})$), we use $\gamma_{x}(x)$, $\gamma_{y}(y)$ to denote its marginals. All probability measures $\gamma(x,y)\in\mathcal{P}(\mathcal{X}\times\mathcal{Y})$ s.t. $\gamma_{y}=\mathbb{Q}$ for a given probability measure $\mathbb{Q}$ are denoted as $\Pi(\mathbb{Q})$.
We use $\gamma(y|x)$ to denote the conditional probability measure.
For a measurable map $T$, $T_{\#}$ denotes the corresponding push-forward operator. The Fenchel conjugate of a function $\psi(u)$ is denoted as $\overline{\psi}(t)\stackrel{{\scriptstyle\text{def}}}{{=}}\sup_{u\in\mathbb{R}
}\{ut-\psi(u)\}$.

## 2 Background

In this section, we give an overview of the theoretical concepts related to our paper. In \wasyparagraph2.1, we state the SUOT problems and its semi-dual formulation. In \wasyparagraph2.2, we formulate the OT/UOT barycenter problem. Then, in \wasyparagraph2.3, we describe our computational setup. For details on OT, we refer to (Santambrogio, ; Villani et al., ; Peyré et al., ), unbalanced OT - (Chizat, ; Liero et al., ; Séjourné et al., ), OT barycenters - (Agueh & Carlier, ; Chizat, ).

To begin with, we introduce the concept of $\psi$-divergences needed for our further derivations.

$\psi$-divergences for positive measures.
Let $\mu_{1},\mu_{2}\in\mathcal{M}_{+}\!(\mathcal{X}^{\prime})$ be two positive measures. Then, for a function $\psi:\mathbb{R}_{+}\rightarrow\mathbb{R}_{+}\cup\{\infty\}$, the $\psi$-divergence between $\mu_{1},\mu_{2}$ is given by:

$$ $\text{D}_{\psi}\left(\mu_{1}\|\mu_{2}\right)\stackrel{{\scriptstyle\text{def}} }{{=}}\int_{\mathcal{X}}\psi\bigg{(}\frac{\mu_{1}(x)}{\mu_{2}(x)}\bigg{)}d\mu_ {2}(x).$ $$

Here the generator function $\psi(t)$ is assumed to be convex, non-negative, lower semi-continuous and attain zero uniquely when $t=1$. Note that $\overline{\psi}(t)$ is non-decreasing since function $\psi(t)$ is non-negative (Choi et al., ). Under these assumptions, $D_{\psi}$ is a valid measure of dissimilarity between two positive measures.
Well known examples of such $\psi$-divergences include Kullback-Leibler divergence $\text{D}_{\text{KL}}$ (Chizat, ; Séjourné et al., ) and $\chi^{2}$-divergence $\text{D}_{\chi^{2}}$. For the extended discussion on admissible divergences, we refer to (Gazdieva et al., , Appendix C).

### 2.1 Optimal Transport

Classic OT problem. Consider two probability measures $\mathbb{P}\!\in\!\mathcal{P}(\mathcal{X})$, $\mathbb{Q}\!\in\!\mathcal{P}(\mathcal{Y})$ and the cost function $c(x,y)\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$. Then the OT problem (Kantorovich, ) between $\mathbb{P}$ and $\mathbb{Q}$ is given by

$$ $\displaystyle\text{OT}_{c}(\mathbb{P},\mathbb{Q})\stackrel{{\scriptstyle\text{ def}}}{{=}}\inf_{\pi\in\Pi(\mathbb{P},\mathbb{Q})}\int_{\mathcal{X}\times \mathcal{Y}}c(x,y)\pi(x,y)dxdy.$ (1) $$

Under the mild assumptions on the distributions $\mathbb{P}$, $\mathbb{Q}$ and cost function $c$, the minimizer $\pi^{*}$ of (1) always exists but is not guaranteed to be unique, see (Villani et al., ). In the special case of the quadratic cost $c(x,y)=\frac{\|x-y\|^{2}}{2}$, problem (1) reduces to the well-known (squared) Wasserstein-2 distance ($\mathbb{W}_{2}^{2}(\mathbb{P},\mathbb{Q})$). Problem (1) admits a semi-dual reformulation:

$$ $\displaystyle\text{OT}_{c}(\mathbb{P},\mathbb{Q})=\sup_{f\in\mathcal{C}( \mathcal{Y})}\big{[}\int_{\mathcal{X}}f^{c}(x)d\mathbb{P}(x)+\int_{\mathcal{Y} }f(y)d\mathbb{Q}(y)\big{]}$ (2) $$

where $f^{c}(x)\stackrel{{\scriptstyle\text{def}}}{{=}}\inf_{\mu\in\mathcal{P}(
\mathcal{Y})}\int_{\mathcal{Y}}\big{(}c(x,y)-f(y)\big{)}d\mu(y)$
is the $c$-transform. Note that the standard $c$-transform in the OT duality is usually defined as $f^{c}(x)=\inf_{y\in\mathcal{Y}}\{c(x,y)-f(y)\}$. Here, for the future needs, we substitute the classic c-transform by the weak c-transform (Backhoff-Veraguas et al., , Theorem 1.3), (Gozlan et al., ). This transition is valid since we work the infimum in weak transform is anyway attained at any $\mu\in\mathcal{P}(\mathcal{Y})$ supported on the $\arg\inf_{y\in\mathcal{Y}}\{c(x,y)-f(y)\}$ set.

Unbalanced & semi-unbalanced OT problems. The standard formulation of the OT problem entails various issues (Balaji et al., ; Séjourné et al., ), such as sensitivity to outliers, inability to deal with class imbalance in the source and target measures and inapplicability to positive measures. Still, we consider the mass transportation between probability measures and focus on the first two issues. Fortunately, they can be overcome by relaxing the hard marginal constraints of the OT problem which yields the Unbalanced OT (Chizat, ; Liero et al., , UOT) problem.
Formally, the UOT problem between the probability measures $\mathbb{P}\!\in\!\mathcal{P}(\mathcal{X})$, $\mathbb{Q}\!\in\!\mathcal{P}(\mathcal{Y})$ is

$$ $\displaystyle\text{UOT}_{c,\psi,\phi}(\mathbb{P},\mathbb{Q})=\inf_{\gamma\in \mathcal{M}_{+}(\mathcal{X}\times\mathcal{Y})}\int_{\mathcal{X}\times\mathcal{ Y}}c(x,y)\gamma(x,y)dxdy+\text{D}_{\psi}\left(\gamma_{x}\|\mathbb{P}\right)+ \text{D}_{\phi}\left(\gamma_{y}\|\mathbb{Q}\right)$ (3) $$

where $c(x,y)\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$ is a continuous cost function and $\text{D}_{\psi}$, $\text{D}_{\phi}$ are $\psi$-divergences over $\mathcal{X}$ and $\mathcal{Y}$ respectively.
Note that the classic OT problem (1) is a particular instance of the UOT problem (3) when the $\psi$-divergences $\text{D}_{\psi}$, $\text{D}_{\phi}$ satisfy $\text{D}_{\psi^{\prime}}(\mathbb{P}^{\prime},\mathbb{Q}^{\prime})=0$ if $\mathbb{P}^{\prime}=\mathbb{Q}^{\prime}$, and $+\infty$ otherwise. Equivalently, it means that the generator functions $\psi$, $\phi$ are the convex indicators of $\{1\}$ and their conjugates are defined as $\overline{\psi}(t)=\overline{\phi}(t)=t$. We can tune the degree of penalizing the marginal distributions mismatch by introducing the unbalancedness parameter $\tau>0$ and considering $\text{D}_{\psi^{\prime}}=\tau\text{D}_{\psi}$, $\text{D}_{\phi^{\prime}}=\tau\text{D}_{\phi}$. Informally, when $\tau\rightarrow+\infty$, the corresponding UOT problem tends to classic OT (1).

In this paper, we focus on the semi-unbalanced OT (SUOT) problem, i.e., the case when only the first marginal constraint in the OT problem (1) is softened:

$$ $\displaystyle\text{SUOT}_{c,\psi}(\mathbb{P},\mathbb{Q})=\inf_{\gamma\in\Pi( \mathbb{Q})}\int_{\mathcal{X}\times\mathcal{Y}}c(x,y)\gamma(x,y)dxdy+\text{D}_ {\psi}\left(\gamma_{x}\|\mathbb{P}\right).$ (4) $$

Note that the minimizer $\gamma^{*}$ of (4) always exists thanks to the compactness of the space $\mathcal{X}\times\mathcal{Y}$ (yielding the compactness of $\Pi(\mathbb{Q})$), continuity of the cost function $c(x,y)$ and lower semi-continuity of the function $\psi$ (yielding the lower-semi-continuity of the optimized functional w.r.t. $\gamma$).

The semi-dual SUOT problem is given by

$$ $\displaystyle\text{SUOT}_{c,\psi}(\mathbb{P},\mathbb{Q})=\!\sup_{f\in\mathcal{ C}(\mathcal{Y})}\!\Bigg{\{}\!\!\!-\!\!\int_{\mathcal{X}}\!\!\overline{\psi}(-f ^{c}(x))d\mathbb{P}(x)\!+\!\!\!\int_{\mathcal{Y}}\!\!f(y)d\mathbb{Q}(y)\Bigg{ \}}\!.$ (5) $$

Note that when $\overline{\psi}(t)=t$, we get dual problem (2) for balanced OT problem (1).

### 2.2 (Semi-)Unbalanced Optimal Transport Barycenter

Consider probability measures $\mathbb{P}_{k}\in\mathcal{P}(\mathcal{X}_{k})$ and continuous cost functions $c_{k}(x,y):\mathcal{X}_{k}\times\mathcal{Y}\mapsto\mathbb{R}$, $k\in\overline{K}$. Given wights $\lambda_{k}\geq 0$ s.t. $\sum_{k=1}^{K}\lambda_{k}=1$, the classic OT barycenter problem consists in finding a minimizer of the sum of OT problems with fixed first marginals $\mathbb{P}_{[1:K]}$:

$$ $\displaystyle\inf_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\sum_{k=1}^{K}\lambda _{k}\text{OT}_{c_{k}}(\mathbb{P}_{k},\mathbb{Q}).\vspace{-2mm}$ (6) $$

To ensure the robustness of the barycenter estimation it is natural to consider the relaxation of the marginals $\mathbb{P}_{[1:K]}$, i.e., substitute the OT problems in (6) with the SUOT problems. Let $\psi_{[1:K]}$ be the set of $\psi$-divergences. Then the SUOT barycenter problem is

$$ $\displaystyle\mathcal{L}^{*}\stackrel{{\scriptstyle\text{def}}}{{=}}\inf_{ \mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\mathcal{B}_{u}(\mathbb{Q})\stackrel{{ \scriptstyle\text{def}}}{{=}}\inf_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\sum_ {k=1}^{K}\lambda_{k}\text{SUOT}_{c_{k},\psi_{k}}(\mathbb{P}_{k},\mathbb{Q}). \vspace{-2mm}$ (7) $$

Clearly, the SUOT barycenter problem (7) subsumes the conventional OT one (6). We also note that in principle one may substitute OT with UOT problem (3) unbalanced from the both sides. However, this is not practically meaningful. Indeed, the ultimate goal is to get rid of the potential outliers and noises in the input data, not in the barycenter. So there is no need for unbalancedness in the barycenter.

Thanks to (Liero et al., , Corollary 2.9), the functionals $\mathbb{Q}\mapsto\text{SUOT}_{c_{k},\psi_{k}}(\mathbb{P}_{k},\mathbb{Q})$ ($k\in\overline{K}$) are convex and lower semi-continuous. Thus, the functional $\mathbb{Q}\mapsto\mathcal{B}_{u}(\mathbb{Q})$ is convex and lower semi-continuous itself.
We note that $\mathcal{P}(\mathcal{Y})$ is weakly compact. Thus, thanks to the Weierstrass theorem (Santambrogio, , Box 1.1), $\mathcal{B}_{u}(\mathbb{Q})$ admits at least one minimizer meaning that the barycenter $\mathbb{Q}^{*}$ exists. Still, we can not claim that it is unique since the functional $\mathbb{Q}\mapsto\mathcal{B}_{u}(\mathbb{Q})$ is not proved to be strictly convex.

### 2.3 Computational setup

In a real-world scenario, the distributions $\mathbb{P}_{k}$ are not available explicitly and can be assessed only via empirical samples.
Assume that we are given $N_{k}$ empirical samples $x^{k}_{[1:N_{k}]}\sim\mathbb{P}_{k},k\in\overline{K}$. Our goal is to find approximations $\widehat{\gamma}_{k}$ of the SUOT plans $\gamma^{*}_{k}$ between the measures $\mathbb{P}_{k}$ and unknown SUOT barycenter $\mathbb{Q}^{*}$ for the given cost functions and $\psi$-divergences $c_{k},\psi_{k},\;k\in\overline{K}$. After that, we can use the learned plans to perform the conditional sampling, i.e., derive new points $y\sim\widehat{\gamma}_{k}(\cdot|x_{k})$ from the (approximate) barycenter taking samples $x_{k}\sim(\widehat{\gamma}_{k})_{x}$ as inputs. Actually, sampling from the left marginals of the learned plans is not an easy task since the marginal can not be easily assessed using the learned plans or potentials. Fortunately, we can deal with this issue using the rejection sampling procedure (Forsythe, ), see our \wasyparagraph4.2 for the additional details.

It is important to note that the learned plans should admit the new input samples, i.e., not necessarily present in the training datasets. This setup is typically called continuous (Li et al., ; Kolesov et al., ; Korotin et al., ) and significantly differs from the discrete one (Peyré et al., ; Cuturi & Doucet, ). The latter is aimed at solving the barycenter problem between the empirical measures which makes the application to new data samples challenging.

## 3 Related Works

The first subsection below is devoted to continuous OT solvers with a specific focus on unbalanced setup. In turn, the second subsection covers the most relevant OT barycenter solvers.

Neural OT/UOT solvers. Neural network-based continuous OT is a popular and fruitful area of recent generative modelling research. Some keynote solvers include: (Makkuva et al., ; Korotin et al., ; Amos, ) (ICNN-based, quadratic cost); (Seguy et al., ; Daniels et al., ; Mokrov et al., ) (Entropic OT); (Vargas et al., ; De Bortoli et al., ; Gushchin et al., ; Tong et al., ; Shi et al., ; Gushchin et al., ) (Schrödinger bridge); (Liu et al., ; Tong et al., ; Kornilov et al., ; Klein et al., ) (Flow matching). Of special importance for our developed method are max-min (adversarial) OT solvers based on (semi-) dual OT formulation (Rout et al., ; Korotin et al., ; ; Fan et al., ). Recently, the adversarial methodology has been extended to unbalanced setup (Choi et al., ; Gazdieva et al., ; Choi et al., ), opening up a new intriguing research direction in the field of robust continuous OT.

OT Barycenter Methods. Below, we start with a brief overview of continuous OT barycenter solvers and then proceed to the current state of robust (unbalanced) barycenter research. The continuous OT barycenter methods could be categorized as follows: (Fan et al., ; Korotin et al., ) stick to ICNN (Amos et al., ) parameterization and work only with quadratic Euclidean costs ($\ell_{2}^{2}$); (Korotin et al., ) also deals with $\ell_{2}^{2}$ but utilizes fixed point algorithm (Álvarez-Esteban et al., ); (Noble et al., ) builds upon Schrödinger bridge; (Li et al., ; Chi et al., ) take the advantage of congruence condition; (Kolesov et al., ; ) empower the congruence condition with Neural OT (Korotin et al., ) and Energy-guided Neural OT (Mokrov et al., ) correspondingly. The latter works are the current SOTA in the continuous OT barycenter domain.

The area of robust barycenter computation is much less explored and, to the best of our knowledge, limited exclusively by discrete (§2.3) solvers (Chizat et al., ; Le et al., ; Wang et al., ; Manupriya et al., ; Yang & Ding, ) based on different principles (MMD regularization/Sinkhorn algorithm etc.). In contrast, we take a significant step forward and propose the first continuous robust barycenter approach with proper theoretical support and practical validation (§4).

## 4 Proposed Method

In \wasyparagraph4.1, we derive and investigate our novel optimization objective for learning SUOT barycenters. In \wasyparagraph4.2, we propose the algorithm to solve this problem. The proofs for all theoretical results are given in Appendix A.

### 4.1 Deriving the Optimization Objective

To develop a procedure for estimating the barycenter, one may simply substitute the dual form of SUOT problem (5) in the barycenter problem (7) and get a $\min$-$\max$ problem. Actually, without additional modifications, it will lead to the $\min$-$\max$-$\min$ problem since the $c$-transforms $f_{k}^{c_{k}}$ also need to be optimized. In our Theorem 1 below, we present the way for solving the optimization problem (7) in $\max$-$\min$ manner and without the optimization over distributions $\mathbb{Q}\in\mathcal{P}(\mathcal{Y})$.

###### Theorem 1 (Semi-dual form of SUOT barycenter problem) .

The dual form of SUOT barycenter problem (7) is given by

$$ $\displaystyle\mathcal{L}^{*}=\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]} \in\mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \big{[}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c_{k}}(x_{k}))d \mathbb{P}_{k}(x_{k})+m\big{]}.$ (8) $$

This Theorem shows that if potentials $f_{[1:K]}$ satisfy the $m$-congruence condition $\sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m$ (for some $m\in\mathbb{R}$) then the optimal value of the SUOT barycenter problem (4) can be derived by solving (8). Substituting the definition of the $c$-transform in (8), we get our final optimization objective.

###### Corollary 1 (Maximin reformulation for the semi-dual problem ( 8 )) .

It holds:

$$ $\displaystyle\mathcal{L}^{*}=\!\!\!\!\!\!\sup_{\begin{subarray}{c}m\in\mathbb{ R},\\ f_{[1:K]}\in\mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\!\!\!\!\!\!\!\inf_{ \gamma(\cdot|x_{k})\in\mathcal{P}(\mathcal{Y})}\underbrace{\sum_{k=1}^{K} \lambda_{k}\Bigg{[}\!\!\int_{\mathcal{X}_{k}}\!\!\!\!-\overline{\psi}_{k}\Big{ (}\!\!-\!\!\!\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)\!-\!f_{k}(y)\big{)}d \gamma_{k}(y|x_{k})\Big{)}d\mathbb{P}_{k}(x_{k})\!+\!m\Bigg{]}}_{\widetilde{ \mathcal{L}}(f_{[1:K]},\{\gamma_{k}(\cdot|x_{k})\}_{[1:K]},m)\stackrel{{ \scriptstyle\text{def}}}{{=}}}.$ (9) $$

where the $\sup$ is taken over potentials $f_{[1:K]}\in\mathcal{C}^{K}(\mathcal{Y})$ and $\inf$ $-$ over conditional plans $\gamma_{k}(\cdot|x_{k})$.

Interestingly, in the special case when at least one of the $\text{SUOT}_{c_{k},\psi_{k}}$ terms in the right-hand-side of (7) reduces into $\text{OT}_{c_{k}}$, the $m$-congruence condition in (9) turns into congruence condition $\sum_{k=1}^{K}\lambda_{k}f_{k}\equiv 0$ which typically appears in balanced settings (Kolesov et al., ; ).

###### Corollary 2 (Congruence Condition of the Special SUOT barycenter problem) .

Suppose that $\overline{\psi_{1}}(t)=t$, i.e., $\text{SUOT}_{c_{1},\psi_{1}}=\text{OT}_{c_{1}}$ in (7). Then, dual form (9) turns into the following dual formulation:

$$ $\displaystyle\mathcal{L}^{*}=\!\!\!\sup_{\begin{subarray}{c}f_{[1:K]}\in \mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv 0\end{subarray}}\inf_{\gamma(\cdot|x_{k}) \in\mathcal{P}(\mathcal{Y})}\sum_{k=1}^{K}\lambda_{k}\Bigg{[}\int_{\mathcal{X} _{k}}\!\!\!-\overline{\psi}_{k}\Big{(}\!\!-\!\!\int_{\mathcal{Y}}\big{(}c_{k}( x_{k},y)\!\!-\!\!f_{k}(y)\big{)}d\gamma_{k}(y|x_{k})\Big{)}d\mathbb{P}_{k}(x_{ k})\Bigg{]}.$ (10) $$

The following Theorem demonstrates the two important properties of the true plans $\gamma^{*}$ solving the SUOT barycenter problem (7). First, the left marginals of these plans $(\gamma^{*}_{k})_{x}$ can be characterized via the optimal potentials $f^{*}_{k}$ which allows us to find a procedure for sampling from the marginal density during the inference, see \wasyparagraph4.2. Second, the Theorem guarantees that the family of optimal plans is indeed one of the solutions to (9).

###### Theorem 2 (SUOT barycenter conditional plans are contained in optimal saddle points) .

Assume that there exist $m$ and $\{f_{k}^{*}(y)\in\mathcal{C}(\mathcal{Y})\}_{k=1}^{K}$ which deliver maximum to the problem (9). Assume that $c_{k}(x,y)$ are continuous cost functions and $\overline{\psi}_{k}$ are continuously differentiable functions. Let $\{\gamma_{k}^{*}\}_{k=1}^{K}$ be a family of optimal SUOT plans between $\mathbb{P}_{k}$ and some barycenter $\mathbb{Q}^{*}$.
Then for every plan $\gamma_{k}^{*}$
(a) the marginal of the plan can be represented as:

$$ $\displaystyle(\gamma_{k}^{*})_{x}(x)=\nabla\overline{\psi}(-(f^{*}_{k})^{c}(x) )\mathbb{P}_{k}(x);$ (11) $$

(b) the corresponding conditional plan $\gamma_{k}^{*}(\cdot|x_{k})$ satisfies

$$ $\displaystyle\gamma^{*}_{k}(\cdot|x_{k})\in\arg\inf_{\gamma_{k}(\cdot|x_{k})} \widetilde{\mathcal{L}}(f_{[1:K]}^{*},\{\gamma_{k}(\cdot|x_{k}))\}_{[1:K]},m).$ (12) $$

Our Theorem (2) shows that for some optimal saddle points $\{(f_{k}^{*},\gamma_{k}^{*})\}_{k=1}^{K}$ of (9) it holds that $\{\gamma_{k}^{*}\}_{k=1}^{K}$ is the family of true SUOT plans between $\mathbb{P}_{[1:K]}$ and $\mathbb{Q}$. Meanwhile, for any family of optimal $f^{*}_{[1:K]}$, the $\arg\inf_{\gamma_{k}({\cdot,x_{k}})}$ sets might contain not only optimal SUOT plans $\{\gamma_{k}^{*}\}_{k=1}^{K}$ but other functions as well which is a known fake solutions issue of neural OT solvers (Korotin et al., ).

###### Theorem 1 (Semi-dual form of SUOT barycenter problem) .

###### Corollary 1 (Maximin reformulation for the semi-dual problem ( 8 )) .

###### Corollary 2 (Congruence Condition of the Special SUOT barycenter problem) .

###### Theorem 2 (SUOT barycenter conditional plans are contained in optimal saddle points) .

### 4.2 Parametrization and Practical Optimization Procedure

Parametrization. To implement the optimization over distributions $\gamma(\cdot|x_{k})\in\mathcal{P}(\mathcal{Y})$ ($k\in\overline{K}$) in (9), we consider parametrizing them with stochastic or deterministic functions, using the strategy defined in (Kolesov et al., , \wasyparagraph4.1) and (Korotin et al., , \wasyparagraph4.1).

To realize the optimization over conditional plans $\gamma_{k}(\cdot|x_{k})$ in (9), we define an auxiliary space $\mathcal{S}\in\mathbb{R}^{D}_{s}$, atomless distribution $\mathbb{S}\in\mathcal{P}(\mathcal{S})$ and measurable maps $T_{k}:\mathcal{X}_{k}\times\mathcal{S}\rightarrow\mathcal{Y}$. For every plan $\gamma_{k}$, we consider the representation using the map $T_{k}$ s.t. $\gamma_{k}(\cdot|x)=T_{k}(x,\cdot)_{\#}\mathbb{S}$. Using the stochastic parametrization of the conditional plans, we can reformulate the optimization objective (9) as:

$$ $\displaystyle\mathcal{L}^{*}=\!\!\!\!\!\!\!\!\!\sup_{\begin{subarray}{c}m\in \mathbb{R},\\ f_{[1:K]}\in\mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}=m\end{subarray}}\!\!\!\!\!\inf_{T_{[1:K]}} \underbrace{\sum_{k=1}^{K}\lambda_{k}\!\Bigg{[}\!\!\int_{\mathcal{X}_{k}}\!\! \!\!-\overline{\psi}\Bigg{(}\!\!\int_{\mathcal{S}}\Big{(}f(T_{k}(x_{k},s))-c(x _{k},T_{k}(x_{k},s))\Big{)}d\mathbb{S}(s)\Bigg{)}d\mathbb{P}_{k}(x_{k})\!+\!m} _{{\mathcal{L}(f_{[1:K]},T_{[1:K]},m)\stackrel{{\scriptstyle\text{def}}}{{=}}} }\Bigg{]}.$ (13) $$

We denote the expression under $\sup\inf$ in (13) as $\mathcal{L}(f_{[1:K]},T_{[1:K]},m)$. In some of the setups, the stochasticity of the conditional plans $\gamma(\cdot|x)$ is not needed. Then we can consider a measurable map $T_{k}:\mathcal{X}_{k}\rightarrow\mathcal{Y}$ which specifies the deterministic conditional plans as $\gamma_{k}(\cdot|x)=\delta_{T_{k}(x)}(\cdot)$.

For solving (13), we parametrize the maps $T_{k}$ and potentials $f_{k}$ ($k\in\overline{K}$) as neural networks $T_{k,\omega}\!:\!\mathbb{R}^{D_{k}}\times\mathbb{R}^{D_{s}}\mapsto\mathbb{R}^{D}$ and $f_{k,\theta}\!:\!\mathbb{R}^{D_{k}}\mapsto\mathbb{R}$ with weights $\omega\stackrel{{\scriptstyle\text{def}}}{{=}}\omega_{[1:K]}\in\Omega$, $\theta\stackrel{{\scriptstyle\text{def}}}{{=}}\theta_{[1:K]}\in\Theta$. Here $\Omega=\Omega_{1}\times\Omega_{2}\times...\times\Omega_{K}$ and $\Theta=\Theta_{1}\times\Theta_{2}\times...\times\Theta_{K}$ are the parameter spaces for the maps and potentials respectively. Note that $\mathbb{R}^{D_{s}}$ denotes the stochastic dimension which should be omitted in the case of deterministic maps. In order to ensure the $m$-congruence condition for the potentials, we parametrize them as $f_{k,\theta}=g_{k,\theta}-\sum_{n\neq k}\frac{\lambda_{n}}{\lambda_{k}(K-1)}g_
{n,\theta}+\frac{m}{\lambda_{k}}$ where $g_{k,\theta}$ denote the auxiliary neural nets.
Here we take inspiration from similar strategies in (Li et al., ; Kolesov et al., ; ).

Figure: Algorithm 1 SUOT Barycenter with Semi-Unbalanced Neural Optimal Transport

Training. Recall that we work in a continuous setup of barycenter problem, i.e., the distributions $\mathbb{P}_{[1:K]}$ are accessible only through the empirical samples of data. Thus, we opt to estimate the objective (13) from samples using the Monte-Carlo method. Specifically, the objective is optimized using stochastic gradient descent-ascent algorithm over random batches of samples from the distributions $\mathbb{P}_{k}$ and auxiliary distribution $\mathbb{S}$(^1^11Sampling from the stochastic distribution is omitted when we deal with the deterministic maps $T_{\omega,[1:K]}$.). For simplicity, we replace the minimization of the objective $\mathcal{L}(f_{\theta,[1:K]},T_{\omega,[1:K]})$ w.r.t. $T_{\omega,[1:K]}$ with the direct minimization of the $c$-transform $\int_{\mathcal{X}_{k}}\int_{\mathcal{S}}\Big{(}c(x_{k},T_{k}(x_{k},s))\!-\!f(T
_{k}(x_{k},s))\Big{)}d\mathbb{S}(s)d\mathbb{P}_{k}(x_{k})$ which has the same set of minimizers.
We detail our optimization procedure in Algorithm 1.

Inference. Our approach for SUOT barycenter estimation relaxes the marginal constraints related to the input distributions $\mathbb{P}_{k}$. It means that during the inference, the input points $x_{k}$ should be sampled not from the distributions $\mathbb{P}_{k}$ but from the left marginals of the learned plans $(\gamma_{k})_{x}$. However, sampling from these marginals is not available explicitly and requires the usage of the established techniques for sampling from the reweighted distributions, e.g., the rejection sampling technique. We apply this technique in the following manner.

Thanks to our Theorem 2 (a), we are able to approximate the fraction of the densities of measures $(\gamma_{k}^{*})_{x}$ and $\mathbb{P}_{k}$: $\frac{(d\gamma_{k}^{*})_{x}}{d\mathbb{P}_{k}}(x)\approx\nabla\overline{\psi}(-
(\widehat{f}_{k})^{c}(x))$ where $\widehat{f}_{k}$ denotes the potential learned during training. Thus, we can (1) generate new samples $X_{k}\sim\mathbb{P}_{k}$ and samples of the same shape from the uniform distribution $u_{k}\sim U$; (2) calculate the constant $c\stackrel{{\scriptstyle\text{def}}}{{=}}\max\big{(}\nabla\overline{\psi}(-(
\widehat{f_{k}})^{c}(x))\big{)}$;
(3) accept the samples $X_{k}^{\prime}$ which satisfy the inequality $u_{k}\leq\nabla\overline{\psi}(-(\widehat{f_{k}})^{c}(X_{k}^{\prime}))$ and reject the others.

## 5 Experiments

In this section, we evaluate our model through various experiments. In \wasyparagraph5.1, we compare the numerical accuracy of our method against baseline methods. In \wasyparagraph5.2, we investigate two key properties of our U-NOTB method: robustness to class imbalance and robustness to outliers. Specifically, we show that these properties can be controlled by adjusting the unbalanced parameter $\tau$. In \wasyparagraph5.3, we evaluate our model for the general costs, leveraging shape and color-invariant cost on the image dataset, other than the quadratic cost. We provide the details about settings and baselines in Appendix B.

### 5.1 Barycenter Evaluation on Synthetic Datasets

Baselines.
We aim to evaluate whether our model accurately estimates the optimal barycenter $\mathbb{Q}$ and the corresponding transport maps $T_{[1:K]}$. Because our work is the first attempt to address the continuous SUOT barycenters, there are no existing baselines for direct comparison. Therefore, we carefully designed
two baseline models for comparison with our approach. These baseline models are motivated by the equivalence between the OT barycenter of two distributions and the interpolation between them (Villani et al., ). Each baseline model is derived from two approaches for learning the unbalanced transport map $T_{uot}$ between $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$: (i) the semi-dual UOT model (UOTM) (Choi et al., ) and (ii) the OT Map model (OTM) (Rout et al., ) combined with the mini-batch UOT sampling (Eyring et al., ). Because the UOT problem is equivalent to the OT problem between the rescaled distributions, we conduct the displacement interpolation using this $T_{uot}$ to approximate the barycenter. Specifically, we consider the following SUOT barycenter problem:

$$ $\displaystyle\inf_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\{\lambda_{1}\text{OT }_{c_{1}}(\mathbb{P}_{1},\mathbb{Q})+\lambda_{2}\text{UOT}_{c_{2}}(\mathbb{P}_ {2},\mathbb{Q})\},$ (14) $$

where $c_{1},c_{2}$ are standard quadratic cost function $c_{i}(x,y)=\frac{\|x-y\|^{2}}{2}$ and $\lambda_{1}=\lambda_{2}=0.5$. Let $T_{1}\stackrel{{\scriptstyle\text{def}}}{{=}}\lambda_{1}Id+\lambda_{2}T$. Then, ${T_{1}}_{\#}\mathbb{P}_{1}$ and $T_{1}$ correspond to the barycenter distribution $\mathbb{Q}$, and the transport map from $\mathbb{P}_{1}$ to $\mathbb{Q}$, respectively. Note that we should set the $\text{OT}_{c_{1}}$ distance for the first distribution $\mathbb{P}_{1}$ to ensure that the estimated barycenter $\mathbb{Q}={T_{1}}_{\#}\mathbb{P}_{1}$ has a unit mass in two baselines. Since the OT problem is a specific case of the SUOT one, our method can be naturally applied to this case.
We test U-NOTB on the synthetic dataset of Spiral-to-Gaussian Mixture (S$\rightarrow$GM) (Fig. 1(a)).

Experimental Results.
Fig. 1 presents a qualitative visualization of the results from our model and two baseline methods.
Because there is no closed-form solution for the continuous SUOT barycenter problems, we consider the SUOT barycenter between two discrete empirical distributions, obtained from the POT (Flamary et al., ) library, as the ground truth solution $\mathbb{Q}^{\star}$ and $T^{\star}$ (Fig. 1(b)).
In Fig. 1, for clarity, we visualized only the spiral distribution $\mathbb{P}_{1}$ and the corresponding barycenter $\mathbb{Q}$. Note that, because we set $\text{OT}_{c_{1}}$ for $\mathbb{P}_{1}$ in Eq. 14, the rejection sampling is not applied in this case. Fig. 1 shows that our model and other baseline methods present decent results.

Table 1 provides a quantitative evaluation of our model for a more detailed assessment.
The approximate barycenter $\mathbb{Q}_{\theta}$ and the transport map $T_{\theta}$ from each method are evaluated from two perspectives. First, the distribution error between $\mathbb{Q}_{\theta}$ and $\mathbb{Q}^{\star}$ is evaluated by measuring the Wasserstein-2 distance $\mathbb{W}^{2}_{2}(\tilde{\mathbb{Q}},\mathbb{Q}^{\star})$ ($\mathbb{W}_{2}$ metric). Second, the optimality error of the transport map is assessed by measuring $\mathcal{L}_{2}$-norm between the transport maps ( $\mathcal{L}_{2}$ metric), i.e., $\int_{\mathcal{X}_{1}}\lVert T_{1}(x_{1})-T^{\star}(x_{1})\rVert^{2}_{2}d
\mathbb{P}_{1}(x_{1})$.
As shown in Table 1, our model achieves the most accurate barycenter estimation across both metrics ($\mathbb{W}_{2}$ and $\mathcal{L}_{2}$) and both distributions ($\mathbb{P}_{1}$ and $\mathbb{P}_{2}$).

Figure: (a) Setup
Refer to caption: extracted/5903048/pics/compa_1.png

**Table 1: Comparison on Benchmarks on Toy datasets.**
| Data | Metric | UOTM | Mini-Batch UOT | U-NOTB (ours) |
| --- | --- | --- | --- | --- |
| Sprial $\rightarrow$ Gaussian Mixture | $\mathcal{L}_{2}$ | 0.65 | 0.31 | 0.25 |
| $\mathbb{W}_{2}$ | 0.29 | 0.15 | 0.12 |  |
| Moon $\rightarrow$ Spiral | $\mathcal{L}_{2}$ | 0.89 | 0.64 | 0.52 |
| $\mathbb{W}_{2}$ | 0.59 | 0.93 | 0.81 |  |

### 5.2 Robust Barycenter Estimation under Outlier and Class Imbalance

In this section, we examine two key properties of our U-NOTB model: (i) robustness to class imbalancedness and (ii) robustness to outliers. The value of OT functional $\text{OT}_{c}(\cdot,\cdot)$ is known to be sensitive to outliers and class imbalance problems. Hence, the OT barycenter is also largely affected by these factors. To address this sensitivity, the UOT problem extends the traditional OT problem by relaxing the constraint on marginal densities. This relaxation introduces two notable characteristics. First, even in the presence of class imbalances, the reweighting process allows the model to appropriately align with the relevant modes in a reasonable manner (Eyring et al., ). Second, the model exhibits robustness to outliers, maintaining reliable transport despite their presence (Balaji et al., ; Choi et al., ).
The goal of this section is twofold: (1) to investigate whether these two properties are also observed in our U-NOTB model by comparing it with the OT counterpart of our model (Korotin et al., ) and (2) to demonstrate that this robustness can be controlled through the unbalancedness parameter $\tau$.

Robustness to Class Imbalance.
The class imbalance issue refers to the case where each class has a different proportion across $\mathbb{P}_{i}$. This class imbalance problem can lead to undesirable behaviour in the standard OT barycenter. For example, in Fig. 2, we consider the upper-modes and bottom-modes for each distribution as class 1 and 2, respectively. In $\mathbb{P}_{1}$, 75% of the data is concentrated in class 1 (upper-left white dots) while only 25% is concentrated in class 1 for $\mathbb{P}_{1}$ (upper-right yellow dots). Therefore, the balanced OT barycenter (NOTB, Kolesov et al. ()) generates a mode between two majority modes of $\mathbb{P}_{i}$, which can be an undesirable phenomenon for the barycenter (Fig. 2(a)).

As shown in Fig. 2, we evaluated our model for diverse unbalancedness parameter $\tau\in\{1,20,100\}$ (Eq. 3). Additionally, for comparison, we conducted experiments under a balanced setting using the OT counterpart. It is important to note that as $\tau$ increases, the UOT problem converges to the OT problem, allowing for less flexibility in marginal distributions (Choi et al., ).
When $\tau=1$, our unbalanced barycenter demonstrates robust results by generating barycenter modes between the pair of the closest modes of the $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$. Our model rejects contour samples from majority modes, reweighting them to better align the corresponding modes.
Moreover, our unbalanced barycenter offers controllability of this robustness through $\tau$. As $\tau$ increases, the flexibility on marginal distributions decreases in the UOT problem. Thus, when $\tau=200$, our model accepts nearly all samples, making the UOT barycenter closely resemble the OT barycenter. We believe this adaptability of the UOT barycenter offers practical benefits in real-world scenarios.

Figure: (a) NOTB
Refer to caption: extracted/5903048/pics/im_1.png

Figure: (a) Marginal distributions
Refer to caption: extracted/5903048/pics/outlier/outlier_marginal_distribution.png

Robustness to Outliers.
To evaluate the robustness to outliers, we conducted experiments on a synthetic dataset that included a small proportion of outliers. As depicted in Fig. 3, the dataset comprises three marginal distributions: $\mathbb{P}_{1}$ (blue dots), $\mathbb{P}_{2}$ (orange dots), and $\mathbb{P}_{3}$ (green dots). Outliers, constituting 5% of the marginal $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$, were added to the outermost point of each marginal.
We evaluated our model for various unbalancedness parameters $\tau$ and compared it with the balanced OT counterpart. Similar to the robustness to class imbalance experiments, we expect that smaller values of $\tau$ will provide more robustness to outliers by allowing higher flexibility in marginal distributions.

We illustrate the results in Fig. 3. Note that for comparative analysis, we also tested the NOTB (Korotin et al., ) approach for OT barycenter estimation under the data without outliers. This experiment aims to demonstrate the desired behaviour of barycenter when there are no outliers in the dataset. When $\tau=1$, our U-NOTB closely aligns with the OT barycenter in this outlier-free scenario. This result shows that our model offers robustness to outliers by rejecting outlier samples at smaller $\tau$. As $\tau$ increases, our model accepts a higher proportion of outliers. Hence, the outliers on the left and right begin to influence the barycenter, making the modes of barycenter lie between the majority modes of distributions $\mathbb{P}_{i}$. In summary, reducing $\tau$ enhances outlier robustness, while increasing $\tau$ yields a more precise barycenter by incorporating all data points. We believe that the flexibility in adjusting $\tau$ offers promising potential for broader applications, where different levels of tolerance to outliers may be required.

Figure: (a) In-distribution Samples from $\mathbb{P}_{0}$. Acceptance Rate: 63%
Refer to caption: extracted/5903048/pics/stylegan/P0_indist.png

### 5.3 Shape-color experiment

In this section, we illustrate one interesting example demonstrating how the UOT barycenter problem can be applied using a general cost.
We designed the problem at general costs when there are some undesirable outliers in the marginal distributions.

- •
Shape distribution $\mathbb{P}_{1}$. The first marginal distribution consists of grayscaled images of digits ‘2’ ($49\%$ of training dataset), ‘3’ ($50\%$) and ‘7’($1\%$ - outliers) of MNIST data.
- •
Color distribution $\mathbb{P}_{2}$. The second marginal distribution consists of three color points: red (probability mass $p_{0}=0.495$), green ($p_{1}=0.495$), white ($p_{2}=0.01$ - outliers).
- •
Manifold. We pretrained the StyleGAN Karras et al. () generator $G:\mathcal{Z}\rightarrow\mathbb{R}^{3\times 32\times 32}$ on the colored MNIST dataset of digits ‘2’ and ‘3’.
Let $E^{i}_{\theta}$ be the network which encodes each marginal $\mathbb{P}_{i}$ to the latent $\mathcal{Z}$. Throughout the experiment, we transport $x_{i}\sim\mathbb{P}_{i}$ to the barycenter point $y_{i}$ by $y_{i}\!=\!G\circ E^{i}_{\theta}(x_{i})$, thus confining the transformed images to the colored MNIST of digits ‘2’ and ‘3’.
- •
General Transport costs. For the first marginal sample $x_{1}\sim\mathbb{P}_{1}$ and its corresponding barycenter point $y_{1}$, we use the following shape-preserving cost: $c_{1}(x_{1},y_{1})=\frac{1}{2}\lVert x_{1}-H_{g}(y_{1})\rVert^{2}$, where $H_{g}$ is a decolorization operator. Moreover, for the second marginal sample $x_{2}\sim\mathbb{P}_{2}$ and the corresponding barycenter point $y_{2}$, we use the following color-preserving cost: $c_{2}(x_{2},y_{2})=\frac{1}{2}\lVert x_{2}-H_{c}(y_{2})\rVert^{2}$, where $H_{c}$ is a color projection operator defined in Kolesov et al. ().

Our results for unbalancedness $\tau=10$ are presented in Figure 4. We demonstrate the examples of learned mapping ($x_{k},T_{k}(x_{k},s)$) from in-distribution points $x_{k}$, sub-figures 4(a), 4(c), and from outliers, sub-figures 4(b), 4(d). Expectedly, the learned in-distribution mappings preserve the shape ($\mathbb{P}_{0}\rightarrow\mathbb{Q}^{*}$) and color ($\mathbb{P}_{1}\rightarrow\mathbb{Q}^{*}$). In turn, the outlier mappings have no reasonable interpretation.

Along with the qualitative performance, we report the acceptance rates of input points $x_{k}$, see the explanation of our inference procedure in §4.2. Importantly, the acceptance of outliers is much smaller compared to in-distribution samples. In particular, if $x_{0}$ is a grayscaled digit ‘7’, it will be accepted (and processed through the mapping $T_{0}$) only with probability $0.19$, while for grayscaled digits ‘2’ and ‘3’ the same figure reads as $0.63$. Thus, the experiment showcases the applicability of our robust barycenter methodology for non-trivial cases with non-Euclidean OT costs.

## 6 Discussion

Our work continues the recent and fruitful branch of OT barycenter research. We present the first attempt to build robust OT barycenters under continuous setup based on unbalanced OT formulation. From now on, the barycenter researchers have a new deep learning tool at their disposal that gives them the ability to deal with imperfect data (with outliers/noise/class imbalance) at a large scale. We believe that our proposed method will further expand the toolbox of deep learning practitioners.

Limitations. In practice, we found that the practical optimization procedure (Algorithm 1) may work unstably. In particular, the training may diverge under improper hyper-parameters selection or the resulting quality might depend on the random seed. We hypothesise that this behaviour is due to the reliance on adversarial training and leave its thorough investigation to future research. From the theoretical side, the recovered semi-unbalanced OT barycenter is not guaranteed to be unique (§2.2); also, the convergence properties of our procedure (13) are not established. All these questions pave the way for further research.

Code of Ethics.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

## Appendix A Proofs

### A.1 Proof of Theorem 1 and Corollary 1

###### Proof of Theorem 1 .

Substituting the dual form (5) into (13), we can formulate the barycenter problem as finding

$$ $\displaystyle\mathcal{L}^{*}=\min_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})} \underbrace{\sup_{f_{1},...,f_{K}\in\mathcal{C(Y)}}\sum_{k=1}^{K}\lambda_{k} \big{[}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c}(x))d\mathbb{P}_{k }(x_{k})+\int_{\mathcal{Y}}f_{k}(y)d\mathbb{Q}(y)}_{\mathcal{F}(\mathbb{Q},f_{ [1:K]})\stackrel{{\scriptstyle\text{def}}}{{=}}}\big{]}$ (15) $$

where we replace $\inf$ with $\min$ thanks to the existence of SUOT barycenter, see \wasyparagraph2.2.

Note that the compactness of the space $\mathcal{Y}$ yields the weak compactness of $\mathcal{P}(\mathcal{Y})$. At the same time, the functional $\mathbb{Q}\mapsto\mathcal{F}(\mathbb{Q},f_{[1:K]})$ is continuous and linear in $\mathbb{Q}$. At the same time, it is concave in $f_{[1:K]}$. Indeed, the convex conjugate of the generator function $\psi$ is convex and non-decreasing by definition, the c-transform operation is convex, see (Kolesov et al., , Proposition A.1 (iii)). Thus, the superposition of a concave, non-increasing function $-\overline{\psi}_{k}(\cdot)$ and concave function $-f^{c}_{k}(\cdot)$, i.e., $-\overline{\psi}_{k}(-f^{c}_{k}(\cdot))$ is a concave function as well (for every $k\in\overline{K}$). Thus, we can apply the minimax theorem (Terkelsen, , Corollary 1) and swap $\min$ and $\sup$ in (15):

$$ $\displaystyle\mathcal{L}^{*}=\sup_{f_{1},...,f_{K}\in\mathcal{C(Y)}}\min_{ \mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\sum_{k=1}^{K}\lambda_{k}\big{[}\int_{ \mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c}(x_{k}))d\mathbb{P}_{k}(x_{k})+ \int_{\mathcal{Y}}f_{k}(y)d\mathbb{Q}(y)\big{]}=$ $\displaystyle\sup_{f_{1},...,f_{K}\in\mathcal{C(Y)}}\Big{[}\sum_{k=1}^{K} \lambda_{k}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c}(x_{k}))d \mathbb{P}_{k}(x_{k})+\min_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\int_{ \mathcal{Y}}\underbrace{\sum_{k=1}^{K}\lambda_{k}f_{k}(y)}_{\overline{f}(y) \stackrel{{\scriptstyle\text{def}}}{{=}}}d\mathbb{Q}(y)\Big{]}=$ $\displaystyle\sup_{f_{1},...,f_{K}\in\mathcal{C(Y)}}\underbrace{\Big{[}\sum_{k =1}^{K}\lambda_{k}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c}(x_{k}) )d\mathbb{P}_{k}(x_{k})+\min_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\int_{ \mathcal{Y}}\overline{f}(y)d\mathbb{Q}(y)}_{\mathcal{G}(f_{1},...f_{K}) \stackrel{{\scriptstyle\text{def}}}{{=}}}\Big{]}.$ (16) $$

Now we use the following fact:
$\inf_{\mathbb{Q}\in\mathcal{P}(\mathcal{Y})}\int_{\mathcal{Y}}\overline{f}(y)d
\mathbb{Q}(y)=\inf_{y\in\mathcal{Y}}\overline{f}(y).$
Assume that the true value $m\stackrel{{\scriptstyle\text{def}}}{{=}}\inf_{y\in\mathcal{Y}}\overline{f}(y)$ is known. Then we can restrict $\sup$ in (16) to the potentials $f_{[1:K]}$ satisfying the $m$-congruence condition: $\sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m$. Indeed, for each tuple $(f_{1},...,f_{K})$, let us consider the tuple $(\widetilde{f}_{1},...,\widetilde{f}_{K})\stackrel{{\scriptstyle\text{def}}}{{
=}}(f_{1},...,f_{K}+\frac{m-\overline{f}}{\lambda_{k}})$ satisfying the congruence condition. For this tuple, $\inf_{y\in\mathcal{Y}}\sum_{k=1}^{K}\lambda_{k}\widetilde{f}_{k}=m$.
Besides, we get:

$$ $\displaystyle\mathcal{G}(\widetilde{f}_{1},...,\widetilde{f}_{K})\!\!-\!\! \mathcal{G}(f_{1},...f_{K})=\lambda_{K}\int_{\mathcal{X}_{K}}[-\overline{\psi} _{K}(-\widetilde{f}_{K}^{c}(x_{K}))+\overline{\psi}_{K}(-f_{K}^{c}(x_{K}))]d \mathbb{P}_{K}(x_{K})\!+\!\cancel{m-m}=$ $\displaystyle\lambda_{K}\int_{\mathcal{X}_{K}}\Big{[}\overline{\psi}_{K}(-f_{K }^{c}(x_{K}))-\overline{\psi}_{K}(-(f_{K}+\frac{m-\overline{f}}{\lambda_{k}})^ {c}(x_{K}))\Big{]}d\mathbb{P}_{K}(x_{K})\geq 0.$ (17) $$

Here the inequality in line (17) follows from the properties of the $c$-transform function and convex conjugate of the divergence’ generator function $\overline{\psi}$. First, $c$-transform is a decreasing function of its argument, see (Kolesov et al., , Proposition A.1 (i)). Thus,

$$ $f_{K}+\frac{m-\overline{f}}{\lambda_{K}}\leq f_{K}+\cancel{\frac{m-m}{\lambda_ {K}}}=f_{K}\Longrightarrow(f_{K}+\frac{m-\overline{f}}{\lambda_{K}})^{c_{K}} \geq f_{K}^{c_{K}}.$ $$

Second, the convex conjugate is a non-decreasing function which yields

$$ $-f_{K}^{c_{K}}\geq-(f_{K}+\frac{m-\overline{f}}{\lambda_{K}})^{c_{K}} \Longrightarrow\psi_{K}(-f_{K}^{c_{K}})\geq\psi_{K}(-(f_{K}+\frac{m-\overline{ f}}{\lambda_{K}})^{c_{K}}).$ $$

From this, the inequality (17) becomes obvious. It means that the transition to congruent tuple is fair for the known value $m$.
It is important to note that in practice this value is not given and should be optimized. Thus, the final optimization objective can be formalized as

$$ $\displaystyle\mathcal{L}^{*}=\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]} \in\mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\mathcal{G}(f_{1},...,f_{ K})=\!\!\!\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]}\in\mathcal{C}^{K}( \mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\Big{[}\sum_{k=1}^{K} \lambda_{k}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c}(x_{k}))d \mathbb{P}_{k}(x_{k})+m\Big{]}$ (18) $$

where (18) coincides with (9).
∎

###### Proof of Corollary 1 .

By substituting the definition of the $c$-transform in (8), we get

$$ $\displaystyle\mathcal{L}^{*}=\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]} \in\mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \Bigg{[}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}\Big{(}-\inf_{\mu(y)\in \mathcal{P}(\mathcal{Y})}\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)-f_{k}(y)\big{ )}d\mu(y)\Big{)}d\mathbb{P}_{k}(x_{k})+m\Bigg{]}=$ (19) $\displaystyle\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]}\in\mathcal{C}^{ K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \Bigg{[}\int_{\mathcal{X}_{k}}\inf_{\mu(y)\in\mathcal{P}(\mathcal{Y})}\Big{\{} -\overline{\psi}_{k}\Big{(}-\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)-f_{k}(y) \big{)}d\mu(y)\Big{)}\Big{\}}d\mathbb{P}_{k}(x_{k})+m\Bigg{]}$ (20) $$

In transition from (19) to (20), we interchange the operation of taking $\inf$ over $\mu(y)\in\mathcal{P}(\mathcal{Y})$ and $\overline{\psi}_{k}$. This is possible since $-\overline{\psi}_{k}$ is monotone decreasing and infimum is attained.

Now we note that the map $(x_{k},\mu)\mapsto\Big{\{}-\overline{\psi}_{k}\Big{(}-\int_{\mathcal{Y}}\big{(
}c_{k}(x_{k},y)-f_{k}(y)\big{)}d\mu(y)\Big{)}\Big{\}}$ is measurable and bounded from below. Indeed, the spaces $\mathcal{X}_{k}$ and $\mathcal{Y}$ are compact, and $\mathcal{P}(\mathcal{Y})$ is weakly compact. At the same time, the potentials $f_{k}(y)$, cost function $c_{k}(x_{k},y)$ are continuous and $-\overline{\psi}_{k}$ is lower semi-continuous. It means that their superposition is also a lower semi-continuous function and, thus, lower bounded on the compact space. Thanks to the properties of this map, we can use the rule of interchange between integral and $\inf$ (Bertsekas & Shreve, , Propositions 7.27 and 7.50) and get

$$ $\displaystyle\eqref{cor-inf-after}=$ $\displaystyle\!\!\!\!\!\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]}\in \mathcal{C}^{K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \Bigg{[}\inf_{\gamma_{k}(\cdot|x_{k}):\mathcal{X}_{k}\mapsto\mathcal{P}( \mathcal{Y})}\int_{\mathcal{X}_{k}}\!\!\!-\overline{\psi}_{k}\Big{(}-\!\!\int_ {\mathcal{Y}}\big{(}c_{k}(x_{k},y)\!-\!\!f_{k}(y)\big{)}d\gamma_{k}(y|x_{k}) \Big{)}\Big{\}}d\mathbb{P}_{k}(x_{k})+m\Bigg{]}$ (21) $$

where the $\inf$ is taken over families of measurable maps $\{\gamma_{k}(\cdot|x_{k})\}_{[1:K]}$. Note that (21) is equal to (9) after the trivial interchange between the summation and $\inf$ operations. This completes the proof.
∎

###### Proof of Corollary 2 .

Since $\text{SUOT}_{c_{1}}=\text{OT}_{c_{1}}$, $\overline{\psi}_{1}=\text{Id}$ by the definition of convex conjugate. By substituting $\tilde{f}_{1}:=f_{1}-\frac{m}{\lambda_{1}}$ and $\overline{\psi}_{1}=\text{Id}$ into (9), we can obtain the congruence condition as follows:

$$ $\displaystyle\mathcal{L}^{*}=\!\!\!\sup_{\begin{subarray}{c}m,f_{[1:K]}\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\!\!\inf_{\gamma_{k}( \cdot|x_{k})}\sum_{k=1}^{K}\lambda_{k}\Bigg{[}\int_{\mathcal{X}_{k}}-\overline {\psi}_{k}\Big{(}-\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)-f_{k}(y)\big{)}d \gamma_{k}(y|x_{k})\Big{)}d\mathbb{P}_{k}(x_{k})+m\Bigg{]}=$ $\displaystyle\!\!\!\sup_{\begin{subarray}{c}m,f_{[1:K]}\\ \lambda_{1}\tilde{f}_{1}+\sum_{k=2}^{K}\lambda_{k}f_{k}\equiv 0\end{subarray}} \inf_{\gamma_{k}(\cdot|x_{k})\in\mathcal{P}(\mathcal{Y})}\Bigg{[}\lambda_{1} \int_{\mathcal{X}_{1}}\int_{\mathcal{Y}}\left(c_{1}(x_{1},y)-\left(\tilde{f}_{ 1}(y)+\cancel{\frac{m}{\lambda_{1}}}\right)\right)d\gamma_{1}(y|x_{1})d\mathbb {P}_{1}(x_{1})+$ $\displaystyle\sum_{k=2}^{K}\lambda_{k}\int_{\mathcal{X}_{k}}-\overline{\psi}_{ k}\Big{(}-\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)-f_{k}(y)\big{)}d\gamma_{k}(y |x_{k})\Big{)}d\mathbb{P}_{k}(x_{k})+\cancel{m}\Bigg{]}=$ $\displaystyle\!\!\!\sup_{\begin{subarray}{c}\lambda_{1}\tilde{f}_{1}+\sum_{k=2 }^{K}\lambda_{k}f_{k}\equiv 0\end{subarray}}\inf_{\gamma_{k}(\cdot|x_{k})\in \mathcal{P}(\mathcal{Y})}\Bigg{[}\lambda_{1}\int_{\mathcal{X}_{1}}\int_{ \mathcal{Y}}\overline{\psi}_{1}\Big{(}-\int_{\mathcal{Y}}\big{(}c_{1}(x_{1},y) -\tilde{f}_{1}(y)\big{)}d\gamma_{1}(y|x_{1})\Big{)}d\mathbb{P}_{1}(x_{1})+$ $\displaystyle\sum_{k=2}^{K}\lambda_{k}\int_{\mathcal{X}_{k}}-\overline{\psi}_{ k}\Big{(}-\int_{\mathcal{Y}}\big{(}c_{k}(x_{k},y)-f_{k}(y)\big{)}d\gamma_{k}(y |x_{k})\Big{)}d\mathbb{P}_{k}(x_{k})\Bigg{]}.$ $$

∎

###### Proof of Theorem 1 .

###### Proof of Corollary 1 .

###### Proof of Corollary 2 .

### A.2 Proof of Theorem 2

Below we recall an auxiliary thereotical result which will be used in our proof of Theorem 2.

###### Theorem 3 (Connection between solutions of dual OT and UOT problems (Choi et al., 2024a ) ) .

Assume that $\overline{\psi}$ is a continuously differentiable function. Then if the optimal potential $f^{*}$ delivering maximum to the dual UOT problem exists, $f^{*}$ is a solution of the following objective

$$ $\displaystyle\text{OT}_{c}(\widetilde{\mathbb{P}},\widetilde{\mathbb{Q}})=\sup _{f}\int_{\mathcal{X}}f^{c}(x)d\widetilde{\mathbb{P}}(x)+\int_{\mathcal{Y}}f(y )d\widetilde{\mathbb{Q}}(y)$ $$

where $\widetilde{\mathbb{P}}(x)=\nabla\overline{\psi}(-(f^{*})^{c}(x))\mathbb{P}(x)$ and $\widetilde{\mathbb{Q}}(y)=\nabla\overline{\psi_{2}}(-f^{*}(y))\mathbb{Q}(y)$.

This Theorem highlights the intriguing connection between the solutions of the classic OT problem and its unbalanced counterpart. Note that in the case of SUOT problem, $\overline{\psi_{2}}(y)=y$, thus, the measure $\widetilde{\mathbb{Q}}$ reduces to $\mathbb{Q}$. Now we are ready to prove our Theorem 2.

###### Proof of Theorem 2 .

Below we show that optimal potentials $f^{*}_{[1:K]}$ delivering maximum to the dual SUOT barycenter problem (8) are also optimal for each of the underlying $\text{SUOT}(\mathbb{P}_{k},\mathbb{Q})$ problems. Recall that the SUOT barycenter problem (7) admits a minimizer which we denote as $\mathbb{Q}^{*}$. Then:

$$ $\displaystyle\sum_{k=1}^{K}\lambda_{k}\text{SUOT}_{c_{k},\psi_{k}}(\mathbb{P}_ {k},\mathbb{Q}^{*})\stackrel{{\scriptstyle\eqref{eq:dual-suot-bary-before}}}{{ =}}\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]}\in\mathcal{C}^{K}( \mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \big{[}\int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c_{k}}(x_{k}))d \mathbb{P}_{k}(x_{k})+m\big{]}=$ (22) $\displaystyle\sup_{\begin{subarray}{c}m\in\mathbb{R},f_{[1:K]}\in\mathcal{C}^{ K}(\mathcal{Y})\\ \sum_{k=1}^{K}\lambda_{k}f_{k}\equiv m\end{subarray}}\sum_{k=1}^{K}\lambda_{k} \int_{\mathcal{X}_{k}}-\overline{\psi}_{k}(-f_{k}^{c_{k}}(x_{k}))d\mathbb{P}_{ k}(x_{k})+\int_{\mathcal{Y}}\underbrace{\sum_{k=1}^{K}\lambda_{k}f_{k}(y)}_{=m }d\mathbb{Q}^{*}(y)\leq$ (23) $\displaystyle\sup_{f_{[1:K]}\in\mathcal{C}^{K}(\mathcal{Y})}\!\sum_{k=1}^{K} \lambda_{k}\int_{\mathcal{X}_{k}}\!\!\!-\overline{\psi}_{k}(-f_{k}^{c_{k}}(x_{ k}))d\mathbb{P}_{k}(x_{k})\!+\!\!\!\int_{\mathcal{Y}}\sum_{k=1}^{K}\lambda_{k} f_{k}(y)d\mathbb{Q}^{*}(y)\!=\!\!\sum_{k=1}^{K}\lambda_{k}\text{SUOT}(\mathbb{ P}_{k},\mathbb{Q}^{*}).$ (24) $$

Here the transition from line (22) to (23) follows from the fact that the optimization is performed over $m$-congruence potentials. If we then remove this restriction on potentials, the value of $\sup$ can become bigger which explains the transition between (23)-(24). Importantly, we get that the final objective in line (24) is equal to the initial objective in line (22). Thus, the inequality in line (23) turns into equality. At the same time, it justifies that the potentials $f^{*}_{[1:K]}$ delivering maximum to the SUOT barycenter (8) problem coincide with the optimal potentials for each of the SUOT problems.

Then, thanks to Theorem 3, the SUOT barycenter problem (24) admits a reformulation:

$$ $\displaystyle\mathcal{L}^{*}=\sum_{k=1}^{K}\lambda_{k}\text{SUOT}(\mathbb{P}_{ k},\mathbb{Q}^{*})=\sum_{k=1}^{K}\lambda_{k}\text{OT}(\widetilde{\mathbb{P}}_{ k},\mathbb{Q}^{*})$ (25) $$

where each of the distributions $\widetilde{\mathbb{P}}_{k}$ ($k\in\overline{K}$) is specified via the optimal potential $f^{*}_{k}$ delivering maximum to the problem $\text{OT}(\widetilde{\mathbb{P}}_{k},\mathbb{Q}^{*})$ in its dual form (2): $d\widetilde{\mathbb{P}}_{k}(x_{k})=\nabla\overline{\psi}(-(f^{*})^{c}(x_{k}))d
\mathbb{P}_{k}(x_{k})$. Thus, the optimal SUOT plans now correspond to the optimal OT plans $\gamma^{*}_{k}(x,y)$ belonging to $\Pi(\widetilde{\mathbb{P}}_{k},\mathbb{Q}^{*})$, i.e., having marginals $(\gamma^{*}_{k})_{x}=\widetilde{\mathbb{P}}_{k}$, $(\gamma^{*}_{k})_{y}=\mathbb{Q}^{*}$.

Now we fix $k\in\overline{K}$. We aim to show that ${\gamma_{k}^{*}(\cdot|x_{k})\in\arg\inf_{\gamma_{k}(\cdot|x_{k}))}\int_{
\mathcal{Y}}(c_{k}(x_{k},y)-f^{*}_{k}(y))d\gamma(\cdot|x_{k})}$ where $\inf$ is taken over measurable maps $\gamma_{k}(\cdot|x_{k})$. Then by repeating the derivations of Corollary 1 we easily get that ${\gamma^{*}_{k}(\cdot|x_{k})\in\arg\inf_{\gamma_{k}(\cdot|x_{k}))}\mathcal{L}(
f_{k}^{*},\gamma_{k}(\cdot|x_{k})))}$. To prove the former, we again use the fact that $\gamma_{k}^{*}$ is a solution of the balanced OT problem between the rescaled marginal $\widetilde{\mathbb{P}}_{k}$ and $\mathbb{Q}^{*}$. Note that $\nabla\overline{\psi}$ is a continuous function since $\overline{\psi}$ is assumed to be a continuously differentiable function. This yields the fact that $\widetilde{\mathbb{P}}_{k}$ is an absolutely continuous distribution. Thus, we can use (Villani et al., , Remark 5.13) which states that $\gamma_{k}^{*}(\cdot|x_{k})\in\arg\inf_{y\in\mathcal{Y}}{\big{(}c(x,y)-f(y)
\big{)}}$. The latter statement is equivalent to $\gamma_{k}^{*}(\cdot|x_{k})\in\int_{\mathcal{Y}}(c_{k}(x_{k},y)-f^{*}_{k}(y))d
\gamma(\cdot|x_{k})$ which completes the proof.
∎

**Table 2: Hyper-parameter settings of Algorithm 1 for various experiments.**
| Experiment | 2D Toy (§5.1) | Class Imbalance (§5.2) | Outlier (§5.2) | Shape-Color (§5.3) |
| --- | --- | --- | --- | --- |
| D | 2 | 2 | 2 | 512 |
| K | 2 | 2 | 3 | 2 |
| $\tau$ | 5 | 1,20,200 | 1,20,200 | 10 |
| batch size | 1024 | 256 | 256 | 64 |
| Epochs | 10K | 10K | 20K | 3K |
| $N_{T}$ | 3 | 3 | 5 | 10 |
| $\lambda_{1}$ | 1/2 | 1/2 | 1/3 | 1/2 |
| $\lambda_{2}$ | 1/2 | 1/2 | 1/3 | 1/3 |
| $\lambda_{3}$ | - | - | 1/3 | - |
| Divergence $\psi_{1}$ | balanced | KL | KL | Softplus |
| Divergence $\psi_{2}$ | KL | KL | KL | Softplus |
| Divergence $\psi_{3}$ | - | - | KL | - |
| $f_{k,\theta}$ | MLP | MLP | MLP | ResNet |
| $T_{k,\phi}$ | MLP | MLP | MLP | ResNet |
| $lr_{f_{k,\theta}}$ | 1e-3 | 1e-3 | 1e-3 | 1e-4 |
| $lr_{T_{k,\phi}}$ | 1e-3 | 1e-3 | 1e-3 | 1e-7 |
| $lr_{m}$ | 1e-3 | 1e-3 | 1e-3 | 1e-2 |

###### Theorem 3 (Connection between solutions of dual OT and UOT problems (Choi et al., 2024a ) ) .

###### Proof of Theorem 2 .

## Appendix B Implementation Details

### B.1 Description of Toy Experiments

#### Moon, Spiral, and 8-Gaussian Datasets

The datasets used in Section 5.1 were implemented by following (Choi et al., ).

Datasets for Class Imbalance Experiments.
For the distribution $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$, we employ the Gaussian mixture of $\frac{1}{4}\mathcal{N}((-5,4),0.4^{2})+\frac{3}{4}\mathcal{N}((-5,-4),0.4^{2})$ and $\frac{3}{4}\mathcal{N}((5,4),0.4^{2})+\frac{1}{4}\mathcal{N}((5,-4),0.4^{2})$, respectively.

Datasets for Outlier Experiments.
We consider three marginal distributions, denoted as $\mathbb{P}_{1}$, $\mathbb{P}_{2}$, and $\mathbb{P}_{3}$. For $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$, we generate datasets consisting of 95% in-distribution data and 5% outliers. In contrast, $\mathbb{P}_{3}$ consists solely of in-distribution data. The in-distribution data for each marginal follows a Gaussian mixture model with four modes and uniform weights, i.e., $\sum_{i=1}^{4}\frac{1}{4}\mathcal{N}(m_{i},\sigma^{2}I)$. The means and variances for $\mathbb{P}_{1}$, $\mathbb{P}_{2}$, and $\mathbb{P}_{3}$ are defined as follows:

- •
$\mathbb{P}_{1}$: $m_{1}=(-5,-1)$, $m_{2}=(5,1)$, $m_{3}=(1,-5)$, $m_{4}=(-1,5)$, with $\sigma=0.1$,
- •
$\mathbb{P}_{2}$: $m_{1}=(-5,1)$, $m_{2}=(5,-1)$, $m_{3}=(1,5)$, $m_{4}=(-1,-5)$, with $\sigma=0.1$,
- •
$\mathbb{P}_{3}$: $m_{1}=(-5,0)$, $m_{2}=(5,0)$, $m_{3}=(0,5)$, $m_{4}=(0,-5)$, with $\sigma=0.1$,

respectively.
For the outlier distributions of $\mathbb{P}_{1}$ and $\mathbb{P}_{2}$, we again use a Gaussian mixture model with four modes, $\sum_{i=1}^{4}\frac{1}{4}\mathcal{N}(m_{i},\sigma^{2}I)$, with the means and variances defined as:

- •
$\mathbb{P}_{1}$ outliers: $m_{1}=(10,2)$, $m_{2}=(10,1)$, $m_{3}=(10,0)$, $m_{4}=(10,-1)$, with $\sigma=0.02$,
- •
$\mathbb{P}_{2}$ outliers: $m_{1}=(-10,1)$, $m_{2}=(-10,0)$, $m_{3}=(10,-1)$, $m_{4}=(10,-2)$, with $\sigma=0.02$,

respectively.

Implementation Details.
For the experiments in Section 5.1 and the class imbalanced experiments, we use the number of iterations of 10K.
For the outlier experiments in Section 5.2, we employ the number of iterations of 20K.
We use the batch size of 1024.
We update $m$ with the Adam Optimizer with learning rate of $10^{-3}$ and $(\beta_{1},\beta_{2})=(0,0.9)$.
For other hyperparameters, we follow the experimental settings of (Kolesov et al., ).

Evaluation metric. As discussed in Section 5.1, we utilize the W2 and L2 metrics, defined as $W^{2}_{2}({T_{1}}_{\#},\mathbb{Q}^{\star})$ and $\lVert{T_{1}}-T^{\star}_{1}\rVert^{2}_{2}$, respectively.
Here, $\mathbb{Q}^{\star}$ and $T^{\star}_{1}$ represent the true barycenter and transport maps.
To approximate $(T^{\star}_{1},\mathbb{Q}^{\star})$, we solve the discrete Unbalanced Barycenter problem.
Specifically, we sample 2000 points and solve the Unbalanced Optimal Transport (UOT) problem from $\mathbb{P}_{1}$ to $\mathbb{P}_{2}$ using the Python Optimal Transport (POT) library (Flamary et al., ).
The obtained discrete UOT map $T$ is then used to define $T_{1}=\lambda_{1}\text{Id}+\lambda_{2}T$, and $\mathbb{Q}^{\star}={T_{1}}_{\#}\mathbb{P}_{1}$.
Using these discrete samples and transport maps, we compute the approximation of the metrics.

### B.2 StyleGAN Experiments

Unless otherwise stated, we follow the implementation of Kolesov et al. ().
Note that we use Adam Optimizers with $(\beta_{1},\beta_{2})=(0,0.9)$.
The exact hyperparameter is described in Table 2.

## Peer Reviews

confidence: 3
contribution: 1
initial_rating: 5
presentation: 2
questions: Q1. How does the time efficiency of U-NOTB compare with other baselines?

Q2. The experiments seem to lack details about the size of the training/testing sets and some training specifics, such as learning rate and other hyperparameters.

Q3. Since the overall training is similar to GANs with a max-min approach, is the training of U-NOTB stable?
soundness: 2
strengths: 1. The proposed method may be the first continuous robust barycenter estimation approach with proper theoretical support and practical validation.

2. The proposed method is robust and can estimate the barycenter based on the data containing the outliers.
summary: This paper proposes a neural network-based method to estimate continuous barycenter via the dual formulation of the semi-unbalanced OT.
weaknesses: My main concern is the novelty of the paper. The primary method seems to be a combination of barycenter calculation, neural optimal transport, and semi-unbalanced optimal transport. In my view, simply combining methods may not be sufficient for publication at ICLR, so I give a negative score. I will consider raising my score if the authors can demonstrate unique contributions that go beyond a straightforward combination of these methods.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Given that the Wasserstein barycenter of Gaussian distributions has a closed-form solution, could the authors provide experimental verification for the Gaussian distribution case?
- Can the robustness parameter $\tau$ be automatically optimized during training rather than manually tuned? If not, for different proportions or types of noise, different $tau$ is usually required. Will this limit the practicality of the scheme?
- How does the method work with high-dimensional data? In other words, the support size of the measures may be large in practical applications. While the support size of the measures in your experiments seems rather small.
- Why don't you compare the methods [1,2] based on robust OT? Please add corresponding experiments.

[1] Nietert S, Goldfeld Z, Cummings R. Outlier-robust optimal transport: Duality, structure, and statistical analysis[C]//International Conference on Artificial Intelligence and Statistics. PMLR, 2022: 11691-11719.

[2] Wang X, Huang J, Yang Q, et al. On Robust Wasserstein Barycenter: The Model and Algorithm[C]//Proceedings of the 2024 SIAM International Conference on Data Mining (SDM). Society for Industrial and Applied Mathematics, 2024: 235-243.
soundness: 3
strengths: - This paper introduces a continuous SUOT barycenter estimation method that addresses the issue of robustness against outliers and imbalances in real-world datasets.
- Rigorous derivation of the SUOT-based framework and solid theoretical support for the proposed model give the approach a solid foundation.
summary: This paper proposes a scalable approach for estimating the robust continuous barycenter by using the dual formulation of the (semi-)unbalanced OT problem. They first attempt to develop an algorithm for robust barycenters under continuous distribution setup. They model this problem as a min-max optimization problem and provide theoretical underpinnings, along with experimental results on synthetic and real-world datasets to demonstrate the robustness and adaptability of their method.
weaknesses: Experiments are not sufficient and the baseline method is not comprehensive.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: See Weaknesses section.
soundness: 3
strengths: 1. The paper provides an important and novel dual theory for the SUOT barycenter problem.

2. The proposed algorthm works well and persistent to the imbalance and outliers. 

3. The paper is well-structured and easy to follow.
summary: In this paper, the authors investigate the problem of finding the barycenter with semi-unbalance optimal cost (SUOT) from perspective of dual theory. In particular:

(1) They derive the theory for the dual form of SUOT barycenter with reformulation, and necessary condition based on marginal plan and conditional plan

(2) Then, from the necessary condition,  they propose the an algorithm (Algorithm 1) to approximate the optimal transport plan from the solution. This algorithm is based on training the neural network of optimal plan.  

(3) Then, they empirically validate the efficacy of the algorithm on both synthetic and real dataset.
weaknesses: 1. The necessary condition seems to be not difficult to obtain, which makes it not convincing enough. So, would it be possible to obtain some sufficient conditions? 

2. It would be better if the authors can highlight the advantages of finding the optimal plan over the optimal solution. 

3. As stated in the limitations, the authors did not provide a rigorous analysis for the convergence of the algorithm.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: A term seems to be missing in the definition of $\psi$-divergence, see Definition 1 in [2].

The sentence line 227 is not clear to me. The $m$-congruence is a constraint of the problem. So I do not understand why it is written "if the potentials satisfy the m-congruence, then the optimal value of the SUOT barycenter can be derived by solving (8)". The SUOT barycenter can always be derived by solving (8) and the potentials necessarily satisfy this constraint when solving (8).

In Corollary 1, it is stated that the sup is taken over $f_{[1,K]}$. Isn't it also taken over $m$?

In Theorem 2, equation (12), shouldn't it be an argmin?

In Section 5.1, it is stated that the UOT problem is equivalent with the OT problem between rescaled distribution. Is it truly equivalent? And could we solve the OT barycenter between the rescaled distributions instead of solving the SUOT barycenter?

Is it proved that $T_{1\\#}\mathbb{P}_1$ gives the UOT barycenter?

I think that the reference [3] is missing. In their experiments, they compute robust OT barycenters via unbalancedness.

Typos:
- Line 11: The first sentence of the abstract feels weird: I don't see what is the "common challenge"
- Line 107: "This transition is valid since we work the infimum in weak"
- Line 142: "wights"
- Line 365: what is $T$ in $T_1 = \lambda_1 Id + \lambda_2 T$
- Legend of Figure 4: $\mathbb{P}_0$ -> $\mathbb{P}_1$ and $\mathbb{P}_1$ -> $\mathbb{P}_2$

[1] Kolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., & Korotin, A. (2024). Estimating Barycenters of Distributions with Neural Optimal Transport. arXiv preprint arXiv:2402.03828.

[2] Séjourné, T., Peyré, G., & Vialard, F. X. (2023). Unbalanced optimal transport, from theory to numerics. Handbook of Numerical Analysis, 24, 407-471.

[3] Séjourné, T., Bonet, C., Fatras, K., Nadjahi, K., & Courty, N. (2023). Unbalanced optimal transport meets sliced-Wasserstein. arXiv preprint arXiv:2306.07176.
soundness: 4
strengths: The paper is overall well written and nice to read.

It provides one of the first method to solve continuous UOT barycenters leveraging a dual formulation and neural networks parametrizations, which is an interesting contribution.

Several convincing experiments are done demonstrating that the method learns well the UOT barycenter, that it is robust to outliers and class imbalance, and that it works with several costs.
summary: In this work, authors propose to solve the continuous unbalanced optimal transport barycenter. To do that, they derive a dual formulation to the problem, which is a min-max problem over potentials and conditional distributions. They solve the problem by parametrizing the potential and conditional distributions with neural networks. Finally, they show that their method works on three experiments, demonstrating that they recover the right barycenters, that the barycenter is robust to outliers and to class imbalance, and that it can handle using different costs.
weaknesses: The main weakness in my opinion is that the method feels really incremental compared to [1]. The main difference is that it it adapted to the UOT problem, which just changes slightly the formulation of the dual, and how to sample from the barycenter for inference.

Another weakness, which is classical with UOT, is that the choice of the unbalancedness parameters does not seem easy.

The method also needs to solve a min-max problem, which is probably unstable and very costly.

Meta-review: In the paper, the authors proposed a new approach to estimate the (semi)-unbalanced barycenter of continuous distributions. It is done via formulated the problem as a min-max optimization problem and the method can be used for any general cost function. All the reviewers agree that the proposed method is novel and the theoretical results are thorough. During the rebuttal period, the authors also included additional theoretical results on the quality bounds on the recovered plans based on the duality gaps, which strengthens the theoretical guarantee for the Algorithm 1 in the paper. 

While there are some concerns about the experiments that are not extensive as well as the method may be slightly incremental compared to the previous works, I believe that the work has sufficient novelty and merit for ICLR. Therefore, I recommend accepting the paper. 

The authors are encouraged to incorporate the suggestions of the reviewers into the revision of their manuscript.

Decision: Accept (Poster)