Light Unbalanced Optimal Transport

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

(1) The optimization objective and Gaussian Mixture approximation in Sec 4 are similar to [1].

Our solver can be considered as the generalization of the one from [LightSB, 1] in the sense that it subsumes LightSB for the specific choice of $f$-divergences. However, this generalization is not straightforward or direct since our objective is built on the completely different principles:

1. Our solver is derived from minimizing $D_{\text{KL}}$ divergence (defined as a discrepancy between positive measures) between ground truth plan $\gamma^*$ and its approximation $\gamma_{\theta,\omega}$. This definition of divergence notably differs from the ordinary definition of $D_{\text{KL}}$ for probability measures used in [1].

2. We parametrize the entire plan using Gaussian mixtures while in [1] it is done only for conditional plans. It is an important difference, since in an unbalanced case the marginals of optimal plan do not coincide with source and target measures. Our parametrization allows sampling from the left marginal of UEOT plan and identifying potential outliers in the source measure.

(2) While this paper provides the universal approximation property for Gaussian Mixure approximation, I have concerns about whether this Gaussian Mixture parametrization can achieve decent results for more complex distributions, such as in generative modeling within the data space on CIFAR-10.

In general, methods based on Gaussian parametrization are usually not appropriate for tasks with complex data, e.g., images. We mention this limitation in our paper (lines 256-257,702-707). Still, our aim was to develop a lightweight unbalanced solver which can serve as a simple and easy-to-use baseline in moderate-dimensional tasks. As expected, we get this ease in exchange for the rich parametrization required for large-dimensional tasks and vice versa.

(3) In the Unpaired Image-to-Image Translation task, Table 2 only presents the accuracy of keeping the attributes of the source images. However, since the goal of this task is semantic translation, the accuracy of the target semantic is also required. For example, in the Young-to-Adult task, the accuracy of whether the generated image is indeed an adult image. Could you provide this target semantic accuracy results?

We are working on this and aim to add the results soon.

(4) In the Appendix, Tables 5 and 7 show the Frechet distance (FD) between the learned and target measures. I believe this FD metric evaluates whether the semantic translation is successful, at the marginal level. Generally, increasing $\tau$ decreases (improves) FD metrics in Table 5 and 7. Could you clarify how the optimal tau is selected? I am curious because when tau is overly large, U-LightOT achieves worse accuracy compared to other models in Table 2.

In Appendix C (Tables 5, 7), we perform an ablation study of our method in order to show that it offers a flexible way to select a domain translation configuration (unbalancedness parameter $\tau$) that either allows for very good level of preserving the properties of the input objects or generation of a distribution which is a very good approximation of the target distribution. In that section, we highlighted the parameter which is optimal in the sense that it provides the best tradeoff between the closeness of the learned translations and target ones (Pareto-optimal), and the ability of the learned latents to keep the features of input latents. However, the final selection of the optimal configuration remains at the discretion of the user.

(5) Could you present the FD results (Tables 5 and 7) for the other models?

We will add the results soon.

(6) For the optimal tau=500 in the Man-to-Woman task, Table 6 shows an accuracy of 83.85 at best, while Table shows an accuracy of 92.85. Could you clarify which result is correct?

The Tables provide results for different number of training steps and parameters $\tau$. Table 2 shows the accuracy for $\tau=100$ and 5K steps of the algorithm which is specified in Appendix B.3. Table 6 presents the results for only 3K steps and different parameters $\tau$. (Note that the value for $\tau=100$ in the Table 6 (88.59 $\pm$ 0.40) is close to 92.85 - up to the difference in number of training steps.)

References.

[1] Korotin, Alexander, Nikita Gushchin, and Evgeny Burnaev. "Light schr" odinger bridge." ICLR 2024.

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

(1) The paper does not discuss how K and L, i.e., the number of Gaussians in the mixtures, affect the results. The generalization error bound mentions that K and L will appear as constants in the error bound, but the practical implications of the choice of K and L are missing from the paper. [...] How does the performance change as a function of K(assuming K=L)?

To address your question, we perform additional experiments (both on Gaussians and in the latent space of ALAE autoencoder) with our solver varying the number of Gaussian modes ($K$, $L$) in potentials.

The setup of this experiment follows the setup introduced in our Section 5.1. We test our solver with diverse number of potentials $K\in\{1,3,5\}$ and $L\in\{1,2,3,4,5\}$. The results are visualized in Fig. 1 of the attached PDF file. It can be seen that for insufficient number of modes in potentials, the solver exhibits issues with convergence and do not correctly solve the task.

(2) It would be beneficial to have an experiment on robustness to outliers, as this is included in the main claim.

Thank you for this valuable suggestion. We conduct the experiment on Gaussian Mixtures with added outliers and visualize the results in Fig. 2 of the attached PDF file. The setup of the experiment, in general, follows the Gaussian mixtures experiment setup described in section 5.2 of our paper. The difference consists in outliers (small gaussians) added to the input and output measures. The results show that our U-LightOT solver successfully eliminates the outliers and manages to simultaneously handle the class imbalance issue. At the same time, the balanced LightSB [4] solver fails to deal with either of these problems.

