Energy-Guided Continuous Entropic Barycenter Estimation for General Costs

Dear reviewer, thank you for spending time reviewing our paper. Please take a look at our general answer to all the reviewers where we address common questions and show additional experimental results. Below we answer to your questions and comments.

(0) The authors acknowledge this, but it is difficult to compare with the true entropic optimal transport barycenter quantitatively. However, it would have been nice to see it in cases where it is known (for example, with Gaussians).

See the general answer to all reviewers for the discussion of this question and newly added comparison with ground-truth unregularized barycenter.

(1) The theorems all seem to have standard proofs that are straightforward extensions of existing results. I am not sure that the theory proposed gives any new techniques or insights.

Thank you for this opinion. However, from our humbly point of view, this is not that right.

(a) "standard proofs [...] straightforward extensions[...]".

While our Theorems 2, 3, 4 are indeed largely based on the similar results from [4], we argue that our Theorem 1 (the key result for our proposed methodology) is original, and its proof is non-trivial. We have even been forced to prove some auxiliary properties of weak $C$-transform. The latter are interesting on its own.

(b) "[...] not sure that the theory gives new techniques or insights."

Our developed theory allows us to build Entropic OT barycenter solver for general cost function based on extensively developed EBMs techniques. We think this is a worthy result of our theoretical deductions. The closest approach [4] develops Energy-based training procedures only for Entropic OT problem itself. It is related but different problem. We do not hide our connection to [4], but we take a considerable step forward from both a theoretical and a practical perspectives.

(2) Furthermore, more important proofs are missing, such as how efficiently one can sample from the barycenter. Furthermore, the finite sample learning guarantees are nice to see but have no practical impact since there are no cases where they could estimate these errors in any practical example.

(a) efficient sampling from the barycenter

Our proposed practical algorithm is based on Energy-based models. It is an autonomous decades-long research direction which comes with its own advantages, techniques and limitations. The question of efficient sampling from barycenters resorts to the similar question for general-purpose EBMs. We point the reviewer's attention to [8], [9], [10] for some attempts to improve the sampling characteristics of EBMs. Note that typically these methods follows from practical observations, not theorems.

In our paper, we also do not make an attempt to equip EBMs with theoretically-grounded efficient sampling method. In fact, we treat the EBM training procedure as the particular solver for our proposed methodology. We leave the question regarding efficient sampling techniques to a separate research.

(b) "finite sample learning guarantees [...] no practical impact"

We agree that the particular statement of our theorem is not directly related to practical procedure, just like many Statistical Machine Learning researches in Generative Modelling. At the same time, our result justifies that our methodology actually permits the particular Statistical and Approximation guarantees. Note that deriving the similar guarantees for other scalable continuous barycenter solvers is complicated. This is because they utilize nontrivial adversarial [2]/iterative [1]/diffusion-based [7] objectives, while our objective (up to MCMC procedure) is the direct minimization of sum of KL divergences (see Theorem 2).

(3) The method learns OT plans but not maps. This is true of the general entropically regularized setting. However, this adds complexity to the usefulness of the method because one doesn’t have access direction to the barycenter distribution itself and rather only has conditional access through the reference measures.

While we partially agree with the reviewer, to be honest, we do not think this is a problem. Indeed, one can learn generative model (e.g., a GAN or a diffusion) on top of samples transported with our solver from marginal distributions to the learned barycenter. This will provide the direct access to samples of barycenter. Moreover, in the manifold-constrained case (Sections 5.2, 5.3), such generative model has to be constructed in the latent space (of a StyleGAN) which is presumably even easier.

Dear reviewer, thank you for spending time reviewing our paper. Please take a look at our general answer to all the reviewers where we address common questions and show additional experimental results. Below we answer to your questions and comments.

(1) The authors only motivated the optimal transport barycenter problem in Section 1, but they did not explain why its entropy-regularization variant is worth discovering.

