SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Given the goal of producing samples directly, we focus on a particular form of loss that is well-studied 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 \in \mathcal{P}_2(\mathbb{R}^J)$  describes the distribution of particles located at specific points in  $\mathbb{R}^J$ . The overall energy of a collection of particles sampled from  $Q$  is decomposed into three parts: (i) the external potential  $V(\theta)$  which acts on each particle individually, (ii) the interaction energy  $\kappa(\theta, \theta')$  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
$$
L ^ {\mathrm {f e}} (Q) := \int V (\theta) d Q (\theta) + \frac {\lambda_ {1}}{2} \iint \kappa \left(\theta , \theta^ {\prime}\right) d Q (\theta) d Q \left(\theta^ {\prime}\right) + \lambda_ {2} \int \log q (\theta) q (\theta) d \theta , \tag {4}
$$
for  $Q \in \mathcal{P}_2(\mathbb{R}^J)$  with Lebesgue density  $q$ ,  $\lambda_1 \geq 0$ ,  $\lambda_2 \geq 0$ . Note that for  $\lambda_2 > 0$  we implicitly assume that  $Q$  has a density. Following Section 9.1 in Villani et al. (2009), its WG is
$$
\nabla_ {W} L ^ {\mathrm {f e}} [ Q ] (\theta) = \nabla V (\theta) + \lambda_ {1} \int (\nabla_ {1} \kappa) (\theta , \theta^ {\prime}) d Q (\theta^ {\prime}) + \lambda_ {2} \nabla \log q (\theta),
$$
where  $\theta \in \mathbb{R}^J$ , and  $\nabla_{1}\kappa$  denotes the gradient of  $\kappa$  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}\in \mathbb{N}$  particles  $\theta_1(0),\ldots ,\theta_{N_E}(0)$  independently from  $Q_{0}\in \mathcal{P}_{2}(\mathbb{R}^{J})$
Step 2: Evolve the particle  $\theta_{n}$  by following the stochastic differential equation (SDE)
$$
d \theta_ {n} (t) = - \left(\nabla V \left(\theta_ {n} (t)\right) + \frac {\lambda_ {1}}{N _ {E}} \sum_ {j = 1} ^ {N _ {E}} \left(\nabla_ {1} \kappa\right) \left(\theta_ {n} (t), \theta_ {j} (t)\right)\right) d t + \sqrt {2 \lambda_ {2}} d B _ {n} (t), \tag {5}
$$
for  $n = 1,\dots ,N_E$ , and  $\{B_n(t)\}_{t > 0}$  stochastically independent Brownian motions.
As  $N_{E} \to \infty$ , the distribution of  $\theta_{1}(t), \ldots, \theta_{N_{E}}(t)$  evolves in  $t$  in the same way as the sequence of densities  $q(t, \cdot)$  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  $\theta_{1}(T), \ldots, \theta_{N_{E}}(T)$  for  $N_{E} \in \mathbb{N}$  and  $T > 0$ . Our intuition and Section 4 tell us that, as desired, the distribution of  $\theta_{1}(T), \ldots, \theta_{N_{E}}(T)$  will be close to the global minimiser of  $L^{\mathrm{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) = \int \ell(\theta)dQ(\theta)$ , corresponding to  $\lambda = 0$  in (1). This is also a special case of (4) with  $\lambda_1 = \lambda_2 = 0$ . As  $\lambda_1 = 0$ , there is no interaction term, and all particles can be simulated independently from one another as
$$
\theta_ {1} (0), \dots , \theta_ {N _ {E}} (0) \sim Q _ {0}, \quad \theta_ {n} ^ {\prime} (t) = - \nabla \ell \big (\theta_ {n} (t) \big), \quad n = 1, \dots , N _ {E}.
$$
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  $\ell$ .
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  $\ell$  with a probability that depends on the domain of attraction and the initialisation distribution  $Q_0$ .
Theorem 1. If  $\ell$  has countably many local minima  $\{m_i : i \in \mathbb{N}\}$ , then it holds independently for each  $n = 1, \ldots, N_E$  that
$$
\theta_ {n} (t) \xrightarrow {\mathcal {D}} \sum_ {i = 1} ^ {\infty} Q _ {0} (\Theta_ {i}) \delta_ {m _ {i}} =: Q _ {\infty}
$$
for  $t\to \infty$ . Here  $\xrightarrow{\mathcal{D}}$  denotes convergence in distribution and  $\Theta_{i} = \{\theta \in \mathbb{R}^{J}:\lim_{t\to \infty}\theta_{*}(t) = m_{i}$  and  $\theta_{*}(0) = \theta \}$  denotes the domain of attraction for  $m_{i}$  with respect to the gradient flow  $\theta_{*}$ .
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  $\theta_{1}(T),\ldots ,\theta_{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\mapsto L(Q)$  for the unregularised case  $\lambda = 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_{\infty}$  strongly depends on the initialisation  $Q_{0}$  . Importantly, the dependence of  $Q_{\infty}$  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 Kullback-Leibler divergence  $(D = \mathrm{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^{*}(\theta)\propto \exp (-\frac{1}{\lambda}\ell (\theta))p(\theta)$  (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^{\mathrm{fe}}$  in (4) by setting  $V(\theta) = \ell(\theta) - \lambda \log p(\theta)$ ,  $\lambda_1 = 0$  and  $\lambda_2 = \lambda$ . Compared to the unregularised objective of the previous section (where  $\lambda = 0$ ), the external potential is now adjusted by  $-\lambda \log p(\theta)$ , 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  $\lambda_1 = 0$ , the corresponding particle method still does not have an interaction and is given as
$$
\theta_ {1} (0), \dots , \theta_ {N _ {E}} (0) \sim Q _ {0} \quad d \theta_ {n} (t) = - \nabla V (\theta_ {n} (t)) d t + \sqrt {2 \lambda} d B _ {n} (t), \quad n = 1, \dots , N _ {E}. \tag {6}
$$
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, \ldots, t_T$ , DLE approximates measures using the end-points of  $N_E$  trajectories given by  $\{\theta_n(t_T)\}_{n=1}^{N_E}$ . In contrast, ULA would use a (sub)set of the samples  $\{\theta_1(t_i)\}_{i=1}^T$  generated from one single particle's trajectory. To analyse DLEs, we build on the Langevin dynamics literature: in Appendix D, we show that  $\theta_n(t) \xrightarrow{\mathcal{D}} Q^*$  as  $t \to \infty$ , independently for each  $n = 1, \ldots, N_E$ . Hence  $\theta_1(T), \ldots, \theta_{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_\infty$  in Theorem 1 which is not the global minimiser. Specifically neither  $Q_\infty$  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 MMD and KL.
For a kernel  $\kappa : \mathbb{R}^J \times \mathbb{R}^J \to \mathbb{R}$ , the squared MMD between measures  $Q$  and  $P$  is
$$
\begin{array}{l} \operatorname {M M D} (Q, P) ^ {2} = \iint \kappa (\theta , \theta^ {\prime}) d Q (\theta) d Q \left(\theta^ {\prime}\right) - 2 \iint \kappa (\theta , \theta^ {\prime}) d Q (\theta) d P \left(\theta^ {\prime}\right) \\ + \iint \kappa (\theta , \theta^ {\prime}) d P (\theta) d P \left(\theta^ {\prime}\right). \\ \end{array}
$$
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 = \mathrm{MMD}^2$  introduces interactions characterised precisely by the kernel  $\kappa$ , and we can show this explicitly by rewriting  $L$  of (1) into the form of  $L^{\mathrm{fe}}$  in (4). In other words, inclusion of  $\mathrm{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  $\mu_P(\theta) \coloneqq \int \kappa(\theta, \theta') dP(\theta')$ , we see that up to a constant not depending on  $Q$ ,  $L(Q) = L^{\mathrm{fe}}(Q)$  for  $V(\theta) = \ell(\theta) - \lambda_1 \mu_P(\theta)$ ,  $\lambda = \frac{\lambda_1}{2}$ , and  $\lambda_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 = \mathrm{MMD}^2$ .
To remedy this problem, we additionally regularise with the KL: since  $\mathrm{KL}(Q,P) = \infty$  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 = \lambda \mathrm{MMD}^2 +\lambda '\mathrm{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^{\mathrm{fe}}(Q)$  up to a constant not depending on  $Q$  by taking  $\lambda = \frac{\lambda_1}{2}$ ,  $\lambda^{\prime} = \lambda_{2}$ , and  $V(\theta) = \ell (\theta) - \lambda_{1}\mu_{P}(\theta) - \lambda_{2}\log p(\theta)$ . 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  $\kappa$ , 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  $\theta_{n}(T)$ ,  $n = 1,\dots ,N_{E}$ , generated via (5). Then  $Q^{n,N_E}(T)\xrightarrow{\mathcal{D}}Q^*$  for each  $n = 1,\ldots ,N_E$  and as  $N_{E}\to \infty$ ,  $T\to \infty$ .
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
Figure 2: We generate  $N_{E} = 300$  particles from DE, DLE and DRLE. The theoretically optimal global minimisers  $Q^{*}$  are depicted with dashed line strokes, and the generated samples are displayed via histograms. We use  $P = Q_{0} = \mathcal{N}(0,1)$  for DLE and DRLE. Notice that the optimal  $Q^{*}$  differs slightly between DLE and DRLE.
demonstrates how impressively powerful the WGF is as tool to derive inference algorithms. Note that this is completely different from sampling methods employed for Bayesian methods, for which the form of  $Q^{*}$  is typically known explicitly up to a proportionality constant.
A notable shortcoming of Theorem 2 is its asymptotic nature. A more refined analysis could quantify how fast the convergence happens in terms of  $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.
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 \mapsto \int \ell(\theta) dQ(\theta)$  over  $\mathcal{P}(\mathbb{R}^J)$  (which is given as Dirac measure located at  $\theta = -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  $\mathcal{P}_2(\mathbb{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  $\ell(\theta) := -|\sin(\theta)|$ ,  $\theta \in [-1000\pi, 1000\pi]$ , which has 2000 local minima. Correspondingly  $Q^{*}$  will have many local modes for all methods. Note that  $\nabla V$  is the same for all approaches since  $\log p$  and  $\mu_{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
Figure 3: We generate  $N_{E} = 300$  particles from DE, DLE, DRLE and FD-GVI with Gaussian parametrisation. The multimodal loss  $\ell$  is plotted in grey and the particles of the different methods are layered on top. The prior in this example is flat, i.e.  $\log p$  and  $\mu_{P}$  are constant. The initialisation  $Q_{0}$  is standard Gaussian.
Figure 4: We generate  $N_{E} = 20$  samples from the three infinite-dimensional gradient descent procedures discussed in Section 4. The  $x$ -axis shows the location of the particles after training. Since the same initialisation  $\theta_{n}(0)$  is chosen for all methods, we observe that particles fall into the same local modes. Further, 16/20 particles are alone in their respective local modes and the location of the particles varies only very little between the different methods (which is why they are in the same bucket in the above histogram). See also Figure 5 in the Appendix for an alteration of Figure 3 with only 4 particles which emphasizes the same point.
their local modes, but not explore other parts of the space. This implies that they will not improve the approximation quality of  $Q^{*}$ . Note that this problem is a direct parallel to multi-modality—a well-known problem for Markov Chain Monte Carlo methods (see e.g. Syed et al., 2022).
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.
KIN8NM	CONCRETE	ENERGY	NAVAL	POWER	PROTEIN	WINE	YACHT
DE	0.33±0.1	6.10±0.3	2.83±0.2	-0.40±0.3	13.70±2.6	11.22±2.2	14.65±1.9	2.20±0.4
DLE	13.25±4.3	5.11±0.2	2.43±0.1	3.46±2.4	13.87±2.3	43.20±12.5	13.73±1.4	1.64±0.1
DRLE	0.46±0.1	8.30±0.6	4.01±0.3	-3.04±0.2	23.21±2.0	48.80±2.1	7.13±0.6	7.80±2.7
Table 1: Table compares the average (Gaussian) negative log likelihood in the test set for the three ID-GVI methods on some UCI-regression data sets (Lichman, 2013). We observe that no method consistently outperforms any of the others.
VW was supported by the scattered scholarship and the EPSRC grant EP/W005859/1. SG was a PhD student of the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1), and also received funding from the Oxford Radcliffe Scholarship and Novartis. JK was funded by EPSRC grant EP/W005859/1.