Title: A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods
Venue: NeurIPS 2023
Accepted: yes
Keywords: Wasserstein gradient flow

generalised variational inference

deep ensembles

Bayesian deep learning

variational Bayes
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

Abstract

We establish the first mathematically rigorous link between Bayesian, variational Bayesian, and ensemble methods. A key step towards this is 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 lens 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 standard 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.

1 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 a class of parameterised distributions (Graves, 2011; Blundell et al., 2015; Gal and Ghahramani, 2016; Louizos and Welling, 2017). An alternative approach is to (approximately) generate samples from the Bayes posterior via Monte Carlo methods (Welling and Teh, 2011; Neal, 2012). Deep ensembles are another approach, and rely on a train-and-repeat heuristic to quantify uncertainty (Lakshminarayanan et al., 2017).

Much ink has been spilled over whether one can see deep ensembles as a Bayesian procedure (Wilson, 2020; Izmailov et al., 2021; D'Angelo and Fortuin, 2021) and over how these seemingly different 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 a 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 a well-studied generally infinitedimensional regularised optimisation problem over the space of probability measures (see e.g. Guedj and Shawe-Taylor, 2019; Knoblauch et al., 2022). As a result, we find that the differences between

37th Conference on Neural Information Processing Systems (NeurIPS 2023).

these algorithms map directly onto different choices regarding this optimisation problem. Key differences between the algorithms boil down to different choices of regularisers, and whether they implement a finite-dimensional or infinite-dimensional gradient descent. Here, finite-dimensional gradient descent corresponds to parameterised VI schemes, whilst the infinite-dimensional case maps onto ensemble methods.

The contribution of this paper is a new theory that generates insights into existing algorithms and provides links between them that are mathematically rigorous, unexpected, and useful. On a technical level, our innovation consists in analysing them as algorithms that target an optimisation problem in the space of probability measures through the use of a powerful technical device: the Wasserstein gradient flow (see e.g. Ambrosio et al., 2005). While the theory is this paper's main concern, its potentially substantial methodological payoff is demonstrated through the derivation of a 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-unlike standard deep ensembles-is derived from a strictly convex objective defined over the space of probability measures. Thus, it targets a 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 convergent; and we hope that it can help plant the seeds for renewed innovations in theory-inspired algorithms for (Bayesian) deep learning.

The paper proceeds as follows: Section 2 discusses the advantages of lifting losses defined on Euclidean spaces into the space of probability measures through a generalised variational objective. Section 3 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 5, where the findings of the paper are illustrated numerically.

2 Convexification through probabilistic lifting

One of the most technically challenging aspects of contemporary machine learning theory is that the losses ℓ : R J → R we wish to minimise are often highly non-convex. For instance, one could wish to minimise ℓ ( θ ) := 1 N ∑ N n =1 ( y n -f θ ( x n ) ) 2 where { ( x n , y n ) } N n =1 is a set of paired observations and f θ a neural network with parameters θ . 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 ℓ has many local (or global) minima (Fort et 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 a finite-dimensional problem into an infinite-dimensional one, which is numerically more difficult to tackle. Figure 1 illustrates our approach:

min

e(0)

min

(0)Cp(o)

min

(0)Op(θ))

+ 入D(Q, P)

AEO

EP(R)

EP(IRJ

First, we transform a non-convex optimisation min θ ∈ Θ ℓ ( θ ) into an infinite-dimensional optimisation over the set of probability measures P ( R J ) , yielding min Q ∈P ( R J ) ∫ ℓ ( θ ) dQ ( θ ) . As an integral, this objective is linear in Q . However, linear functions are not strictly convex. We therefore need to add a strictly convex regulariser to ensure uniqueness of the minimiser. 2 We prove in Appendix A that indeed-for a regulariser D : P ( R J ) × P ( R J ) → [0 , ∞ ] such that ( Q,P ) ↦→ D ( Q,P ) is strictly convex and P ∈ P ( R J ) a fixed measure-existence and uniqueness of a global minimiser can be guaranteed. Given a scaling constant λ > 0 , we can now put everything together to obtain

the loss L and the unique minimiser Q ∗ as

Throughout, whenever Q and Q ∗ have an associated Lebesgue density, we write them as q and q ∗ . Moreover, all measures, densities, and integrals will be defined on the parameter space R J of θ . Similarly, the gradient operator ∇ will exclusively denote differentiation with respect to θ ∈ R J .

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 θ -averaged loss θ ↦→ ℓ ( θ ) 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 a corresponding weight. The choices for λ , D and P determine the trade-off between the initial loss ℓ and reference measure P and therefore the weights we assign to different solutions.

Yet, (1) is not a new problem form: it has various interpretations, depending on the choices for D,λ,ℓ and the framework of analysis (see e.g. Knoblauch et al., 2022, for a discussion). For example, if ℓ ( θ ) is a negative log likelihood, D is the Kullback-Leibler divergence ( KL ), and λ = 1 , Q ∗ is the standard Bayesian posterior, and P is the Bayesian prior. This interpretation of P as a prior carries over to generalised Bayesian methods, in which we can choose ℓ ( θ ) to be any loss, D to be any divergence on P ( R J ) , and λ to regulate how fast we learn from data (see e.g. Bissiri et al., 2016; Jewson et al., 2018; Knoblauch et al., 2018; Miller and Dunson, 2019; Knoblauch, 2019; Alquier, 2021a; Husain and Knoblauch, 2022; Matsubara et al., 2022; Wild et al., 2022; Wu and Martin, 2023; Altamirano et al., 2023). 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).

A second line of research featuring objectives as in (1) are PAC-Bayes methods, whose aim is to construct generalisation bounds (see e.g. Shawe-Taylor and Williamson, 1997; McAllester, 1999b,a; Gr¨ unwald, 2011). Here, ℓ is a general loss, but P only has the interpretation of some reference measure that helps us measure the complexity of our hypotheses via Q ↦→ λD ( Q,P ) (Guedj and Shawe-Taylor, 2019; Alquier, 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; B´ egin 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 (1) as a 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 min Q ν ∈Q L ( Q ) for a set of measures Q := { Q ν : ν ∈ Γ } ⊂ P ( R J ) parameterised by a parameter ν ∈ Γ ⊆ R I . To find Q ν , one now simply performs (finite-dimensional) gradient descent with respect to the function ν ↦→ L ( Q ν ) . For the special case where P is a Bayesian prior, λ = 1 , ℓ ( θ ) is a negative log likelihood parametrised by θ , and D = KL , this recovers the well-known 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-GVI) 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

choose Q and P to be simple distributions such as Gaussians to ensure that L ( Q ν ) is a tractable function of ν . This often results in a Q that is unlikely to contain a good approximation to Q ∗ ; raising doubt if Q ν ∗ can approximate Q ∗ in any meaningful sense. Secondly, even if Q ↦→ L ( Q ) is strictly convex on P ( R J ) , the parameterised objective ν ∈ Γ ↦→ L ( Q ν ) is usually not. Hence, there is no guarantee that gradient descent leads us to Q ν ∗ . This point also applies to very expressive variational families (Rezende and Mohamed, 2015; Mescheder et al., 2017) which may be sufficiently rich that Q ∗ ∈ Q , but whose optimisation problem ν ∈ Γ ↦→ 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 a rigorous theoretical analysis outside of narrowly defined settings.

FD-GVI in function space: A collection of approaches formulated as infinite-dimensional problems are GVI methods on an infinite-dimensional function space (Ma et al., 2019; Sun et al., 2018; Ma and Hern´ andez-Lobato, 2021; Rodriguez-Santana et al., 2022; Wild et al., 2022). Here, the loss is often convex in function space. In practice 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 often challenging), and lead to a typically highly non-convex variational optimization problem in the parameterised space.

