A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods

Summary

This paper provides a theoretical framework based on generalized variational inference [1] and Wasserstein gradient flows (WGF) for analyzing deep ensemble methods and their regularized versions. The authors demonstrate that deep ensembles and other variational Bayesian methods can be cast as instances of an infinite dimensional variational inference problem and the WGF of different instantiations of a free energy functional. The authors additionally use their theoretical framework to derive a new algorithm for generating samples from a target distribution. [1] Knoblauch, Jeremias, Jack Jewson, and Theodoros Damoulas. "An optimization-centric view on Bayes’ rule: Reviewing and generalizing variational inference." Journal of Machine Learning Research 23.132 (2022): 1-109.

Strengths

The paper is well-organized and written, and the benefits of the unifying theoretical framework are compelling. The discussion of how deep ensemble methods can be viewed through the lens of WGF and the use of this lens to prove theoretical guarantees on the limiting behavior of particle estimations is useful and insightful. This work also holds the promise of deriving new algorithms, as demonstrated by the deep repulsive Langevin ensembles presented in Section 4.3.

Weaknesses

#### **Experiments section is difficult to follow** While the details corresponding to the various figures in Section 5 are fully provided in Appendix G, this section is currently difficult to follow as a stand-alone section in the main paper. Without (even high level) details on the experimental setup and the general motivation of each experiment it is difficult to dive right into the results and Figures as they are currently presented. I recommend moving some details from Appendix G into the main text and providing the context for each experiment before diving into the results. --- #### **Motivation for convexification is unclear** While the authors prove that convexity of the infinite-dimensional variational form of the learning objective guarantees uniqueness of a minimizer, this is somewhat disconnected from the presented goal of optimizing $\ell(\theta)$ via probabilistic lifting. For example, in footnote 1 on page 2, the authors argue that the unregularized variational objective has non-unique optimum. However, the local optima all have equivalent values of the objective and are simply weighted averages of equivalent optima of $\ell$, hence it is not clear why uniqueness is a desiderata here. Additionally, Figure 2 in Section 5 demonstrates how deep ensembles (DE) do not converge to $Q^*$. However, although deep Langevin ensembles (DLE) and deep repulsive Langevin ensembles (DRLE) provable converge $Q^*_{DLE}$ and $Q^*_{DRLE}$, respectively, these optimal distributions are also not equal to $Q^*$. Hence a clearer exposition as to why regularized optima are preferred to the unregularized ones is needed. --- #### **Motivation for DRLE is lacking** While DRE / DRLE is indeed interesting as a new algorithm that can be derived from the presented theoretical framework, it would be great if the authors also provided some intuition / motivation as to why MMD is perhaps a better suited divergence regularizer than KL.

Questions

In Lines 188-191, the authors state that > In theory, the PDE in (4) provides us with a direct way of implementing infinite-dimensional gradient descent for (1): simply follow the WGF. In practice however, this is impossible: **numerical solutions to PDEs become computationally infeasible for the high-dimensional parameter spaces which are common in deep learning applications.** I am unclear what is meant by this. Why is the high-dimensionality of deep learning parameterizations relevant here? Isn’t this problem simply impossible since it involves an infinite-dimensional space?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Authors can potentially elaborate on the future directions of this work, specifically around analyzing the approximation errors of approximating WGF with finite number of particles over a finite time horizon.

Summary

The paper established theoretical connections between ensembling an old and established method of deriving uncertainty estimates They use theory from iteraction of particles in a thermodynamic system to generalise and connect seemingly different ways of ensembling and Variational Bayes(Inference) methods. This is done by formulating the original non-convex optimization problem ubiqitious in ML and stats as infinite dimensional convex optimization problem in the space of probability measures. The addition of a regularization quantity ensures the strict convexity of the problem and by choosing different forms of this quantity lead to derivation of various inference algorithms.

Strengths

1. The paper is well written, theory heavy and addresses important topic of deep ensembling and its connection with variational Bayes methods 2. I am not so good with theory, but the theorems and equations looked ok to me without obvious mistakes. 3. Although this is a theory paper, the theoretical claims are well supported by the experiments and where they are not the authors they explain it well. 4. The distinction between IDGVI and FDGVI is well drawn out and explained. Also the limitations with FDGVI that the approximation family is limited by construction serves as a motivation for using IDGVI methods.

