Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

The main Prop 3.1-1, which provides the justification for the proposed general framework, was proved in [1] (The proof is different).

➤ Thanks for pointing out the similarities to the concurrent work [Choi+'2023]. However, while we agree that these results appear to be related we respectfully disagree that Point 1 of our Prop 3.1 was proved in [Choi+'2023]. To highlight that point, we analyze their theory. Since [Choi+'2023, Theorem 3.4] presents a bound on the dual sub-optimality gap, hence dealing with aspects not covered in our paper, it seems that the reviewer mentions [Choi+'2023, Theorem 3.3]. For this reason, we focus on the latter Theorem. In the following, we want to lay out the differences between our and their results.

• [Choi+'2023]'s theory is developed under stronger assumptions than ours, which are given in the first paragraph of [Choi+'2023, Appendix A]. On the one hand, they assume the cost to be the squared Euclidean distance $c(\mathbf{x}, \mathbf{y}) = \tfrac{\tau}{2}\|\mathbf{x} - \mathbf{y}\|_2^2$ $(\tau > 0)$, while our results hold for any costs $c(\mathbf{x}, \mathbf{y}) = h(\mathbf{x} - \mathbf{y})$ with $h$ strictly convex. We hence recover the squared Euclidean distance with $h = \tfrac{\tau}{2} \|\cdot\|_2^2$. On the other hand, they assume that $\mu, \nu \in \mathcal{M}_1^+(\Omega)$ are probability measures of same total mass 1, while our results hold for any positive measures $\mu, \nu \in \mathcal{M}^+(\Omega)$ of arbitrary total mass, that might be different.

• We then clarify the correspondence between our and theirs notations. They denote $(v^{\star, c}, v^{\star})$ the optimal potentials solving the Unbalanced Kantorovich Dual Problem (UDKP), whereas we denote them by $(f^\star, g^\star)$. Furthermore, they denote $\Psi$ the entropy function associated to the divergence used to measure the deviation from the marginals, whereas we denote it $\phi$. Finally, in the same way as us, they note $\tilde{\mu} = \bar{\phi}(-f^\star)\cdot\mu$ and $\tilde{\nu} = \bar{\phi}(-g^\star)\cdot\mu$.

• We now highlight that point 1 of Prop 3.1 is not shown in [Choi+'2023]. Indeed, [Choi+'2023, Theorem 3.3] proves that if we let $T^\star$ be the Monge map between $\tilde{\mu}$ and $\tilde{\nu}$, which exists because of the assumptions on the costs and the fact that $\tilde{\mu} \ll \mu \ll \mathcal{L}_d$, then one has:

$$\forall \mathbf{x} \in \mathcal{X}, \textrm{ } T(\mathbf{x}) \in \arg \inf_{\mathbf{y} \in \mathcal{Y}} c(\mathbf{x}, \mathbf{y}) - g^\star(\mathbf{y})$$

However, they do not show that $\pi_\mathrm{UOT}$, the unbalanced UOT plan between $\mu$ and $\nu$, is supported on the graph of that map, i.e. $\pi_\mathrm{UOT} = (\mathrm{Id}, T^\star)\sharp\tilde{\mu}$ . Therefore, their statement do not show that $p_1 \sharp \pi_\mathrm{UOT} = \tilde{\mu}$ and $p_2 \sharp \pi_\mathrm{UOT} = \tilde{\nu}$, i.e. that $\pi_\mathrm{UOT}$ has marginals $\tilde{\mu}$ and $\tilde{\nu}$, which is the result stated in Point 1 of our Prop. 3.1. Moreover, in their proof, we don't understand how they obtain that $\tilde{\mu}$ and $\tilde{\nu}$ have the same total mass, which must be justified before considering the Monge map $T^\star$ between $\tilde{\mu}$ and $\tilde{\nu}$. In our proof, we obtain it by showing that $\tilde{\mu}$ and $\tilde{\nu}$ are the marginals of $\pi_\mathrm{UOT}$, so the equality of masses follows from Fubini's Theorem.

• Proposition 3.1 (Point 1 and 2) is in fact a more general version of their [Choi+'2023, Theorem 3.3]. Indeed, for all costs $c(\mathbf{x}, \mathbf{y}) = h(\mathbf{x} - \mathbf{y})$, we show that

$$\forall \mathbf{x} \in \mathcal{X}, \textrm{ } T^\star(\mathbf{x}) = \mathrm{Id} - \nabla h^\star \circ \nabla f^\star(\mathbf{x})$$

and that $\pi_\mathrm{UOT}$ is unique, and characterized by $\pi_\mathrm{UOT} = (\mathrm{Id}, T^\star)\sharp\tilde{\mu}$. This result then allow us to extend their result for $c(\mathbf{x}, \mathbf{y}) = \tfrac{\tau}{2} \|\mathbf{x} - \mathbf{y}\|_2^2$. As we mentioned in Remark. 3.2, in this case:

$$T^\star(x) = \mathbf{x} - \nabla f^\star(\mathbf{x}) = \nabla \varphi^\star(\mathbf{x})$$

where $\varphi^\star : \mathbf{x} \mapsto \|\mathbf{x}\|_2^2 - f^\star(\mathbf{x})$ is convex and unique. We then obtain an unbalanced generalization of [Brenier'1987]'s Theorem, stating that as in the balanced OT case, when the source measure is a density, the unbalanced OT plan is unique supported on the graph of a convex function.

We would be happy to answer any questions that might arise and clarify if some arguments were not clear.

Based on the reweighting solution mentioned above, can the imbalanced mass problem be addressed by detecting the degree of shifting mass (i.e., estimating the weights in Sec. 3.1) and solving the reweighting (balanced) OT? If it is feasible, what will be the advantages and weaknesses of the proposed methods?

➤ As discussed in more detail above, there are three limitations to this approach: First, there might not be sufficient information to rescale the marginals (e.g. due to poor or non-existent labels). Second, unbalancedness also compensates for imbalances within one batch, which is also likely to happen in class balance on the full dataset. Third, outliers cannot be automatically discarded in this way. An advantage would be the independence of the unbalancedness parameters, which, as discussed, is not necessarily straightforward.

As the basic problem is similar to label shift, the experiments on related problems are highly appreciated. For example, the generalized label shift scenario [r1] and partial domain adaptation [r2] (which can be taken as an extreme scenario for label shift), where the comparison methods for corresponding tasks should also be carefully considered.

➤ As mentioned above, (unbalanced) neural Monge maps in general cannot be applied to the domain adaptation task straight-forwardly.

Point 2 in Prop. 3.1, the definition of $\mathcal{L}_d$ is not provided.

➤ We define $\mu \ll \mathcal{L}_d$ at the start of Section 2 as $\mu$ absolutely continuous w.r.t. the Lebesgue measure.

---

[Srivatsan+'2020] Srivatsan, Sanjay R., et al. "Massively multiplex chemical transcriptomics at single-cell resolution." Science 367.6473 (2020): 45-51.

[Fartras+'2021] Kilian Fatras, Thibault S'ejourn'e, Nicolas Courty, and R'emi Flamary. "Unbalanced minibatch optimal transport; applications to domain adaptation" (2021).

Some related concepts and fields are omitted, which should also be discussed appropriately.

➤ As far as we understand, this comment is mainly about the field of domain adaptation. We first give definitions domain of adaptation and translation, which the remaining response builds upon, and would like to ask the reviewer for confirmation:

• Domain adaptation is concerned with learning domain-invariant representations or transformations to train robust classification models that can generalize across domains.
• Domain translation, on the other hand, concentrates on training a generative map to translate the source domain data into the style and distribution of the target domain, e.g. inspired by reducing the total cost incurred by this translation. The goal of domain translation is to learn the mapping between domains rather than performing a prediction task.

How to leverage UOT in the domain adaptation task is thoroughly discussed in [Fartras+'2021], which we also refer to as an application of UOT in Section 1 ("The practical significance of unbalancedness in discrete OT has been demonstrated in various applications, e.g. in video registration (Lee et al., 2020), computer vision (Plaen et al., 2023), or domain adaptation (Fatras et al., 2021)"). However, our work is concerned with learning Monge maps for domain translation, and hence our framework is not directly applicable to domain adaptation. For domain translation, we believe we cover all important concepts and related fields.

The limitations of the proposed method should be discussed.

➤ As mentioned in the general comment, we now discuss these in Appendix B.3 .

The experiments can be enhanced by considering more practical and challenging problems.

➤ As mentioned in the general comment, we have added another real-world application in a different perturbation task on the single-cell data measured in [Srivatsan 2020] as well as a benchmark on CIFAR-10 image generation. What more challenging datasets would you be interested in seeing?

As far as I understand this paper, the basic goal is to address the potential mismatch induced by the imbalanced weights of different domains; more precisely, the imbalance means the mass proportion of the ideal transport pairs is different. Such a scenario is also analogous to the label shift problem, which is usually considered in OT methodology. To address this problem, there are two natural and common solutions: relaxation (i.e., UOT) or reweighting (i.e., adjusting the marginal distribution) [r1, r2]. Since this work adopts relaxation as a solution (while also noting that the idea of reweighting is also implicitly shown in Sec. 3.1), I think more justification and discussion on the related fields are highly expected.

➤ In the following, we reply to the question in three parts. For the first part of the question regarding imbalanced weights of different domains (in the following referred to as label imbalance), we assume we have access to labeled data. We agree that unbalancedness helps to mitigate class imbalances, as demonstrated in multiple examples, e.g. in the EMNIST experiments. Yet, we would like to highlight that unbalancedness acts on batches, and hence there does not have to be a label imbalance on the full dataset for unbalancedness to have a positive effect, as within one batch, the label distribution might not reflect the label distribution of the full dataset. Moreover, we would like to highlight that unbalancedness is not only able to adapt label imbalances but also automatically discard outliers.

For the second part of the question, we fully agree that rescaling the marginal distribution is an option to address class imbalance. Yet, this assumes access to labels. In many real-world scenarios, data is not labeled, or the quality of the labels is insufficient. An example of the latter is cell-type labeling in single-cell genomics, which clusters cells into disjoint classes. In fact, cells are known to evolve continuously between clusters. With respect to image labels, the labels considered in the CelebA experiments (e.g. glasses vs. no glasses) might be clear. In contrast, labels for hair color might be less obvious, and hence reweighting the distribution accordingly is difficult.

Lack of External Benchmarks: Without a comparison to existing methods or frameworks, it's difficult to assess how much of an improvement the proposed approach offers.

➤ We compare to competing methods for single-cell trajectory inference (Appendix C.5) and image translation (Table 3, Appendix C.1). For the perturbation modeling task on 4i data, the Monge Gap is the current state-of-the-art to the best of our knowledge. Hence, we continue this line of research by comparing the balanced Monge Gap to our proposed unbalanced version as stated in Section 5.2. ("Since Uscidda & Cuturi (2023) show that OT-MG outperforms OT-ICNN in this task, we continue this line of research by applying unbalanced OT-MG (UOT-MG) to predict cellular responses to 35 drugs from data profiled with 4i technology (Gut et al., 2018).")

We agree with the reviewer that competing unbalancedness estimators were not considered. We hope to overcome this limitation with the newly added benchmark for the sciPlex dataset.

Complexity: Introducing unbalancedness into neural Monge maps could add computational or conceptual complexity, which may not be fully addressed in the paper.

➤ We agree with the reviewer and as mentioned in the general comment have now added a discussion of the computational complexity introduced by our framework in Appendix B.3. We discuss precisely the cost of incorporating unbalancedness in each Monge map estimator.

Dependency on Empirical Validation: While the paper does include empirical validation, its contributions could be further strengthened with more diverse datasets or a broader set of real-world applications.

➤ As mentioned in the general comment, we have added another real-world application in a different perturbation task on single-cell data measured in [Srivatsan+'2020] as well as a benchmark on CIFAR-10 image generation. What other datasets or tasks would you be interested in seeing?

What are the limitations of introducing unbalancedness into Monge maps, and are there scenarios where this approach may not be beneficial?

➤ Thanks for this question. As mentioned in the general comment, we've added a section on the limitations of our framework.

How does the proposed unbalanced neural Monge map approach compare with existing methods in terms of efficiency and accuracy?

➤ In the following, we assume that the question addresses the comparison of efficiency and accuracy to other unbalancedness estimators. The comparison of UOT-ICNN with [Lübeck+'2022] is discussed in length in section 4. In particular, we state "Third, our approach is computationally more efficient [than the approach by Luebeck et al.], as our method requires the computation of at most one discrete unbalanced OT plan, while the algorithm proposed in [Lübeck+'2022] includes the computation of two discrete unbalanced OT plans." A comparison of UOT-ICNN with other approaches that model unbalancedness (e.g. [Yang+'2019] and [Choi+'2023]) with respect to efficiency is not possible due to the different modeling assumptions. Similarly, comparisons with respect to the efficiency (complexity) of UOT-MG and UOT-FM are not possible. Regarding accuracy, as mentioned in the general comment we've now added a new benchmark comparing UOT-ICNN to related unbalancedness estimators.

---

[Yang+'2019] Karren D. Yang and Caroline Uhler. "Scalable unbalanced optimal transport using generative adversarial networks", (2019).

[Srivatsan+'2020] Srivatsan, Sanjay R., et al. "Massively multiplex chemical transcriptomics at single-cell resolution." Science 367.6473 (2020): 45-51.

[Lübeck+'2022] Frederike L"ubeck, Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. "Neural unbalanced optimal transport via cycle-consistent semi-couplings", arXiv preprint arXiv:2209.15621 (2022).

In the single-cell trajectory inference and image translation experiments in Sections 5.1 and 5.3 (Table 1, 2, and 4), the article compared the performance of OT and unbalanced OT. Were there comparisons with other state-of-the-art methods?

➤ Yes, for single-cell trajectory we compare to competing methods in Appendix C.5 as mentioned at the end of Section 5.1 . For EMNIST, we chose not to compare to any competing methods, as this is meant as an illustrative example to show the difference between FM, OT-FM, and UOT-FM. For CelebA, we compare to competing methods in Table 3 (latent space) and Appendix C.1 (pixel-space). For Table 4, we chose to evaluate the learned transport cost only for the FM-based methods as the difference between them is the main result we want to emphasize.

The experiments in Table 3 on the CelebA dataset show that the proposed method performs suboptimally FID score in comparison to UVCGAN [1], which was introduced in 2022. Further explanation is required to address this disparity.

➤ We argue that UOT-FM should be viewed as a principled method for unpaired image translation and compared to other such methods like CycleGAN [Zhu+'2017] or [Korotin+'2022], which UOT-FM significantly outperforms in all tasks as detailed in Table 3 and Appendix C.1 .

As mentioned in Appendix E.3 UVCGAN [Torbunov+'2022] is a highly optimized version of CycleGAN that uses an improved architecture (UNet + Vision Transformer at the bottleneck layer) and self-supervised pre-training. Due to computational reasons, we pre-compute all latent embeddings with the Stable Diffusion encoder and thus are also not benefitting from random cropping during training, which UVCGAN leverages as stated in Appendix F.2.1 . We believe a similar kind of optimization could also be done for UOT-FM to significantly improve its performance.

OT-FM [2] has conducted numerous experiments in image generation tasks and outperforms many GAN-based and diffusion-based methods. Have you considered experimenting with UOT-FM on the same tasks for further validation of the method's effectiveness?

➤ Originally, we chose not to benchmark on generative modeling tasks as the main focus of the paper is domain translation. However as mentioned in the general comment, we have now added an experiment on CIFAR-10 image generation.

---

[Zhu+'2017] Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. "Unpaired image-to-image translation using cycle-consistent adversarial networks". In IEEE International Conference on Computer Vision 2017 (ICCV), 2017.

[Korotin+'2022] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. "Neural optimal transport". 2022. URL https://arxiv.org/abs/2201.12220.

[Torbunov+'2022] Dmitrii Torbunov, Yi Huang, Haiwang Yu, Jin Huang, Shinjae Yoo, Meifeng Lin, Brett Viren, and Yihui Ren. "Uvcgan: Unet vision transformer cycle-consistent gan for unpaired image-to-image translation", 2022.