(3) [...] speed comparison for either experiment. It would be great to have a wall-clock comparison of competing methods.

Thank you for the suggested idea. We compared the running time of our algorithm and unbalanced competitors on the image translation task (Adult$\rightarrow$Young), its setup is described in section 5.2 of our paper. The results for all of the methods (wall-clock times for 10k updates) are presented in the Table below. We omit the results for other variants of translations since they are quite similar to the results in the Table.

As you can see, our proposed solver outperforms its competitors (unbalanced methods) in terms of convergence time.

(4) The Gaussian mixture assumption limits the method's applicability to only low-dimensional problems, and it is not clear whether this limitation can be overcome.

In general, methods based on Gaussian parametrization are usually not appropriate for tasks with complex data, e.g., images. We mention this limitation in our paper (lines 256-257,702-707). Still, our aim was to develop a lightweight unbalanced solver which can serve as a simple and easy-to-use baseline in moderate-dimensional tasks. As expected, we get this ease in exchange for the rich parametrization required for large-dimensional tasks and vice versa.

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.

[1] L. Eyring et al. Unbalancedness in neural monge maps improves unpaired domain translation. ICLR, 2024

[2] J. Choi et al. Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.

[3] K. D. Yang et al. Scalable unbalanced optimal transport using generative adversarial networks. ICLR, 2018.

[4] Korotin et al. "Light schrodinger bridge", ICLR 2024.

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

(1) Appears to be specific only to the case of having $D_{\text{KL}}$ divergence penalties for the mass constraints. [...] Do you have any intuition if one were to use other penalties beyond $D_{\text{KL}}$ to enforce the mass constraint?

Our solver admits different divergences except for the $D_{\text{KL}}$ one. From the theoretical point of view, we describe the set of admissible divergences in our Appendix C (lines 549-672). Besides, we provide a numerical example illustrating the performance of our solver with $\mathcal{D}_{\chi^2}$ divergence, see Fig. 3 and description in lines 559-677.

(2) Paper can appear a bit difficult and dense to read.

We are upset that you found our work difficult to read. We will try to improve this aspect if you indicate in more detail which points were difficult to understand.

(3) The reduction you get appears to have (perhaps even superficially) some relationship to the way WAE decomposes the coupling into conditionals. More coincidentally, WAE also uses conditional Gaussians to parametrize the encoder distribution, although for different purposes. Do you have any comments if there is any deeper link here?

Thanks for asking. We think that there is no direct link. Indeed, in WAE, the encoder for each input $x$ outputs some Gaussian, while in our case, all the conditional Gaussians (more precisely, Gaussian mixtures) are tight together. This means that given one conditional distribution $\gamma_{\theta}(y|x=x_0)$, one can immediately express all the other $\gamma_{\theta}(y|x=x_{\text{other}})$. In fact, the densities of all these conditional distributions are parameterized by a single scalar-valued function $v$; see eq. (8) in our paper. This is achieved because of the properties of the entropic optimal transport solutions which we exploited to construct our algorithm.

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 thorough feedback. Please find the answers to your questions below.

(1) Performance of the method on the UEOT plan between Gaussian distributions, see Janati et al.,2020 [5]

Thank you for your suggestion. Unfortunately, a comparison of our method's solutions with the analytical solutions proposed in [5] is not relevant, since this paper consideres a different setup of the UEOT problem. Namely, this paper derives solutions for the UEOT problem (between Guassian measures) with $D_{\text{KL}}$ as entropy regularization instead of the differential entropy used in our paper. We noted the difference between the problem we are considering and the one considered in [5] in our paper, see lines 91-92 and corresponding footnote.

(2) The Gaussian mixture approach is likely to work only in low dimensions. It would be interesting to see when it fails, e.g. using the benchmark above.

As we explained in the previous answer the benchmark provided in [5] is not relevant for us as it considers another UEOT problem.

(3) The authors state that for OT-FM and UOT-FM in the FFHQ dataset, they use a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. Where does this parameterization come from? It seems to be relatively small for a flow matching architecture on images, and does not seem to be the architecture used in the original papers.

It is important to understand here that we run our experiments in the latent space of the ALAE autoencoder, and not on the images directly. Accordingly, we adapted the architectures of the neural networks used in OT-FM and UOT-FM to work with latent codes. In this case, architectures such as fully connected neural networks are relevant.

(4) Why are FID scores not reported for the image translation tasks?

We conduct the image translation experiment in the latent space of ALAE autoencoder. For this reason, we did not report the FID metrics assessing the quality of the generated images but rather focus on assessing the quality of the generated latents and focus on the Frechet distance (FD) defined as the difference in means and covariances of the generated and target latents.

However, to fully address the raised question, we report the FID scores between the generated images (produced by ALAE decoder from the generated latent codes) and target images distributions in the Table below. The results show that FID is nearly the same for all of the models under consideration. It supports our intuition that FID is indeed not a representative metric for assessing the performance of models performing the translation of latent codes.

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

[1] Korotin et al. "Light schrodinger bridge", ICLR 2024.

