Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport

(9) Is it proved that $T\_1\sharp\mathbb{P}\_1$ gives the UOT barycenter?

As shown in Theorem 2, the optimal plan $T^\star\_1$ lies in the saddle point solution of our max-min optimization problem (equation 12). Moreover, as shown in equation 11, \\mathbb{Q} = T^{\\star}_{1}\\# \widetilde{\\mathbb{P}}_1, where $\widetilde{\mathbb{P}}\_1 = \nabla \bar{\psi} (-(f^\star\_1)^c (x) ) \mathbb{P}\_1$. Thus, to obtain a point of the barycenter, we first sample $x\sim \widetilde{\mathbb{P}}$ by rejection sampling (line 327-338). Then, we pass $x$ through our learned transport map $T_1$. Could you please reply, does this answer your question?

(10) I think that the reference (Séjourné et al., 2023a) is missing. In their experiments, they compute robust OT barycenters via unbalancedness.

Thank you for suggesting the relevant paper. We added the reference in Related Work section of the revised paper. Based on our understanding of the paper, we attribute it as the discrete method.

(11) Typos.

Thank you for your carefullness, we have fixed all the typos which you have mentioned.

Concluding remarks. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

References.

J. Choi, J. Choi, and M. Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. In Advances in Neural Information Processing Systems, volume 36, 2024.

M. Gazdieva, A. Asadulaev, E. Burnaev, and A. Korotin. Light Unbalanced Optimal Transport. In Advances in Neural Information Processing Systems, volume 36, 2024.

Kolesov et. al., Estimating Barycenters of Distributions with Neural Optimal Transport, Proceedings of the 41st International Conference on Machine Learning, PMLR 235:25016-25041, 2024.

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

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

K. D. Yang and C. Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2018.

Korotin et al. Wasserstein iterative networks for barycenter estimation. Advances in Neural Information Processing Systems 35, 15672-15686, 2022.

Kolesov et. al., Energy-Guided Continuous Entropic Barycenter Estimation for General Costs, NeurIPS, 2024.

Thank you for your detailed feedback. Please find the answers to your questions below.

(1) The main weakness in my opinion is that the method feels really incremental compared to (Kolesov et al., 2024, Estimating). The main difference is that it is adapted to the UOT problem, which just changes slightly the formulation of the dual, and how to sample from the barycenter for inference.

We would like to clarify the novelty of our paper. It would be most correct to position our method as a generalization of the recent SOTA approach to estimating continuous balanced OT barycenters - NOTB (Kolesov et al., 2024, Estimating), to the case of semi-unbalanced OT barycenters. However, this generalization is not straightforward as it is based on different principles.

First, the congruence condition on the potentials, which is used in the NOTB approach, ceases to be true in the case of semi-unbalanced OT. In our paper, we developed a completely different condition on the potentials, which required the usage of an additional parameter $m$ in the optimization algorithm.

Second, the procedure of sampling from the barycenter during the inference has major differences from the one used in NOTB. In the semi-unbalanced case, to get samples from the barycenter, we first need to define the procedure of sampling from the left marginals of the plans. Defining such a sampling procedure is tricky even for the continuous unbalanced OT (UOT) solver by its own. As far as we know, only one of the existing continuous UOT solvers (Gazdieva et al., 2024) propose such a strategy using the restricted Gaussian mixture parametrization of potentials. Other scalable and theoretically justified continuous UOT solvers (Choi et al., 2024, Yang et al., 2018) do not offer such a precise sampling procedure and consider only sampling from the input measure. In contrast to previous approaches, we suggest a precise sampling procedure that utilizes equation (11) from our paper. Using the learned potentials, our solver allows for precise sampling of barycenter from the left marginals of learned plans without any restrictions on the underlying potentials. In the case of previous works in the field such as NOTB, the situation is much easier since these marginals coincide with input measures and sampling is defined a priori.

Third, during the rebuttal period, we conducted a novel high-dimensional experiment to further demonstrate the practical advantages of our method. Specifically, as detailed in Appendix C.3, we performed interpolation between two image distributions within the FFHQ 256×256 dataset, focusing on transitioning from the young to elderly categories. This experiment highlights the potential of our approach to manipulate images by learning barycenters between distinct image distributions, enabling controlled transitions across semantic attributes. Moreover, our model also demonstrates robustness to class imbalancedness in this practical task compared to other baselines.

