Energy-Guided Continuous Entropic Barycenter Estimation for General Costs

Dear reviewer, thank you for spending time reviewing our paper. Your kind words that our paper well suits NeurIPS are encouraging for us. Below you can find the answers to your comments and questions.

(W1) Difficulties with proper evaluation.

We agree that the out-of-the-box evaluation is typically unavailable for the problem we aim to tackle (entropic barycenter). This is due to the inherent complexity of the barycenter problem. In particular, there seem to be no easy-to-calculate statistics to monitor to control the quality of the recovered barycenter (like the loss gap $L - L^*$). Surprisingly, we were able to derive a solution to the notoriously challenging entropic barycenter problem. On the other hand, we did our best to validate our methodology as accurately as possible. In particular, we consider several experimental setups where the unregularized barycenter is known: (a) Toy 2D Twister (Section 5.1), (b) Ave Celeba experiment (Section 5.3), Gaussian case (Appendix C.4). These setups allow us quantitatively/qualitatively estimate the recovered EOT barycenter.

A comment on small $\epsilon$ regime. From the theoretical point of view (approximation bounds etc.), the regime $\epsilon \rightarrow 0$ is indeed not so favourable. However, in terms of method evaluation, it is the regime that allows more-or-less accurate ways to "monitor" the quality. This is thanks to the aforementioned practical setups, where the true unregularized barycenters are known, and we can directly compare them with the generated ones. Practically speaking, in our paper, we consider rather small values of $\epsilon$ (up to $10^{-4}$), and our approach treats such choices.

(W2) Questions about the experimental section and $\ell^2$ barycenters.

Your understanding is correct: when we consider $\ell^{2}$ barycenters of images, the $\ell^{2}$ between two particular images is computed by flattening each image into a vector (e.g., $32\times 32\rightarrow 1024$). In turn, when there are just two distributions ($K=2$) of images considered, their $\ell^{2}$ barycenter is indeed a certain McCahn interpolant. However, one needs to recover somehow the true $\ell^{2}$ OT map (plan) between the distributions to access the barycenter through this interpolant.

(W3) Our work vs. [77]. Theoretical and practical efforts upon [77] which are required to end up with our approach.

Please find the answer to your comment in the general answer to all the reviewers.

(Q1) Non-convexity of the energy potentials. Extending the discussion and practical observations.

The problem with the non-convex energy function and its effect on MCMC is a hallmark of energy-based methods (and even score-based approaches). The non-convexity may cause some MCMC chains to become trapped by local minima of the energy landscape, making the corresponding samples somewhat different from the expected samples from $\pi^f_k(y \vert x_k)$. By taking a closer look at our results, e.g., Figure 7 (conditional samples from MNIST 0/1 barycenter), one may distinguish that some generated pictures "fall out of the line", i.e., noticeably differ from the others. We think this phenomenon is inherently caused by the non-convexity. We underline that non-convexity does not stop EBMs from successfully modelling high-dimensional distributions (see, e.g., references [34], [112] from our paper), which makes our methodology practically appealing. In the revised version of our manuscript, we will expand the discussion on the non-convexity.

(Q2) The type of the considered Entropic Barycenter problem. "Blurring" effect. Entropic barycenter of $\delta$-distributions.

Please note that we directly mention the blurring bias in our experiment with MNIST01 in data space (section 5.2, Figure 5 and lines 323-325) in our paper. However, we do not focus on it too much because in the manifold-constrained setups (MNIST01 in Section 5.2, Ave, Celeba in Section 5.3) this bias in fact disappears. More generally, we believe that manifold-constrained setups are generally more interesting and probably even necessary to obtain meaningful barycenters for downstream practical problems (see Appendix B.2 for a discussion).

Regarding the $\epsilon$-barycenters of $\mathbb{P}\_{k}=\delta_{x_{k}}$, it seems like it is some Gaussian centered at $\sum_{k=1}^{K}\lambda_k x_{k}$ with variance proportional to $\epsilon$ (in the case of $\ell^{2}$ barycenters). Namely, the higher $\epsilon$ is, the more "blurry" (noisy) the $\epsilon$-barycenter is compared to the unregularized one. Following your question, we will add an extended discussion of this aspect to B.3.

(Q3) Missed citation. (Possible) closed-form expression for entropic barycenters in 1D.

Thanks for noting, we were not aware of this paper. It indeed seems like their $\alpha_{\text{OT}^{\mathcal{L}}}$ barycenter coincides with our considered Schrodinger barycenter. We will surely include the citation and proper discussion, including their formula for the closed-form barycenter for 1D Gaussians and $\ell^{2}$ cost.