The logic behind the introduction of entropy is analogous to its introduction in case of discrete OT. Considering entropy-regularized problem in discrete case allows to

(a) provide improved theoretical properties such as the strict convexity of the OT/barycenter problem, uniqueness of the solution, derive the tractable OT dual form [1] which, in turn, allows to

(b) establish convenient algorithms such as the Sinkhorn-like OT/barycenter methods with guarantees of performance.

The situation in the continuous case is the same. Point (a) still holds, and, as we show in our paper, the tractable OT dual form (Eq. 3) for EOT (b) gives us the opportunity to apply the well-known techniques from Energy-based modeling (EBM) to solve the barycenter problem under consideration.

Furthermore, our established solver allows to use general cost functions $c_k$ and, in particular, permits employing pre-trained generative models (e.g., StyleGAN) to restrict the sought-for barycenter to the image manifold. This aspect contributes to scalability of the proposed algorithm and expands practical applications of the barycenter problem.

(2) The MCMC algorithm leveraged in Algorithm 1 is not efficient. [...] If we replace the MCMC algorithm with other sampling methods, namely importance sampling, would the running time be improved? If not, what sampling methods are possible to achieve an improved running time. The authors should discuss more about this in Section 6.

We would like to note that we develop a generic methodology for solving the EOT barycenter problem and demonstrate that it does work even with the simplest sampling procedures. In principle, importance sampling might be used as such procedure (instead of ULA), however, it scales poorly with dimensionality (see Section 3.3 of [2]) and demands the usage of advanced improvements. Thus, the result of generating potentially might be improved with utilization either more efficient MCMC algorithms [5] or advanced techniques for importance sampling procedures [2, 6], however, these considerations are beyond the scope of our paper.

Following your suggestion, we mentioned corresponding papers in Section 6.

(3) Does the result in equation (6) hold true for any continuous function $f \in \mathcal{C}(\mathcal{Y})$ or only for the optimal function $f$ in equation (3)?

The results holds for every $f \in \mathcal{C}(\mathcal Y)$. If $f$ is optimal ($f=f^{*}$), then Equation 6 turns to the equality $0 = 0$.

(4) If the input distributions $\mathbb{P}_1,..,\mathbb{P}_k$ are not probability distributions and do not share the same mass, would we be able to extend the energy-guided approach to approximate the barycenter of those distributions?

Adapting our approach to distributions of unequal mass (unbalanced setup) is beyond the scope of our current paper. It would require significant changes in methodology and theoretical foundations. Exploring this direction is an interesting avenue for future research.

(6) paper organization is not good: [...] move the notation paragraph in Section 2 to the end of Section 1. [...] In Appendix A, the authors should separate the proofs of theoretical result into different subsections.

Up to the reviewer's request, we moved the notation and separated the proofs in the revised version of our paper.

(7) The paper violates the 9-page rule of ICLR 2024.

Please note that according to ICLR 2024 rules, it is allowed to position the reproducibility statement on an additional 10th page, see the reproducibility section in https://iclr.cc/Conferences/2024/AuthorGuide

(8) The authors should cite more relevant papers. In Section 2.2, the entropic barycenter problem was also previously studied in papers [3, 4].

We thank the reviewer for pointing to the relevant papers [3, 4]. We added the citations in Section 2.2, Appendix B.1 and C.2.

(9) Typos: Below equation (8), the term 'which differs from (8)' should be 'which differs from (7)', [...] clarify the abbreviation l.s.c although I can guess that it stands for lower semi-continuous.

We thank the reviewer for indicating the typos, we fixed them in the revised version of our paper.

Concluding remarks. Please reply to our post and inform us if the clarifications provided adequately address your concerns regarding our work. We are more than willing to discuss any remaining points during the discussion phase. If the responses offered meet your satisfaction, we kindly request you to consider raising your score.

Dear reviewer, thanks for your positive feedback.

