Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport

Authors: Milena Gazdieva, Jaemoo Choi, Alexander Kolesov, Jaewoong Choi, Petr Mokrov, Alexander Korotin
Primary area: generative models
Keywords: unbalanced optimal transport, barycenter, generative modeling

Abstract:
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, in real-world scenarios, 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$-$\max$ optimization problem and is adaptable to *general* cost functions. 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. Our source code is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters.

Decision: accept (poster)
Reviewer scores: 6, 6, 5, 8

Reviews and discussion:
<thread>
creation date:
1730020681800

thread_type:
official_review

<thread_object>
creator: reviewer_niqu

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730353523863

thread_type:
official_review

<thread_object>
creator: reviewer_hyxj

<content>
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.

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.

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.

questions:
See Weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730608819166

thread_type:
official_review

<thread_object>
creator: reviewer_mxnh

<content>
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.

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.

weaknesses:
Experiments are not sufficient and the baseline method is not comprehensive.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706550299

thread_type:
official_review

<thread_object>
creator: reviewer_nqfu

<content>
summary:
This paper proposes a neural network-based method to estimate continuous barycenter via the dual formulation of the semi-unbalanced OT.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732312557397

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
Dear reviewers, 

thank you for your thoughtful reviews! We appreciate that you positively highlight our theoretical insights (Reviewers nQfu, MXHN), proper practical validation (Reviewers nQfu, niQu) and overall quality of the text (Reviewer hyxj).

We have uploaded an updated version of the paper. The newly added content is highlighted with the **blue** color. **The changes include**:

- [nQfu] New Table with training/inference times of our solver and baselines in **Appendix B.1**;
- [MXNH] New **Appendicies C.1/C.2** showing the performance of our solver and baselines in OT/SUOT barycenter problem for Gaussian distributions with known ground-truth solutions. The experiments reveal the poor performance of classic discrete method (Cuturi & Doucet, 2014) in continuous OT barycenter problem;
- [MXHN] New **Appendix C.3** with high-dimensional experiment demonstrating the practical advantages of our solver;
- [MXHN, niQu] Added references to papers (Nietert et al., 2022), (Wang et al., 2024), (Séjourné et al., 2023), see **Section 3** and **Section 6**;
- [niQu] Fixed the unclear aspects/typos in the **main text**.

**References.**

Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International
conference on machine learning, pp. 685–693. PMLR, 2014

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.

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

Séjourné, T., Bonet, C., Fatras, K., Nadjahi, K., \& Courty, N. (2023). Unbalanced optimal transport meets sliced-Wasserstein. arXiv preprint arXiv:2306.07176.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732825075394

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
New theoretical result

comment:
Dear Reviewers,

as per the request of the Reviewer hyxj, we have prepared an additional revision of our paper which includes a **new theoretical result** (Theorem 3 in Section 4) establishing the **quality bounds** on the recovered plans based on the duality gaps, i.e., the errors for solving inner and outer optimization problems in our objective (9). Theorem 3 shows that when our Algorithm 1 optimizing this $\max$-$\min$ objective converged nearly to the optimum, its solutions are close to the true conditional plans.

Best regards,
the Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.