A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods
Authors: Veit David Wild
Sahra Ghalebikesabi
Dino Sejdinovic
Jeremias Knoblauch
Venue: NeurIPS 2023 oral (2023)
Keywords: Wasserstein gradient flow
generalised variational inference
deep ensembles
Bayesian deep learning
variational Bayes
TL;DR: We establish the first unified theory connecting Bayesian, variational Bayesian, and ensemble methods for deep learning by leveraging Wasserstein Gradient Flows.

Abstract: We establish the first mathematically rigorous link between Bayesian, variational Bayesian, and ensemble methods. A key step towards this it to reformulate the non-convex optimisation problem typically encountered in deep learning as a convex optimisation in the space of probability measures. On a technical level, our contribution amounts to studying generalised variational inference through the lense of Wasserstein gradient flows. The result is a unified theory of various seemingly disconnected approaches that are commonly used for uncertainty quantification in deep learning---including deep ensembles and (variational) Bayesian methods. This offers a fresh perspective on the reasons behind the success of deep ensembles over procedures based on parameterised variational inference, and allows the derivation of new ensembling schemes with convergence guarantees. We showcase this by proposing a family of interacting deep ensembles with direct parallels to the interactions of particle systems in thermodynamics, and use our theory to prove the convergence of these algorithms to a well-defined global minimiser on the space of probability measures.