Weaknesses

1. There is a lot of content that has been compressed in 9 pages which can be challenging for a reader, and a journal might have been more appropriate for this work.

Questions

1. in practice, deep networks use stochastic gradients, does the theory hold for stochastic gradients, esp. as Section 2 explicitly uses gradients for motivation. 2. For case of FDGVI, the authors do not consider normalizing flows or SIVI as methods to overcome the limited capacity of approximating family problem when comparing it with IDGVI.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The limitations or practical challenges with the derived inference algorithms can be addressed. How practical is the result from Theorem 2, is it something that will only work asymptotically or will this work practically and if so how fast ?

Summary

The authors propose to unify existing theory on Bayesian (variational) inference (VI) by addressing a generalized objective, which is obtained from standard parameterized loss minimization by “probabilistic lifting” (re-casting in a space of probability measures over the parameter) and “convexification” (ensuring the existence of a global minimizer by regularization), with infinite-dimensional gradient flows in 2-Wasserstein space. A general recipe is provided to implement such a Wasserstein gradient flow (WGF) via an energy objective and a system of interacting particles. In the key contribution of the paper, the authors study WGF with different types of regularization–most notably, the unregularized version corresponding to deep ensembles (DE). It is shown that DE do not conduct a Bayesian learning procedure and systematically fail to generate samples from the optimal distribution, yet perform competitively thanks to the flexibility of the infinite-dimensional inference they realize (as opposed to, e.g., classical parametric VI).

Strengths

* [S1] **Unifying framework**. After much discussion in the past few years, the authors are–to the best of my knowledge–the first to establish a comprehensive theory that encompasses (finite-dimensional) VI and DE. * [S2] **Clarity**. Despite the rather abstract subject, the authors present a coherent and easy-to-follow sequence of arguments. Complexity is strictly limited to the necessary extent. * [S3] **Rigor**. Mathematical concepts and notation are sound. Extensive proofs and/or references to prior work underline every claim (though I did not check every proof in detail).

Weaknesses

* [W1] **Analysis of DE behavior** (see Questions) * It is not entirely clear if the paper studies arbitrary variants of DE or only a very narrowly defined version (see Q5). * The authors conjecture that the number of samples being vastly smaller than the number of local minima is responsible for D(R)LE not outperforming DE consistently and point to Fig. 4. This evidence seems rather anecdotal and could benefit from a more detailed investigation. Also see Q7--Q8. * [W2] **Omissions in notation**. While the notation is consistent and comprehensible overall, the authors tend to omit integration domains, objects of differentiation etc. (e.g., Eq. 1, l. 150, l. 177, l. 179). With the shifting of integration spaces and various gradients involved, it would be helpful to be as explicit as possible in this regard.

Questions

* [Q1] Eq. 6: What does the index $j$ relate to? * [Q2] l. 218: Are there any convergence results in the respective limits of $T$ and $N_E$? * [Q3] l. 223: The experiments are promised to confirm small approximation errors due to finite samples and time, in particular in comparison to finite-dimensional methods. Where, exactly, do I find evidence for this claim? * [Q4] l. 236: Just for the sake of clarity, is $\theta^\prime_n(t) = - \nabla \ell(\theta_n(t))$ equivalent to $d \theta_n(t) - \nabla V(\theta_n(t)) dt$? * [Q5] l. 237: Do I understand correctly that you interpret DE as training with no regularization whatsoever (weight decay, batch normalization etc. – let alone variations with weight sharing and the like)? I doubt that many researchers actually apply such a decidedly naive approach. * [Q6] l. 253: Can DE implementing infinite-dimensional GD be understood as taking a non-parametric/functional approach as to what the distribution of the generated samples looks like (as opposed to, e.g., mean-field VI with a finite parameter vector)? * [Q7] Fig. 2: Is there an explanation why DE exhibits this precise 50/50 spread of the probability mass? * [Q8] Fig. 4: I’m not sure I understand the key message here. Why do the particles end up in the same modes despite different $Q^\ast$? What would be the expected behavior? — Minor remarks * l. 53: Redundant “i” in “probability” * l. 72: Shouldn’t $D$ provide the mapping $(Q, P) \mapsto D(Q, P)$ according to the stated domain? * l. 170: I feel enough space can be freed to include the definition of the 2-Wasserstein metric, it seems an odd choice to omit this arguably relevant information. * l. 292: Is there a $d$ missing in front of $Q(\theta)$ in the first double integral? * l. 297: $L^{FE}$ with capitalized superscript ($L^{fe}$ otherwise). * l. 329: White space after “minimisers” * l. 330: Remove either “which” or “that” * l. 337: “matters” * Fig. 3 (caption): White space in “FD-GVI” * Table 1: I would recommend removing the “Boston” dataset due to its racism issues. Also: “methods… outperform” or “method… outperforms”. * l. 355: “lens”

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Given that the main contribution is a unifying framework for existing theories, this point doesn’t apply as usual. However, the authors should state more clearly that the evidence shown in Section 5 for findings in Section 4 is quite limited.