Infinite-dimensional GVI: Instead of minimising the (non-convex) problem ν ↦→ L ( Q ν ) , we want to exploit the convex structure of Q ↦→ L ( Q ) . Of course, a priori it is not even clear how to compute the gradient for a function Q ↦→ L ( Q ) defined on an infinite-dimensional nonlinear space such as P ( R J ) . 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 (1) 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 infinite-dimensional gradient descent as infinite-dimensional GVI (ID-GVI) . Infinite-dimensional gradient descent methods have recently gained attention in the 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 ∈ P ( R J ) ↦→ E ( Q, ̂ Q ) where E ( · , · ) is a discrepancy measure. Some methods apply the Wasserstein gradient flow (WGF) for different choices of E (Arbel et al., 2019; Korba et al., 2021; 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 a 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 et 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 a 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 (1), and not known explicitly.

3 Gradient flows in finite and infinite dimensions

Before we can realise our ambition to solve (1) with an ID-GVI scheme, we need to cover the relevant bases. To this end, we will discuss gradient flows in finite and infinite dimensions, and explain how they can be used to construct infinite-dimensional gradient descent schemes. In essence, a gradient 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 finds local minima of losses ℓ : R J → R by iteratively improving an initial guess θ 0 ∈ R J through the update θ k +1 := θ k -η ∇ ℓ ( θ k ) , k ∈ N , where η > 0 is a step-size and ∇ ℓ denotes the gradient of ℓ . For sufficiently small η > 0 , this update can equivalently be written as

Gradient flows formalise the following logic: for any fixed η , we can continuously interpolate the corresponding gradient descent iterates { θ k } k ∈ N 0 . To do this, we simply define a function θ η : [0 , ∞ ) → R as θ η ( t ) := θ t/η for t ∈ η N 0 := { 0 , η, 2 η, ... } . For t / ∈ η N 0 we linearly interpolate 3 . As η → 0 , the function θ η converges to a differentiable function θ ∗ : [0 , ∞ ) → R called the gradient flow of ℓ , because it is characterised as solution to the ordinary differential equation (ODE) θ ′ ∗ ( t ) = -∇ ℓ ( θ ∗ ( t )) with initial condition θ ∗ (0) = θ 0 . Intuitively, θ ∗ ( t ) is a 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-time gradient descent as though it were a continuous gradient flow-with the hope that for sufficiently small η , the behaviour of both will essentially be the same.

Our results in the infinite-dimensional 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 infinite-dimensional setting.

3.1 Gradient flows in Wasserstein spaces

Let P 2 ( R J ) be the space of probability measures with finite second moment equipped with the 2-Wasserstein metric given as

where C ( P, Q ) ⊂ P ( R J × R J ) denotes the set of all probability measure on R J × R J such that π ( A × R J ) = P ( A ) and π ( R J × B ) = Q ( B ) for all A,B ⊂ R J (see also Chapter 6 of Villani et al., 2009). Further, let L : P 2 ( R J ) → ( -∞ , ∞ ] be some functional-for example L in (1). In direct analogy to (2), we can improve upon an initial guess Q 0 ∈ P 2 ( R J ) by iteratively solving

for k ∈ N 0 and small η > 0 (see Chapter 2 of Ambrosio et al., 2005, for details). Again, for η → 0 , an appropriate limit yields a continuously indexed family of measures { Q ( t ) } t ≥ 0 . If L is sufficiently smooth and Q 0 = Q (0) has Lebesgue density q 0 , the time evolution for the corresponding pdfs { q ( t ) } t ≥ 0 is given by the partial differential equation (PDE)

with q (0 , · ) = q 0 (Villani, 2003, Section 9.1). Here ∇· f := ∑ J j =1 ∂ j f j denotes the divergence operator and ∇ W L [ Q ] : R J → R J the Wasserstein gradient (WG) of L at Q . For the purpose of this paper, it is sufficient to think of the WG as a gradient of the first variation; i.e. ∇ W L [ Q ] = ∇ L ′ [ Q ] where L ′ [ Q ] : R J → R J is the first variation of L at Q (Villani et al., 2009, Exercise 15.10). The Wasserstein gradient flow (WGF) for L is then the solution q † to the PDE (3). If L is chosen as in (1), our hope is that the logic of finite-dimensional gradient descent carries over; and that lim t →∞ q † ( t, · ) is in fact the density q ∗ corresponding to Q ∗ .

Following this reasoning, this paper applies the WGF for (1) to obtain an ID-GVI algorithm. In Section 4 and Appendices C-F, we formally show that the WGF indeed yields Q ∗ (in the limit as t →∞ ) 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 infinite-dimensional gradient descent for (1). 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. Rather than trying to first approximate the q solving (3) and then sampling from its limit in a second step, we will instead formulate equations which replicate how the

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 is wellstudied in the context of thermodynamics (Santambrogio, 2015, Chapter 7), and which recovers various forms of the GVI problem in (1) (see Section 4). In thermodynamics, Q ∈ P 2 ( R J ) describes the distribution of particles located at specific points in R J . The overall energy of a collection of particles sampled from Q is decomposed into three parts: (i) the external potential V ( θ ) which acts on each particle individually, (ii) the interaction energy κ ( θ, θ ′ ) describing pairwise interactions 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

for Q ∈ P 2 ( R J ) with Lebesgue density q , λ 1 ≥ 0 , λ 2 ≥ 0 . Note that for λ 2 > 0 we implicitly assume that Q has a density. Following Section 9.1 in Villani et al. (2009), its WG is

where θ ∈ R J , and ∇ 1 κ 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 a nonlinear Fokker-Planck equation associated with a stochastic process of McKean-Vlasov type (see Appendix B for details). Fortunately for us, it is well-known that such processes can be approximated through interacting particles (Veretennikov, 2006) generated by the following procedure:

Step 1: Sample N E ∈ N particles θ 1 (0) , . . . , θ N E (0) independently from Q 0 ∈ P 2 ( R J ) .

Step 2: Evolve the particle θ n by following the stochastic differential equation (SDE)

for n = 1 , . . . , N E , and { B n ( t ) } t> 0 stochastically independent Brownian motions.

As N E →∞ , the distribution of θ 1 ( t ) , . . . , θ N E ( t ) evolves in t in the same way as the sequence of densities q ( t, · ) solving (3). This means that we can implement infinite-dimensional 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 θ 1 ( T ) , . . . , θ N E ( T ) for N E ∈ N and T > 0 . Our intuition and Section 4 tell us that, as desired, the distribution of θ 1 ( T ) , . . . , θ N E ( T ) will be close to the global minimiser of L fe .

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 (1). Since N E and T are finite, and since we need to discretise (131), there will be an approximation error. Given this, how good are the samples produced by such methods? As we 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 it yields ID-GVI algorithms to solve (1). We put particular focus on the analysis of the regulariser D ; providing new perspectives on heuristics for uncertainty quantification in deep learning in the process. Specifically, we establish formal links explaining how they may (not) be understood as a Bayesian procedure. Beyond that, we derive the WGF associated with regularisation using the maximum mean discrepancy, and provide a theoretical analysis of its convergence properties.

4.1 Unregularised probabilistic lifting: Deep ensembles

We start the analysis with the base case of an unregularised functional L ( Q ) = ∫ ℓ ( θ ) dQ ( θ ) , corresponding to λ = 0 in (1). This is also a special case of (4) with λ 1 = λ 2 = 0 . As λ 1 = 0 , there is no interaction term, and all particles can be simulated independently from one another as

This simple algorithm happens to coincide exactly with how deep ensembles (DEs) are constructed (see e.g. Lakshminarayanan et al., 2017). In other words: the simple heuristic of running gradient descent algorithm several times with random initialisations sampled from Q 0 is an approximation of the WGF for the unregularised probabilistic lifting of the loss function ℓ .

Following the WGF in this case does not generally produce samples from a global minimiser of L . Indeed, the fact that L generally does not even have a unique global minimiser was the motivation for regularisation in (1). Even if L had a 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 Q 0 .

Theorem 1. If ℓ has countably many local minima { m i : i ∈ N } , then it holds independently for each n = 1 , . . . , N E that

for t → ∞ . Here D -→ denotes convergence in distribution and Θ i = { θ ∈ R J : lim t →∞ θ ∗ ( t ) = m i and θ ∗ (0) = θ } denotes the domain of attraction for m i with respect to the gradient flow θ ∗ .

A proof with technical assumptions-most importantly a version of the famous Lojasiewicz inequality-is in Appendix C. Theorem 1 derives the limiting distribution for θ 1 ( T ) , . . . , θ N E ( T ) , which shows that-unless all local minima are global minima-the WGF does not generate samples from a global minimum of Q ↦→ L ( Q ) for the unregularised case λ = 0 . Note that the conditions 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-that Q ∞ strongly depends on the initialisation Q 0 . Importantly, the dependence of Q ∞ on Q 0 remains true for all losses constructed via deep learning architectures. However, despite these theoretical shortcomings, DEs remain highly competitive in practice and typically beat FD-GVI methods like standard VI (Ovadia et al., 2019; Fort et al., 2019). This is perhaps not surprising: DEs implement an infinite-dimensional 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 5, and find that it is a consequence of the fact that in deep learning, N E is small compared to the number of local minima (cf. Figure 4).

4.2 Regularisation with the Kullback-Leibler divergence: Deep Langevin ensembles

In Section 2, we argued for regularisation by D to ensure a unique minimiser Q ∗ . The KullbackLeibler divergence ( D = KL ) is the canonical choice for (generalised) Bayesian and PAC-Bayesian methods (Bissiri et al., 2016; Knoblauch et al., 2022; Guedj and Shawe-Taylor, 2019; Alquier, 2021b). Now, Q ∗ has a known form: if P has a pdf p , it has an associated density given by q ∗ ( θ ) ∝ exp( -1 λ ℓ ( θ )) p ( θ ) (Knoblauch et al., 2022, Theorem 1).

Notice that the KL -regularised version of L in (1) can be rewritten in terms of the objective L fe in (4) by setting V ( θ ) = ℓ ( θ ) -λ log p ( θ ) , λ 1 = 0 and λ 2 = λ . Compared to the unregularised objective of the previous section (where λ = 0 ), the external potential is now adjusted by -λ log p ( θ ) , 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 λ 1 = 0 , the corresponding particle method still does not have an interaction and is given as

Clearly, this is just the Langevin SDE and we call this approach the deep Langevin ensemble (DLE) . While the name may suggest that DLE is equivalent to the unadjusted Langevin algorithm (ULA) (Roberts and Tweedie, 1996), this is not so: for T discretisation steps t 1 , t 2 , . . . t T , DLE approximates measures using the end-points of N E trajectories given by { θ n ( t T ) } N E n =1 . In contrast, ULAwoulduse a (sub)set of the samples { θ 1 ( t i ) } T i =1 generated from one single particle's trajectory . To analyse DLEs, we build on the Langevin dynamics literature: in Appendix D, we show that θ n ( t ) D -→ Q ∗ as t → ∞ , independently for each n = 1 , . . . , N E . Hence θ 1 ( T ) , . . . , θ N E ( 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 ∗ as defined per (1) is unique, has the form of a Gibbs measure, and can be sampled from using (6). In contrast, unregularised DE produces samples from Q ∞ in Theorem 1 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 is 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-mean discrepancy (MMD) (see e.g. Gretton et al., 2012) only, and then a combination of MMDand KL.

For a kernel κ : R J × R J → R , the squared MMD between measures Q and P is

MMD measures the difference between within-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 (1) with D = MMD 2 introduces interactions characterised precisely by the kernel κ , and we can show this explicitly by rewriting L of (1) into the form of L fe in (4). In other words, inclusion of MMD 2 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 µ P ( θ ) := ∫ κ ( θ, θ ′ ) dP ( θ ′ ) , we see that up to a constant not depending on Q , L ( Q ) = L fe ( Q ) for V ( θ ) = ℓ ( θ ) -λ 1 µ P ( θ ) , λ = λ 1 2 , and λ 2 = 0 . While we can show that a 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 2 .

To remedy this problem, we additionally regularise with the KL : since KL( Q,P ) = ∞ if P has a Lebesgue density but Q has not, this now guarantees that Q ∗ has a density q ∗ . In terms of (1), this means that D = λ MMD 2 + λ ′ KL . Adding regularisers like this has a long tradition, and is usually done to combine the different strengths of various regularisers (see e.g. Zou and Hastie, 2005). Here, we 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 fe ( Q ) up to a constant not depending on Q by taking λ = λ 1 2 , λ ′ = λ 2 , and V ( θ ) = ℓ ( θ ) -λ 1 µ P ( θ ) -λ 2 log p ( θ ) . Using the same algorithmic blueprint as before, we evolve particles according to (5). As these particles follow an augmented Langevin SDE that incorporates repulsive particle interactions via κ , 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 Q n,N E ( T ) be the distribution of θ n ( T ) , n = 1 , . . . , N E , generated via (5) . Then Q n,N E ( T ) D -→ Q ∗ for each n = 1 , . . . , N E and as N E →∞ , T →∞ .

This is remarkable: we have constructed an algorithm that generates samples from the global minimiser Q ∗ -even though a formal expression for what exactly Q ∗ looks like is unknown! This

2.5

0.0

2.5

0

2

Frequency

DE

2.5

0.0

2.5

0.0

0.5

DLE

2.5

0.0

2.5

0.0

0.5

DRLE

2.5

0.0

2.5

0

20

Loss function

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 a proportionality constant.

Anotable shortcoming of Theorem 2 is its asymptotic nature. A more refined analysis could quantify how fast the convergence happens in terms of N E , 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 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). That being said, for the case of deep learning, FD-GVI approaches fail to provide even the most basic asymptotic 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, and-by virtue of Theorem 1-over DEs as well.