A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods
Veit D. Wild University of Oxford
Sahra Ghalebikesabi University of Oxford
Dino Sejdinovic University of Adelaide
Jeremias Knoblauch
University College London
Abstract
We establish the first mathematically rigorous link between Bayesian, variational Bayesian, and ensemble methods. A key step towards this is reformulate the non-convex optimisation problem typically encountered in deep learning as a convex optimisation in the space of probability measures. On technical level, our contribution amounts to studying generalised variational inference through the lens of Wasserstein gradient flows. The result is unified theory of various seemingly disconnected approaches that are commonly used for uncertainty quan- deep ensembles and (variational) Bayesian
tification in deep methods. This offers a fresh perspective on the reasons behind the success of deep ensembles over procedures based on standard variational inference, and allows the derivation of new ensembling schemes with convergence guarantees. We show- case this by proposing a family of interacting deep ensembles with direct parallels to the interactions of particle systems in thermodynamies, and use our theory to prove the convergence of these algorithms to a well-defined global minimiser on the space of probability measures.
Introduction
A major challenge in modern deep learning is the accurate quantification of uncertainty. To develop trustworthy AI systems, it will be crucial for them to recognize their own limitations and to convey the inherent uncertainty in the predictions. Many different approaches have been suggested for this. In variational inference (VI), a prior distribution for the weights and biases in the neural network is assigned and the best approximation to the Bayes posterior is selected from class of parameterised distributions (Graves) (011]) Blundell et al. 2013 Gal and Ghahramani] 2011) Louizos and Welling,
2017). An alternative approach is (approximately) generate samples from the Bayes posterior via Monte Carlo methods (Welling and Teh 2011 Neal 2012). Deep ensembles are another approach, and rely train-and-repeat heuristic quantify uncertainty ([aakshminrrayanan et al. 2017).
Much ink has been spilled over whether one can see deep ensembles as a Bayesian procedure (Wil- son 2020; Izmailov al. 2021: D' Angelo and Fortuin 2021) and over how these seemingly dif- ferent methods might relate to each other. Building on this discussion, we shed further light on the connections between Bayesian inference and deep ensemble techniques by taking different vantage point. In particular, we show that methods as different as variational inference, Langevin sampling (Ermak 1975), and deep ensembles can be derived from well-studied generally infinite- dimensional regularised optimisation problem over the space of probability measures (see e.g. Guede and Shawe-Taylor 2019 Knoblauch et al. 2022). As result, find that the differences between
37th Conference on Neural Information Processing Systems (NeurIPS 2023).
veit. wild @ destats.ox.ac.uu

---

these algorithms map directly onto different choices regarding this optimisation problem. Key dif- ferences between the algorithms boil down to different choices of regularisers, and whether they im- plement
or infiniteddimensional gradient descent. Here, finite-dimensional gra- dient descent corresponds to parameterised VI schemes, whilst the infniie-dimensiooall case maps onto ensemble methods.
The contribution of this paper is new theory that generates insights into existing algorithms and provides links between them that are mathematically rigorous, unexpected, and useful. On a techni- cal level, our innovation consists in analysing them as algorithms that target an optimisation problem in the space of probability measures through the use of powerful technical device: the Wasser- stein gradient flow (see e.g. Ammrosio et al. 2005). While the theory is this paper's main concern, its potentially substantial methodological payoff is demonstrated through the derivation of new inference algorithm based on gradient descent in infinite dimensions and regularisation with the maximum mean discrepancy. We use our theory to show that this algorithm- ensembles-is derived from a strictly convex objective defined over the space of probability mea- sures. Thus, it targets unique minimum, and is capable of producing samples from this global minimiser in the infinite particle and time horizon limit. This makes the algorithm provably con- vergent; and we hope that it can help plant the seeds for renewed innovations in theory-inspired algorithms for (Bayesian) deep learning.
standard deep
The paper proceeds as follows: Section 2 discusses the advantages of lifting losses defined on Eu- clidean spaces into the space of probability measures through generalised variational objective. Section [] introduces the notion of Wasserstein gradient flows, while Section 4 links them to the aforementioned objective and explains how they can be used to construct algorithms that bridge Bayesian and ensemble methods. The paper concludes with Section [] where the findings of the paper are illustrated numerically.
Convexification through probabilistic lifting
One of the most technically challenging aspects of contemporary machine learning theory is that the losses f: R + R we wish to minimise are often highly non-convex, For instance, one could wish
to minimise E()) 1/2 Σn= (yn- foo--f(rrn) fo(#x) 2 where {(zn. Vo.) bn)) J is set of paired observations and fe a neural network with parameters 0. While deep learning has shown that non-convexity is often a negligible practical concern, it makes it near-impossible to prove many basic theoretical results that a good learning theory is concerned with, as E has many local (or global) minima (Fort eet al., 2019 Wilson and Izmailov 2020). We reintroduce convexity by lifting the problem onto a computationally more challenging space. In this sense, the price we pay for the convenience of convexity is the transformation of finite-dimensional problem into an nnfniteedimennionall one, which is numerically more difficult to tackle. Figure illustrates our approach:
/ ((0))QQ0))
{ (10))Q10)D((000
min ((0) AE
min QEP(R-)
min @EP(R))
Step 1: probabilistic lifting
Step 2: convesification through regularisation
Figure 1: Illustration of convexification through probabilistic lifting.
First, we transform a non-convex optimisation mingee ((0) into an infinitediimessional optimisa- tion over the set of probability measures P(R'), yielding mingep(r)) J (()) As integral, this objective is linear in Q. However, linear functions are not strictly convex. We therefore need to add a strictly convex regulariser ensure uniqueness of the minimiser We prove in Appendix A that indeed-for a regulariser D : P(R') x P(R')) + [0, oo) such that (Q, + D(Q, P) is strictly convex and P P(R") fixed measure- existence and uniqueness of global minimiser can be guaranteed. Given scaling constant À > 0, we can now put everything together to obtain
TTo illustrate why this is necessary, assume that there are Aa and 0g in Rd so that E(0x) E(OD) minace ((0). Due to linearity, each measure Q, E P(R'), t E [0, 1]. defined as Q == (1 t)60,, + tōe provides one (of the infinitely many) global minima in the set arg minge [(() dQ(().

---

the loss L and the unique minimiser Q" as
L(Q) : f AD(Q)
Q == argmin L(Q). QEP(3))
(I)
Throughout, whenever Q and Q" have an associated Lebesgue density, we write them as and g² Moreover, all measures, densities, and integrals will be defined on the parameter space RJ of 0. Similarly, the gradient operator V will exclusively denote differentiation with respect to 0 E RR",
2.1 One objective with many interpretations
In the current paper, our sole focus lies on resolving the difficulties associated with non-convex optimisation of on Euclidean spaces. Through probabilistic lifting and convexification, we can identify a unique minimiser Q" in the new space, which minimises the 0-averaged loss 0 + (()) without deviating too drastically from some reference measure P. In this sense, Q* summarises the quality of all (local and global) minimisers of {(() by assigning them corresponding weight. The choices for X, D and P determine the trade-off between the initial loss & and reference measure P and therefore the weights we assign different solutions.
Yet, (T) is not a new problem form: it has various interpretations, depending on the choices for D, X, and the framework of analysis (see e.g. Knoblauch et al., 2022, for discussion). For ex- ample, if {(0) is a negative log likelihood,
is the Kullback-Leible divergence (KI.), and Q' is the standard Bayesian posterior, and P is the Bayesian prior. This interpretation of P as prior carries over to generalised Bayesian methods, which we can choose {()) to be any loss, D to be any divergence on (( ). and À to regulate how fast we learn from data (see e.g. Bissiri
1,
ct al. 2016; Jewson et al.) 2018: Knoblauch et al. 2018; Miller and Dunson 2019 Knoblauch
2019 Alquier 202 la Husain and Knoblauch 2022 Matsubara et al. 2022 Wild et al. 2022] Wu
and Martin [023] Altamirano et al. (023). In essence, the core justification for these general-
isations that the very assumptions justifying application of Bayes' Rule violated in modern machine learning. In practical terms, this results in a view of Bayes' posteriors as one-of of many possible-measure-eee veluued estimators Q' of the form in (I. Once this vantage point is taken, it is not clear why one should be limited to using only one particular type of loss and regulariser for every possible problem. Seeking parallel with optimisation on Euclidean domains, one may then compare the orthodox Bayesian view with the insistence on only using quadratic regularisation for any problem. While it is beyond the scope of this paper to cover these arguments in depth, we refer the interested reader to Knoblauch et al. (2022).
A second line of research featuring objectives as in ( are PAC-Bayes methods, whose aim is to construct generalisation bounds (see c.g. [Shawe-Tayll or and Williamson 1997 McAllester 1999b a Grünwald 201T Here, l is general loss, but 1 only has the interpretation of some reference measure that helps us measure the complexity of our hypotheses via Q + AD(Q, P) (Guedj and Shawe-Taylon Taylor 2019 AAlquerr (2021b). Classic PAC-Bayesian bounds set D to be KL, but there has been a recent push for different complexity measures (Alquier and Guedj] 2018 18 Bégin et al. 2016
Haddouche and Guedj] 2023).
2.2 Generalised variational inference (GVI) in finite and infinite dimensions
In line with the terminology coined in Knoblauch et al. (2022), we refer to any algorithm aimed at solving ① generalised variational inference (GVI) method. Broadly speaking, there are two ways one could design such algorithms: in finite or infinite dimensions. Finite-dimensional GVI: This is the original approach advocated for in Knoblauch et al. (2022): instead of trying to compute Q', approximate it by solving Q.- = arg m.n L(Q) for a set of measures @ == {Q v E r} C P(R') parameterised by a parameter v E 1 C IR'. To find Q.. one now simply performs dinee-diiensional) gradient descent with respect the function v + L(Qz). For the special case where P is Bayesian prior, > >= 1, {()) is negative log likelihood parametrised by 0, and D = KL, this recovers the well-knnwn standard VI algorithm. To the best of our knowledge, all methods that refer to themselves as VI or GVI in the context of deep learning are based on this approach (see e.g. Graves 2011 Blundell et al. 2015 Louizos and Welling 2017 Wild et al. 2022). Since procedures of this type solve a finite-dimensional version of (1), we refer to them as finite-dimensional GVI (FD-GVT) methods throughout the paper. While such algorithms can perform well, they have some obvious theoretical problems: First of all, the finite-dimensional approach typically forces us to
3

---

choose Q and P to be simple distributions such as Gaussians to ensure that L(Q,) is tractable function of This often results in a Q that is unlikely to contain a good approximation to Q"; raising doubt if Qu- can approximate Q" in any meaningful sense. Secondly, even if Q + [[)) is strictly convex on P(TR.)) the parameterised objective v E r -> T(Qx)) is usually not. Hence, there is no guarantee that gradient descent leads us to Qu. This point also applies to very expressive variational families (Rezende and Mohamed| 2015 Mescheder et al., 2017) which may be sufficiently rich that Q' E Q. but whose optimisation problem v E T + L(Q,) is typically non-convex and hard to solve, so that no guarantee for finding Q" can be provided. While this does not necessarily make FD-GVI impractical, it does make it exceedingly difficult to provide rigorous theoretical analysis outside of narrowly defined settings.
FD-GVI in function space: A collection of approaches formulated as dimenie-dimeniional prob- lems are GVI methods on an ddinite-dimensional function space (Ma et al. 2019 Sun et al. 2018 Ma and Hernández-L is often convex in function space. In practice however, the variational stochastic process still re- quires parameterization to be computationally feasible-and and in this sense, function space methods are FD-GVI approaches. The resulting objectives require good approximation of the functional KL-divergence (which is often challenging), and lead to typically highly non-convex variational optimization problem in the parameterised space.
2021 Rodrriuezz-Santana et al. 2022 Wild et al. (022). Here, the loss
Infiniteddimensional GVI: Instead of minimising the (non-conven, problem V L(Qu). we want to exploit the convex structure of Q + L(Q). Of course, a priori it not even clear how to compute the gradient for function Q / IQQ) defined on an aifinite-dimensiona nonlinear space such as ((RR).. However, in the next part of this paper we will discuss that it is possible to implement a gradient descent in infinite dimensions by using gradient flows on a metric space of probability measures (Ambrosio et al., 2005). More specifically, one can solve the optimisation problem ( by following the curve of steepest descent in the 2-Wasserstein space. As it turns out, this approach is not only theoretically sound, but also conceptually elegant: it unifies existing algorithms for uncertainty quantification in deep learning, and even allows us to derive new ones. We refer to algorithms based on some form of drinite-dimensione gradient descent as infinite-dimensional gradient descent methods have recently gained attention in the
GVI (ID-GVI). machine learning community. For existing methods of this kind, the goal is to generate samples from a target Q that has a known form (such as the Bayes posterior) by applying a gradient flow to Q E P(RR') + E(Q, @) where E(.,,) is discrepancy measure. Some methods apply the Wasserstein gradient flow (WGF) for different choices of E (Arbel et al. 2019 Korba et al. [00TT] Glaser et al. 2021). whilst other methods like Stein variational gradient descent (SVGD) (Liu and
Wang 2016 stay within the Bayesian paradigm (E KL, Q Bayes posterior) but use gradient
flow other than the WGF (Liu 2017). D'Angelo and Fortuin ( (2021) exploit the WGF in the standard Bayesian context and combine it with different gradient estimators (Li and Turner 2017 Shi al., 2018) to obtain repulsive deep ensembling schemes. Note that this is different from the repulsive approach we introduce later in this paper: Our repulsion term is the consequence of regulariser, not a gradient estimator. Since our focus is on tackling the problems associated with non-convex optimisation in Euclidean space. the approach we propose is inherently different from all of these existing methods: our target Q" is only implicitly defined via (I. and not known explicitly.
3 Gradient flows in finite and infinite dimensions
Before we can realise our ambition to solve ( with an ID-GVI scheme, we need to cover the rele- vant bases. To this end, we will discuss gradient flows in finite and infinite dimensions, and explain how they can be used to construct dient flow is the limit of a gradient descent whose step size goes to zero While the current section introduces this idea for the finite-dimensional case for ease of exposition, its use in constructing algorithms within the current paper will be for the infinite-dimensional case.
gradient descent schemes. In essence, gra-
Gradient descent finds local minima of losses E RJ_ - R by iteratively improving an initial guess 8 E ' through the update 0 := 0g ((xx), & E N, where 17 > 0 is a step-size and V denotes the gradient of l. For sufficiently small 17 > 0, this update can equivalently be written as
0x+1 = arg min {((0) + 2n
1
0 - 0x1
(2)
BER

---

Gradient flows formalise the following logic: for any fixed 1,, we can continuously interpolate the
corresponding gradient descent iterates {0x ) KENo To do this, we simply define a function f [0, oo) + IR as (*(1)= Ou/n for & E 1yN :- {0. 17. 21,,...). For i & n/oo we linearly interpolate As 7 - 0, the function 0 converges to differentiable function 0. [0, oo) + I called the gradient flow of f, because is characterised as solution to the ordinary differential equation (ODE)
(c) = - [t)) with initial condition 0_ (0) = t-- Intuitively, 0, (t) is continuous-time version of discrete-time gradient descent; and navigates through the loss landscape so that at time t, an infinitesimally small step in the direction of steepest descent is taken. Put differently: gradient descent is nothing but an Euler discretisation of the gradient flow ODE (see also Santambrogio 2017). The result is that for mathematical convenience, one often analyses discrete- gradient descent as though it were continuous gradient flow with the hope that for sufficienlly small f), the behaviour of both will essentially be the same.
Our results in the infiniteddimensional case follow this principle: we propose an algorithm based on discretisation, but use continuous gradient flows to guide the analysis. To this end, the next section generalises gradient flows to the nonlinear
setting.
3.1 Gradient flows in Wasserstein spaces
Let P2(R') be the space of probability measures with finite second moment equipped with the 2- 2-Wasserstein metric given
W2(P,Q)² = inf { 11-0012 :E
where C(P.Q) Q) C P(R x Ra) denotes the set of all probability measure on TR x RR such that A(A R
((A) and T (H B) = Q(B) for all A, C R (see also Chapter 6 of Villani
et al. 2009
Further, let L: (P())) (-oc, oo] be some functional-for example L in ( In
direct analogy to (2). we can improve upon an initial guess Qu E P2((1) ) by iteratively solving
Qk == arg min (B*)(
{L(Q) + (2(Q2QQQ Q+)²)
21
for k E No and small >> 0 (see Chapter 2 of Ambrosio et al. 2005 for details). Again, for 7) + O. an appropriate limit yields continuously indexed family of measures {Q(1))22) If L is sufficiently smooth and Qo (0) has Lebesgue density qu. the time evolution for the corresponding pdfs {g(t))12o is given by the partial differential equation (PDE)
Oig(t, 0) V (o(t, VW
with ((0,-) = go (Villani 2003 Section 9.1). Here v- Ei= 0,f, denotes the divergence
operator and W I.
[Q]: - RJ the Wasserstein gradient (WG) of T. at Q. For the purpose of this
paper, it is sufficient to think of the WG as a gradient of the first variation; i.c. Lw = V
where L'[Q]= IR + I is the first variation of L at Q (Villani et al. 2009 Exercise 10. The Wasserstein gradient flow (WGF) for L is then the solution q¹ to the PDE B If L is chosen as in ①I. our hope that the logic of dinen--immniio gradient descent carries over, and that lim, q (t. is in fact the density 4² corresponding to Q"
Following this reasoning, this paper applies the WGF for (1) obtain ID-GVI algorithm. In Section 4 and Appendices C-F, we formally show that the WGF indeed yields Q" (in the limit as +- oo) for a number of regularisers of practical interest.
3.2 Realising the Wasserstein gradient flow
In theory, the PDE in 3) could be solved numerically in order to implement the infinitediimensional gradient descent for (I. In practice however, this is impossible: numerical solutions to PDEs be- come computationally infeasible for the high-dimensional parameter spaces which are common in deep learning applications. Rather than trying first approximate the q solving (3) and then sam- pling from its limit in a second step, we will instead formulate equations which replicate how the
"This means (t) 9=11-9 (t sn) 0s, for & E [sn, (s + 1)m) and No.
5

---

samples from the solution to (3) evolve in time. This leads to tractable inference algorithms that can be implemented in high dimensions.
Given the goal of producing samples directly, we focus on a particular form of loss that well- studied in the context of thermodynamics (Santambrogio 2015 Chapter 7), and which recovers var- ious forms of the GVI problem in (I) (see Section 4p. In thermodynamics, Q E P2 (2 describes the distribution of particles located at specific points in R. The overall energy of collection of particles sampled from Q is decomposed into three parts: (i) the external potential ( which acts on each particle individually, (ii) the interaction energy k(0, 6²) describing pairwise interac- tions between particles, and (iii) an overall entropy of the system measuring how concentrated the distribution Q is. Taking these components together, we obtain the so called free energy
(,)( - (Q() + A
/ʃ (Q()))Q(1)))) + X2 / /logg(0)((0) ((0)a(0)) d0,
(4)
for Q E P2( with Lebesgue density A 2 0, A 2 0. Note that for X2 > 0 we implicitly assume that Q has density. Following Section 9.1 in Villani et al. (2009), its WG is
V V 10))(0) = VV( + A1 (VV)))0 dQ(p)) + V log (0))
where f E RR, and VIK denotes the gradient of к with respect to the first variable. We plug this into (3) to obtain the desired density evolution. Importantly, this time evolution has the exact form of nonlinear Fokker-Planck equation associated with a stochastic process of McKean- Vlasov type (see Appendix for details). Fortunately for us, it well-known that such processes can be approx- imated through interacting particles (Veretennikov 2006) generated by the following procedure:
Step 1:
Sample Ng E N particles 01 (0), IN. ee (0) independently from Qu E P.
Step 2: Evolve the particle fn by following the stochastic differential equation (SDE)
NE
(1) = (O.(())) +
(0),
A
Σ (()))(((())00000)) + /ooddd++///aa// 222dBa((), (5)
NE
for L......... NE. and (t), stochastically independent Brownian motions.
As Ng -> oo, the distribution of 0, (.) ON (t) evolves in f in the way as the sequence of densities g(t, solving (3. This means that we can implement infiniteddimension gradient descent by following the WGF and simulating trajectories for infinitely many interacting particles according to the above procedure. In practice, we can only simulate finitely many trajectories over a finite time horizon. This produces samples 0, (T), ONg (T) for NE N and T > O. Our intuition and Section 4 tell us that, as desired, the distribution of 0, (T),... ONE (T) will be close to the global minimiser of T.fo.
In the next section, we will use the above algorithm to construct an ID-GVI method producing samples approximately distributed according to Q' defined in (I). Since Ng and T are finite, and since we need to discretise (T3T), there will be an approximation error. Given this, how good are the samples produced by such methods? As shall demonstrate in Section 5 the approximation errors are small, and certainly should be expected to be much smaller than those of standard VI and other FD-GVI methods.
4 Optimisation in the space of probability measures
With the WGF on thermodynamic objectives in place, we can now finally show how yields ID-GVI algorithms to solve (T). We put particular focus on the analysis of the regulariser D; providing new perspectives on heuristics for uncertainty quantification in deep learning the process Specifically, we establish formal links explaining how they may (not) be understood yond that, we derive the WGF associated with regularisation using the maximum mean discrepancy, and provide theoretical analysis of its convergence properties.
Bayesian procedure. Be-

---

N), then it holds independently for 4.1 Unregularised probabilistic lifting: Deep ensembles
We start the analysis with the base case of an unregularised functional L(Q) = (()) corre- sponding to À = 0 in (T). This is also a special case of (4) with A x2 = 0. As X = 0. there is no interaction term, and all particles can be simulated independently from one another
0 (0) NN (0) ~ Qo. (t) = Ve(( n n= I NE.
This simple algorithm happens to coincide exactly with how deep ensembles (DEs) are constructed (see e.g. Laaksmiiarayaaaan et al., 2017). In other words: the simple heuristic of running gradient descent algorithm several times with random initialisations sampled from Qo is an approximation of the WGF for the unregularised probabilistic lifting of the loss function L.
Following the WGF in this case does not generally produce samples from global minimiser of L. Indeed, the fact that L generally does not even have unique global minimiser was the motivation for regularisation in ( Even if L had unique minimiser however, a DE would not find it. The result below proves this formally: unsurprisingly. deep ensembles simply sample the local minima of with a probability that depends on the domain of attraction and the initialisation distribution Qo- Theorem 1. If has countably many local minima {m: each n 1,. 1...., NE that
D0
On(() D
Σ
Qo(O:) 8m; =: Qoo
for + 00. Here 2, denotes convergence in distribution and O₂ {0 E RR, limm-+ (t) my and (0) = 0} denotes the domain of attraction for m, with respect to the gradient flow 0+.
A proof with technical assumptions- most importantly version of the famous Lojasiewicz inequality-is in Appendix C. Theorem | derives the limiting distribution for (T),,.... ONg (T), which shows that-unless all local minima are global minima-the WGF does not generate samples from global minimum of Q + L(Q) for the unregularised case À = 0. Note that the condi- tions of this result simplify the situation encountered in deep learning, where the set of minimisers would typically be uncountable (Liu et al. 2022). While one could derive a very similar result for the case of uncountable minimisers, this becomes notationally cumbersome and would obscure the main point of the Theorem- thhorem-thatt Qoe strongly depends on the initialisation Qu- Importantly, the dependence of Qoe on Qo remains true for all losses constructed via deep learning architec- tures. However, despite these theoretical shortcomings. DEs remain highly competitive in practice and typically beat FD-GVI methods like standard VT (Ovadia et al.) 2019 Fort et al. 2019). This is perhaps not surprising: DEs implement an anfnite-dimeesionall gradient descent, while FD-GVI methods are parametrically constrained. Perhaps more surprisingly. we observe in Section 5 that DEs can even easily compete with the more theoretically sound and regularised ID-GVI methods that will be discussed in Section 4.2 and 4.3 We study this phenomenon in Section [S] and find that it is a consequence of the fact that deep learning, Ng is small compared to the number of local minima (cf. Figure 4
4.2 Regularisation with the Kullback-Leible divergence: Deep Langevin ensembles
In Section 2 we argued for regularisation by D) to ensure a unique minimiser Q*. The Kullback-
Leibler divergence (D KL) is the canonical choice for (generalised) Bayesian and PAC- -ay-Bayn
methods (Bissiri et al. 2016 Knoblauch et al., 2022 Guedj and Shawe-Taylo Taylor 2019: Alquier
2021b). Now, Q" has known form: if has a pdf p. has an associated density given by ('()) exp(- x e(0) )p(O) (Knoblauch et al. 2022 Theorem 1).
Notice that the KL-regularised version of L in T can be rewritten in terms of the objective L in (4) by setting V(0)= = {(6) Alogp( A and A2 = X. Compared to the unregularised objective of the previous section (where X = 0), the external potential is now adjusted by -X -()), forcing Q' to allocate more mass in regions where p has high density. Beyond that, the presence of the negative entropy term has three effects: it ensures that the objective is strictly convex, that Q' is more spread out, and that it has a density qº Since \ = O. the corresponding particle method still does not have an interaction and is given as
01 (O),..., (0) Qu dt (t) = - (O. + n n= 1,.... Ng.
(6)
7

---

Clearly, this is just the Langevin SDE and we call this approach the deep Langevin ensemble (DLE). While the name may suggest that DLE equivalent to the unadjusted Langevin algorithm (ULAA) (Roberts and Tweedie| 1996), this is not so: for T discretisation steps 4 L2, ... Lr, DLE approximates measures using the end-points of Ng trajectories given by (trr) Ne 1- In contrast,
ULA would use a (sub)set of the samples (81(t) [t:)]]= generated from one single particle's 's trajectory. To analyse DLEs, we build on the Langevin dynamics literature: in Appendix D
we show that (l(()) 2x Q" as E + oc, independently for each = 1,... NE. Hence 0 (T), ...... ON (T) will for large T > 0 be approximately distributed according to Q-. Comparing DE and DLE in this light, we note several important key differences: Q* defined per (I) is unique, has the form of Gibbs measure, and can be sampled from using (6). In contrast, unregularised DE produces samples from Qx in Theorem M which is not the global minimiser. Specifically neither Q nor Q* for DEs correspond to the Bayes posterior. It is therefore not a Bayesian procedure in any commonly accepted sense of the word.
4.3 Regularisation with maximum mean discrepancy: Deep repulsive Langevin ensembles
Regularising with KL attractive because Q" has a known form. However, in our theory, there is no reason to restrict attention to a single type of regulariser: we introduced D to convexify our objective. It is therefore of theoretical and practical interest to see which algorithmic effects are induced by other regularisers. We illustrate this by first considering regularisation using the squared maximum-mes mean discrepancy (MMD) (see e.g. Gretton et al. 2012) only, and then a combination of MMD and KL.
For kernel i: : RR R + R. the squared MMD between measures Q and P is
MMD(Q, P)2 = /] k(0, ((0.8))(())QQ10)) - 2 /f k(0, ((10))()))))))(0
[f k(0, 0¹ (dP())))))
MMD measures the difference between sithln-sample similarity and across-sample similarity, so it is smaller when samples from P are similar to samples from Q. but also larger when samples within Q are similar to each other. This means that regularising ① with D = MMD² introduces interactions characterised precisely by the kernel k, and we can show this explicitly by rewriting of (T) into the form of I,fe in 4 other words, inclusion of MMD² makes particles repel each other, making it more likely that they fall into different (rather than the same) local minima. Writing the kernel mean embedding as ppp == x(0, 8') dP(('). we see that up to a constant not depending on Q, L(Q) = (Q) for ((0) = {() - (0). X - and 22 = O. While we can show that global minimiser Q" exists, and while we could produce particles using the algorithm of Section 3.2 we cannot guarantee that they are distributed according to Q" (see Appendix (F.. Essentially, this is because in certain situations, we cannot guarantee that Q" has a density for D = MMD*.
To remedy this problem, we additionally regularise with the KL: since KL P) = oc if P has a Lebesgue density but Q has not, this now guarantees that Q' has a density q² terms of (I. this means that D > MMD2 +X' KT. Adding regularisers like this has long tradition, and is usually done to combine the different strengths of various regularisers (see e.g. Zou and Hastic (005). Here, follow this logic: the KL ensures that Q" has a density, and the MMD makes particles repel each other. With this, we can rewrite L(Q) in terms of L up to a constant not depending on Q by taking À = 21 X' = X2. and V(0)- E(O) X1 (0) d2 log p(0). Using the same algorithmic blueprint as before, we evolve particles according to (5). As these partieles follow an augmented Langevin SDE that incorporates repulsive particle interactions via K, we call this method the deep repulsive Langevin ensemble (DRLE). We show in Theorem 2 (cf. Appendix E for details and assumptions) that DRLEs generate samples from the global minimiser Q' in the infinite particle and infinite time horizon limit.
Theorem 2. Let Qu.Ne (T) be the distribution of 0₂ (T) 1, 1,...,,N Ne. generated via Then Ne (T) D, Q" for each 1,..., .gg and as Ng xo, T x..
This is remarkable: we have constructed an algorithm that generates samples from the global min- imiser ev-nvven though a formal expression for what exactly Q* looks like is unknown! This
8

---

Loss function
DE
DLE
DRLE
Frequency
20
0.5
0.5
0
0.0
0.0
-255
0.0
2.5
-2.5
0.0
2.5
-2.5
0.0
2.5
-2.5
0.0
0
0
Figure 2: We generate Ng =00 particles from DE, DLE and DRLE. The theoretically optimal global minimisers Q' are depicted with dashed line strokes, and the generated samples are displayed via histograms. We use P = Qo NO 1) for DLE and DRLE. Notice that the optimal Q* differs slightly between DLE and DRLE.
demonstrates how impressively powerful the WGF is as tool to derive inference algorithms. Note that this is completely different from sampling methods employed for Bayesian methods, for which the form of Q" is typically known explicitly up to proportionality constant.
A notable shortcoming of Theorem 2 is its asymptotic nature. A more refined analysis could quantify how fast the convergence happens in terms of NE. T, the SDE's discretisation error, and maybe even the estimation errors due to sub-sampling of losses for constructing gradients. While the existing literature could be adapted to derive the speed of convergence for DRLE in T (Ambrosio et al, 2005 Section 11.2), this would require a strong convexity assumption on the potential V. which will not be satisfied for any applications in deep learning. This is perhaps unsurprising: even for the most thoroughly analysed algorithm in this literature no convergence rates have been derived that are applicable to the highly multi-modal target measures encountered in Bayesian deep learning (Wibisono) 2019 Chewi et al., [002). That being said, for the case of deep learning, FD-GVI approaches fail to provide even the most basic asymptotic
the Langevin
convergence guarantees. Thus, the fact that it is even possible for us to provide any asymptotic guarantees derived from realistic assumptions marks a significant improvement over the available theory for FD-GVI methods, by virtue of Theorem over DEs as well.
5 Experiments
Since the paper's ' primary focus is on theory, we use two experiments to reinforce some of the pre- dictions it makes in previous sections, and third experiment that shows why-in direct contradiction to a naive interpretation of the presented theory-it is typically difficult to beat simple DEs. More details about the conducted experiments can be found in Appendix G The code is available bttps://githue com/ wtt.coo/////o.cc sss.lehhalbbbiaeee kesabi
Global minimisers: Figure 2 illustrates the theory of Sections 4 and Appendices C-F: DLE and DRLE produce samples from their respective global minimisers, while DE produces a distribution which -in accordance with Theorem
! does not correspond to the global minimiser of Q + J ((0) (0) over ((R)) (which is given as Dirac measure located at 0 = -2).
FD-GVI vs ID-GVI: Figure 3 jilustraate two aspects. First, the effect of regularisation for DLE and DRLE is that particles spread out around the local minima. In comparison, DE particles fall directly into the local minima. Second, FD-GVI (with Gaussian parametric family) leads to qualitatively poorer approximations Q".. This is because the D0-GVI explore the whole space (2(RR)), whilst FD-GVI is limited to learning a unimodal Gaussian.
DEs D(R)EEs, and why finite NE matters: Table compares DE, DLE and DRLE on number of real world data sets, and finds a rather random distribution of which method performs best. This seems contradict our theory, and suggests there essentially no difference between regularised and unregularised ID-GVI. What explains the discrepancy? Essentially, it is the fact that Ng is not only finite, but much smaller than the number of minima found in the loss landscape of deep learning. In this setting, each particle moves into the neighbourhood of a well- separated single local minimum and typically never escapes, even for very large T. We illustrate this in Figure with a toy example. We choose a uniform prior P and initialisation Qu and the loss E(O) := sin (+),, E [-10007, 1000m), which has 2000 local minima. Correspondingly Q" will have many local modes for all methods. Note that VV is the same for all approaches since log and jup are constant. The difference between the methods boils down to repulsive and noise effects. However, these noise effects are not significant if each particle is stuck in a single mode: the particles will bounce around

---

DE DLE DRLE
FD-GVI
5
5
5
0o
Figure 3: We generate Næ 300 particles from DE, DLE, DRLE and FD-GVI with Gaussian parametrisation. The multimodal loss E is plotted in grey and the particles of the different methods layered on top. The prior in this example is flat, i.c. log p and uP constant. The initialisation Qo is standard Gaussian.
DE
DLE
Count
4
DRLE
0
2000
0 θ
2000
Figure 4: We generate Ng 20 samples from the three infinitediimensional gradient descent
procedures discussed in Section 4 The æ-axis shows the location of the particles after training. Since the same initialisation O, (0) is chosen for all methods, we observe that particles fall into the same local modes. Further, 16/20 particles are alone in their respective local modes and the location of the particles varies only very little between the different methods (which is why they are in the same bucket in the above histogram). See also Figure 5 in the Appendix for an alteration of Figure 3 with only 4 particles which emphasizes the same point.
their local modes, but not explore other parts of the space. This implies that they will not improve the approximation quality of Q". Note that this problem is a direct parallel to multi-modality- well-known problem for Markov Chain Monte Carlo methods (see e.g. [Sydd et al. 2022
Conclusion
In this paper, we used infiniteddimensional gradient descent via Wasserstein gradient flows (WGFs) (see e.g. Ambrosio et al. 2005) and the lens of generalised variational inference (GVI) (Knoblauch al., 2022) unify a collection of existing algorithms under common conceptual roof. Arguably, this reveals the WGF to be a powerful tool to analyse ensemble methods in deep learning and beyond. Our exposition offers a fresh perspective on these methodologies, and plants the seeds for new ensemble algorithms inspired by our theory. We illustrated this by deriving a new algorithm that includes a repulsion term, and use our theory to prove that ensembles produced by the algorithm converge to a global minimum. A number of experiments showed that the theory developed in the current paper is useful, and showed why the performance difference between simple deep ensembles and more intricate schemes may not be numerically discernible for loss landscapes with many local minima.
KIN8NM CONCRETE
ENERGY
NAVAL
POWER
PROTEIN
WINE
YACHT
DE DLE
0.33
6.10g0.3
2.83+0.2
0.40+ao
13.70+2.6
11.2242.2
14.65
2,20+0.4 1.64
13.25143
5.11
2.43+01
3.46124
13.87123
43.20
13.73
DRLE
0.464011
8.30
4.01
-3.044..2
23.21
48.80421
7.13
7.80
Table I: Table compares the average (Gaussian) negative log likelihood in the test set for the three ID-GVI methods on some UCI-regression data sets (Lichman) 2013). consistently outperforms any of the others.
We observe that no method
10

---

Acknowledgements
VW was supported by the scatchered scholarship and the EPSRC grant EP/W005859/1. SG was PhD student of the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1), and also received funding from the Oxford Radcliffe Scholarship and Novartis. JK was funded by EPSRC grant EP/W0058591.
References
Ahmed, N. and Ding, X. (1993). On invariant measures of nonlinear Markov processes. Journal of Applied Mathematics and Stochastic Analysis, 6(4):385-406.
Alquier, P. (2021a). Non-exponentially weighted aggregation: regret bounds for unbounded loss
functions. In International Conference on Machine Learning, pages 207-218. PMLR.
Alquier, (2021b). arXiv:2110.112166
User-friendly introduction to PAC-Bayes bounds.
arXiv preprint
Alquier, P. and Guedj, B. (2018). Simpler PAC-Bayesian bounds for hostile data, Machine Learning. 107(5):877-90..
Altamirano, M., Briol, F.-X,, and Knoblauch, J. (2023).
Robust and scalable Bayesian online
changepoint detection. arXiv preprint arXiv:2302.04759.
Ambrosio, L., Gigli, N., and Savaré, G. (2005). Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media.
Arbel, M., Korba, A., Salim, A., and Gretton, A. (2019). Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, 32.
Barbu, V. and Röckner, M. (2020). From nonlinear Fokker-Planck equations solutions of distri- bution dependent sde. arXiv preprint IrXiv:8808 10706.
Bégin, L., Germain, P., Laviolette, F., and Roy, J.-F. (2016). PAC-Bayesian bounds based on the Rényi divergence. In Artificial Intelligence and Statistics, pages 435-444-
Bissiri, P. G., Holmes, C. C., and Walker, S. G. (2016). A general framework for updating be- lief distributions. Journal of the Royal Statistical Sociery: Series (Statistical Methodology). 78(5): 100-1130.
Blundell, C.,, Comeebise, J., Kavukcuoglu, K., and Wierstra, D. (2015). Weight uncertainty in neural network. International conference on machine learning, pages 1613-1622. PMLR
Bressan, A. (2003). Tutorial on the center manifold theorem. Hyperbolic systems of balance laws, 1911:327-344.
Chewi, S., Erdogdu, M. A., Li, M., Shen, R., and Zhang, S. (2022). Analysis of langevin monte carlo from poincare to log-sobolev. In Conference Learning Theory, pages 1-2. PMLR.
Chiang, T.-S., Hwang, C.-R., and Sheu, S. J. (1987). Diffusion for global optimization in R". SIAM Journal on Control and Optimization, 25(3):737-753.
Colding, T. H. and Minicozzi II, W. P. (2014). Lojasiewicz inequalities and applications. arXiv preprint arXiv:1402.5087,
D' Angelo, F. and Fortuin, V. (2021). Repulsive deep ensembles are Bayesian. Advances in Neural Information Processing Systems, 34:3451-3465.
Ermak, D. L. (1975). A computer simulation of charged particles in solution. i. technique and equilibrium properties. The Journal of Chemical Physics, 629-49)4189-4196
Fort, S., Hu, H., and Lakshminarayanan, B. (2019). Deep ensembles: loss landscape perspective. arXiv preprint

---

Gal, Y. and Ghahramani, Z. (2016). Dropout as Bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059. PMLR.
Garreau, D., Jitkrittum, W. and Kanagawa, M. (2017). Large sample analysis of the median heuris- tic. arXiv preprint wwrXiv::707.0777
Glaser, P., Arbel, M., and Gretton, A. (2021), Kale flow: A relaxed kl gradient flow for probabilities with disjoint support. Advances in Neural Information Processing Systems, 11:0018-031.
Graves, A. (2011). Practical variational inference for neural networks. Advances in neural informa- tion processing systems, 24.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel two- sample test. Journal of Machine Learning Research, 13(25):723-773.
Grünwald, P. (2011). Safe learning: bridging the gap between Bayes, MDL and statistical learn- ing theory via empirical convexity. In Proceedings of the 24th Annual Conference on Learning Theory, pages 397-420.
Guedj, B. and Shawe-Taylos J. (2019). primer pac-Bayesian learning. In ICML 2019-Thirty- sixth International Conference on Machine Learning.
Haddouche, M. and Guedj, B. (2023). Wasserstein PAC-Bayes learning: A bridge between general- isation and optimisation. arXiv preprint arXiv:2304.0000
Husain, H. and Knoblauch, (2022). Adversarial interpretation Bayesian inference. Interna- tional Conference on Algoritimic Learning Theory, pages 553-572. PMLR.
Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. G. (2021). What are Bayesian neural network posteriors really like? International conference on machine learning, pages 4629- 4640. PMLR.
Jewson, J., Smith, J., and Holmes, C. (2018). Principles of Bayesian inference using general diver- gence criteria. Entropy, 20(6):442.
Knoblauch, J. (2019). Frequentist consistency of generalized variational inference. arXiv preprint 19X:::991..0496.
Knoblauch, J. (2021). Optimization-centricc generalizations of Bayesian inference. PhD thesis, University of Warwick.
Knoblauch, J., Jewson, J., and Damoulas, T. (2018). Doubly robust Bayesian inference for non- stationary streaming data using B-divergences. In Advances in Neural Information Processing Systems (NeurIPS), pages 64-75.
Knoblauch, J., Jewson, J., and Damoulas, T. (2022). An optimivation-ccntric view on Bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research, 23(132):1-109.
Kolokoltsov, V. N. (2010). Nonlinear Markov processes and kinetic equations, volume 182, Cam- bridge University Press.
Korba, A., Frain-Frankowwki,, P.-C., Majewski, S., and Ablin, P. (2021). Kernel Stein discrepancy descent. In International Conference on Machine Learning, pages 5719- PMLR.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017). Simple and scalable predictive un- certainty estimation using deep ensembles. Advances neural information processing systems, 30.
Lee, D., Simchowitz, M., Jordan, M. I., and Recht, B. (2016). Gradient descent only converges to minimizers. Conference on learning theory, pages 1246-1257. PMLR.
Li, Y. and Turner, R. E. (2017). Gradient estimators for implicit models. arXiv preprint 1705:1100..07107
12

---

Lichman, M. (2013). UCI machine learning repository.
Liggett, T. M. (2010), Continuous time Markov processes: an introduction, volume 113, American Mathematical Soc.
Liu, C., Zhu, I.., and Belkin, M. (2022). Loss landscapes and optimization in over-parameterrized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85- 116.
Liu, Q. (2017). Stein variational gradient descent gradient flow. Advances in neural information processing systems, 30.
Liu, Q. and Wang, D. (2016). Stein variational gradient descent: A general purpose Bayesian inference algorithm. Advances in neural information processing systems, 29.
Louizos, C. and Welling, M. (2017). Multiplicative normalizing flows for variational Bayesian neural networks. In International Conference on Machine Learning, pages 2218-2227. PMLR.
Lu, J., Lu, Y., and Nolen, (2019). Scaling limit of the stein variational gradient descent: The mean field regime. SIAM Journal Mathematicul Analysis, 51(2):648-67 6711
Ma, C. and Hemánnee--Loato. J. M. (2021). Functional variational inference based on stochastic process generators. Advances in Neural Information Processing Systems,
Ma, C., Li, Y., and Hernández-LLobtto, M. (2019). Variational implicit processes. In International Conference on Machine Learning, pages 4222-4233. PMLR.
Matsubara, T., Knoblauch, J., Briol, F.-X., and Oates, C. J. (2022). Robust generalised Bayesian inference for intractable likelihoods. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(3):997-1022.
McAllester, D. A. (1999a). PAC-Bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pages 164-170. ACM.
McAllester, D. A. (1999b). Some PAC-Bayesian theorems. Machine Learning, 37(3):355-363.
Mescheder, L., Nowozin, S., and Geiger, A. (2017). Adversarial variational bayes: Unifying varia- tional autoencoders and generative adversarial networks. In International conference on machine learning, pages 2391-2400 PMLR.
Miller, J. W. and Dunson, D. B. (2019). Robust Bayesian inference via coarsening. Journal of the American Statistical Association, 114(527):1133-1222
Muandet, K., Fukumizu, K., Sriperumbudur, B.. Schölkopf, B., al. (2017). Kernel mean embed- ding of distributions: A review and beyond. Foundations and Trends@ in Machine Learning, 10(1-2):1-141.
Neal, R. M. (2012). Bayesian learning for neural networks, volume 118. Springer Science & Business Media.
Ovadia, Y., Fertig, E.,, Ren, J., Nado, Z., Sculley, D., Nowozin, S., Dillon, J., Lakshminarayanan, B., and Snock, J. (2019). Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in neural information processing systems, 32.
Polyanskiy, Y. and Wu, Y. (2014). Lecture notes on information theory. Lecture Notes for ECF563 (UIUC) and, 6(2012-2016)77.
Rezende, and Mohamed, S. (2015). Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR.
Roberts, G. O. and Tweedie, R. L. (1996). Exponential convergence Langevin distributions and their discrete approximations, Bernoulli, pages 341-363.
Rodriguez-Santann, S., Zaldivar, B., and Hermandez-Lobato, D. (2022). Function-space inference with sparse implicit processes. International Conference on Machine Learning, pages 18723- 18740. PMLR.
13

---

Santambrogio, F. (2015). Optimal transport for applied mathematicians. Birkäuser, NY, 55.
Santambrogio. F. (2017). {Euclidean, metric, and Wasserstein) gradient flows: overview. Bulletin of Mathematical Sciences, 7:87-154.
Shawe-Taylor, J. and Williamson, R. C. (1997). A PAC analysis of a Bayesian estimator. In Annual Workshop on Computational Learning Theory: Proceedings of the tenth annual conference Computational learning theory, volume 6, pages 2-9.
Shi, J., Sun, S., and Zhu, J. (2018). A spectral approach to gradient estimation for implicit distribu- tions. In International Conference on Machine Learning, pages 4644-4653, PMLR.
Sun. S., Zhang, G., Shi, J., and Grosse, R. (2018). Functional variational bayesian neural networks. In International Conference on Learning Representations.
Syed, S., Bouchard-Côté, A., Deligiannidis, G., and Doucet, A. (2022). Non-reversible parallel tempering: a scalable highly parallel meme scheme. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(2):321-355
Veretennikov, A. Y. (2006). On ergodic measures for McKean-Vlasov stochastic equations. In Monte Carlo and Quasi-Monte Carlo Methods 2004, pages 471-488 186.. Springer Berlin Heidelberg.
Villani, C. (2003). Topics in optimal transportation, volume 58. American Mathematical Soc.
Villani, C. et al. (2009). Optimal transport: old and new, volume 338. Springer.
Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pages 681- 688.
Wibisono, A. (2019). Proximal langevin algorithm: Rapid convergence under isoperimetry. arXiv preprint ar Xiv: 191004499
Wild, V. D., Hu, R., and Sejdinovic, D. (2022). Generalized variational inference in function spaces: Gaussian measures meet Bayesian deep learning. Advances in Neural Information Processing Systems, 35:3716-3730.
Wilson, A. G. (2020). The case for Bayesian deep learning. arXiv preprint arXiv:2001.10995.
Wilson, A. G. and Izmailov, P. (2020). Bayesian deep learning and probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-4478..
Wu, P.-S. and Martin, R. (2023). A comparison of learning rate selection methods in generalized Bayesian inference. Bayesian Analysis, 18(1):105-132.
Zou, H. and Hastie, T. (2005). Regularization and variable selection via the elastic net. Journal of the royal statistical society: series B (statistical methodology). 67(2)3301-320.
14

---

A Existence and uniqueness of global minimiser
In this section, we discuss assumptions under which the global minimiser of the optimisation prob- lem
L(Q) - (0)f(00)000) (()))(())))))))))))))))) AA(Q.)))
over P(TR-) exists and is unique. We assume throughout that the optimisation problem not patho- logical, the sense that there exists a measure Q E P(R²) such that L(Q) < co. This in
applications often trivial to verify. A good candidate for Q is typically the reference measure P.
Loss assumptions Let 8: RJ -> IR be loss satisfying the following assumptions:
(L1) The loss is bounded from below which means that
inf {E(0) E IR') > -C.
(8)
(L22) The loss is norm-coereive which means that
(()) + 00
(9)
if |0|| + oo.
(L3) The loss is lower semi-continuous which means that lim inf e(0) 2 ((0o)
(10)
0-10g
for all Bu E R
Regulariser assumptions Let D P(R)) P(R)) - [0, [o] be regulariser and P E P(R)) a reference measure. We define Dp(.) = D(, P) for notational convenience. We assume the following for Dp:
(DI) The function Dp is lower semi-continuous w.r.t. to the topology of conv-cone,,gence, i.e. for all sequences (Qn) NEN C ((R)) and all Q with (p < co, holds that Qn D, Q implics
lim inf (r(Q.) (Q(
(11)
Here, Dg denotes convergence in distribution.
(D2) Dr is strictly convex, i.e. for all Qı # Q2 E P(R)) with Dr(Qi) < and Dr(Q2)
xo, it holds that
+ <a
<aDp(Q1) (1-
(12)
with (0,
The next theorem provides an existence result for the optimisation problem 7 in spirit Lemma 2.1 Knoblauch (202T) with the important difference that our assumptions are easier to verify, since they are formulated in terms of l and Dr.
The result is similar
Theorem (Existence of global minimiser). Under the assumptions (L.)--13) and (DI) there exists probability measure Q' (RR)) with
I(Q") inf (L(Q) Q QEP(R'))
(13)
Proof. Let c > -00 be the lower bound for l. It follows immediately that L(Q) > for all Q E P(R') since D(P,Q) Q) 2 0. As consequence we know that
xe > L* == inf (L(Q) Q E P(R)) <c> -00.
(14)
By definition of the infimum we can construct a sequence I, L(Qz) the image of L such
I, + L
(15)
15

---

for n + oc. We now show by contradiction that the corresponding sequence (Qn) C P(R') is Assume that (Qn) is not tight. By definition can then find an e > such that for each
tight f
k N there exists ni E N with Qn (I-k, & 1 We set A = | kj Rd and
obtain
Ans I(Qm))
(16)
J ((0)dmm(())
1 JE'LA ((0) + (0)
(17)
2 E(0) (0) + SEMA e(0)
(18)
An
2 Qn, (Ax) + inf {e(0) EE AR (TR (RE)
(19)
2 Qnn (Ax) + cinf {{(0) E (R))
(20)
Due to the coerciveness
we know that inf {E(0) : O [ 00 for + xc and there- fore in - oo for k: + co. However, this a contradiction: The sequence (I..) is convergent and therefore in particular bounded. As consequence, it cannot contain the unbounded sub-sequence (Ln,). It follows that the sequence (Q.) is tight. By Prokhorov's 's theorem we can now extract sub sequence (Qnx) of (Qn) and measure Q" P(R)) such that
nn
Q"
(21)
for k - co. Due to Lemma 5.1.7 in Ambrosio et al. (2005) the lower semi-continuity of implies that Q + J E(O) dQ(() lower semi-continuous. This combined with the lower semi-continuity of Dp gives
lim inf L(Q L(Q").
(22)
00000
From this it immediately follows that
L(Q") < lim inf L(Qnz) L"
(23)
but by definition L is the global minimum of which implies L < I(Q"). We therefore conclude that I(Q") T.*
Theorem 3 only shows the existence of a global minimiser. In order to show uniqueness we use the convexity assumption (D2). The proof is the same as in finite dimensions and only included for completeness.
Theorem (Uniqueness of global minimiser). Assume that (D2) holds. Then, the global minimiser of L is unique (whenever it exists).
Proof. Assume there exits two probability measures Q1, Q2 E P(R-) such that I(Q1) = I," = T((22)- where 0c L inf (L(Q) Q P(R')) -c. We define the probability measure
(24)
Q3 = Q1 Q3. By strict convexity we obtain
((Qa) < (Q))))
-
(25)
which is contradiction to Qı and Q2 being global minimisers.
Note that in the literature on GVI (Knoblauch et al. 2022 it common to assume that the regulariser is definite, i.e.
P=Q
(26)
for all P.Q E P(R))). We did not use this assumption in neither Theorem nor Theorem However, the next lemma shows that is basically implied by strict convexity.
"A sequence probability measures (Qx) called tight if and only if for every 0 there exists a compact set K _ such that for all N holds: (n(K)
16

---

Lemma 1. Let Dp : P(R") + [0, oo] be strictly convex and assume further D(Q.Q)
for all Q E P(R)). Then it follows that D(Q, implies Q.
Proof. We prove the claim by contradiction. Assume that there exists P o Q such that D(P.Q) = The strict convexity and [PP, P) = 0 imply combined that
1
1 P)Q P(O 1 00 P) < D(P, P) +
D P
(27)
2
2
0.
(28)
However, know that D + QQ 2 0 by assumption. This is contradiction.
L
Discussion on loss assumptions The assumptions on the loss & in (L1) and (L3) are rather weak. Typically loss functions in machine learning are bounded from below and continuous (and there- fore in particular lower semi-continuous). However, norm-coercivity can be violated. Consider for example the squared loss
{(0) Σ (yn fo())))
(29)
where fo is the parametrisation of neural network with one hidden layer, i.e. (wc, A) and
fo(z) = w (Ar),
(30)
where a: R + R is an activation function which is applied pointwise to the vector Az and has the property that (0) = O. is now possible to find a sequence of parameters (0x (N)Nen C IR' with
||0=|| + oc such that (0x) does not converge to infinity. Define we == k(1...1). A == and 8g (wy Ax) for E N. Then we obviously have that
||0x|| = ||=a|| + 00
(31)
for 1- - 0o but
N
E(Ox))
Σ
(Un. -ox((n) (Ix))²
(32)
n=l
Σ (yn w w- o(0))²
(33)
(34)
n=1
which is constant and therefore does converge co. similar, but notationally more involved, construction can be made for neural networks with more than one hidden layer. However, this is an issue that can be easily resolved by adding what is known as weight decay to the loss. For example, consider for >> 0 the loss
N
((0)
= Σ (vn - (o(x)) +110012
(35)
==1
with weight decay. This loss is by construction norm-coercive and therefore the previous existence proof applics.
Discussion on regulariser assumptions The assumptions (DI) and (D2) are quite weak. The KL- divergence for example is known to be lower semi-continuous (Polyanskiy and Wu. 2014 Theorem 3.7) and strictly convex (Polyanskiy and Wu 2014 Theorem 4.1). This immediately implies lower semi-continuity and convexity of KL(-, P) for any fixed P. The MMD is also known to be strictly convex (Arbel et al.) 2019, Lemma 25), whenever it is well-defined, which can be guaranteed under weak assumptions also depend on the kernel k. However, for bounded kernels trivial to verify. We include the proof for completeness, but assume this has been shown before elsewhere.
(Muandet et al. 2017 Lemma 3.1). The lower semi-continuity properties
17

---

Lemma 2. Let the kernel R X RJ be continuous and bounded: ||<|
ou |k(0, 8'] < oo and P be fixed. Then MMD(.,, continuous and therefore, partic-
ular, lower semi-continuous.
Proof. Let (Qn)neN and Q" such that
Qn _ Q"
(36)
for n 00. This immediately implies that
& 2 Q
(37)
for where Q @ Q" denotes the product measure of Q' with itself. Further, note that the kernel mean embedding μP is continuous as integral with respect to the second component of continuous function and bounded since
|pr(0)) =1 / k(0,0') dP()')
(38)
$ / In(0,0')aP((
(39)
||1|
(40)
By the definition of weak convergence for measures, we therefore have
lf x0006") d(Qn, @ Qn)(0,0) - /ʃ x(0,0') (Q)(0,0))
(41)
f pr((0) (Q. - (ur(0)dQ(())
(42)
for n - oc. This immediately implies continuity of MMD(;, P) with respect to the topology of weak convergence.
Notice that most kernels common in machine learning, such as the squared exponential or the Matérn kernel, are continuous and bounded and therefore Lemma applies.
The astute reader may have noticed that our existence proof only guarantees the ex- istence of measure Q" E (R.). However, the Wasserstein gradient flow by definition only formulated in the space of probability measures with finite second moment, denoted P2 (R2).. As- sumptions which guarantee that Q* (R-)) are easy formulate. For example, we can require that there exists C > 0 and R 0 such that the loss & satisfies
Remark 1.
[((0)] > C/
(43)
for all ||0|| > R. This immediately implies that Q" E (R( since otherwise
/ [1(00) dQ' (0) 0
(44)
gives a contradiction to the finiteness of L(Q"). However, even (43 is violated, the reference measure P may still guarantee that Q" E (2))). For example, if P (R then (p(Q") will typically be large if Q* & (2 and the global minimiser is therefore in sense unlikely to have fat tails. We therefore assume Q" E (R)))) throughout the paper and consider it to be a minor practical concern.
B Realising the Wasserstein gradient flow
In this section, we identify suitable stochastic process that allows to follow the WGF.
Let L P(RR') - (-o0, oo] be the free energy discussed in Section 3.2 given as
L (Q) = / v(0))Q(0)+ 2 I (Q())())()))) + Aa / log (1())) d0, (
A
(45)
18

---

where A1, A2 2 0 are constants, V R. + is the potential, &. R'' x R - R is symmetric. We will write L for L from now on simplify notation. The Wasserstein gradient of L is given as (cf. Chapter 9.1 Villani 2003 Equation 9.4)
V w L(Q VV + À ( I Q)(0) + X log (q(0)),
(46)
where VIK: R X RR' is the (vector-valued) derivative of with respect to the first compo- nent, V denotes the euclidean gradient with respect to and (Vin* ()) := J V ve 6*) dQ((")
for 0 E RJ. The corresponding Wasserstein gradient flow is therefore given (cf. Chapter 9.1 Villani 2003
Equation 9.3)
0,g(t,0)= - V. V (VV)) + +1
Q)(0) + A2 V log (a)(0))).
(47)
In general the probability density evolution of a stochastic process is-via the Tokker-Planck equation-associaane associated with the adjoint of the (infinitesimal) generator of the stochastic process. We will therefore try to identify the generator associated to the density evolution in (47). To this end let h CRR,, where (R R) denotes the space of twice continuously differentiable func- tions with compact support. We multiply both sides of (47) with h, integrate, and apply the partial integration rule to obtain
d di h(0)((10)d0
/ Vw L(())() (0) Vh(00) a(t, do.
(48)
/ (VV(0) + \1 V Q.) (Q:))
(49)
A2 / V Vog log (9.(0) (h( (Q).
(50)
By chain-rule and partial integration, (50) can be rewritten as
->2 f Vlog (0) Vh(0)dQ.(0 ->2 Vq. VV(()) do
(51)
A2 / Ah(0)dQ10)
(52)
Putting everything together, we obtain d
dt / ((O))(t,0)00 (41Q()))4(()ddQ.(0))
(53)
where {A[Q]] QEP(Ri) a family of operators defined as
(AIQ))) = (V + (V Q)(0)) Vh(0) + A2hh.
(54)
for h E (R RR. The reader may recognize this operator family as the generator of so called nonlinear Markov processes (Kolokoltsov) 2010, Chapter 1.4). The nonlinearity in this case refers to the dependency on the measure Q. Linear Markov processes have no measure-dependency. This family of generators corresponds to McKean- -elsen-Vlasov process of the form
dt(t) (OV))()) (V Q.)((11)) dt + 22)2dB(1).
(55)
where (B(t)) t>0 is a Brownian motion and Qe the law of (t). In other words: The solution to (5) has the time marginals Q(t) such that (53) holds for every h CR R). Furthermore, the corresponding pdfs (1(1) satisfy the nonlinear Fokker-Plande equation given
0gg. A'
(56)
where A*( denotes the adjoiit of the operator A[Q] and is given as
(A) v (MO)(VV(0)) A (V K Q)( + X2 log (N())))
(57)
19

---

for h E (RR',R)(Brbb and Röckner 2020 cf. equation (1.1)-(1.4). Note that (56) corresponds
exactly to the Wasserstein gradient flow equation in (47). We can therefore follow the WGF by simulating solutions to (55).
The standard approach simulate solutions to interacting particles. Formally, we replace Q(t) by
55 NE
Veretennikov
2006) is ensemble of n=1 (0, and obtain
Ng (0,,000)
(1) - (VV + NE Σ
222ddB..1)
(58)
for I...., NE where Ng N denotes the number of particles. The Euler-Maruyama approxi- mation (58) leads to the final algorithm:
Step 1:
Initialise Ng E N particles 01,0. ,... from a use chosen initial distribution Qo-
Step 2: Evolve the particles forward time according to
A1 NE
No (On)) 0/.)) - 22r122Zn (59)
0n.441 On, 7 VV (On,k)
(On.))
for n = 1, NE, k: = 0,..., 7 with Zn,k ~ Nx()
Note that On. is thought of as approximation of (t) at position = kn. Furthermore, as dis- cussed in Section various choices of V, and 22 allow to implement the WGF for different regularised optimisation problems in the space of probability measures. This is summarised below:
Deep ensembles: V(0) - {()), Å1 = 0, A2 = 0
Deep Langevin ensembles: V(0) = =E(0) > Alogp(0), A1 = 0, == A2
Deep repulsive Langevin ensembles: V(0) - {()) - X logp(0) - (2)
C
Asymptotic distribution of particles: unregularised objective
In this section, we investigate the asymptotic distribution of the WGF for the objective
L(Q) = f E(0) dQ())
f
(60)
for Q P(RR)). The associated particle method is:
Sample (0), ... ONS (0) independently from Qo-
Simulate (deterministically) 0, = - VE(O,, for = I,..., Ng.
We start by introducing some notation for the deterministic gradient system. Let o' (6n) denote the solution to the ordinary differential equation (ODE)
0(0) = 0n E RJ 6' (t) = - ve(((t))
(61)
(62)
at time t > O. In a first step, we show the following lemma, which is a simple application of the famous Lojasiewicz theorem (Colding and Minicozzi II, 2014). and the fact that Lebesgue almost every initialisation leads to a local minimum (Lee et al. 2016).
Lemma 3. Assume & R + IR is norm-coercive and satisfies the Lojasiewicz inequality, i.e. for every 0 E Rdd exists an environment U of 0 and constants 0 < < and C >> 0 such that
[((0) {(ō) < C(V)()))
(63)
for all 0 E U. Then we know that o² (0o) converges for oo to a local minimum of E for Lebesgue almost every 0g J
20

---

Proof. First we show that t + 6² (60) is bounded. We proof this by contradiction. Assume that ('(6) is unbounded. Then there exists subsequence (taJDeN C [0, oc) with t, - oc for n oo such that
|6' (66) oc
(64)
for 7t + oo. The norm-coercivity immediately implies that E( ((o"'(00) (0v)) +
(65)
for n + oo. However, this contradicts
((d'(0o)) (00)) = E(6o) < co, where the first inequality follows from the fact that / + { (0o)) is decreasing, which is conse- quence of
(66)
d di (()((0)) = ve((o'(oo) _ ø¹ (0o)
d
(67)
dt
(V 12 0.
(68)
Hence t + (80) is bounded. By the Bolzano- Weieess-sessssssss theorem we can find sequence (In)NEN C [0, 00) with in + oo and a point O0 E RR' such that
6²n (6o) + 000
(69)
for n + co. Hence (6'(6)) e>-o has the accumulation point 0x. The Lojasiewicz theorem (Colding and Minicozzi II 2014) allows us to deduce that
o' (0n) 0xx
(70)
for f. + oo, and that 0x satisfies (((0~))= =0 0.
It remains to show that fxx is not saddle point for Lebesgue almost every initial value 9y. However, this is very similar to the proof in Lee et al. (2016). The only difference is that one would need to use continuous-time version of the stable manifold theorem, which is readily available, for example in Bressan (2003).
Let {m: Jien denote the local minima of l which are by assumption countable. Denote further by
+ == { RR''
lim ('())) - m.} 1-200
(71)
the domain of attraction for the minimum m. The next theorem then an easy consequence of Lemma 3
Theorem S. Assume that the loss function only has countably many local minima, is norm coer- cive, and satisfies the Lojasiewicz inequality. Let further Og Qo for some Qo P(R.) such that E @o(O.) = 1. Then,
8'(
D
Σ
Qo(O:) 8m; =: Qoo
(72)
for co. Here 2 denotes convergence in distribution.
Proof. Let 0o E IR be fixed. Due to Lemma
3 we know that
('(D))
Σ m; (0o O,}
(73)
for Lebesgue almost every Bg for + oo. Here, 1f.} denotes the indicator function. Let Y now be random variable with law Qo By assumption, know that Y O for some i E N with probability 1. Hence,
80 Σ
m, e
(74)
i=1
21

---

almost surely for t + oo. Since almost sure convergence implies convergence in distribution, we conclude that
(M) 2, L ml E
(75)
where L() denotes the law of random variable. However, the law of the RHS is easily recognised
as
C m,ffY e0,)) - Σ
Qo(O.))m.
(76)
==1
which concludes the proof.
Remark 2. Note that the condition
Qo(O.)
(77)
in Theorem is easy satisfy. According to Lemma the set
R U
(78)
has Lebesgue measure zero. Therefore, any Qo which has density w.t.t the Lebesgue measure will satisfy (77).
D Asymptotic distribution for deep Langevin ensembles
In this section, we analyse the objective
I(Q) = (()) dQ(0) +KKT(Q P)
(79)
for Q P(R)). The corresponding particle method is given as:
Sample (0), ONS (0) independently from Qo.
Simulate the SDE de,, (t) (V. (t for each 1,..., Ng-
Recall that V()) (()) A p()) This case is well-studied in the literature and known as
Langevin diffusion. Under mild assumptions (Chiang al. 1987 Roberts and Tweedie 1996).
(t) D Qoo
(80)
for f - oo and each particle n = 1,..., NE independently. The probability measure Q has the density
Z exp (- V(0),
1
Jo0 (f)
(81)
Z exp
1
)))(()
{(0) x )(()),
(82)
where Z > 0 is the normalising constant. As a consequence, the WGF asymptotically produces samples from Qoo. However, is priori unclear that Qoo is in fact the same as the global minimiser Q* of L.
We investigate this question by relating invariant measures stationary points of the Wasserstein gradient.
Definition 1. (Liggett 2010 Thm. 3.3.7) A measure Q is called an invariant measure (for given
Feller-process)
f
Ah(0) dQ(() =
(83)
for all & (IR)). Here A is the infinitesimal generator of the corresponding Feller-process.
22

---

Recall that the infinitesimal generator of the Langevin diffusion for h (IR) is given as
Ah VV Vh + Aåh.
(84)
Definition 2. A measure Q E P2(R²) is called a stationary point of the Wasserstein gradient if
V LQ) =0
(85)
for Q alaamost every 0 E IR..
In finite dimensions, This carries over to the well-known that local minimiser is a stationary point of the gradient. case, with similar proof. Since we could not find this result anywhere in the literature we included for completeness.
Lemma 4. Let Q be a local minimiser of L, i.e. there exits and c > 0 such that
l(Q)= L(Q)
(86)
for all Q with Q)(Q,Q) S Then Q is stationary point of the Wasserstein gradient in the sense of Definition 2
Proof. Let h E C2 be arbitrary and Q E P2(())) be local minimum of Further, let Øi (fg)
be the solution the initial value problem
{(0) = 8o
(87)
(1() V)(((())
(88)
for E (-e,e) for some e > We now define Q(1) := d'# for I (-e,e) where f denotes the push-forward of the measures through the function f. In the Riemannian interpretation of the
Wasserstein space, (Q(t)) rE(-r.)) is curve in P22(RR) with tangent vector h at point Q Ambrosio
et al. 2005 Chapter 8). We, further, define f: (-c, - TR as f(t) = I(Q(1)). Application of the
chain-rule (Ambrosio et al. 2005 P. 233) gives
f() L(Q(t)) dt
(89)
(Vw L[QO). V))
(90)
V L Vh(0) dQ()).
(91)
We know that f has local minimum 0 and, therefore. f((0) 0 which gives
0
Vw L Vh(0) dQ()).
(92)
Since 92 holds for arbitrary test functions h E (IR and as CR(R)) is dense L²((),, we obtain that L() = 0 for Q 0 e R..
The next lemma relates invariant measures and stationary points of the Wasserstein gradient for infinitesimal generators of the form (84). It will prove extremely useful to translate between the Langevin diffusion literature and our optimisation perspective.
Lemma 5. Let Q Pa PR) be such that Q has a density with respect to the Lebesgue measure. Then, the following two statements are equivalent:
Q a stationary point of the Wasserstein gradient.
. Q is an invariant measure.
Proof. Let Q be a measure with density q. Recall that the generator of the Langevin diffusion is for h (Rd) given
Ah VV. Vh + AAh.
(93)
23

---

By partial integration, it is easy to verify that the - adjoint (w.r.t the Lebesgue measure) is given as
4 v. (h- VV) AAh.
(94)
We, therefore, conclude that
J Ah(0) dQ(o)
/ Ah(0)g(0) do
(95)
/ (00))'((0)0) do
(96)
/ (O)(V (g(0) VV(()) + AAq(O)) do
+
(97)
Furthermore, we have Lw = VV + AV log g, and therefore
/ Vw vh(0) = / Vw vh(0)g(0) do
(98)
/ (VV Vh(0) d0
(99)
- (V) (V- + AAq(0)) d8,
(100)
where the last line follows from applying partial integration. This allows us to conclude that
f Ah(o) dQ(0) - LQQ(0) Vh(0) dQ(0)
(101)
whenever Q has density. As consequence we have that Q is invariant if and only if it is stationary point of the Wasserstein gradient.
Lemma 5 allows to move between the optimisation and stochastic differential equation perspec-
tive. In Appendix A we discussed the existence and uniqueness of global minimiser Q" of L. We know that Q" has density since the Kullback-Leibler divergence would be infinite otherwise (assuming P has Lebesgue-density which we assume throughout the paper). Lemma 4 that Q" stationary point of the Wasserstein gradient. Due to Lemma [5] we infer that Q' must be an invariant measure. However, due to the uniqueness of the invariant measure under the previously mentioned mild assumptions (Chiang al. 1987 Roberts and Tweedie 1996), we can conclude that Q* Q
E Asymptotic distribution of deep repulsive Langevin ensembles
In this section, we consider
((Q)= / (Q())))) + MMD(Q, - + K2
\ 2
(102)
/ V(0) (Q))) 2
\ /
(Q) const,
(103)
as optimisation objective. Here, (QQ l (0) (0) d0 denotes the differential entropy.
Recall that in this case V(0) = E(O) (p)r(0) X2 log pl (). We already discussed in Appendix that the AcKean-Vlason process of the form
8(0) Qu
(104)
()))) VV(O())) + (V Q.)((1))) di + 222dB(())
(105)
with (B(t)) 120 being Brownian motion achieves the desired density evolution. Furthermore, the
particle approximation of [04) is given as
No (t). - (VV + Ng Σ
(VV)))(((0)))) dt 2222dB...1)
(106)
j=l
24

---

for nt = ,,....NNg where Ng E N denotes the number of particles.
The approach follows the same procedure as in Appendix D We show the notions of invariant
measures and stationary points of the Wasserstein gradient are the same for measures with Lehesgue density. We start by introducing the concept of invariant measure for nonlinear Markov process (Ahmed and Ding 1993 Definition 1).
Definition 3. A measure Q is called an invariant measure for nonlinear Markov process with the family of infinitesimal generators {A[Q]] QEP(W) if
/ A[Q]A(0) (Q(() =
(107)
for all h E (RR)).
Recall that the family of infinitesimal generators in our case is given as
(AIQI))
VV + iV 2)(0) Vh(o) A2
(108)
for h E CRRR,, R). In analogy to Lemma 5 we obtain the following result.
Lemma 6. Let Q E (2)))) be such that Q has a density with respect to the Lebesgue meusure. Then, the sollowing two statements are equivalent:
a stationary point of the Wasserstein gradient for in (102 in the sense of Def.
Q is invariant measure for the McKean-Vlasov process with infinitesimal generator defined in (108)
Proof. First, notice that
f A(Q)M())Q(() / A(Q))(0)q(0)) do
(109)
h(0) (0) d0.
(110)
Recall, that A*[ denotes the -2-adjoint of the operator A(Q] and that it is given as
(A)(Q)) v. (N()(VV(O) (V ()) + X2 V log (h())))
(111)
for h E CR R) with compact support. This implies
(A)))) = v (1()) - (V () + +^24q(0).
(112)
We plug this into 110 to obtain
I A(Q](())QQQ
(113)
= -JMOOV- (2(0) +AVI Q)())) dh \eh(0))a((0)d0. d0.
(114)
On the other hand, we have that
()) = vV 1 Q)(() + logg()
(115)
and therefore
VI(Q(() Vh(O)dQ(A)
(116)
(VV())+X(V + +2 logg(0)) VQ(()
(117)
a( + 2)()) Vh(0) + A / Va(0)- Vh(0)-Vh(())) do
(118)
=- - (110) (V)V(()((0))))))))++++(++)))))() ((0))) A / (00))h(0) d0,
(119)
25