(1)The results demonstrate generalization error estimates and universal approximation via neural network results for their algorithm. However, this part is a bit too general. I wonder if there are more specific results for special kind of neural network towards their approximation quality.

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. The key challenge here is the analysis of weak $C_{\text{EOT}}$-transformed functions $f^{C}$. This is much more difficult than, say, analysis of just functions (Neural Networks) $f$. Therefore, we expect that the derivation of the desired bounds will require substantial theoretical and numerical efforts: definitions, assumptions, theorems, proofs and supporting experiments. We leave all this stuff to a follow-up research.

Concluding remarks. Please reply to our post and inform us if the clarifications provided adequately address your concerns regarding our work. We are more than willing to discuss any remaining points during the discussion phase.

Dear reviewer, thank you for spending time reviewing our paper.

(1) Trivial theoretical results. "Theorem 2 appears to follow from the definitions".

At first, we highlight that the key theoretical result of our work is Theorem 1. From our point of view, it is not a trivial result. Secondly, we agree that proving Theorem 2 is not a difficult task. At the same time, we argue that Theorem 2 is important and worth presenting. It states that the gap between optimal value of (dual) training objective $\mathcal{L}$ and an actual value directly bounds the quality of the recovered barycentric plans $\pi^{f_k}$. It leads to the straightforward conclusion that the closer is our optimized objective $\mathcal{L}$ to the optimal one, the more accurate are the recovered barycentric plans $\pi^{f_k}$. Combined with Theorem 1, this result forms the foundation of our approach.

(2) Convergence speed/rates discussions. "[...] more discussion on the convergence behavior of the proposed algorithm say convergence rates under simple settings but that is lacking in the paper."

We thank the reviewer for this interesting question. Let us formalize the problem a bit. In Section $\S$ 4.3 we introduced the set of perfect potentials $(\widehat{f_1}, \dots \widehat{f_K})$ which solve the entropic barycenter problem given limited number of empirical samples and limited class of available potentials. If we correctly understand, your question is about how fast the optimized potentials $(\widehat{f_{1,opt}} , \dots \widehat{f_{K,opt}})$ converge to the perfect ones. Let us call the discrepancy between optimized and perfect potentials the optimization error. There are multiple sources of optimization error and there are multiple factors which affect its dynamics: inexact sampling due to MCMC and stochastic gradient ascent (Eq. 12). The analysis of these quantities is a completely different domain in Machine Learning and is out of scope of our work. As with most generative modeling research, we do not attempt to analyze optimization errors.

(3) "Could you answer how relevant Theorems 4 and 5 are towards the understanding of the proposed algorithm?"

Continuous the previous answer, Theorem 4 and 5 provide a statistical/approximation analysis for our learning objective assuming there is no optimization error (as it is usually done in statistical ML).

The theorems tell us that our empirical learning objective is valid from the statistical perspective (error could be bounded by the well-celebrated Rademacher complexity) and the approximation perspective (neural networks can achieve arbitrary small error). Obtaining more precise bounds requires the detailed analysis of weak $C$-transform's statistical properties. This requires significant work for establishing particular theorems, assumptions, definitions. We leave this for the future work.

(4) In addition, from my knowledge of the literature, it does not appear to be common to use neural networks for parameterization of the potentials (correct me if I am wrong)

Please note that we consider the continuous OT barycenter problem (see Section 2.1) but not discrete. Here the dual variable is a (continuous) function, not just a finite-length vector. Hence, using the neural-network parameterization for dual potentials is common [1,2]. It not only allows to do the out of sample estimation (which is tricky to do in the discrete case) but also improves the scalability of solvers.

(5) How much computational overhead this will incur on the computation, especially if K is large?

In accordance with the proposed algorithm on page 6 (see Section 4.2), one can see that the time of the computation will grow linearly. At the same time, we note that some parts of our algorithm could be easily run in parallel. In particular, the MCMC sampling procedures for each $k \in \{1, 2, \dots, K\}$ and computing the $k$-th components of our optimized loss $\widehat{L}_k$ are completely independent. This may significantly improve the computation rates.