5 Experiments

Since the paper's primary focus is on theory, we use two experiments to reinforce some of the predictions it makes in previous sections, and a 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 on https://github.com/sghalebikesabi/GVI-WGF .

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 1-does not correspond to the global minimiser of Q ↦→ ∫ ℓ ( θ ) dQ ( θ ) over P ( R J ) (which is given as Dirac measure located at θ = -2 ).

FD-GVI vs ID-GVI: Figure 3 illustrates 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 of Q ∗ . This is because the ID-GVI methods explore the whole space P 2 ( R J ) , whilst FD-GVI is limited to learning a unimodal Gaussian.

DEs vs D(R)LEs, and why finite N E matters: Table 1 compares DE, DLE and DRLE on a number of real world data sets, and finds a rather random distribution of which method performs best. This seems to contradict our theory, and suggests there is essentially no difference between regularised and unregularised ID-GVI. What explains the discrepancy? Essentially, it is the fact that N E 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 4 with a toy example. We choose a uniform prior P and initialisation Q 0 and the loss ℓ ( θ ) := -| sin( θ ) | , θ ∈ [ -1000 π, 1000 π ] , which has 2000 local minima. Correspondingly Q ∗ will have many local modes for all methods. Note that ∇ V is the same for all approaches since log p and µ P 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