---

where the last line follows from partial integration. Comparing (14 to 119 gives
/ A(Q])(0) dQ ()) Vh(0) dQ(o)
(120)
for all h E (R²) whenever Q has a density. This immediately implies that is invariant iff is a stationary point.
Again, we leverage this correspondence between stationary point and invariant measures. There is a rich literature on ergodicity of nonlinear Markov processes. For example, Theorem 2 of Vereten- nikov 2006) specifies conditions on K. and such that
(n) 2,Qxx
(121)
for Ng. ł > oc. Here QT.NE (f) denotes the law of fixed particle 0, (t), n 1..... .N......,.. whose distribution is characterised by the SDE (58) The measure Qxx is the unique invariant measure of the nonlinear Markov process. By Lemma 6 every invariant measure
stationary point of the Wasserstein gradient and vice versa. Hence, existence and uniqueness of the stationary point of the Wasserstein gradient is immediately implied. However, since the global minimiser Q" is a stationary point of the Wasserstein gradient (cf. Lemma 4. conclude by uniqueness that Qoo Q"
F Asymptotic analysis of deep repulsive ensembles
In this section, we consider the objective
I(Q)= == f (Q()) + MMD(Q, P)
(122)
for Q E P(R.). The corresponding McKean-Vlasov process is of the form
(t (VV()()) A(VI * Q.) dt.
(123)
where Qe denotes the distribution of 0t) and (0)) {(0) (0) with (0) x(0, the kernel mean-embedding of P. We call the particle method in this case deep repulsive ensembles (DRE).
k(0, 0*) dP)
The existence the global minimiser Q* is still guaranteed under the assumptions in Appendix Lemma guarantees that Q" is a stationary point of the Wasserstein gradient, i.e.
(V))) + I(VIx= (0) 0
(124)
for -ae E R.. Recall that the infinitesimal generator in this case is given as
(A[Q])) == - (V) I(V ()) Vh(0)
(125)
for Q P(R'), h CRRR). It immediately follows from the definition that
(AIQ))) Vw (() (h(())))(()
(126)
for all h CR 0 R As in Lemma & 6 this implies that each stationary point of the Wasserstein gradient is an invariant measure of the McKean-Vlasov process and vice versa. In Appendix D & E we cite relevant literature that guarantees uniqueness of the invariant measure, which a necessary (but not sufficient) condition for convergence to the invariant measure. The next theorem shows that uniqueness will in general not hold without the presence of the diffusion term.
Theorem 6. The invariant measure for the McKean- Vlasov process with the family of generators (A[Q)) ( defined in (126) is (in general) not unique.
Proof. Let Ng E N and define T. (R.J) NE
- RR
No.
N 160,,
I(e
AN =
Σ
V(O.)
A
0,)-
(127)
2N g
iJ=1
26

---

Assume that V is bounded from below and norm-coercive. Then is bounded from below and norm-coercive and therefore we can find a global minimiser 0 == (....... Ng E (Ry) A of T.. Since L differentiable, we know that 6" is a stationary point of the gradient which implies
A
VV
Σ V (K) 07)
(128)
NE
for all i 1,.... NE- Here, MMD literature. Note that (128) is equivalent to
that the kernel is symmetric, which is standard in the
VV(0) + A(VIN ()) =
(129)
for Q-a.e. 0 E R. where
NE 00; (d0).
(d0) = NE
(130)
This means that Q is a stationary point of the Wasserstein gradient, and therefore an invariant mea- sure for the McKean- Vlasov process. Since Ne E N was arbitrary, have constructed countably many invariant measures and therefore uniqueness can't t hold in general.
The reason that non-uniqueness of the invariant measure is an immediate contradiction conver- gence is the following: If we initialise with any of the invariant measures constructed in the proof of Theorem 6 then the particle distribution of the McKean- McKean-Vlasov process will remain unchanged over time. Convergence to the global minimiser can therefore surely not hold for arbitrary initialisation Qo It may be possible construct conditions example, for Stein variational gradient descent a similar issue occurs. However, in this case one can guarantee convergence (Lu al. 2019] Theorem 2.8) if Qo has Lebesgue-denss (and if the ker- nel satisfies further restrictive assumptions). The existence of conditions that guarantee convergence for DRE remains an open problem.
Qo under which convergence still holds. For
G Implementation details
In Appendix we derived the following algorithm:
Step 1:
Simulate Ng E N particles 01,0- ... NN.0 from a use chosen initial distribution Qo-
Step 2: Evolve the particles forward in time according to
N 0n. = On. 17( (VV(On.) + ((Vi))(On + 22
A1
(131)
N j=1
for 1, NE. k: = .,..., K with Zn,k (10,Isx.)
We can generate samples from DE, DLE and DRLE by setting the potential and regularisation parameters as described below:
Deep ensembles: V(0) = (00), A1 = 0, A2 =0
Deep Langevin ensembles: V(0) {(0) X llgg((0), \ = =0 d2
Deep repulsive Langevin ensembles: V(0) = {(0) - A log - (2)
we can think of oI. I global minimiser Q for DLE and DRLE if K is large enough. All experiments the SE kernel given as
Due to Appendix
as approximately sampled from the
|0 011
x(0,0') exp
(132)
20
with lengthscale parameter dx > 0. The kernel mean embedding uo can easily be approximated as
(p = Σ x(0,0,), 0 0 E
(133)
where 0 01,...,0x 0M independently. We chose M = 20.
27

---

G.1 Toy example: global minimiser
We describe details regarding the experiments conducted to produce Figure 2 below. We generate NE = 300 particles and make the following choices:
Loss: {(0) = 0 + _ 02)
Prior: P~ (0, and therefore log p(0) = 02
Initialisation: Qo = P
Reg. parameter: ADLE = 1, ADRLE = 1, XDRLE =1
Step size: 10 10- Iterations: K = 100, 000
Kernel lengthscale, xx is chosen according to the median heuristic (Garreau et al. 2017) based on samples from the prior P'
The loss is constructed such that we have global minimum at 0 = -2, turning point at 0
0,
and local minimum
Deep ensembles: The optimal Q" is Dirae measure located the global minimiser 0 -2. However, proved in Theorem | the WGF produce samples from
(d0) = 2 2 + 2 81 (d0),
(134)
(-o0, 0) is the region of attraction for the global minimum and (0, oo) for the local minimum which both have probability 0.5 under Qo P = -N 1). In particular, Qoo + Q" as expected.
Deep Langevin ensembles: The optimal measure has the pdf
(0) exp ( {(O) p(()
(135)
for A E R. As expected the WGF produces samples from Q"
Deep repulsive Langevin ensembles: The optimal q for deep repulsive ensembles is harder to determine. From the condition that qº is a stationary point of the Wasserstein gradient, we can derive that u(0) log q"' (0) satisfies the differeeiif equation
1 (')
A
( (VIK))(, 0') exp(u(f')) de' X2
n' (0)
(136)
22
with some initial value (0) = un- In principle,
could choose uo such that (())
exp(u(0)) integrates to 1. However, since we do not know the appropriate initial condition priori, we choose an arbitrary uo and normalise the pdf afterwards. We use an numerical solver to evaluate u0) on a fixed grid. As expected, the WGF produces samples from Q* in this case.
G.2 Toy example: multimodal loss
The details below correspond to the experimental results presented in Figures 3 and Figure 5 is an alteration of Figure 3 with only 4 particles with the goal of stressing the importance of the number of particles relative to the number of local minima.
DE, DLE, DRLE We generate NØ = 300 particles and make the following choices:
Loss: {(O) log Σ - 1 N p,, T2), 0 R2 (+3, +33
Prior: P flat and therefore log (0) =
Initialisation: Qo (0,
Reg. parameter: ÀDLE 0.2, ADRLE 0.2, X DRLE 0.6
Step size: 77 0.1, Iterations: K 10, 000
28

---

DE DLE DRLE
FD-GVI
5
5
5
5
00
Figure 5: We generate NØ particles from DE, DLE, DRLE and FD-GVI with Gaussian parametrisation. The multimodal loss E is plotted in grey and the particles the different methods layered on top. The prior in this example is flat, i.c. log and pp constant. The initialisation Qo is standard Gaussian.
Kernel lengthscale, σR is chosen according to the median heuristic (Garreau et al. 2017) based samples from the prior P'
Note that for translation-invariant kernel such the SE kernel obtain for the flat prior P that
(p(
k(0, 0') do'
(137)
Φ(0 0') da
(138)
((E) dE,
(139)
where the second line follows from the fact that we can write any translation-invariant kernel as k(0, 8') 6(0 8') for some function R - R and the second line is simple variable substitu- tion. If 139) is finite, the above expression well-defined and therefore (4p constant. Note that in
particular for the SE kernel, we have (bE) = exp(-I|<|I (2a and therefore (139) is finite. As consequence, we have that for a flat prior P the gradient of the potential V is the same for all three methods. This means that the loss l isn't adjusted and the only difference between the three methods is the presence of repulsion and noise effects.
Remark 3. The astute reader may have noticed that flat prior P is in fact not covered by theory in Appendix A The problem is that (L L), where L denotes the Lebesgue measure, is not positive (and not even bounded from below). To see this, choose Q (O, 2) with diag(oñ, oï) and note that
K) - / log
do = -H(Q),
(140)
where H(Q) denotes the differential entropy. For a Gaussian, it is known that
H(Q)- = 2 log ((2Te)* det det(E')
2 (oog(2me)* + log(oo) + log(o?))
(141)
and therefore if either of + oo or of + oo then KL(Q,C) + -0. However, note that this difficulty is rather technical in nature and can easily be remedied. Instead of C, we could have chosen the uniform prior P ~ U(-10100 10200). In this case, the positivity of KL(,P)) is guaranteed by Jensen's inequality. This choice of P gives up to additive constant- the same objective as a flat prior and up to machine precision the same kernel mean embedding μp. is, therefore, algorithmically irrelevant if P is flat or uniform
very large set.
FD-GVI We use the same prior and loss for DE, DLE and DRLE. We parameterise the variational family as independent Gaussian, i.e.
o = (N(,,, E)Ip TR?, Σ diag( cxp(3))) (2), ER?
(142)
29

---

We learn the variational parameters v = (p. E R by minimising
Σ(v)
((O) dQu(0) + AKL(Qu, P)
(143)
/ ((0) (Q) - (H(N(m,)))
(144)
200
X
4( IU. 2,) -
2 log ((2Te)² exp(31 exp($2))
(145)
200
200
X (3:
Σ j=1
E(p + 2.. 2j) (3: + 32) + const,
(146)
200
where 7.,..., 7200 NN and (H, Σ)) denotes the differential entropy of the normal
distribution. For the regularisation parameter, we chose À 0.5.
G.3 Toy example: more modes than particles
We generate Ng 20 particles and make the following choices:
Loss: {(0) - sin 0 [-M MT], with M 1000
Prior: P flat and therefore p(0) = 0 and Hr const. (cf. Appendix G.2
Initialisation: Qo U-M Mm)
Reg. parameter: ADLE 0.001, ADRLE 0.001, XDRLE 0.6
Step size: 7) 0.01, Iterations: = 1.000
Kernel lengthscale, ox is chosen according to the median heuristic (Garreau et based on samples from the prior P
2017
Note that & has 2M = 2000 local minima at locations
m, 2 in, i (-M O, ...... (M 1)}-
(147)
Due to the flat prior VV = Ve for all three methods. We observe that it is hard to distinguish the methods since most particles are their local modes by themselves.
G.4 UCI Regression
The UCI data sets are licensed under Creative Commons Attribution 4.0 International license (CC BY 4.0). Following Lakshminarayanas et al. (2017). train 5 one-hidden-layer neural networks fo with 50 hidden nodes for 40 epochs. We split each data set into train (81 % of samples), validation (9% of samples), and test set (10% of samples). Based on the best hyperparameter runs (according to Gaussian NLL) found via grid search on validation data set, we make the following choices:
(fe(zn) - yn) where (In In n=1 are paired observations. Loss: E(0) N Em-n(folz,
Prior: P P~
Initialisation: Kaiming intilisation, iee for each layer I E {1,...L that maps features with dimensionality n4-1 into dimensionality ny. sample Qi, (0, 2/n))
Reg. parameter: ADLE 10-4 ÀDRLE 10 XDRLE 10
Step size: 7 0.1, Iterations:
10, 000
Kernel lengthscale, σw. is chosen according to the median heuristic (Garreau et 2017 based on samples from the prior P
G.5 Compute
While the final experimental results can be run within approximately an hour on a single GeForce RTX 3090 GPU, the complete compute needed for the final results, debugging runs, and sweeps amounts to around 9 days.
30


Reviews:
[{"review_id": "mb1EZyWpzq", "reviewer": "Reviewer_6hMj", "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.  ", "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. \n\n2) 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?  \n\n3) 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? ", "limitations": "All limitations are addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper is technically sound and well-written. Overall it was easy to follow.\n2) 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. \n3) 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. \n\n2) 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). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "#### Weaknesses\n\n1. We understand the reviewer’s concern and agree that the paper does not provide a full analysis of empirical performances. One reason for this is the limited amount of space that 9 pages allow in order to comprehensively present our firmly grounded theoretical framework and a number of its technical aspects. That being said, the more important reason is that we do not intend the paper to be a thorough study of any particular ‘market-ready’ methodology that can compete with the fine-tuned algorithms prevalent in industrial scale deep learning. Instead, our focus is on theoretical insight into what existing methods for uncertainty quantification in deep learning are actually doing—and this is reflected by our emphasis on connecting Bayesian and non-Bayesian methods through the analytical lens the paper proposes. As part of this, we conducted a range of experiments in the experimental section (including on UCI regression tasks) to showcase the limitations of naively adopting the current framework for designing new algorithms. In this sense, we believe that our empirical investigation serves the main purposes of the paper: to highlight how the derived theory aligns with reality, and to explain any divergence between theory and reality to help follow-up research with exploiting our ideas in order to develop more efficient and more effective algorithms.\n\n2.  We thank the reviewer for their comment. Indeed, we do not want to exaggerate our proposal’s novelty. While we agree with the reviewer that many of our contribution’s building blocks have been floating around in the literature, we are unaware of any work that has combined them in the way the current manuscript has. Bayesian (and generalized Bayesian) procedures typically focus on justifications that are centered around updating prior to posterior knowledge. Our focus is different: we instead focus on regularization in the space of probability measures, and on algorithms that allow us to solve the resulting problems. On a technical level, we also innovate by introducing WGFs for solving these types of problems. To the best of our knowledge, WGF has not previously been used to study how different regularisers translate into different properties of the inference algorithms; and how this in turn places inference algorithms as different as deep ensembles and variational Bayesian methods under the same overarching framework \n\n\n#### Questions\n1. Thank you for pointing this out; the sentence is indeed misleading in a certain sense: The initial space here is the Euclidean space and the ‘more challenging’ space refers to the space of all probability measures on the Euclidean space. From a point grounded in analysis, the latter is indeed more challenging as it provides less structure: Unlike Euclidean spaces, it is infinite-dimensional, non-linear and does not have an inner-product structure, which makes a thorough analysis more challenging. And yet, the reviewer is also correct to point out that the optimization problem itself is much simpler thanks to its convexity on this (more challenging-to-analyse) space. We intended to describe precisely this trade-off: We trade the simple Euclidean space with the more complicated space of probability measures in order to obtain an ‘easier’ convex objective function. We will slightly rephrase this sentence in the new version of the manuscript to stress this point more clearly.\n\n2. Thank you very much for making us aware of this—you are absolutely correct: this sentence should not have been in the final manuscript. We will ensure that this part of the paper is fixed by providing the adequate context. If you believe there are additional relevant papers related to these types of problems that we should be citing but that are currently not contained in the manuscript, we would be very appreciative of you letting us know so that we can include them.\n\n3. This is an excellent question. Two assumptions prevent this behavior: First, we assume that every saddle point has at least one strictly negative eigenvalue. Second, we assume that the Lojasiewicz inequality is satisfied (Lemma 3 in Appendix). Intuitively, the first assumption guarantees that locally around a saddle-point, the domain of attraction has Lesbegue measure zero. Hence, if our initialisation measure $Q_0$ has a Lesbegue density, we will almost surely not be attracted to the saddle point and leave its domain of attraction. The Lojasiewicz condition then guarantees that once we are in the domain of attraction of a local minimum, we will stay there and converge eventually to it. For more details see [1]. Regarding the reviewer’s specific scenario: A loss with a flat local maximum that has non-zero measure under the initialization would violate the Lojasiewicz inequality—as the reviewer correctly hypothesizes, this setting is therefore not covered by the theory.\n\n\n[1] Lee, J. D., Simchowitz, M., Jordan, M. I., and Recht, B. (2016). Gradient descent only converges to minimizers. In Conference on learning theory, pages 1246–1257. PMLR.\n"}, {"review_id": "lVFOHTu90m", "reviewer": "Reviewer_NdiF", "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.\n\nThe 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).  \n\n", "questions": "\n* 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]\n\n* Since now inference is conducted without the guarantees provided by the Bayesian method, how should the distributions obtained be interpreted or used? \n  \n* 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.    \n\n(_References included in the **\"Limitations\"** section_)", "limitations": "\n* 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). \n\n* 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. \n\n--- \n**References**:\n\n[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.\n\n[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.\n\n[3] Santana, S. R., & Hernández-Lobato, D. (2022). Adversarial α-divergence minimization for Bayesian approximate inference. Neurocomputing, 471, 260-274.\n\n[4] Ma, C., Li, Y., and Hernández-Lobato, J. M. (2019). “Variational implicit processes”. In: International Conference on Machine Learning, pp. 4222–4233.\n\n[5] Sun, S., Zhang, G., Shi, J., and Grosse, R. (2019). “Functional variational Bayesian neural networks”. In: International Conference on Learning Representations.\n \n[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.\n\n[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.\n\n[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.\n\n[9] Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477), 359-378.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Good idea, could be interesting to the community were it proven in some other contexts. \n\n* 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.\n", "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.\n\n* Some literature relevant to the topic at hand seems to be missing from the discussion, or at least should be discussed more thoroughly:\n  \n  *  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].\n\n  * 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]. \n\n* 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.\n\n* (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. \n\n* (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.\n\n(_References included in the \"**Limitations**\" section_)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "#### Weaknesses\n\n[This is meant as a response to the first 2 bullet points]\nWe thank the reviewer for their helpful suggestions. We will include the suggested references and give a thorough discussion in the final version of the manuscript. It allows us to contrast infinite-dimensional gradient-flow methods with infinite-dimensional parameter-space methods. The functions space literature focuses on a formulation of the problem in infinite-dimensional function spaces—but still uses a finite-dimensional gradient flow based on parameterisations to implement its algorithms. In contrast, our method operates on a finite-dimensional parameter space, but implements an infinite-dimensional gradient flow. The function space view has the benefit that the functional loss  is often convex and consequently the target $Q^*$ can be unimodal. However, the variational stochastic process still requires parameterization to be computationally feasible—and in this sense, function space methods are FD-GVI approaches. The resulting objectives require a good approximation of the functional KL-divergence (which is challenging), and lead to a typically highly non-convex variational optimization problem in the parameterised space. While good initialization strategies may well be able to overcome some of these issues, this type of tuning requirement is common amongst FD-GVI approaches, a direct result of non-convexity, and typically not grounded in theory. Moving away from relying on these types of tuning strategies (whose effect is often poorly understood) serves as another motivation to attempt an infinite-dimensional gradient flow procedure: in principle, it allows us to exploit the convexity in the space of probability measures directly. The derived inference algorithms and the resulting asymptotic guarantees give some evidence to this conjecture: they hold regardless of the chosen initialisation.\n\nThat being said, we want to clarify that it is not our intention to question the validity of FD-GVI procedures or advocate for their abandonment: they are the dominant Bayesian deep learning paradigm for a reason, have considerable practical merits, and often work well in practice. We have mentioned and stressed this point more clearly in the new version of the manuscript. In spite of their practical utility, we do also believe that the mathematical challenges associated with FD-GVI serve as an additional motivation to investigate infinite-dimensional GD procedures. Indeed, doing this reproduced popular competing deep learning algorithms as diverse as DE and DLE—and even allowed us to derive DRLE, which provides a template for how our theory not only draws links btw existing deep learning algorithms, but can also inspire new ones.\n\n#### Questions\n1. LLA relies heavily on a function space perspective, which we did not adopt in our work. It is possible that a thorough investigation of function space methods would lead to further connections but we consider this beyond the scope of this paper. \n2. [This paragraph answers Question 2 and Limitation 1] Thank you for raising this point—this is indeed a very important point. We have included the following in the final version of the manuscript: \n> In essence, the core justification for these generalisations is that the very assumptions justifying application of Bayes’ Rule are violated in modern machine learning. In practical terms, this results in a view of Bayes’ posteriors as one—of many possible—measure-valued estimators $Q^∗$ of the form in (1). Once this vantage point is taken, it is not clear why one should be limited to using only one particular type of loss and regulariser for every possible problem. Seeking a parallel with optimisation on Euclidean domains, one may then compare the orthodox Bayesian view with the insistence on only using quadratic regularisation for any problem. While it is beyond the scope of this paper to cover these arguments in depth, we refer the interested reader to Knoblauch et al. (2022).\n\n3. Again, the reviewer is pointing out a very interesting avenue for future research. In principle, it is possible to use other divergences such as the ones mentioned by the reviewer (and indeed all f-divergences). However, often the implementation of the gradient flow then requires us to have access to the evolved pdf of the samples at time $t$. This can generically be replaced with a kernel density estimator based on the samples available at time $t$. However, it is well-known that kernel density estimators suffer greatly from the curse of dimensionality. Therefore, it is practically infeasible to use them in the context of deep learning (as the parameter space over which we need to compute said kernel density estimators) is huge. \n\n#### Limitations:\n1. See Question 2.\n2. We agree that the paper does not provide a full analysis of empirical performances. One reason for this is the limited amount of space. That being said, the more important reason is that we do not intend the paper to be a thorough study of any particular ‘market-ready’ methodology that can compete with the fine-tuned algorithms prevalent in industrial scale deep learning. Instead, our focus is on theoretical insight into what existing methods for uncertainty quantification in deep learning are actually doing—and this is reflected by our emphasis on connecting Bayesian and non-Bayesian methods through the analytical lens the paper proposes. As part of this, we conducted a range of experiments in the experimental section (incl. on UCI regression tasks) to showcase the limitations of naively adopting the current framework for designing new algorithms. In this sense, we believe that our empirical investigation serves the main purposes of the paper: to highlight how the derived theory aligns with reality, and to explain any divergence between theory and reality to help follow-up research with exploiting our ideas in order to develop more efficient and more effective algorithms.\n"}, {"review_id": "NVWAhmOLS2", "reviewer": "Reviewer_AtWA", "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).", "questions": "* [Q1] Eq. 6: What does the index $j$ relate to?\n* [Q2] l. 218: Are there any convergence results in the respective limits of $T$ and $N_E$?\n* [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?\n* [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$?\n* [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.\n* [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)?\n* [Q7] Fig. 2: Is there an explanation why DE exhibits this precise 50/50 spread of the probability mass?\n* [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?\n\n—\n\nMinor remarks\n\n* l. 53: Redundant “i” in “probability”\n* l. 72: Shouldn’t $D$ provide the mapping $(Q, P) \\mapsto D(Q, P)$ according to the stated domain?\n* 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.\n* l. 292: Is there a $d$ missing in front of $Q(\\theta)$ in the first double integral?\n* l. 297: $L^{FE}$ with capitalized superscript ($L^{fe}$ otherwise).\n* l. 329: White space after “minimisers”\n* l. 330: Remove either “which” or “that”\n* l. 337: “matters”\n* Fig. 3 (caption): White space in “FD-GVI”\n* Table 1: I would recommend removing the “Boston” dataset due to its racism issues. Also: “methods… outperform” or “method… outperforms”.\n* l. 355: “lens”", "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.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n* [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.\n* [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).\n", "weaknesses": "* [W1] **Analysis of DE behavior** (see Questions)\n  * It is not entirely clear if the paper studies arbitrary variants of DE or only a very narrowly defined version (see Q5). \n  * 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.\n* [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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "#### Weaknesses\n1.  One of the aims of this paper is to show that the naive strategy of train-and-repat can be understood as Wasserstein gradient flow of the probabilistic lifting of the loss function. Which strategies are covered depends on a case-by-case basis. Weight-decay for example is covered by our theory, as it just means that $\\ell$ is given as MSE+$\\lambda |\\theta|^2$. Batch-normalization is also covered in the sense that an ensemble of neural-networks which implement batch-normalization can be understood through the same WGF-lense. However, we did not want to give the impression that DEs are a bad strategy (even in their most basic form). Quite the opposite: we stress throughout the paper that they can be interpreted as implementing a rather sophisticated version of gradient descent and that especially in the multi-modal loss-landscapes of deep learning it is hard to come up with a better strategy. \n2.  Thank you for pointing this out, we will make a more concerted effort at explaining this in the updated manuscript. As our theory shows, a key difference between DE and D(R)LE is the form the solutions take: with infinitely many particles, D(R)LE produces a unique probability measure that has a continuous density (see Figure 2). In contrast, Theorem 1 shows that DEs produce probability measures that have atomic support: they have probability measure zero almost everywhere—except at a few individual points. In a sense, DEs are an extremely sparse representation of parameterisations of the neural network we care about. This is different for the measures produced by D(R)LEs: (with infinitely many particles) because they are densities, they try to assign probability mass to all regions of the parameter space. It is reasonable to believe that this would lead to a less sparse/better representation of parameterisations of the neural network we care about. Figure 2 supports this idea: in the setting where there are many more particles than minima, we can cover the parameter space of the neural network well—and in these settings, the purely atomic nature of DEs would generally be a drawback. To show this, we sampled the initial values for the parameters for the DE uniformly between -2.5 and 2.5. Since there are two minima in Figure 2 whose basins of attraction for gradient descent are around [-2.5, 0] and [2.5,0], a uniform initialization leads to the exact 50:50 spread. \nHowever, in practice, neural nets don’t look like Figure 2—they have lots and lots of minima, and we have very few particles to try and approximate the densities that look so clean and regular in Figure 2. That’s why we include Figure 4: it shows what actually happens in this scenario. If the number of particles is far lower than the number of minima, the particles in D(R)LE  are not enough to be a good approximation to a density of this type of space. Instead, they drift to the closest minima (exactly like they would in DEs!) and get stuck there for a long time. To show this phenomenon even more clearly, we have attached the original picture in Figure 3 together with a variant of the same experiment that uses only four particles to approximate the D(R)LE measures to this review. As is clear from that picture, when the number of particles is small relative to the number of minima, a naive interpretation of the theory is misguided; and D(R)LE behaves very similar to DEs. This is consistent with how the underlying equations evolve: When we discretise the evolution equations, what we get are basically slightly modified and randomized gradient descent schemes (see e.g. eq (7)). Theory (and experience from Langevin sampling for Bayesian posteriors) shows that as one keeps evolving the particle ensemble long enough, they will eventually escape and behave sufficiently differently from DEs—but in deep learning, we will not have the resources to keep the processes running forever. The naive expectation that we would get the type of behavior we see in Figure 2 in neural networks is therefore voided—instead, we should expect that D(R)LEs will act very similarly to DEs. \n3. Thank you for pointing this out; we will improve the paper and make these points clearer. Specifically we will explain that all integrals are over $R^J$ and all gradients are with respect to $\\theta$.\n\n#### Questions\n1. The number of ensemble members is $N_E$ (e.g. $N_E=10$). This means we train 10 neural networks where each neural network has its own set of parameters. The index $n$ therefore would run from 1 to 10 and denote the set of parameters corresponding to the n-th neural network.\n2. Yes, indeed. In Section 4.2 and Theorem 2 we show that under certain conditions our intuitions are validated and we obtain convergence results.\n3. In Figure 2 we illustrate that the samples we obtain from implementing our method for DLE and D(R)LE are close to the optimal measures.\n4. Yes, this is indeed the same. \n5. See answer to Weakness 1.\n6. Yes. Commonly, ‘non-parametric’ methods refer to the setting where the number of parameters can be arbitrarily high, and where in the limit of infinitely many parameters, one can recover an (infinite-dimensional) truth. In the context of our paper, the non-parametric object is indeed the targeted distribution itself, while the number of particles (of DE or D(R)LE) can be seen as the (arbitrarily large) number of parameters. As our theory shows, if the number of particles/parameters goes to infinity, we can recover the infinite-dimensional object we are targeting. This is different from parametric/parameterised problems, where there is not a natural way of making the parameter space arbitrarily large. E.g., typical variational families such as the family of normal distributions is parametric—but the collection of normal mixture distributions (with arbitrarily many and potentially infinitely many mixture components) would be non-parametric.\n7. See Weakness 1.\n8. See Weakness 2.\n\n#### Limitations\nSee global response point in 4."}, {"review_id": "10GLwalfSx", "reviewer": "Reviewer_4EE1", "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.   ", "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. \n2. 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. \n", "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 ? ", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper is well written, theory heavy and addresses important topic of deep ensembling and its connection with variational Bayes methods\n2. I am not so good with theory, but the theorems and equations looked ok to me without obvious mistakes.\n3. 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.\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "#### Questions\n1. This is indeed an excellent question. It is true that in practice, we will need to replace full gradients by mini-batch versions and for example the kernel mean embedding with its monte carlo estimator. The reviewer correctly observes that the theoretical results in Section 4 do not account for this type of approximation and sub-sampling. However, this is not unique to our analysis, and a rather common simplification. More importantly, we believe that the results built by ignoring this added complication already provide the most important insights into the differences between various algorithms and can be used as a starting point for further theoretical investigations.\nMore precisely, we think of the Wasserstein gradient flow as a powerful and principled tool to derive new inference algorithms for different types of regulariseres (MMD/KL and maybe others in the future). We can then combine it with standard plug-in estimators (like mini-batch gradient estimators) to immediately obtain an algorithm that approximately performs gradient descent in infinite dimensions.\n2. We thank the reviewer for pointing us to these alternatives for VI. We have included references for VI via normalizing flows and SIVI  in the main text. To the more general point: It is indeed true that one can make the class of variational inference more expressive, but this comes with a trade-off: If we use a more expressive approximating family, the KL-divergence term is typically not available in closed form anymore, which means we need to introduce approximations for the regulariser. Furthermore, the resulting optimization problems for the variational parameters are still highly non-convex and therefore often depend heavily on good initialisations. We are not aware of any works in the FD-GVI literature that obtain results competitive with deep ensembles. We believe that this is a direct consequence of the above problems.\n\n#### Limitations\n\nThe reviewer raises an excellent point. Results such as the one presented in Theorem 2 are only asymptotic in nature, and consequently we have no guarantee that the convergence is fast enough. These results should therefore be seen as necessary but not sufficient conditions for a good inference algorithm. Yet, stronger results which quantify the speed of convergence, would require us to make extremely strong assumptions that will be violated for deep learning. For example, we could trivially obtain quantitative results by citing the relevant literature albeit under very strong assumptions that would never be satisfied in the context of deep learning. Section 11.2 of  Ambrosio et al. (2005) shows that if $L$ is lambda-convex along generalized geodesics, we obtain exponentially fast convergence. Lambda-convexity of $L$ can be guaranteed if the potential $V$ is strictly convex (cf. Section 9.3 of Ambrosio et al. (2005)), which is surely never satisfied in deep learning applications. \n\nThe final version will make this difference between qualitative and quantitative results clearer, and discuss the strong assumptions required to obtain quantitative results.\n\nHowever, we want to point out that standard parameterized FD-GVI comes with no guarantees of any type. Although qualitative results have their shortcomings, they at least show that the inference algorithm in principle is powerful enough to solve the optimization problem at hand.\n\n\n\n#### References\n\nAmbrosio, L., Gigli, N., and Savare, G. (2005). Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media.\n"}, {"review_id": "lax8UOG0rg", "reviewer": "Reviewer_9edE", "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.\n\n[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.\n", "questions": "In Lines 188-191, the authors state that\n> 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.**\n\nI 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?\n", "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.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n", "weaknesses": "#### **Experiments section is difficult to follow**\nWhile 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.\n\n---\n\n#### **Motivation for convexification is unclear**\nWhile 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.\n\nAdditionally, 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^*$.\n\nHence a clearer exposition as to why regularized optima are preferred to the unregularized ones is needed.\n\n---\n\n#### **Motivation for DRLE is lacking**\nWhile 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. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "#### Details on experimental section\n\nAs the reviewer notes themselves, the paper is already rather densely packed. Because of space limitations, we were not able to include all relevant experimental details in the paper. We chose to focus on methodological details, but will move as many details into the main part in accordance with the available space in the final version. \n\n#### Motivation for convexification\n\nWe thank the reviewer for their comment. One needs to distinguish between two questions. The first one relates to the uniqueness of the minimiser of the optimisation problem. The point that we try to make with the footnote is simple: without regularization, the optimization problem can in general not have a unique minimiser. This relates to the bigger question that the reviewer raises: Why should one care about uniqueness? The main problem with non-uniqueness is that this typically translates into undesirable properties of the inference algorithm—for example, unless the solution is unique, it is not even clear which solution (of the many possible ones) an algorithm should target. This is illustrated by Theorem 1: the asymptotic distribution of Deep Ensembles depends strongly on initializations and how large the domain of attraction for a given local minimum is. Therefore, if we initialize poorly, we may end up putting most particles close to a comparatively poor local minimum.\n\nIdentifying a unique target to aim for, which in our case is $Q^*$, is typically seen as the first step in designing an inference algorithm with theoretical guarantees—after all, unless the target is clearly defined, it is not generally possible to assess an algorithm’s performance. Once the target has been identified, the second step is to build an inference algorithm that finds it.\n\nRegarding the interpretation of Figure 2: There may be some confusion about what $Q^*$ is in the plot. The dotted line in the background is the (unique and optimal) $Q^*$ of the optimisation problem; and the histogram clearly shows that DLE and DRLE do generate samples from this global minimiser (in line with our theory). On the other hand, DE does not (!). While it would be possible to find an initialisation distribution $Q_0$ such that DE converges to its global minimum, we would have to know where the global minimum (and its region of attraction) are located to do this—and this is precisely the kind of dependence on initialisation that is undesirable! DLRE and DLE converge to their respective optimal $Q^*$ regardless of initialization, thereby losing this undesirable trait. Importantly, note that this would be impossible to achieve if we hadn’t guaranteed the existence of a unique minimizer through convexification.\n\n#### Motivation for DRLE \n\nWe thank the reviewer for this great question. It is a point that we will discuss more thoroughly in the final version of the manuscript. The MMD is based on the kernel $\\kappa$ which introduces a repulsive effect between different sets of particles. This means that we can essentially choose our own metric that compares different parameter vectors (i.e. particles) with each other. It now depends heavily on the application, which sets of parameters should be considered ‘similar’. We chose a very basic kernel, the squared exponential, which essentially just measures the euclidean distance between the parameter vectors. However, in certain applications one might have a very good understanding of the type of diversity one would want to encourage and this knowledge could be embedded in the kernel $\\kappa$. We considered a thorough investigation of the effects of the kernel to be beyond the scope of this paper, but it is an interesting avenue for future research.\n\n\n#### Questions\n[PDE solvers] We are grateful the reviewer pointed out that this sentence is misleading. We mean the following: Once we have a closed form expression for the Wasserstein gradient (as for example in the KL and MMD case) the PDE in (4) gives us a way to determine the pdf of $Q(t)$ at any time $t$. We can simply apply numerical PDE solvers to find an approximation of $q(t,\\theta)$. However, PDE solvers quickly become computationally infeasible if the input space is larger than 2 or 3. As the parameter space in deep learning is typically at least in the thousands and more often than not in the millions or billions, it is computationally infeasible to deploy numerical PDE solvers. \n\n#### Limitations\nThis is an excellent point. We have written a response to this in the general rebuttal point 3. "}]

Author rebuttal:
### General Response: 

We want to thank all the reviewers for taking the time to read our manuscript so carefully and for providing valuable feedback that we believe will significantly improve the manuscript further. Overall, we have obtained a median score of 8, which is a wonderful reward for the countless hours that we have spent on this project and its technical content in the past year.  
We summarize the reviewer’s main feedback below, and explain how we have addressed the points in the uploaded and updated version of the manuscript which implements some of the most called-for changes.

1. **Reviewer NdiF** has expressed a wish for a more thorough discussion of FD-GVI in our manuscript, and has raised concerns that our representation of parameterised variational methods may fall short; and **Reviewer 4EE1** raised a similar concern. We thank both reviewers for pointing out this oversight, and have added to the discussion on related literature in Section 2.2, which now includes a paragraph on the state-of-the-art on function-space inference, implicit variational inference strategies, and normalizing flows. We believe this discussion will further highlight that the manuscript’s intent is not to imply that FD-GVI methods are impractical; and that rather, its intent is to thoroughly highlight their conceptual shortcomings relative to ID-GVI methods (which are the main subject of our investigations).
**Reviewer 6hMj** has rightfully pointed out that a sentence regarding the origin of the studied optimization problem in the introduction needs some reshaping to avoid overstating the contribution in Knoblauch et al. (2022). We agree with the reviewer, and have adjusted this sentence.

2. **Reviewer NdiF** has raised some valid points regarding the interpretability and meaning of the non-Bayesian nature of the measure-valued estimators which we derive. We agree that such a discussion is useful and necessary—we have added it to Section 2.1.

3. Both **Reviewer 9edE and 4EE1** have asked questions regarding a quantitative analysis of the approximation error caused by finite time, finite samples and the use of unbiased estimators. We thank the reviewers for raising this, and have added the following discussion in Section 4.3: 
> A notable shortcoming of Theorem 2 is its asymptotic nature. A more refined analysis would quantify how fast the convergence happens in terms of $N_E$, $T$, the SDE's discretisation error, and potentially even the use of unbiased estimators for the loss based on sub-sampling. 
While the existing literature could be adapted to derive the speed of convergence for DRLE in $T$ (Ambrosio et al.,
2005, Section 11.2), this would require a strong convexity assumption on the potential $V$, which will not be satisfied for any applications in deep learning. 
This is perhaps unsurprising: even for the Langevin algorithm—probably the most thoroughly analysed algorithm in this literature—no convergence rates have been derived that are applicable to the highly multi-modal target measures encountered in Bayesian deep learning  (Wibisono, 2019; Chewi et al., 2022).

4. **Several reviewers (6hMj, NdiF, AtWA)** pointed out that relative to the rest of the manuscript, its experimental section is comparatively limited. We agree that the paper does not provide a full comparison of empirical performances between different deep learning methods. The most important reason for this is that we do not intend the paper to be a thorough study of any particular ‘market-ready’ methodology that can compete with the fine-tuned algorithms prevalent in industrial scale deep learning. Instead, our focus is on theoretical insight into what existing methods for uncertainty quantification in deep learning are actually already doing. This is reflected by the title and by our emphasis on connecting Bayesian and non-Bayesian methods through the analytical lens the paper proposes. So while we did conduct a range of experiments in the experimental section (including on UCI regression tasks), this was done to showcase the usefulness of our framework as well as the limitations of naively adopting it for designing new algorithms—rather than as evidence that the proposed framework can outcompete prevalent deep learning paradigms. In this sense, we believe that our empirical investigation serves the main purposes of the paper: to highlight how the derived theory aligns with reality, and to explain any divergence between theory and reality to help follow-up research with exploiting our ideas in order to develop more efficient and more effective algorithms.

We have also included a number of smaller changes listed below:

1. Adding to why the MMD can encode useful properties as a regulariser as contrasted with the KLD (included in 4.3 already; **Reviewer 9edE**).

2. Being more explicit about what we mean by ‘challenging space’ (included in 2 already;  **Reviewer 6hMj**).

3. Explaining that all integrals are over the parameter space $R^J$ of $\theta$, and that similarly, all gradients are with respect to $\theta$ (included just before 2.1;  **Reviewer AtWA**)

4. Inclusion of a further setting for one of the experiments (the one for Figure 3) to make the messaging around the number of particles 
relative to the number of local minima clearer & more explicit (**Reviewer AtWA**; the resulting experiments are can be found in the attached PDF).  

In addition to changes that are already implemented, we will include further adjustments in the camera-ready version as listed below:

1. Adjustments to make the experimental section more readable and to state the purpose of the experimental section more clearly; with scope depending on space available (**Reviewer 9edE**)
2. Including a definition of the Wasserstein distance (**Reviewer NdiF, AtWA**), 
3.  Mentioning the two key assumptions required for Theorem 1. (**Reviewer 4EE1**)
4. More explicitly mention why it is hard to solve the PDE in (4) (**Reviewer 9edE**)



Meta-review: There has been a prevailing question as to why seemingly naive deep ensembles have surprisingly good properties for uncertainty quantification.  This paper tackles that question using an analysis through the lens of Wasserstein gradient flow and optimization in the space of measures.  Several variations of ensembling are derived within this framework.

The reviewers were unanimously in favor of acceptance, saying the work indeed provides a fresh and informative perspective about the behavior of deep ensembles.  The most common criticism amongst reviewers was that the experiments were too small scale and a bit confusing---which I think is a minor flaw (but one that the authors should look to improve, of course, especially wrt clarity).  There were also questions about the appropriateness / usefulness of Theorem 2, since it is an asymptotic result.

Notes for camera-ready version: D’Angelo & Fortuin [2021] also discuss deep ensembles with a repulsive term, obtained through the lens of Stein variational gradient descent / gradient flow in both parameter and function space.  Thus I was surprised to see their work mentioned only once and in passing.  I believe this paper goes beyond their work with a more rigorous treatment; yet I think the draft would be much improved to add more context about the relationship to their work.  This discussion should also be included along with the other related work mentioned in point 1 of the general rebuttal.
Decision: Accept (oral)