Concluding remarks. Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. Thank you one more time for your efforts.

Dear reviewer. Thank you for spending time reading and evaluating our work. We were greatly encouraged by your positive feedback on our theory and the readability of our manuscript. Below you can find the answers to your questions and comments.

(W1) The role and impact of entropic regularizer's weight. Robustness test of the method to the regularizer's weight.

We agree that considering too small $\epsilon$ within our proposed method may cause some difficulties with optimizing our energy-based barycenter objective. The heart of the matter is the answer to: "can our approch permit practically-reasonable choices for $\epsilon$"? We believe that our considered practical use cases (Figure 5, Figure 7, Figure 8) testify that this is the case. The images sampled from our learned conditional plans $\pi^k(y \vert x)$ are of reasonable diversity. We want to emphasize that for Ave Celeba! experiment (manifold constrained) the considered value of $\epsilon$ was $10^{-4}$. Also, as per request, we additionally deployed our methodology on MNIST 0/1, manifold setup, with an $\epsilon$ of $10^{-4}$ (in addition to the previously considered value of $\epsilon=10^{2}$). The results are demonstrated in Figure 1 in the attached pdf.

Regarding the high-dimensional applications, our comment is as follows. It seems that learning an (especially entropic) barycenter directly in the image data space is, in most cases, not very meaningful - we will obtain just blurred pixel-wise averages of input images. This is exactly where our StyleGAN trick comes into play. In turn, the latent space of StyleGAN is moderate dimensional and allows learning meaningful barycenters in high-dimensions. And our experimental results demonstrate that our approach successfully tackles this setup.

(W2) NNs architectures to parameterize the potential functions.

One of the key insights of our method in practice is the reliance of our approach on already well-established, built-on efficient architectures (and hyper-parameter selections) constructed for EBMs and EOT problems [1]. Since we trust the architectural design and hyperparameter specifications of those authors, we perform no additional find-tuning, using their neural networks as pre-trained foundation models. We think that doing this, in addition, is out of the scope of the current work. Please note that our primal contribution is methodological, not architectural, but has a primarily theoretical intent.

(W3) Reliance on MCMC and gradient descent. Quantitative analysis of convergence and computational efficiency in, e.g., high-dimensional case.

We adopt MCMC and (stochastic) gradient descent as the standard tools for Energy-based Models. It is a generic knowledge that these tools allow for generating high-dimensional images (when carefully tuned); see references [34], [112] from our paper.

Following your request, we quantitatively analyze the performance of our method as a function of the number of discretized Langevin dynamic steps, see the charts (Figure 4) in our attached pdf. In this experiment, we try to learn the barycenter of Gaussian distributions; which, as explained in Appendix C.4, is known in-closed form for the unregularized case, which gives us access to a theoretical object by which we may compare and validate our computational pipeline. We pick the dimensionality $D = 64$ (which is the highest-dimensional case among the considered). As we can see, performance drops when the number of steps is insufficient. Overall, in all our experiments, the number of Langevin steps is chosen to achieve reasonable qualitative/quantitative results.

(Minor W4) Practical application.

For an intriguing practical application, we pick an intriguing Single-Cell Data analysis experiment, where we need to predict the distribution of cell population at some intermediate time point $t_i$ given cell distributions at $t_{i + 1}$ and $t_{i - 1}$. Please see a brief explanation (and the relevant links) of the problem in our General response, our results are in Table 1 in the attached pdf. Our approach achieves competitive performance.

(Minor W5) Call for conclusion section.

Section E contains our conclusions and summary. Currently, this section is located in the Appendix. However, following your suggestion, we will move it to the main part of our manuscript in the final version.

Concluding remarks. Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.

References

[1] Mokrov et. al., Energy-guided entropic neural optimal transport. ICLR'2024.

Dear reviewer. Thank you for your thoughtful review. We are pleased that you noted the theoretical and practical strengths of our work. Below we answer your comments and questions.

(W1)+(Q2). Too many details in the background section. Improving the readability by moving some secondary assumptions/definitions to the Appendix. Shortening Section 2.2.

Following the reviewer's request, in the revised version of our manuscript, we will shorten the Background section by leaving only vital definitions/assumptions in the main part. In particular, we will remove/move to the appendix Eq. (1) (definition of $\text{EOT}^{(2)}$). Indeed, only Eq. (2) actually needed. Correspondingly, Eq. (7) will also disappear from the main part. Following your request, we will also shorten the explanations after Eq. (8) (Section 2.2) and eliminate introducing empirical measure $\widehat{\mathbb{P}}$, line 121.