5

0

5

5

0

5

DE

5

0

5

DLE

5

0

5

DRLE

5

0

5

FD-GVI

0

1

2000

0

2000

0

2

4

6

Count

DE

DLE

DRLE

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-a well-known problem for Markov Chain Monte Carlo methods (see e.g. Syed et al., 2022).

6 Conclusion

In this paper, we used infinite-dimensional gradient descent via Wasserstein gradient flows (WGFs) (see e.g. Ambrosio et al., 2005) and the lens of generalised variational inference (GVI) (Knoblauch et al., 2022) to unify a collection of existing algorithms under a 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.

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, P. (2021b). User-friendly introduction to PAC-Bayes bounds. arXiv preprint arXiv:2110.11216 .

Alquier, P. and Guedj, B. (2018). Simpler PAC-Bayesian bounds for hostile data. Machine Learning , 107(5):887-902.

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´ e, 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¨ ockner, M. (2020). From nonlinear Fokker-Planck equations to solutions of distribution dependent sde. arXiv preprint arXiv:1808.10706 .

B´ egin, L., Germain, P., Laviolette, F., and Roy, J.-F. (2016). PAC-Bayesian bounds based on the R´ enyi 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 belief distributions. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 78(5):1103-1130.

Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015). Weight uncertainty in neural network. In 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 on Learning Theory , pages 1-2. PMLR.

Chiang, T.-S., Hwang, C.-R., and Sheu, S. J. (1987). Diffusion for global optimization in R n . 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 , 62(10):4189-4196.

Fort, S., Hu, H., and Lakshminarayanan, B. (2019). Deep ensembles: A loss landscape perspective. arXiv preprint arXiv:1912.02757 .

Gal, Y. and Ghahramani, Z. (2016). Dropout as a 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 heuristic. arXiv preprint arXiv:1707.07269 .

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 , 34:8018-8031.

Graves, A. (2011). Practical variational inference for neural networks. Advances in neural information processing systems , 24.

Gretton, A., Borgwardt, K. M., Rasch, M. J., Sch¨ olkopf, B., and Smola, A. (2012). A kernel twosample test. Journal of Machine Learning Research , 13(25):723-773.

Gr¨ unwald, P. (2011). Safe learning: bridging the gap between Bayes, MDL and statistical learning theory via empirical convexity. In Proceedings of the 24th Annual Conference on Learning Theory , pages 397-420.

Guedj, B. and Shawe-Taylor, J. (2019). A primer on pac-Bayesian learning. In ICML 2019-Thirtysixth International Conference on Machine Learning .

Haddouche, M. and Guedj, B. (2023). Wasserstein PAC-Bayes learning: A bridge between generalisation and optimisation. arXiv preprint arXiv:2304.07048 .

Husain, H. and Knoblauch, J. (2022). Adversarial interpretation of Bayesian inference. In International Conference on Algorithmic 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? In International conference on machine learning , pages 46294640. PMLR.

Jewson, J., Smith, J., and Holmes, C. (2018). Principles of Bayesian inference using general divergence criteria. Entropy , 20(6):442.

Knoblauch, J. (2019). Frequentist consistency of generalized variational inference. arXiv preprint arXiv:1912.04946 .

Knoblauch, J. (2021). Optimization-centric generalizations of Bayesian inference . PhD thesis, University of Warwick.

Knoblauch, J., Jewson, J., and Damoulas, T. (2018). Doubly robust Bayesian inference for nonstationary streaming data using β -divergences. In Advances in Neural Information Processing Systems (NeurIPS) , pages 64-75.

Knoblauch, J., Jewson, J., and Damoulas, T. (2022). An optimization-centric 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. Cambridge University Press.

Korba, A., Aubin-Frankowski, P.-C., Majewski, S., and Ablin, P. (2021). Kernel Stein discrepancy descent. In International Conference on Machine Learning , pages 5719-5730. PMLR.

Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30.

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.

Li, Y. and Turner, R. E. (2017). Gradient estimators for implicit models. arXiv preprint arXiv:1705.07107 .

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, L., and Belkin, M. (2022). Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis , 59:85116.

Liu, Q. (2017). Stein variational gradient descent as 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, J. (2019). Scaling limit of the stein variational gradient descent: The mean field regime. SIAM Journal on Mathematical Analysis , 51(2):648-671.

Ma, C. and Hern´ andez-Lobato, J. M. (2021). Functional variational inference based on stochastic process generators. Advances in Neural Information Processing Systems , 34:21795-21807.

Ma, C., Li, Y., and Hern´ andez-Lobato, J. 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 variational 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):1113-1125.

Muandet, K., Fukumizu, K., Sriperumbudur, B., Sch¨ olkopf, B., et al. (2017). Kernel mean embedding 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 Snoek, 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 ECE563 (UIUC) and , 6(2012-2016):7.

Rezende, D. 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 of Langevin distributions and their discrete approximations. Bernoulli , pages 341-363.

Rodriguez-Santana, S., Zaldivar, B., and Hernandez-Lobato, D. (2022). Function-space inference with sparse implicit processes. In International Conference on Machine Learning , pages 1872318740. PMLR.

Santambrogio, F. (2015). Optimal transport for applied mathematicians. Birk¨ auser, NY , 55.

Santambrogio, F. (2017). { Euclidean, metric, and Wasserstein } gradient flows: an 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 on Computational learning theory , volume 6, pages 2-9.

Shi, J., Sun, S., and Zhu, J. (2018). A spectral approach to gradient estimation for implicit distributions. 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ˆ ot´ e, A., Deligiannidis, G., and Doucet, A. (2022). Non-reversible parallel tempering: a scalable highly parallel mcmc scheme. Journal of the Royal Statistical Society Series B: Statistical Methodology , 84(2):321-350.

Veretennikov, A. Y. (2006). On ergodic measures for McKean-Vlasov stochastic equations. In Monte Carlo and Quasi-Monte Carlo Methods 2004 , pages 471-486. 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-11) , pages 681688.

Wibisono, A. (2019). Proximal langevin algorithm: Rapid convergence under isoperimetry. arXiv preprint arXiv:1911.01469 .

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 a probabilistic perspective of generalization. Advances in neural information processing systems , 33:4697-4708.

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):301-320.

A Existence and uniqueness of global minimiser

In this section, we discuss assumptions under which the global minimiser of the optimisation problem

over P ( R J ) exists and is unique. We assume throughout that the optimisation problem is not pathological, in the sense that there exists a measure ̂ Q ∈ P ( R J ) such that L ( ̂ Q ) < ∞ . This is in applications often trivial to verify. A good candidate for ̂ Q is typically the reference measure P .

Loss assumptions Let ℓ : R J → R be a loss satisfying the following assumptions:

(L1) The loss ℓ is bounded from below which means that

(L2) The loss is norm-coercive which means that

if ∥ θ ∥ → ∞ .

(L3) The loss ℓ is lower semi-continuous which means that

for all θ 0 ∈ R J .

Regulariser assumptions Let D : P ( R J ) × P ( R J ) → [0 , ∞ ] be a regulariser and P ∈ P ( R J ) a reference measure. We define D P ( · ) := D ( · , P ) for notational convenience. We assume the following for D P :

(D1) The function D P is lower semi-continuous w.r.t. to the topology of weak-convergence, i.e. for all sequences ( Q n ) n ∈ N ⊂ P ( R J ) and all Q with D P ( Q ) < ∞ , it holds that Q n D -→ Q implies