[2] L. Eyring et al. Unbalancedness in neural monge maps improves unpaired domain translation. ICLR, 2024

[3] J. Choi et al. Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.

[4] K. D. Yang et al. Scalable unbalanced optimal transport using generative adversarial networks. ICLR, 2018.

[5] H. Janati et al. Entropic optimal transport between unbalanced gaussian measures has a closed form. NeurIPS, 2020.

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

(1) The method of parameterizing the potential function using GMMs was already proposed in LightSB [1]. The only change here is the switch to a UOT objective, making the methodological contribution minimal.

Our solver can be considered as the generalization of the one from [LightSB, 1] in the sense that it subsumes LightSB for the specific choice of $f$-divergences. However, this generalization is not straightforward or direct since our objective is built on the completely different principles:

1. Our solver is derived from minimizing $D_{\text{KL}}$ divergence (defined as a discrepancy between positive measures) between ground truth plan $\gamma^*$ and its approximation $\gamma_{\theta,\omega}$. This definition of divergence notably differs from the ordinary definition of $D_{\text{KL}}$ for probability measures used in [1].

2. We parametrize the entire plan using Gaussian mixtures while in [1] it is done only for conditional plans. It is an important difference, since in an unbalanced case the marginals of optimal plan do not coincide with source and target measures. Our parametrization allows sampling from the left marginal of UEOT plan and identifying potential outliers in the source measure.

(2) Aside from the universal approximation result, the theoretical contributions are also limited.

We partially agree with the reviewer that the proof of our Universal Approximation Theorem (UAT) is the most difficult and tricky among the results obtained in our paper. However, the theoretical contributions of our paper are not limited to this theorem. Our other results include: (1) Theorem 4.1 $-$ the derivation of the tractable optimization objective in terms of the $D_{\text{KL}}$ divergence between positive measures; (2) Proposition 4.2 $-$ the derivation of the bound for the estimation error of our solver; (3) Theorem A.4 $-$ derivation of the dual form of UEOT problem with the potentials belonging to the space $C_{2,b}(x)$ of continuous functions, bounded by the quadratic polynom (from the both sides) and additionally bounded by constant from above. Proof of each of these results is non-trivial and requires highly specialized knowledge in the diverse fields of mathematics and statistics.

(3) In the I2I experiments, the authors use the ALAE autoencoder, while some of comparison methods are implemented directly in the image space. All of the comparisons should be implemented in the latent space for fairness.

All of the methods included in comparison in the image-to-image translation task were implemented in the latent space of the ALAE autoencoder which is mentioned in the paper, see Section 5.2, line 276. We agree that it might written more clearly and will additionally emphasize this aspect in the final version of our paper.

(4) Since U-LightOT are implemented on a latent space that captures attributes well, the attribute accuracy is expected to be high. Other than accuracy, more general metrics such as c-FID or FID should be used for comparison.

We conduct the image translation experiment in the latent space of ALAE autoencoder. For this reason, we did not report the FID metrics assessing the quality of the generated images but rather focus on assessing the quality of the generated latents and focus on the Frechet distance (FD) defined as the difference in means and covariances of the generated and target latents.

However, to fully address the raised question, we report the FID scores between the generated images (produced by ALAE decoder from the generated latent codes) and target images distributions in the Table below (Adult → Young translation). The results show that FID is nearly the same for all of the models under consideration. It supports our intuition that FID is indeed not a representative metric for assessing the performance of models performing the translation of latent codes.

(5a) Lack of ablation studies on the number of Gaussian modes N and M.*

To address the reviewer's concern, we perform additional experiments with our solver varying the number of Gaussian modes ($N$, $M$) in potentials.

Gaussians mixtures. The setup of this experiment follows the setup introduced in our Section 5.1. We test our solver with diverse number of potentials $N\in\{1,3,5\}$ and $M\in\{1,2,3,4,5\}$. The results are visualized in Fig. 1 of the attached PDF file. It can be seen that for insufficient number of modes in potentials, the solver exhibits issues with convergence and do not correctly solve the task.

(5b) How does the performance change when parameterizing very high-dimensional and multi-modal GMM data with fewer or more N and M?

Unfortunately, to assess the performance of our solver in such an experiment with multi-model GMM data, we need to have some kind of ground-truth solutions. However, for the multi-modal GMM data the solutions are not available making it hard to perform such an experiment. Following your comment, we qualitatively demonstrated the performance of our solver with varying number of potential modes $N,M$ for 2-dimensional Gaussian mixtures and quantitavely assess its performance in Image-to-Image translation task, see the answer above.

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.

Since the end of the rebuttal period is approaching, we want to thank the reviewers for their time spent on their reviews and subsequent discussion. We are grateful for your interesting and valuable suggestions and will add the changes to the final version of our paper. In addition to the changes listed in the general comment, we will include

1. (Addition to Appendix C) Table with OT cost between the source and generated distributions and Wasserstein-2 distance between the generated and target distributions in Gaussian mixture experiment (with varying parameter $\tau$);
2. (Addition to Appendix E) Tables with additional metrics for comparing our method and its competitors (accuracy, FD, FID between the generated latents and the target ones).