(W2. Part 1) The difference between this work and [1].

Please, take a look at our General response (1) to all the reviewers, where we get into details of the connection with [1] and our novelty over [1].

(W2. Part 2) Comparison with [2].

We highlight that the authors of [2] acknowledge that their work is significantly based on our paper, see the Contribution subsection in [2]. Interestingly, they also seem to conduct a comparison of their approach with ours on Ave Celeba! barycenter benchmark. According to their results, on the manifold-constrained data setup, our method achieves FID scores $\approx 9$, while their approach yields a much bigger number $\approx 30$. Therefore, our method demonstrates a significant quantitative boost. This is a clear indicator of the strength of our energy-guided methodology.

(W3) "Most of the experimental results are qualitative, which makes it difficult to attest if the method improves on existing ones." Call for more experiments with quantitative outcomes.

Following your kind suggestion, we conducted a novel experiment with predicting the interpolation between distributions at multiple timepoints following the setup identical to [3]. Please refer to the general answer to all the reviewers. According to the results, our method works on par with competitive methods.

For completeness, we also note that we have a purely quantitative experiment in the Gaussian case in Appendix C.4.

(Q1) "What is $\mathcal{B}(\mathbb{Q})$ in eq. 8?"

The term $\mathcal{B}(\mathbb{Q})$ is defined directly in equation (8), i.e., $\mathcal{B}(\mathbb{Q}):= \sum\limits_{k = 1}^{K} \lambda_k \text{EOT}_{c_k, \varepsilon}(\mathbb{P}_k, \mathbb{Q}).$ We agree that this makes (8) to be a bit overcomplicated, which was done due to space constraints. Sorry for this. In the final version, we will have one extra page and we will separately define $\mathcal{B}(\mathbb{Q})$ or even get rid of this notion completely.

Minor comments and typos

Thank you for pointing them out! We will fix them in the revised version of the manuscript. The FIDs for our method with std. deviations (running the inference with different seeds) look as:

For the competitors, we grab the results from WIN paper where no std. deviations are provided.

Concluding remarks. Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.

References

[1] Mokrov et. al., Energy-guided entropic neural optimal transport. ICLR'2024.

[2] Kolesov et. al., Estimating barycenters of distributions with neural optimal transport. ICML'2024.

[3] Tong et. al. Simulation-free Schrödinger bridges via score and flow matching. AISTATS'2024

Dear reviewer. Thank you for your efforts in reviewing our paper and for your valuable comments. We are delighted that you appreciated both the theoretical and practical contributions of our paper. Below you can find the answers to your questions and comments:

(W1) "The use of the Style GAN seems to be pivotal in the experimental validation, yet it is only explained briefly while reporting the experiments." Proper attention to the idea with StyleGAN manifold.

In the revised version of our manuscript, we will devote a specific subsection to the idea with StyleGAN manifold. We will explain the technique itself and clarify why it is significant and valuable in the context of entropic barycenters. Our reasoning here is approximately as follows:

• At first, solving the entropic barycenter problem directly in the image data space leads to noisy images, see, e.g., our MNIST 0/1 experiment. Such "deterioration" follows from the theory - in fact, the Schrödinger barycenter problem we consider (see Appendix B.3) is known to have blurring bias, see work [15] from the references in our manuscript. Our practical reliance on MCMC also makes it obvious that the generated images from the barycenter are noisy.
• Secondly, solving barycenter problems in image data space is not very practical (for standard costs $c_k$ like $\ell^2$ or around). Indeed, our paper (MNIST 0/1 experiment in data space) and previous research (see, e.g., [1, Figure 7]) witness that $\ell^2$ barycenter of image distributions $A$, $B$, $C$ etc.\ is a bizarre distribution which includes "averaged" images of $a \in A$, $b \in B$, $c \in C$ etc. So, we think that direct manipulation with image data space is more about benchmarking barycenter solvers rather than real-world applicability.
• In contrast, operating with images with the assistance of our proposed StyleGAN-inspired manifold technique seems to be much more practically appealing, see the examples of potential practical applications in our Appendix B.2.

(W2) Computational complexity of the method.

As per request, we report the (approximate) time for training and inference (batch size = $128$) of our method on different experimental setups, see the table below (the hardware is a single V100 gpu):

For Ave Celeba experiment, we additionally report the computational complexity of the competitors (partially grabbed from the original papers, partially obtained through our own evaluations):

As we can see, all the methods in this experiment require a comparable amount of time for training. The inference with our approach is expectedly costlier than competitors due to the reliance on MCMC. In the revised version of our paper, we will add the information about the computational complexity.