Here, D -→ denotes convergence in distribution.

̸

(D2) D P is strictly convex, i.e. for all Q 1 = Q 2 ∈ P ( R J ) with D P ( Q 1 ) < ∞ and D P ( Q 2 ) < ∞ , it holds that

with α ∈ (0 , 1) .

The next theorem provides an existence result for the optimisation problem (7). The result is similar in spirit to Lemma 2.1 in Knoblauch (2021) with the important difference that our assumptions are easier to verify, since they are formulated in terms of ℓ and D P .

Theorem 3 (Existence of global minimiser) . Under the assumptions (L1)-(L3) and (D1) there exists a probability measure Q ∗ ∈ P ( R J ) with

Proof. Let c > -∞ be the lower bound for ℓ . It follows immediately that L ( Q ) ≥ c for all Q ∈ P ( R J ) since D ( P, Q ) ≥ 0 . As a consequence we know that

By definition of the infimum we can construct a sequence l n = L ( Q n ) ∈ R in the image of L such

for n → ∞ . We now show by contradiction that the corresponding sequence ( Q n ) ⊂ P ( R J ) is tight 4 . Assume that ( Q n ) is not tight. By definition we can then find an ϵ > 0 such that for each k ∈ N there exists n = n k ∈ N with Q n k ([ -k, k ] J ) ≤ 1 -ϵ . We set A k := [ -k, k ] J ⊂ R J and obtain

Due to the coerciveness of ℓ , we know that inf { ℓ ( θ ) : θ ∈ R J \ A k } →∞ for k →∞ and therefore l n k → ∞ for k → ∞ . However, this is a contradiction: The sequence ( l n ) is convergent and therefore in particular bounded. As a consequence, it cannot contain the unbounded sub-sequence ( l n k ) . It follows that the sequence ( Q n ) is tight. By Prokhorov's theorem we can now extract a sub sequence ( Q n k ) of ( Q n ) and a measure Q ∗ ∈ P ( R J ) such that

for k →∞ . Due to Lemma 5.1.7 in Ambrosio et al. (2005) the lower semi-continuity of ℓ implies that Q ↦→ ∫ ℓ ( θ ) dQ ( θ ) is lower semi-continuous. This combined with the lower semi-continuity of D P gives

From this it immediately follows that

but by definition L ∗ is the global minimum of L which implies L ∗ ≤ L ( Q ∗ ) . We therefore conclude that L ( Q ∗ ) = L ∗ .

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 4 (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 Q 1 , Q 2 ∈ P ( R J ) such that

where ∞ > L ∗ := inf { L ( Q ) : Q ∈ P ( R J ) } > -∞ . We define the probability measure Q 3 := 1 2 Q 1 + 1 2 Q 3 . By strict convexity we obtain

which is a contradiction to Q 1 and Q 2

being global minimisers.

Note that in the literature on GVI (Knoblauch et al., 2022) it is common to assume that the regulariser is definite, i.e.

for all P, Q ∈ P ( R J ) . We did not use this assumption in neither Theorem 3 nor Theorem 4. However, the next lemma shows that it is basically implied by strict convexity.

Lemma 1. Let D P : P ( R J ) → [0 , ∞ ] be strictly convex and assume further D ( Q,Q ) = 0 for all Q ∈ P ( R J ) . Then it follows that D ( Q,P ) = 0 implies P = Q .

̸

Proof. We prove the claim by contradiction. Assume that there exists P = Q such that D ( P, Q ) = 0 . The strict convexity and D ( P, P ) = 0 imply combined that

However, we know that D ( 1 2 P + 1 2 Q,P ) ≥ 0 by assumption. This is a contradiction.

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 therefore in particular lower semi-continuous). However, norm-coercivity can be violated. Consider for example the squared loss

where f θ is the parametrisation of a neural network with one hidden layer, i.e. θ = ( w,A ) and

where σ : R → R is an activation function which is applied pointwise to the vector Ax and has the property that σ (0) = 0 . It is now possible to find a sequence of parameters ( θ k ) k ∈ N ⊂ R J with ∥ θ k ∥ → ∞ such that ℓ ( θ k ) does not converge to infinity. Define w k := k ( 1 . . . 1 ) , A k := 0 and θ k = ( w k A k ) for k ∈ N . Then we obviously have that

which is constant and therefore does not converge to ∞ . A 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

with weight decay. This loss is by construction norm-coercive and therefore the previous existence proof applies.

Discussion on regulariser assumptions The assumptions (D1) and (D2) are quite weak. The KLdivergence 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 on κ (Muandet et al., 2017, Lemma 3.1). The lower semi-continuity properties also depend on the kernel κ . However, for bounded kernels it is trivial to verify. We include the proof for completeness, but assume this has been shown before elsewhere.

for k →∞ but

Lemma 2. Let the kernel κ : R J × R J be continuous and bounded: ∥ κ ∥ ∞ := sup θ,θ ′ ∈ R J | k ( θ, θ ′ ) | < ∞ and P be fixed. Then MMD( · , P ) is continuous and therefore, in particular, lower semi-continuous.

Proof. Let ( Q n ) n ∈ N and Q ∗ be such that for n →∞ . This immediately implies that

for n → ∞ , 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 a continuous function and bounded since

By the definition of weak convergence for measures, we therefore have

for n → ∞ . 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´ ern kernel, are continuous and bounded and therefore Lemma 2 applies.

Remark 1. The astute reader may have noticed that our existence proof only guarantees the existence of measure Q ∗ ∈ P ( R J ) . However, the Wasserstein gradient flow is by definition only formulated in the space of probability measures with finite second moment, denoted P 2 ( R J ) . Assumptions which guarantee that Q ∗ ∈ P 2 ( R J ) are easy to formulate. For example, we can require that there exists C > 0 and R > 0 such that the loss ℓ satisfies

for all ∥ θ ∥ > R . This immediately implies that Q ∗ ∈ P 2 ( R J ) since otherwise

gives a contradiction to the finiteness of L ( Q ∗ ) . However, even if (43) is violated, the reference measure P may still guarantee that Q ∗ ∈ P 2 ( R J ) . For example, if P ∈ P 2 ( R J ) , then D P ( Q ∗ ) will typically be large if Q ∗ / ∈ P 2 ( R J ) and the global minimiser is therefore in a sense unlikely to have fat tails. We therefore assume Q ∗ ∈ P 2 ( R J ) throughout the paper and consider it to be a minor practical concern.

B Realising the Wasserstein gradient flow

In this section, we identify a suitable stochastic process that allows us to follow the WGF.

Let L fe : P ( R J ) → ( -∞ , ∞ ] be the free energy discussed in Section 3.2 given as

where λ 1 , λ 2 ≥ 0 are constants, V : R J → R is the potential, κ : R J × R J → R is symmetric. We will write L for L fe from now on to simplify notation. The Wasserstein gradient of L is given as (cf. Chapter 9.1 Villani, 2003, Equation 9.4)

where ∇ 1 κ : R J × R J → R J is the (vector-valued) derivative of κ with respect to the first component, ∇ denotes the euclidean gradient with respect to θ and ( ∇ 1 κ ∗ Q )( θ ) := ∫ ∇ 1 κ ( θ, θ ′ ) dQ ( θ ′ ) for θ ∈ R J . The corresponding Wasserstein gradient flow is therefore given as (cf. Chapter 9.1 Villani, 2003, Equation 9.3)

In general the probability density evolution of a stochastic process is-via the Fokker-Planck equation-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 ∈ C 2 c ( R J , R ) where C 2 c ( R J , R ) denotes the space of twice continuously differentiable functions with compact support. We multiply both sides of (47) with h , integrate, and apply the partial integration rule to obtain