Summary

To improve the accuracy of the uncertainty quantification, the authors aim to provide a mathematically rigorous link between Bayesian inference, Variational Bayes methods and ensemble methods. In this work, methods s.a. variational inference, Langevin sampling and deep ensembles can be seen as particular cases of an infinite-dimensional regularised optimization problem formulated via Wasserstein gradient flows. They also provide a novel inference algorithm based on MMD and gradient descent in infinite dimensions plus regularisation. The procedure takes place by reframing the usual finite-dimensional loss function problem into an infinite-dimensional one. This is done rewriting the original optimization problem using an infinite-dimensional problem over the set of probability measures $\mathcal{P}(R^J)$ and introducing a strictly convex regulariser to induce a unique global solution. This solution is assumed to not be too different from the solution to the original problem, which is controlled by a reference measure $P$. This leads to an interpretation of many different inference setups as particular cases of the optimization problem proposed. Therefore, this approach results as a combination of the proposals in Knoblauch et al. (2022) and Ambrosio et al. (2005).

Strengths

* Good idea, could be interesting to the community were it proven in some other contexts. * The formulation is clear and elegant thanks to the gradient flows and the usage of Wasserstein space. The usage of the thremodynamical formulation of free energy is very attractive as well.

Weaknesses

* Altough the proposal is interesting and elegant, I think the experimental part of the paper does not provide enough evidence of the benefits related to this framework change. Results such as those present in Figure 3 could, in principle, be rivaled by previous methods s.a. [1] and [4], neither of which are discussed here. The authors maybe could provide a stronger motivation in this regard, and maybe try to encompass these other methods inside their framework. * Some literature relevant to the topic at hand seems to be missing from the discussion, or at least should be discussed more thoroughly: * Regarding the definition of infinite-dimensional GVI methods, I consider that other methods based on samples are left out and should be considered, such as [1,2,3]. These works may seem specially relevant due to the interest in implicitly-defined target $Q^*$, and in particular those that make use of the function-space formulation s.a. [1] or [4]. * I think finite-dimensional GVI methods are misrepresented as they can be much more expressive than the selection made in Section 2.2 may lead to believe. I consider that this point should be addressed, and the discussion must be readjusted accordingly in order to highlight the benefits of the proposed approach without relying on this fact. As examples of this matter, please see references [4,5,6]. * The writing can be generally improved, since the paper can be at times hard to follow. This is just a consequence of the amount of information provided, which is a positive point, although in sections 3 and 4 could be polished further. * (minor) The presentation could be improved, for example, by convering images to formulas s.a. in Figure 1 or the layout on the final page. * (minor) Since Wasserstein spaces are such a crucial point of this, I would suggest devoting a bit more time to explain the basics of the concept in the main text itself and not fully depend on the sources. (_References included in the "**Limitations**" section_)

Questions

* Can this framework be used to describe methods that obtain an a-posteriori approximation of the predictive distribution via Laplace approximation or similar methods? As examples, please see [7,8] * Since now inference is conducted without the guarantees provided by the Bayesian method, how should the distributions obtained be interpreted or used? * I think other possible interesting regularisation choices would be Renyi divergences and also any proper scoring rule, defined in [9], which may solve issues related to the MMD and KL divergence. (_References included in the **"Limitations"** section_)

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