More detailed explanation regarding this experiment is given in Appendix C.3 in the revised version of our paper.

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

Indeed, the choice of unbalancedness parameter is a tricky point of any method related to UOT problem. In our paper, we followed recent works in Unbalanced Optimal Transport (UOT), such as (J. Choi et al., 2023, J. Choi et al., 2024), and adopted a manual selection of the unbalancedness parameter $\tau$. Still, we believe that this flexibility in choosing $\tau$ can be interpreted as an advantage because $\tau$ can be tailored to suit the specific properties of the distributions, e.g., control the amount of samples treated as outliers. However, we also believe that developing a method to automatically determine $\tau$ based on the proportion of outliers or class imbalance would be a highly interesting direction for future research.

Yes, the method is specifically developed for probability distributions. We believe that extending our method to arbitrary positive distributions represents a promising direction for future work.

Moreover, to establish the connection between OT and SUOT barycenters in the case of probability measures we employ the connection between OT and UOT problems established in (Choi et al., 2024). Since the latter result holds true for the case of probability measures, the connection between OT and SUOT barycenters in the case of arbitrary positive measures also needs further investigations.

We are delighted to hear that your concerns have been addressed. Once again, we sincerely appreciate the time and effort you have put into reviewing our paper.

Thank you for your detailed feedback. Please find the answers to your questions below.