By chain-rule and partial integration, (50) can be rewritten as

Putting everything together, we obtain

where { A [ Q ] } Q ∈P ( R J ) is a family of operators defined as

for h ∈ C 2 c ( R J , R ) . The reader may recognize this operator family as the generator of a 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 a McKean-Vlasov process of the form

where ( B ( t ) ) t> 0 is a Brownian motion and Q t the law of θ ( t ) . In other words: The solution to (55) has the time marginals Q ( t ) such that (53) holds for every h ∈ C 2 c ( R J , R ) . Furthermore, the corresponding pdfs ( q ( t ) ) satisfy the nonlinear Fokker-Planck equation given as

where A ∗ [ Q ] denotes the L 2 -adjoint of the operator A [ Q ] and is given as

for h ∈ C 2 c ( R J , R ) (Barbu and R¨ ockner, 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 to simulate solutions to (55) (Veretennikov, 2006) is to use an ensemble of interacting particles. Formally, we replace Q ( t ) by 1 N E ∑ N E n =1 δ θ n ( t ) and obtain

for n = 1 , . . . , N E where N E ∈ N denotes the number of particles. The Euler-Maruyama approximation of (58) leads to the final algorithm:

Step 1: Initialise N E ∈ N particles θ 1 , 0 , . . . , θ N E , 0 from a use chosen initial distribution Q 0 .

Step 2: Evolve the particles forward in time according to

Note that θ n,k is thought of as approximation of θ n ( t ) at position t = kη . Furthermore, as discussed in Section 4, various choices of V , λ 1 and λ 2 allow us to implement the WGF for different regularised optimisation problems in the space of probability measures. This is summarised below:

Deep ensembles: V ( θ ) = ℓ ( θ ) , λ 1 = 0 , λ 2 = 0

Deep Langevin ensembles: V ( θ ) = ℓ ( θ ) -λ log p ( θ ) , λ 1 := 0 , λ := λ 2

Deep repulsive Langevin ensembles: V ( θ ) = ℓ ( θ ) -λ 1 log p ( θ ) -λ 2 µ P ( θ )

C Asymptotic distribution of particles: unregularised objective

In this section, we investigate the asymptotic distribution of the WGF for the objective

for Q ∈ P ( R J ) . The associated particle method is:

Sample θ 1 (0) , . . . , θ N E (0) independently from Q 0 .

Simulate (deterministically) θ ′ n ( t ) = -∇ ℓ ( θ n ( t ) ) for n = 1 , . . . , N E .

We start by introducing some notation for the deterministic gradient system. Let ϕ t ( θ 0 ) denote the solution to the ordinary differential equation (ODE)

at time t > 0 . 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 J → R is norm-coercive and satisfies the Lojasiewicz inequality, i.e. for every θ ∈ R J exists an environment U of θ and constants 0 < γ < 1 and C > 0 such that

for all ¯ θ ∈ U . Then we know that ϕ t ( θ 0 ) converges for t →∞ to a local minimum of ℓ for Lebesgue almost every θ 0 ∈ R J .

Proof. First we show that t ↦→ ϕ t ( θ 0 ) is bounded. We proof this by contradiction. Assume that ϕ t ( θ 0 ) is unbounded. Then there exists a subsequence ( t n ) n ∈ N ⊂ [0 , ∞ ) with t n →∞ for n →∞ such that

for n →∞ . The norm-coercivity immediately implies that

for n →∞ . However, this contradicts

where the first inequality follows from the fact that t ↦→ ℓ ( ϕ t ( θ 0 ) ) is decreasing, which is a consequence of

Hence t ↦→ ϕ t ( θ 0 ) is bounded. By the Bolzano-Weierstrass theorem we can find a sequence ( t n ) n ∈ N ⊂ [0 , ∞ ) with t n →∞ and a point θ ∞ ∈ R J such that

for n →∞ . Hence ( ϕ t ( θ 0 ) ) t> 0 has the accumulation point θ ∞ . The Lojasiewicz theorem (Colding and Minicozzi II, 2014) allows us to deduce that

for t →∞ , and that θ ∞ satisfies ∇ ℓ ( θ ∞ ) = 0 .

It remains to show that θ ∞ is not a saddle point for Lebesgue almost every initial value θ 0 . However, this is very similar to the proof in Lee et al. (2016). The only difference is that one would need to use a continuous-time version of the stable manifold theorem, which is readily available, for example in Bressan (2003).

Let { m i } i ∈ N denote the local minima of ℓ which are by assumption countable. Denote further by

the domain of attraction for the minimum m i . The next theorem is then an easy consequence of Lemma 3.

Theorem 5. Assume that the loss function ℓ only has countably many local minima, is norm coercive, and satisfies the Lojasiewicz inequality. Let further θ 0 ∼ Q 0 for some Q 0 ∈ P ( R J ) such that ∑ ∞ i =1 Q 0 (Θ i ) = 1 . Then,

for t →∞ . Here D -→ denotes convergence in distribution.

Proof. Let θ 0 ∈ R J be fixed. Due to Lemma 3, we know that

for Lebesgue almost every θ 0 for t → ∞ . Here, ✶ {·} denotes the indicator function. Let Y now be a random variable with law Q 0 . By assumption, we know that Y ∈ Θ i for some i ∈ N with probability 1 . Hence,

almost surely for t → ∞ . Since almost sure convergence implies convergence in distribution, we conclude that

where L ( · ) denotes the law of a random variable. However, the law of the RHS is easily recognised as which concludes the proof.

Remark 2. Note that the condition

in Theorem 5 is easy to satisfy. According to Lemma 3 the set

has Lebesgue measure zero. Therefore, any Q 0 which has a density w.r.t. the Lebesgue measure will satisfy (77).

D Asymptotic distribution for deep Langevin ensembles

In this section, we analyse the objective

for Q ∈ P ( R J ) . The corresponding particle method is given as:

Sample θ 1 (0) , . . . , θ N E (0) independently from Q 0 .

Simulate the SDE dθ n ( t ) = -∇ V ( θ n ( t ) ) dt + √ 2 λdB n ( t ) for each n = 1 , . . . , N E .

Recall that V ( θ ) = ℓ ( θ ) -λ log p ( θ ) . This case is well-studied in the literature and known as Langevin diffusion. Under mild assumptions (Chiang et al., 1987; Roberts and Tweedie, 1996),

for t → ∞ and each particle n = 1 , . . . , N E independently. The probability measure Q ∞ has the density

where Z > 0 is the normalising constant. As a consequence, the WGF asymptotically produces samples from Q ∞ . However, it is a priori unclear that Q ∞ is in fact the same as the global minimiser Q ∗ of L .

We investigate this question by relating invariant measures to stationary points of the Wasserstein gradient.

Definition 1. (Liggett, 2010, Thm. 3.3.7) A measure Q is called an invariant measure (for a given Feller-process) if

for all h ∈ C 2 c ( R J ) . Here A is the infinitesimal generator of the corresponding Feller-process.

Recall that the infinitesimal generator of the Langevin diffusion for h ∈ C 2 c ( R J ) is given as

Definition 2. A measure Q ∈ P 2 ( R J ) is called a stationary point of the Wasserstein gradient if

for Q -almost every θ ∈ R J .

In finite dimensions, it is well-known that a local minimiser is a stationary point of the gradient. This carries over to the infinite-dimensional case, with a similar proof. Since we could not find this result anywhere in the literature we included it for completeness.

Lemma 4. Let ̂ Q be a local minimiser of L , i.e. there exits and ϵ > 0 such that

for all Q with W 2 ( ̂ Q,Q ) ≤ ϵ . Then ̂ Q is a stationary point of the Wasserstein gradient in the sense of Definition 2.

Proof. Let h ∈ C 2 c ( R J ) be arbitrary and ̂ Q ∈ P 2 ( R J ) be a local minimum of L . Further, let ϕ t ( θ 0 ) be the solution to the initial value problem

for t ∈ ( -ϵ, ϵ ) for some ϵ > 0 . We now define Q ( t ) := ϕ t # ̂ Q for t ∈ ( -ϵ, ϵ ) where f # µ denotes the push-forward of the measures µ through the function f . In the Riemannian interpretation of the Wasserstein space, ( Q ( t ) ) t ∈ ( -ϵ,ϵ ) is a curve in P 2 ( R J ) with tangent vector h at point ̂ Q (Ambrosio et al., 2005, Chapter 8). We, further, define f : ( -ϵ, ϵ ) → R as f ( t ) := L ( Q ( t )) . Application of the chain-rule (Ambrosio et al., 2005, p. 233) gives

We know that f has a local minimum at t = 0 and, therefore, f ′ (0) = 0 which gives

Since (92) holds for arbitrary test functions h ∈ C 2 c ( R J ) and as C 2 c ( R J ) is dense in L 2 ( ̂ Q ) , we obtain that ∇ W L [ ̂ Q ]( θ ) = 0 for ̂ Q -a.e θ ∈ R J .

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 ∈ P 2 ( R J ) be such that Q has a density q with respect to the Lebesgue measure. Then, the following two statements are equivalent:

Q is 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 ∈ C 2 c ( R J ) given as

By partial integration, it is easy to verify that the L 2 - adjoint (w.r.t the Lebesgue measure) is given as

We, therefore, conclude that

Furthermore, we have ∇ W L [ Q ] = ∇ V + λ ∇ log q , and therefore

where the last line follows from applying partial integration. This allows us to conclude that

whenever Q has a density. As a consequence we have that Q is invariant if and only if it is a stationary point of the Wasserstein gradient.

Lemma 5 allows us to move between the optimisation and stochastic differential equation perspective. In Appendix A, we discussed the existence and uniqueness of a global minimiser Q ∗ of L . We know that Q ∗ has a density since the Kullback-Leibler divergence would be infinite otherwise (assuming P has a Lebesgue-density which we assume throughout the paper). Lemma 4 guarantees that Q ∗ is a stationary point of the Wasserstein gradient. Due to Lemma 5, we can infer that Q ∗ must be an invariant measure. However, due to the uniqueness of the invariant measure under the previously mentioned mild assumptions (Chiang et 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

as optimisation objective. Here, H ( Q ) = -∫ log q ( θ ) q ( θ ) dθ denotes the differential entropy.

Recall that in this case V ( θ ) = ℓ ( θ ) -λ 1 µ P ( θ ) -λ 2 log p ( θ ) . We already discussed in Appendix B that the McKean-Vlasov process of the form

with ( B ( t ) ) t ≥ 0 being a Brownian motion achieves the desired density evolution. Furthermore, the particle approximation of (104) is given as

for n = 1 , . . . , N E where N 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 Lebesgue density. We start by introducing the concept of an invariant measure for a nonlinear Markov process (Ahmed and Ding, 1993, Definition 1).

Definition 3. A measure Q is called an invariant measure for a nonlinear Markov process with the family of infinitesimal generators { A [ Q ] } Q ∈P ( R J ) if

for all h ∈ C 2 c ( R J ) .

Recall that the family of infinitesimal generators in our case is given as

for h ∈ C 2 c ( R J , R ) . In analogy to Lemma 5, we obtain the following result.

Lemma 6. Let Q ∈ P 2 ( R J ) be such that Q has a density q with respect to the Lebesgue measure. Then, the following two statements are equivalent:

Q is a stationary point of the Wasserstein gradient for L in (102) in the sense of Def. 2.

Q is an invariant measure for the McKean-Vlasov process with infinitesimal generator defined in (108)

Proof. First, we notice that

Recall, that A ∗ [ Q ] denotes the L 2 -adjoint of the operator A [ Q ] and that it is given as

for h ∈ C 2 ( R J , R ) with compact support. This implies

We plug this into (110) to obtain

On the other hand, we have that

and therefore

where the last line follows from partial integration. Comparing (114) to (119) gives

for all h ∈ C 2 c ( R J ) whenever Q has a density. This immediately implies that Q is invariant iff it 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 Veretennikov (2006) specifies conditions on κ and V such that

for N E , t → ∞ . Here Q n,N E ( t ) denotes the law of a fixed particle θ n ( t ) , n = 1 , . . . , N E , whose distribution is characterised by the SDE (58). The measure Q ∞ is the unique invariant measure of the nonlinear Markov process. By Lemma 6 every invariant measure is a 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), we conclude by uniqueness that Q ∞ = Q ∗ .

F Asymptotic analysis of deep repulsive ensembles

In this section, we consider the objective

for Q ∈ P ( R J ) . The corresponding McKean-Vlasov process is of the form

where Q t denotes the distribution of θ ( t ) and V ( θ ) = ℓ ( θ ) -µ P ( θ ) with µ P ( θ ) = ∫ κ ( θ, θ ′ ) dP ( θ ) the kernel mean-embedding of P . Wecall the particle method in this case deep repulsive ensembles (DRE) .

The existence of the global minimiser Q ∗ is still guaranteed under the assumptions in Appendix A. Lemma 4 guarantees that Q ∗ is a stationary point of the Wasserstein gradient, i.e.

. Recall that the infinitesimal generator in this case is given as for Q ∗ -a.e. θ ∈ R J

for Q ∈ P ( R J ) , h ∈ C 2 c ( R J ) . It immediately follows from the definition that

for all h ∈ C 2 c ( R J ) , θ ∈ R J . As in Lemma 5 & 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 is 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 ] ) Q ∈P ( R J ) defined in (126) is (in general) not unique.