(L1) Identification and discussion of the limitations apart those alleviated by StyleGAN.

In Appendix E we have a paragraph called "Methodological limitations" which discusses the limitations not related to StyleGAN. When preparing the manuscript, we accidentally leave the word "limitations" in the "Methodological limitations" to be normal, not bold. It is our bug, sorry for this. Will be fixed.

Formatting issues.

Thank you for pointing them out! We will go through the text carefully and do our best to fix the formatting problems in the final revision.

Concluding remarks. Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.

Dear reviewer, thank you for spending time reviewing our paper and for a positive evaluation of the work. We are happy that you found our work to be interesting and well-presented. Below we answer your questions and comments.

(W1) Convergence bounds discussions. "[...] it is not [...] a measurable convergence diagnostic". Significance of Theorem 4.6.

It seems that there is a tiny misunderstanding. The results from Section 4.3 (Prop. 4.4, Theorem 4.5, Theorem 4.6) exhibit the generalization and approximation properties of our proposed solver. Namely, we give a complete (statistical) picture of what happens when we move from real reference distributions $\mathbb{P}_k$ to their empirical counterparts. To the best of our knowledge, the ability to carry out such the analysis is exceptional in the existing neural OT literature, which strengthens the contribution of our work. While our Prop. 4.4 indeed does not give particular figures on how different the true $\pi_k^*$ and recovered $\pi^{\widehat{f}_k}$ plans are, it provides the rates of convergence (w.r.t. number of empirical samples). Indeed, if we assume the approximation error to be zero (note Theorem 4.6), then Theorem 4.5. gives us the desired rates.

While our Theorem 4.6. indeed looks simply like "we have a loss function, and can approximate it arbitrarily good with NNs", this result is not that easy. Indeed, the considered loss $\mathcal{L}$ operates with (a) congruent potentials and (b) deals with entropic $C_k$-transforms of the potentials, which complicate the overall analysis. In particular, only specific properties of the entropic $C_k$-transforms allow deriving the desired universal approximation statement, see the proof of the Theorem.

(W2) Discussions on intrinsic difficulty of multimarginal problems.

The problem we aim to tackle is indeed significantly complex by nature. The fact that our method theoretically (and practically) recovers entropic barycenters is quite encouraging for us. Having in mind the complexity of the barycenter problem, in our paper we pay special attention to the experimental setups where our method could be evaluated qualitatively or even quantitatively, see, e.g.: Toy 2D Twister (Section 5.1), Ave Celeba (Section 5.3), Gaussian case (Appendix C.4).

At the same time, we humbly think the discussions on the NP-hardness of the multimarginal problems are out of the scope of our paper, since our approach is based on parametric models (Neural Networks) learned with gradient descent/ascent, which eliminates such questions coming from the Theory of Algorithms.

(Q1) Practical aspects of optimizing Energy-based training objectives. Variance of the loss estimate. Reliance on MCMC, stochastic gradient descent, and neural networks.

Even in the presence of variance due to stochasticity induced by MCMC and gradient procedures, the proper tuning of the hyperparameters makes the training procedure to be successful. To practically demonstrate it, we provide the behaviour of $\mathcal{L}_2$-UVP metric (between the unregularized ground truth barycenter mapping and our learned entropic barycenter mapping $\pi^{f_k}(y \vert x_k)$) for our Ave Celeba experiment, see Figure 3 in the attached pdf. We emphasize that $\mathcal{L}_2$-UVP directly measures (by computing pointwise MSE) how the learned mapping differs from the true mapping. So, the results showcase that the presence of stochasticity in the loss does not invalidate the optimization.

(Q1.1) Benchmarking the performance

In order to benchmark the performance, one either needs to know the true barycenter (it holds true for the majority of experiments we consider) or utilize some problem-specific metrics to check whether the recovered barycenter satisfies some expected problem-specific properties. The particular loss curves during the training are indeed not that informative.

Remark 1. "ULA is surely not the simplest MCMC algorithm."

We agree that we were a little hasty with this statement. Will be fixed in the revised version. Thank you!

Remark 2. "[...] the notation feels a bit heavyhanded and can be distracting from the main story of the paper."

We will simplify the background section in the final version of our manuscript by removing/moving to the appendix some not vital definitions/results. Please see our first answer to reviewer unyK.

Concluding remarks. Kindly let us know if the clarifications provided address your concerns about our work. We're happy to discuss any remaining issues during the discussion phase. If you find our responses satisfactory, we would appreciate it if you could consider raising your score.