(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?

We understood the term "sufficient condition" as follows: for every optimal saddle point $\{(f\_k^*, \gamma\_k^*)\}\_{k=1}^K$ of our optimization objective (9), it holds that $\{\gamma\_k^*\}\_{k=1}^K$ is the family of true SUOT plans between $\mathbb{P}\_{[1:K]}$ and $\mathbb{Q}$.

Actually, the latter seems to be true under some additional assumptions related to the strong convexity of $c_k(x_k,y)-f_{k}(y)$, ($k\in[1,K]$). These assumptions were already established and studied in the previous works on balanced OT, see (Fan et al., 2021) or (Makkuva et al., 2020), and balanced OT barycenter problem, see (Kolesov et al., 2024, Estimating). In principle, we think that one can make the same kind of assumptions and presumably derive analogous bounds for our semi-unbalanced barycenter setup. However, these assumptions are far from practice since they usually require the usage of restricted class of neural network architectures, i,e., Input Convex Neural Networks (ICNNs). At the same time, the methods exploiting this type of networks are known to provide quite fair results in the tasks of generative modeling, see, e.g., Fig. 4 of (Korotin et al., 2021) for visualization of their performance.

Another way for ensuring the "sufficient condition" consists in considering more general OT formulations, e.g., using different regularizations such as entropic (Gushchin et al., 2024) or kernel (Korotin et al., 2023). We believe that our ideas for basic OT presented in the paper can be generalized to such cases, but leave the investigation of this aspect for future work.

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

We interpreted this question as asking: Why do we focus on finding an optimal plan instead of directly computing the barycenter distribution? What are the practical advantages of this approach?

There exist a lot of practical tasks which require finding a shared representation of the data which comes from different sources. This task can be positioned as the problem of finding the barycenter of distributions. Then, in order to translate new data from each of the input distributions to the shared representation (barycenter), the practitioner needs to have access to the corresponding translation maps (conditional plans). It explains the importance of finding the optimal plans and not only barycenters by their own. More details on the specific applications where such kind of tasks appear, e.g., finding shared representatons for scans from different MRI scanners or mixing geological simulators, can be found in (Appendix B.2, Kolesov et al., 2024).

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

Deriving theoretical results regarding the convergence of the obtained algorithm for estimating continuous unbalanced OT barycenters is a difficult and even might be not a solvable task under the reasonable conditions. Indeed, even existing results of this kind for the UOT problem itself seem to require restrictive assumptions which are not feasible in practice. For example, (Choi et al., 2023) in their Theorem 3.4 assume strong convexity of the potentials, which is not applicable in practice, as we explained in the first answer to you.

Concluding remarks. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

Thank you for your detailed feedback. Please find the answers to your questions below.

(1) Given that the Wasserstein barycenter of Gaussian distributions has a closed-form solution, could the authors provide experimental verification for the Gaussian distribution case?

Thank you for this valuable suggestion. In the revised version of our paper, we test our solver in the balanced/semi-unbalanced OT barycenter problem for Gaussian distributions with computable ground-truth solutions.

In the balanced OT barycenter problem for Gaussian distributions, the ground-truth barycenter is known to be Gaussian and can be estimated using the fixed point iteration procedure (Álvarez-Esteban et al., 2016). We tested our solver in this balanced setup keeping in mind that for large unbalancedness parameters $\tau$, it should provide a good approximation of balanced OT barycenter problem solutions. Ultimately, we show that while our solver is designed to tackle an unbalanced OT barycenter problem, its performance in the different balanced OT barycenter problem is comparable to the current SOTA solvers.

A recent preprint (Nguyen et al., 2024) provides an iteration procedure for calculating unbalanced SUOT barycenter for Gaussian distributions using the quadratic cost and KL divergences. We test our solver in this unbalanced setup for different parameters $\tau$ and show that it consistently outperforms the SOTA balanced solver (Kolesov et al., 2024a).

We include these new requested experiments in Appendices C.1, C.2, of the revised version of our paper.

(2) 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?

We appreciate the reviewer for highlighting an important aspect of our approach. Following recent works in Unbalanced Optimal Transport (UOT), such as (J. Choi et al., 2023, J. Choi et al., 2024), our model adopts a manual selection of the robustness parameter $\tau$. Meanwhile, we believe that this flexibility in choosing $\tau$ can be interpreted as an advantage because $\tau$ can be tailored to suit the specific properties of the distributions, e.g., control the amount of samples treated as outliers. However, we also believe that developing a method to automatically determine $\tau$ based on the proportion of outliers or class imbalance would be a highly interesting direction for future research.

(3) 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.

We are not quite sure what do you exactly mean by "support" in this question because this notion is usually treated differently by researchers from discrete and continuous OT field. In the case of discrete OT, "support" usually corresponds to the size of the dataset. In the case of continuous OT, it is treated as an ambient or intrinsic data dimension. Thus, we provide the answers for both sides of this question.

Dataset size. The sizes of the datasets do not matter for our algorithm. It uses the stochastic optimization and can handle arbitrarily large datasets.

Data space dimension. Searching for the barycenters directly in the data space may not be very meaningful because, as a result of averaging, some practically meaningless objects may appear. For example, averaging images of '0' and '1' with $\ell_2$ cost directly in the image space will result in straightforward $\ell_2$ interpolation of these images. Thus, dealing with a barycenter problem, it is necessary to restrict the space where we search for this barycenter. For example, this space can be specified using the pretrained generative models, e.g., StyleGAN (Karras et al., 2019), as it was done in the previous related papers (Kolesov et al., 2024a,b) and our experiment with MNIST dataset, see our Section 5.3. Here the data space dimension usually means an intrinsic dimension, i.e., the dimension of the manifold where this data lies. In the context of generative models, this intrinsic dimension is usually treated as a dimension of the model's latent space. Our experiments in Section 5.3 are conducted in standard StyleGAN latent space of the dimension 512.

(See the next comment for continuation of this answer.)

(4) Since the overall training is similar to GANs with a max-min approach, is the training of U-NOTB stable?

Thanks for raising this point. Indeed, the instability of training is a well-known issue of adversarial optimization objectives. Our approach is not an exception and we noted this aspect in Discussion section of our paper, see lines 534-537. Still, we emphasize that the majority of the existing approaches for barycenter computation suffer from some kind of computational issues. Among the recent approaches: (Kolesov et. al., 2024, Estimating), (Korotin et. al., 2022) also resort to min-max optimization; (Kolesov et. al., 2024, Energy) utilizes Energy-based training with Langevin simulation which is time costly and may be unstable under improper hyperparameters setting (e.g., too large Langevin step size, too small number of Langevin steps, etc.). Among the other approaches: while there may be some that do not present serious computational challenges, this comes at the cost of limited applicability (e.g., reliance exclusively to $\ell_2$ OT) and scalability, see Table 1 in (Kolesov et. al., 2024, Energy).

Concluding remarks. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

References.

J. Choi, J. Choi, and M. Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. In Advances in Neural Information Processing Systems, volume 36, 2024.

M. Gazdieva, A. Asadulaev, E. Burnaev, and A. Korotin. Light Unbalanced Optimal Transport. In Advances in Neural Information Processing Systems, volume 36, 2024.

Kolesov et. al., Estimating Barycenters of Distributions with Neural Optimal Transport, Proceedings of the 41st International Conference on Machine Learning, PMLR 235:25016-25041, 2024.

Korotin et al. Wasserstein iterative networks for barycenter estimation. Advances in Neural Information Processing Systems 35, 15672-15686, 2022.

Kolesov et. al., Energy-Guided Continuous Entropic Barycenter Estimation for General Costs, NeurIPS, 2024.

K. D. Yang and C. Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2018.

Thank you for your detailed feedback. Please find the answers to your questions below.

(1) My main concern is the novelty of the paper. <...> if the authors can demonstrate unique contributions that go beyond a straightforward combination of these methods.

We would like to clarify the novelty of our paper. It would be most correct to position our method as a generalization of the recent SOTA approach to estimating continuous balanced OT barycenters - NOTB (Kolesov et al., 2024, Estimating), to the case of semi-unbalanced OT barycenters. However, this generalization is not straightforward as it is based on different principles.

First, the congruence condition on the potentials, which is used in the NOTB approach, ceases to be true in the case of semi-unbalanced OT. In our paper, we developed a completely different condition on the potentials, which required the usage of an additional parameter $m$ in the optimization algorithm.

Second, the procedure of sampling from the barycenter during the inference has major differences from the one used in NOTB. In the semi-unbalanced case, to get samples from the barycenter, we first need to define the procedure of sampling from the left marginals of the plans. Defining such a sampling procedure is tricky even for the continuous unbalanced OT (UOT) solver by its own. As far as we know, only one of the existing continuous UOT solvers (Gazdieva et al., 2024) propose such a strategy using the restricted Gaussian mixture parametrization of potentials. Other scalable and theoretically justified continuous UOT solvers (Choi et al., 2024, Yang et al., 2018) do not offer such a precise sampling procedure and consider only sampling from the input measure. In contrast to previous approaches, we suggest a precise sampling procedure that utilizes equation (11) from our paper. Using the learned potentials, our solver allows for precise sampling of barycenter from the left marginals of learned plans without any restrictions on the underlying potentials. In the case of previous works in the field such as NOTB, the situation is much easier since these marginals coincide with input measures and sampling is defined a priori.

Third, during the rebuttal period, we conducted a novel high-dimensional experiment to further demonstrate the practical advantages of our method. Specifically, as detailed in Appendix C.3, we performed interpolation between two image distributions within the FFHQ 256×256 dataset, focusing on transitioning from the young to elderly categories. This experiment highlights the potential of our approach to manipulate images by learning barycenters between distinct image distributions, enabling controlled transitions across semantic attributes. Moreover, our model also demonstrates robustness to class imbalancedness in this practical task compared to other baselines.

More detailed explanation regarding this experiment is given in Appendix C.3 in the revised version of our paper.

(2) How does the time efficiency of U-NOTB compare with other baselines?

In the revised manuscript, we reported the training/inference time of our method, see new Tables 3,4 in Appendix B.1. As shown in Table 3, the time efficiency of U-NOTB is comparable to SOTA balanced solver NOTB in terms of training time. In terms of inference time, U-NOTB is slower, taking 2 to 10 times longer than NOTB. This gap is due to the additional computational complexity introduced by rejection sampling, which requires calculating the $c$-transform of the potential function $\hat{f}_k$. This process involves a forward pass through the learned potential function. However, we would like to highlight that thanks to our proposed sampling method, our solver gains robustness to outliers and class imbalance. A detailed discussion of rejection sampling can be found in lines 333–338 of our manuscript.

(3) 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.

The implementation details for our solver are given in Appendix B. Specifically, it includes Table 2 with all types of hyperparameters which we used in the experiments. We refer to this Appendix section in the main text of our paper, see lines 346-347. However, we would be happy to include the additional details in a revised version of our paper if you can point out which information is missing. Could you please clarify which extra details should be included in our paper?