Proof. Let N E ∈ N and define ˜ L : ( R J ) N E → R as

Assume that V is bounded from below and norm-coercive. Then ˜ L is bounded from below and norm-coercive and therefore we can find a global minimiser θ ∗ := ( θ ∗ 1 , . . . , θ ∗ N E ) ∈ ( R J ) N E of ˜ L . Since ˜ L is differentiable, we know that θ ∗ is a stationary point of the gradient which implies

for all i = 1 , . . . , N E . Here, we assume that the kernel κ is symmetric, which is standard in the MMDliterature. Note that (128) is equivalent to

This means that ̂ Q is a stationary point of the Wasserstein gradient, and therefore an invariant measure for the McKean-Vlasov process. Since N E ∈ N was arbitrary, we have constructed countably many invariant measures and therefore uniqueness can't hold in general.

The reason that non-uniqueness of the invariant measure is an immediate contradiction to convergence 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-Vlasov process will remain unchanged over time. Convergence to the global minimiser can therefore surely not hold for arbitrary initialisation Q 0 . It may be possible to construct conditions on Q 0 under which convergence still holds. For example, for Stein variational gradient descent a similar issue occurs. However, in this case one can guarantee convergence (Lu et al., 2019, Theorem 2.8) if Q 0 has a Lebesgue-density (and if the kernel satisfies further restrictive assumptions). The existence of conditions that guarantee convergence for DRE remains an open problem.

G Implementation details

In Appendix A, we derived the following algorithm:

Step 1: Simulate N E ∈ N particles θ 1 , 0 , . . . , θ N E , 0 from a use chosen initial distribution Q 0 .

