Energy-guided Entropic Neural Optimal Transport

Dear reviewer! Thank you for your positive feedback and profound questions. Below we answer your questions and comments.

(1) Scalability of the proposed approach. Direct applicability of the approach to image spaces

To begin with, we want to highlight that there exist modern EBM models which are capable of generating high-fidelity image data, see, e.g., [Zhao], [Du]. This fact serves as an indirect witness that our method is scalable to such complex data setups as well. We assume that the combination of the aforementioned advanced techniques with our proposed methodology may yield Energy-based EOT solvers which operates in high-dimensional image spaces directly without StyleGAN proxy. However, in our paper we stick more to the theory and methodology rather than technical stuff with hyperparameters/regularization methods/technical tricks selection and leave these adaptations to a future work.

At the same time, we agree with the reviewer that it is really important to showcase the applicability for the method on the image data. That is why we conduct an experiment on ColoredMNIST "2"$\rightarrow$"3" $3 \times 32 \times 32$ digits transfer and demonstrate the results in Appendix C.1. The results are not so good. However, please note that for this experiment we adapt the code from [Nijkamp] - it is probably the simplest EBM setup for learning image data. In particular, we parameterize the learned potential $f_{\theta}$ as the trivial Conv2d -> LeakyReLU -> Conv2d -> LeakyReLU -> $\dots$ combination. See our code and Appendices C.1 and D.4 for the details. We humbly think that this experiment already demonstrates the applicability of our methodology, as "it is", without engineering improvements, to image data.

[Nijkamp] Nijkamp et. al., On the Anatomy of MCMC-Based Maximum Likelihood Learning of Energy-Based Models.

(1.1) "[...]it is worth noting that this approach couples with a pretrained GAN model and conducts training in latent spaces.

We want to leave a comment regarding our unpaired I2I setup in the latent space of StyleGAN. We want to underline that our Energy-based approach is not a usual method for learning something in a latent space. To the best of our knowledge, a typical method operating in a latent space requires not only a decoder (a model which transforms latent embeddings to data) but also an encoder (which transforms data of interest to the latent embeddings). At the same time, our method does not need any encoder.

(2) Clarifications regarding Eqs. 13 and 18

We are not sure if we catch the question of the reviewer exactly, but we make a try. At first, we propose to compare objectives (18) and (11). In fact, the potential $f_{\theta}$ in (18) serves as a test function which compares the second marginal of distribution $d \pi^{f_\theta}(x, y) = d \mu_x^{f_\theta}(y) d \mathbb{P}(x)$ and $\mathbb{Q}$. Similarly, the potential $E_{\theta}$ in (11) plays the same role by comparing the generative distribution $\mu_\theta$ and GT distribution $\mu$. In other words, our Energy-based methodology for solving EOT is the combination of three principles. We look for a plan $\pi \in \mathcal{P}(\mathcal{X} \times \mathcal{Y})$,

• with the first marginal $\pi(x)$ (projection to $\mathcal{X}$) given by the source distribution $\mathbb{P}$. This condition is satisfied by construction (as we build our plan as $d \pi(x, y) = d \mathbb{P}(x) \cdot d \pi(y \vert x)$).
• with the second marginal $\pi(y)$ (projection to $\mathcal{Y}$) given by the target distribution $\mathbb{Q}$. In order to satisfy this condition, we need the optimization given by Eq. (18).
• whose conditional distributions $\pi(y \vert x)$ satisfy Eq. (13), i.e., $\pi(y \vert x) = \mu_x^f(y)$ for some potential $f$. This condition is also satisfied by construction. Note, that it is the condition which ensures that plan $\pi$ is EOT plan between its first and second marginals.

Combining principle 3.) with 1.) and 2.), our methodology results in Energy-based EOT solver. Does our comments address your question?

Dear reviewer, thanks for your efforts and your thoughtful review. We appreciate that you have pointed out the methodological and theoretical advantages of our work. Please find answers to your questions and comments below.

(1) Problems with MCMC training usage

We agree with the reviewer that our EBMs-based methodology does not provide a "silver bullet" for solving the Entropic Optimal Transport. It indeed may suffer from the difficulties which are common to EBMs. In particular, MCMC procedure can be time-consuming and requires some technical tricks and hyperparameters tuning. Note that we have taken these complications into account in the limitations section. At the same time, EBM is actively developing research direction, and recent works from this area show that EBM can successfully model high-dimensional $512\times 512$ image data in ambient (not latent) space, see, e.g., [Gao], [Zhao]. Adapting advanced EBM techniques from such papers is a fruitful direction for future research.

Second, we would like to draw the reviewer's attention to our high-dimensional I2I translation setup (Section 5.3). In this experiment, we do not experience any difficulties with MCMC stability/convergence, and we use only basic and well-established EBM training techniques, e.g., we use a replay buffer. Nevertheless, we achieve a good quality of learned EOT stochastic mappings. Although this experiment is conducted in the latent space of a pre-trained StyleGAN, we dare to hope that it demonstrates the usefulness of our developed approach as it is, even without substantial technical improvements.

Finally, we would like to emphasise that the main contribution of our work is theoretical and methodological. Here we have managed to establish the link between EBM and EOT and to show that the proposed methodology can be successfully applied in practice. We humbly believe that this is already a good contribution to the ML community.

[Gao] Gao et. al., Learning Energy-Based Models by Diffusion Recovery Likelihood.

