SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Veit D. Wild*  
University of Oxford
Sahra Ghalebikesabi  
University of Oxford
Dino Sejdinovic  
University of Adelaide
Jeremias Knoblauch  
University College London
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.
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 infinite-dimensional regularised optimisation problem over the space of probability measures (see e.g. Guedj and Shawe-Taylor, 2019; Knblauch et al., 2022). As a result, we find that the differences between
*veit.wild@stats.ox.ac.uk
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  $\ell : \mathbb{R}^J \to \mathbb{R}$  we wish to minimise are often highly non-convex. For instance, one could wish to minimise  $\ell(\theta) := \frac{1}{N} \sum_{n=1}^{N} (y_n - f_\theta(x_n))^2$  where  $\{(x_n, y_n)\}_{n=1}^N$  is a set of paired observations and  $f_\theta$  a neural network with parameters  $\theta$ . 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  $\ell$  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:
Figure 1: Illustration of convexification through probabilistic lifting.
First, we transform a non-convex optimisation  $\min_{\theta \in \Theta} \ell(\theta)$  into an infinite-dimensional optimisation over the set of probability measures  $\mathcal{P}(\mathbb{R}^J)$ , yielding  $\min_{Q \in \mathcal{P}(\mathbb{R}^J)} \int \ell(\theta) dQ(\theta)$ . 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: \mathcal{P}(\mathbb{R}^J) \times \mathcal{P}(\mathbb{R}^J) \to [0, \infty]$  such that  $(Q, P) \mapsto D(Q, P)$  is strictly convex and  $P \in \mathcal{P}(\mathbb{R}^J)$  a fixed measure—existence and uniqueness of a global minimiser can be guaranteed. Given a scaling constant  $\lambda > 0$ , we can now put everything together to obtain
2To illustrate why this is necessary, assume that there are  $\theta_{A}$  and  $\theta_B$  in  $\mathbb{R}^J$  so that  $\ell (\theta_A) = \ell (\theta_B) = \min_{\theta \in \Theta}\ell (\theta)$ . Due to linearity, each measure  $Q_{t}\in \mathcal{P}(\mathbb{R}^{J})$ ,  $t\in [0,1]$ , defined as  $Q_{t}\coloneqq (1 - t)\delta_{\theta_{A}} + t\delta_{\theta_{B}}$  provides one (of the infinitely many) global minima in the set  $\arg \min_{Q\in \mathcal{P}(\mathbb{R}^J)}\int \ell (\theta)dQ(\theta)$ .
the loss  $L$  and the unique minimiser  $Q^{*}$  as
$$
L (Q) := \int \ell (\theta) d Q (\theta) + \lambda D (Q, P), \quad Q ^ {*} := \underset {Q \in \mathcal {P} \left(\mathbb {R} ^ {J}\right)} {\arg \min } L (Q). \tag {1}
$$
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  $\mathbb{R}^J$  of  $\theta$ . Similarly, the gradient operator  $\nabla$  will exclusively denote differentiation with respect to  $\theta \in \mathbb{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  $\ell$  on Euclidean spaces. Through probabilistic lifting and convexification, we can identify a unique minimiser  $Q^{*}$  in the new space, which minimises the  $\theta$ -averaged loss  $\theta \mapsto \ell(\theta)$  without deviating too drastically from some reference measure  $P$ . In this sense,  $Q^{*}$  summarises the quality of all (local and global) minimisers of  $\ell(\theta)$  by assigning them a corresponding weight. The choices for  $\lambda$ ,  $D$  and  $P$  determine the trade-off between the initial loss  $\ell$  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, \lambda, \ell$  and the framework of analysis (see e.g. Knoblauch et al., 2022, for a discussion). For example, if  $\ell(\theta)$  is a negative log likelihood,  $D$  is the Kullback-Leibler divergence (KL), and  $\lambda = 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  $\ell(\theta)$  to be any loss,  $D$  to be any divergence on  $\mathcal{P}(\mathbb{R}^J)$ , and  $\lambda$  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ünwald, 2011). Here,  $\ell$  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 \mapsto \lambda 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égin et al., 2016; Haddouche and Guedj, 2023).
2.2 Generalised variational inference (GVI) in finite and infinite dimensions
In line with the terminology coined in Knoblauch et al. (2022), we refer to any algorithm aimed at solving (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_{\nu^{*}} = \arg \min_{Q_{\nu}\in \mathcal{Q}}L(Q)$  for a set of measures  $\mathcal{Q} := \{Q_{\nu}:\nu \in \Gamma \} \subset \mathcal{P}(\mathbb{R}^J)$  parameterised by a parameter  $\nu \in \Gamma \subseteq \mathbb{R}^{I}$ . To find  $\mathcal{Q}_{\nu}$ , one now simply performs (finite-dimensional) gradient descent with respect to the function  $\nu \mapsto L(Q_{\nu})$ . For the special case where  $P$  is a Bayesian prior,  $\lambda = 1$ ,  $\ell (\theta)$  is a negative log likelihood parametrised by  $\theta$ , and  $D = \mathrm{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  $\mathcal{Q}$  and  $P$  to be simple distributions such as Gaussians to ensure that  $L(Q_{\nu})$  is a tractable function of  $\nu$ . This often results in a  $\mathcal{Q}$  that is unlikely to contain a good approximation to  $Q^{*}$ ; raising doubt if  $Q_{\nu^{*}}$  can approximate  $Q^{*}$  in any meaningful sense. Secondly, even if  $Q \mapsto L(Q)$  is strictly convex on  $\mathcal{P}(\mathbb{R}^J)$ , the parameterised objective  $\nu \in \Gamma \mapsto L(Q_{\nu})$  is usually not. Hence, there is no guarantee that gradient descent leads us to  $Q_{\nu^{*}}$ . This point also applies to very expressive variational families (Rezende and Mohamed, 2015; Mescheder et al., 2017) which may be sufficiently rich that  $Q^{*} \in \mathcal{Q}$ , but whose optimisation problem  $\nu \in \Gamma \mapsto L(Q_{\nu})$  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ández-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  $\nu \mapsto L(Q_{\nu})$ , we want to exploit the convex structure of  $Q \mapsto L(Q)$ . Of course, a priori it is not even clear how to compute the gradient for a function  $Q \mapsto L(Q)$  defined on an infinite-dimensional nonlinear space such as  $\mathcal{P}(\mathbb{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  $\widehat{Q}$  that has a known form (such as the Bayes posterior) by applying a gradient flow to  $Q \in \mathcal{P}(\mathbb{R}^J) \mapsto E(Q, \widehat{Q})$  where  $E(\cdot, \cdot)$  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 = \mathrm{KL}$ ,  $\widehat{Q} = \mathrm{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  $\ell : \mathbb{R}^J \to \mathbb{R}$  by iteratively improving an initial guess  $\theta_0 \in \mathbb{R}^J$  through the update  $\theta_{k+1} := \theta_k - \eta \nabla \ell(\theta_k)$ ,  $k \in \mathbb{N}$ , where  $\eta > 0$  is a step-size and  $\nabla \ell$  denotes the gradient of  $\ell$ . For sufficiently small  $\eta > 0$ , this update can equivalently be written as
$$
\theta_ {k + 1} = \underset {\theta \in \mathbb {R} ^ {J}} {\arg \min } \left\{\ell (\theta) + \frac {1}{2 \eta} \| \theta - \theta_ {k} \| _ {2} ^ {2} \right\}. \tag {2}
$$
Gradient flows formalise the following logic: for any fixed  $\eta$ , we can continuously interpolate the corresponding gradient descent iterates  $\{\theta_k\}_{k\in \mathbb{N}_0}$ . To do this, we simply define a function  $\theta^{\eta}:[0,\infty)\to \mathbb{R}$  as  $\theta^{\eta}(t)\coloneqq \theta_{t / \eta}$  for  $t\in \eta \mathbb{N}_0\coloneqq \{0,\eta ,2\eta ,\ldots \}$ . For  $t\notin \eta \mathbb{N}_0$  we linearly interpolate<sup>3</sup>. As  $\eta \rightarrow 0$ , the function  $\theta^{\eta}$  converges to a differentiable function  $\theta_{*}:[0,\infty)\to \mathbb{R}$  called the gradient flow of  $\ell$ , because it is characterised as solution to the ordinary differential equation (ODE)  $\theta_{*}'(t) = -\nabla \ell (\theta_{*}(t))$  with initial condition  $\theta_{*}(0) = \theta_{0}$ . Intuitively,  $\theta_{*}(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  $\eta$ , 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  $\mathcal{P}_2(\mathbb{R}^J)$  be the space of probability measures with finite second moment equipped with the 2-Wasserstein metric given as
$$
W _ {2} (P, Q) ^ {2} = \inf  \left\{\int | | \theta - \theta^ {\prime} | | _ {2} ^ {2} d \pi (\theta , \theta^ {\prime}): \pi \in \mathcal {C} (P, Q) \right\}
$$
where  $\mathcal{C}(P,Q)\subset \mathcal{P}(\mathbb{R}^J\times \mathbb{R}^J)$  denotes the set of all probability measure on  $\mathbb{R}^J\times \mathbb{R}^J$  such that  $\pi (A\times \mathbb{R}^{J}) = P(A)$  and  $\pi (\mathbb{R}^J\times B) = Q(B)$  for all  $A,B\subset \mathbb{R}^{J}$  (see also Chapter 6 of Villani et al., 2009). Further, let  $L:\mathcal{P}_2(\mathbb{R}^J)\to (-\infty ,\infty ]$  be some functional--for example  $L$  in (1). In direct analogy to (2), we can improve upon an initial guess  $Q_{0}\in \mathcal{P}_{2}(\mathbb{R}^{J})$  by iteratively solving
$$
Q_{k + 1}:= \operatorname *{arg  min}_{Q\in \mathcal{P}_{2}(\mathbb{R}^{J})}\left\{L(Q) + \frac{1}{2\eta} W_{2}(Q,Q_{k})^{2}\right\}
$$
for  $k \in \mathbb{N}_0$  and small  $\eta > 0$  (see Chapter 2 of Ambrosio et al., 2005, for details). Again, for  $\eta \to 0$ , an appropriate limit yields a continuously indexed family of measures  $\{Q(t)\}_{t \geq 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 \geq 0}$  is given by the partial differential equation (PDE)
$$
\partial_ {t} q (t, \theta) = \nabla \cdot \left(q (t, \theta) \nabla_ {W} L [ Q (t) ] (\theta)\right), \tag {3}
$$
with  $q(0,\cdot) = q_0$  (Villani, 2003, Section 9.1). Here  $\nabla \cdot f\coloneqq \sum_{j = 1}^{J}\partial_{j}f_{j}$  denotes the divergence operator and  $\nabla_W L[Q]:\mathbb{R}^J\to \mathbb{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.  $\nabla_W L[Q] = \nabla L'[Q]$  where  $L^{\prime}[Q]:\mathbb{R}^{J}\rightarrow \mathbb{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^{\dagger}$  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\to \infty}q^{\dagger}(t,\cdot)$  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\to \infty$ ) 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
<sup>3</sup>This means  $\theta^{\eta}(t) := \frac{\theta_{s+1} - \theta_s}{\eta}\big(t - s\eta\big) + \theta_s$ , for  $t \in [s\eta, (s+1)\eta)$  and  $s \in \mathbb{N}_0$ .
samples from the solution to (3) evolve in time. This leads to tractable inference algorithms that can be implemented in high dimensions.