Step 2: Evolve the particles forward in time according to

We can generate samples from DE, DLE and DRLE by setting the potential and regularisation parameters as described below:

Deep ensembles: V ( θ ) = ℓ ( θ ) , λ 1 = 0 , λ 2 = 0

Deep Langevin ensembles: V ( θ ) = ℓ ( θ ) -λ log p ( θ ) , λ 1 = 0 , λ = λ 2

Deep repulsive Langevin ensembles: V ( θ ) = ℓ ( θ ) -λ 1 log p ( θ ) -λ 2 µ P ( θ )

Due to Appendix D & E, we can think of θ 1 ,K , . . . , θ N E ,K as approximately sampled from the global minimiser Q ∗ for DLE and DRLE if K is large enough. All experiments use the SE kernel given as

with lengthscale parameter σ κ > 0 . The kernel mean embedding µ P can easily be approximated as

where θ 1 , . . . , θ M ∼ P independently. We chose M = 20 .

for ̂ Q -a.e. θ ∈ R J where

G.1 Toy example: global minimiser

We describe details regarding the experiments conducted to produce Figure 2 below.

We generate N E = 300 particles and make the following choices:

Loss: ℓ ( θ ) := 3 2 ( 1 4 θ 4 + 1 3 θ 3 -θ 2 ) -3 8

Prior: P ∼ N (0 , 1) and therefore log p ( θ ) = -1 2 θ 2

Initialisation: Q 0 = P

Reg. parameter: λ DLE = 1 , λ DRLE = 1 , λ ′ DRLE = 1

Step size: η = 10 -4 , Iterations: K = 100 , 000

Kernel lengthscale, σ κ , 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 a global minimum at θ = -2 , a turning point at θ = 0 , and a local minimum at θ = 1 .

Deep ensembles: The optimal Q ∗ is a Dirac measure located at the global minimiser θ = -2 . However, as we proved in Theorem 1, the WGF produce samples from

̸

as ( -∞ , 0) is the region of attraction for the global minimum and (0 , ∞ ) for the local minimum which both have probability 0 . 5 under Q 0 = P = N (0 , 1) . In particular, Q ∞ = Q ∗ as expected.

Deep Langevin ensembles: The optimal measure has the pdf

for θ ∈ 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 ( θ ) := log q ∗ ( θ ) satisfies the integro-differential equation

with some initial value u (0) = u 0 . In principle, we could choose u 0 such that q ( θ ) := exp( u ( θ )) integrates to 1 . However, since we do not know the appropriate initial condition a priori, we choose an arbitrary u 0 and normalise the pdf afterwards. We use an numerical solver to evaluate u ( θ ) 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 5. 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 E = 300 particles and make the following choices:

Loss: ℓ ( θ ) = - log ∑ 4 i =1 1 4 N ( θ ; µ i , I 2 ) , θ ∈ R 2 , µ i = ( ± 3 , ± 3) T , i = 1 , . . . , 4

Prior: P flat and therefore log p ( θ ) = 0

Initialisation: Q 0 ∼ N (0 , I 2 )

Reg. parameter: λ DLE = 0 . 2 , λ DRLE = 0 . 2 , λ ′ DRLE = 0 . 6

Step size: η = 0 . 1 , Iterations: K = 10 , 000

5

0

5

5

0

5

DE

5

0

5

DLE

5

0

5

DRLE

5

0

5

FD-GVI

0

1

Kernel lengthscale, σ κ , is chosen according to the median heuristic (Garreau et al., 2017) based on samples from the prior P

Note that for a translation-invariant kernel such as the SE kernel we obtain for the flat prior P that

where the second line follows from the fact that we can write any translation-invariant kernel as κ ( θ, θ ′ ) = ϕ ( θ -θ ′ ) for some function ϕ : R J → R and the second line is simple variable substitution. If (139) is finite, the above expression is well-defined and therefore µ P constant. Note that in particular for the SE kernel, we have ϕ ( ξ ) = exp( -∥ ξ ∥ 2 / (2 σ 2 κ )) and therefore (139) is finite. As a 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 ℓ 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 a flat prior P is in fact not covered by our theory in Appendix A. The problem is that KL( · , L ) , where L denotes the Lebesgue measure, is not positive (and not even bounded from below). To see this, choose Q = N (0 , Σ) with Σ = diag ( σ 2 1 , σ 2 2 ) and note that

where H ( Q ) denotes the differential entropy. For a Gaussian, it is known that

and therefore if either σ 2 1 → ∞ or σ 2 2 → ∞ then KL( Q, L ) → -∞ . However, note that this difficulty is rather technical in nature and can easily be remedied. Instead of L , we could have chosen the uniform prior P ∼ U ( -10 100 , 10 100 ) . In this case, the positivity of KL( · , P ) is guaranteed by Jensen's inequality. This choice of P gives-up to an additive constant-the same objective as a flat prior and up to machine precision the same kernel mean embedding µ P . It is, therefore, algorithmically irrelevant if P is flat or uniform on a very large set.

FD-GVI Weuse the same prior and loss as for DE, DLE and DRLE. We parameterise the variational family as independent Gaussian, i.e.

We learn the variational parameters ν := ( µ, β ) ∈ R 4 by minimising

where Z 1 , . . . , Z 200 ∼ N (0 , I 2 ) and H ( N ( µ, Σ) ) 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 N E = 20 particles and make the following choices:

Loss: ℓ ( θ ) = -| sin( θ ) | , θ ∈ [ -Mπ,Mπ ] , with M = 1000

Prior: P flat and therefore log p ( θ ) = 0 and µ P = const. (cf. Appendix G.2)

Initialisation: Q 0 ∼ U ( -Mπ,Mπ )

Reg. parameter: λ DLE = 0 . 001 , λ DRLE = 0 . 001 , λ ′ DRLE = 0 . 6

Step size: η = 0 . 01 , Iterations: K = 1 . 000

Kernel lengthscale, σ κ , is chosen according to the median heuristic (Garreau et al., 2017) based on samples from the prior P

Note that ℓ has 2 M = 2000 local minima at locations

Due to the flat prior ∇ V = ∇ ℓ for all three methods. We observe that it is hard to distinguish the methods since most particles are in 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 Lakshminarayanan et al. (2017), we train 5 one-hidden-layer neural networks f θ 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 a Gaussian NLL) found via grid search on a validation data set, we make the following choices:

Loss: ℓ ( θ ) = 1 N ∑ N n =1 ( f θ ( x n ) -y n ) 2 where { x n , y n } N n =1 are paired observations.

Prior: P ∼ N (0 , 1)

Initialisation: Kaiming intilisation, i.e. for each layer l ∈ { 1 , ...L } that maps features with dimensionality n l -1 into dimensionality n l , we sample Q l, 0 ∼ N (0 , 2 /n l )

Reg. parameter: λ DLE = 10 -4 , λ DRLE = 10 -4 , λ ′ DRLE = 10 -2

Step size: η = 0 . 1 , Iterations: K = 10 , 000

Kernel lengthscale, σ κ , is chosen according to the median heuristic (Garreau et al., 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\n#### **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\nQUESTIONS:\nIn 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\nLIMITATIONS:\nAuthors 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe 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 ?", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\n* [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\nWEAKNESSES:\n* [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\nQUESTIONS:\n* [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”\n\nLIMITATIONS:\nGiven 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo 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\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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_)\n\nQUESTIONS:\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_)\n\nLIMITATIONS:\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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1) 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.\n\nWEAKNESSES:\n1) 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).\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nAll limitations are addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}