[Zhao] Zhao et. al., Learning Energy-Based Generative Models via Coarse-to-Fine Expanding and Sampling.

(2) Unpaired I2I task with other OT methods

Our work is not the only approach which can solve Unpaired image-to-image problem following OT formulation. In fact, tackling Unpaired I2I is the primary goal of the vast majority of continuous OT-based solvers. We refer the reviewer to Appendix A.2 (subsection "Continuous OT") for the corresponding survey and to Subsection 4.2 which covers the methods specifically designed for continuous EOT. Among them, the work [Korotin] probably demonstrates the most promising results for Unpaired I2I. This work is based on the semi-dual formulation of the OT problem regularised by a negative variance term. It proposes to optimize a max-min adversarial objective. This is a drawback of [Korotin], as adversarial learning is tricky and can be unstable. In particular, we couldn't get their algorithm to work properly on the AFHQ Cat$\rightarrow$Dog setup, even downscaled to $64\times 64$ images. Probably, one reason for this is the insufficient amount of data for this problem (5K samples for the source and target domains), since the original paper demonstrates good results on $128 \times 128$ setup of other unpaired I2I problems ($\approx$ 60K samples for the source and target domains). We added the examples of unpaired mapping learned by NOT in the newly updated supplementary materials. In conclusion, we want to underline, that in our paper we consider original AFHQ'$512\times 512$ Cat$\rightarrow$Dog setup, but utilize a pre-trained StyleGAN.

[Korotin] Korotin et. al., Neural Optimal Transport.

Dear reviewer! Thank you for your kind words. We are very pleased that you find our work to be interesting and valuable from both theoretical and practical points of view. Below we comment on your questions.

(1) Overall rate for the bound from Theorem 4, using explicit rates of NNs approximation, e.g., [1, 2]

Improving Theorem 4 with explicit numerical bounds is indeed a good question coming from Statistical Learning Theory. However, it seems to be rather non-trivial and deserves a separate research, probably, a separate paper similar to [2]. The key challenge here is the analysis of weak $C_{\text{EOT}}$-transformed functions. This is much more difficult than, say, analysis of functions (Neural Networks) $f$ and their compositions with a pre-definite measurement function $\gamma$ : $\gamma \circ f$. Note that the latter is the setup which is considered in [2]. Therefore, we expect that the derivation of the desired bounds will require substantial theoretical and numerical efforts: definitions, assumptions, theorems, proofs and supporting experiments. It is tantamount to carrying out a distinct research. At the same time, following your suggestions, we cite [1, 2] and indicate the importance of deriving numerical estimation rates for Theorem 4 in the new revision of our manuscript.

We emphasise that our Theorem 4, as given in our paper, is already important and worth presenting. It provides a way for estimating the quality of generative model $\pi^{\hat{f}}$ approximating EOT plan from the Statistical Learning Theory perspectives. As we point out in our paper, the majority of works establishing similar statistical bounds deals with other quantities, e.g., the OT (EOT) cost between fixed distributions or the barycentric projections. Note that work [2] also handles discrepancies (several $f$-divergences) between fixed distributions. This is far from the generative modelling setup.

(2) "The paper seems to focus more on applicability of EBM, [...]." EBM computational and space-time complexity analysis

Probably, here there is a tiny misunderstanding. Indeed, one of the main aims and contributions of our work is to show that EBMs could be applied (with some modifications) for solving the EOT problem. And we humbly hope that our theoretical findings as well as supporting experiments properly support this thesis. At the same time, we do not set out to answer the question regarding the applicability of EBMs itself, for general-purpose generative modelling problem. We do not answer the question to which extent the existing EBM procedures are efficient or not efficient. In fact, we treat the EBM training procedure as the particular solver for our proposed methodology. We resort to decade-long EBM research. And we are not that heroes who managed to equip EBM with solid theoretical grounds. Note that our Theorem 4 is free of assumptions about the particular, e.g. EBM-advised, method used to optimise the weak dual objective (9).

(3) Sampling error in Theorem 4

The sampling error (error due to inexact sampling from EBM due to MCMC) does not enter the bounds in Theorem 4. In fact, Theorem 4 just compares the true EOT plan $\pi^*$ with approximate EOT plan $\pi^{\hat{f}}$ and does not assume any sampling. Note that $\hat{f}$ which appears in Theorem 4 is (by definition) the perfect potential which solves weak dual objective (9) given a limited number of empirical samples from the source and target distributions $\mathbb{P}$ and $\mathbb{Q}$ and restricted class of available potentials $\mathcal{F}$. In turn, the sampling error may cause the potential $\hat{f}_{\text{practical}}$ resulting from our practical algorithm to differ from the perfect potential $\hat{f}$. Let us call this optimization error. Note that MCMC sampling error is not the only source of optimization error. It is also due to stochastic gradient ascent (Eq. 18), usage of batches, etc. The analysis of these quantities is a completely different domain in Machine Learning and out of the scope of our work. As with most generative modelling research, we do not attempt to analyse optimisation errors.

I thank the authors for the response, which addresses all of my questions. The authors are right that I misunderstood a bit -- I was expecting that the training objective uses not samples from Q but the MCMC samples (one comment is that it seems using EBM samples avoids exponentiation in the training, thus adding training stability?). I thus raise confidence and contribution score, and leave the rating unchanged.