* Since the Bayesian framework is abandoned, I fear there are no guarantees about the properties for the distributions obtained in the same sense as with Bayesian inference. Although can be somewhat justified by results, a lot more work is needed in this regard in methods that rely on this extensions (which is a problem for this paper, although definitely not exclusive to it). * The paper is centred on theoretical developments, and as such, the theoretical discussion and argumentation is really interesting. However, and although it is not the core of the paper, the experimental phase leaves a lot to be desired in terms of justifying why this formulation change is needed. --- **References**: [1] Rodrı́guez-Santana, S., Zaldivar, B., & Hernandez-Lobato, D. (2022, June). Function-space Inference with Sparse Implicit Processes. In International Conference on Machine Learning (pp. 18723-18740). PMLR. [2] Mescheder, Lars, Sebastian Nowozin, and Andreas Geiger. "Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks." International Conference on Machine Learning. PMLR, 2017. [3] Santana, S. R., & Hernández-Lobato, D. (2022). Adversarial α-divergence minimization for Bayesian approximate inference. Neurocomputing, 471, 260-274. [4] Ma, C., Li, Y., and Hernández-Lobato, J. M. (2019). “Variational implicit processes”. In: International Conference on Machine Learning, pp. 4222–4233. [5] Sun, S., Zhang, G., Shi, J., and Grosse, R. (2019). “Functional variational Bayesian neural networks”. In: International Conference on Learning Representations. [6] Ma, C., & Hernández-Lobato, J. M. (2021). Functional variational inference based on stochastic process generators. Advances in Neural Information Processing Systems, 34, 21795-21807. [7] Deng, Z., Zhou, F., & Zhu, J. (2022). Accelerated Linearized Laplace Approximation for Bayesian Deep Learning. Advances in Neural Information Processing Systems, 35, 2695-2708. [8] Antorán, J., Janz, D., Allingham, J. U., Daxberger, E., Barbano, R. R., Nalisnick, E., & Hernández-Lobato, J. M. (2022, June). Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning (pp. 796-821). PMLR. [9] Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477), 359-378.

Summary

The paper offers a viewpoint on deep ensembles as a (unregularized) Wasserstein gradient flow in the space of probability measures. This viewpoint enables new algorithms for deep ensembles (Langevin and repulsive via MMD), which are evaluated on some small datasets.

Strengths

1) The paper is technically sound and well-written. Overall it was easy to follow. 2) While many similar ideas have been floating around in the literature, the precise presented view on deep ensembles seems novel, and I found Theorem 1 to be interesting. 3) While experiments on larger neural networks are missing, the effect of the proposed algorithms is clearly demonstrated in some controlled experiments and small data sets.

Weaknesses

1) Perhaps the main weakness of the paper is the lack of a comparison of the new methods on large neural networks. 2) Many of the introduced tools (convexification via probabilistic lifting, Bayes with general divergence function, Wasserstein flows, etc.) are well-known. But I believe Theorem 1 and Theorem 2 offer some new insights (in case they are really correct, see Questions).

Questions

1) In section 2, it is written that one lifts the problem to a more "challenging space" -- but I would instead say that this space is much simpler. The problem suddenly has a closed-form solution (Gibbs measure) that can be written down, one has convexity, etc. 2) The claim that the "infinite-dimensional regularised optimisation problem over the space of probability measures first introduced in Knoblauch et al. (2022)" seems rather quite bold -- the optimization problem (1) is a very fundamental one -- as discussed later in the paper (Section 2.1) there are many references. So perhaps this sentence in the introduction should be reshaped a bit? 3) In Theorem 1, is it really local minima or could it also be saddle-points? Couldn't a gradient flow in a nonconvex objective also get stuck at a saddle-point or local-maximum (when initialized at the maximum)? Imagine a landscape where we have a very large flat local maximum which has non-zero measure under that initial distribution. Is this somehow excluded in the assumptions?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

All limitations are addressed.

Thank you

Thank you for the thorough response. I do not have any additional comments or questions at this time.

Brief response to the rebuttal

I want to thank the authors for their insightful responses and detailed comments on the reviews, including mine. After reading the rebuttals and going over parts of the article again, I'm really happy with what's been presented. I'm now even more convinced that we should accept this submission and I will update my review to reflect this. I consider this submission to be an interesting piece of work with important implications for future research. Thanks again for the good work!

Response to the rebuttal

Thanks to the authors for replying to my questions. I am quite satisfied with their detailed comments to my questions and other reviewers' questions. I recommend a strong acceptance.