Repulsive Deep Ensembles are Bayesian
Year: 2021

Abstract:
Deep ensembles have recently gained popularity in the deep learning community for their conceptual simplicity and efficiency. However, maintaining functional diversity between ensemble members that are independently trained with gradient descent is challenging. This can lead to pathologies when adding more ensemble members, such as a saturation of the ensemble performance, which converges to the performance of a single model. Moreover, this does not only affect the quality of its predictions, but even more so the uncertainty estimates of the ensemble, and thus its performance on out-of-distribution data. We hypothesize that this limitation can be overcome by discouraging different ensemble members from collapsing to the same function. To this end, we introduce a kernelized repulsive term in the update rule of the deep ensembles. We show that this simple modification not only enforces and maintains diversity among the members but, even more importantly, transforms the maximum a posteriori inference into proper Bayesian inference. Namely, we show that the training dynamics of our proposed repulsive ensembles follow a Wasserstein gradient flow of the KL divergence to the true posterior. We study repulsive terms in weight and function space and empirically compare their performance to standard ensembles and Bayesian baselines on synthetic and real-world prediction tasks.

heading: 1 Introduction
text: There have been many recent advances on the theoretical properties of sampling algorithms for approximate Bayesian inference, which changed our interpretation and understanding of them. Particularly worth mentioning is the work of Jordan et al. [38], who reinterpret Markov Chain Monte Carlo (MCMC) as a gradient flow of the KL divergence over the Wasserstein space of probability measures. This new formulation allowed for a deeper understanding of approximate inference methods but also inspired the inception of new and more efficient inference strategies. Following this direction, Liu and Wang [51] recently proposed the Stein Variational Gradient Descent (SVGD) method to perform approximate Wasserstein gradient descent. Conceptually, this method, which belongs to the family of particle-optimization variational inference (POVI), introduces a repulsive force through a kernel acting in the parameter space to evolve a set of samples towards high-density regions of the target distribution without collapsing to a point estimate.
Another method which has achieved great success recently are ensembles of neural networks (socalled deep ensembles), which work well both in terms of predictive performance [42, 80] as well as uncertainty estimation [65], and have also been proposed as a way to perform approximate inference in Bayesian neural networks [82, 36]. That being said, while they might allow for the averaging of predictions over several hypotheses, they do not offer any guarantees for the diversity between those hypotheses nor do they provably converge to the true Bayesian posterior under any meaningful limit. In this work, we show how the introduction of a repulsive term between the members in the ensemble, inspired by SVGD, not only naïvely guarantees the diversity among the members, avoiding their
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
collapse in parameter space, but also allows for a reformulation of the method as a gradient flow of the KL divergence in the Wasserstein space of distributions. It thus allows to endow deep ensembles with convergence guarantees to the true Bayesian posterior.
An additional problem is that BNN inference in weight space can lead to degenerate solutions, due to the overparametrization of these models. That is, several samples could have very different weights but map to the same function, thus giving a false sense of diversity in the ensemble. This property, that we will refer to as non-identifiability of neural networks (see Appendix A), can lead to redundancies in the posterior distribution. It implies that methods like MCMC sampling, deep ensembles, and SVGD waste computation in local modes that account for equivalent functions. Predictive distributions approximated using samples from these modes do not improve over a simple point estimate and lead to a poor uncertainty estimation. Following this idea, Wang et al. [76] introduced a new method to extend POVI methods to function space, overcoming this limitation. Here, we also study an update rule that allows for an approximation of the gradient flow of the KL divergence in function space in our proposed repulsive ensembles.
We make the following contributions:
• We derive several different repulsion terms that can be added as regularizers to the gradient updates of deep ensembles to endow them with Bayesian convergence properties.
• We show that these terms approximate Wasserstein gradient flows of the KL divergence and can be used both in weight space and function space.
• We compare these proposed methods theoretically to standard deep ensembles and SVGD and highlight their different guarantees.
• We assess all these methods on synthetic and real-world deep learning tasks and show that our proposed repulsive ensembles can achieve competitive performance and improved uncertainty estimation.

heading: 2 Repulsive Deep Ensembles
text: In supervised deep learning, we typically consider a likelihood function p(y|f(x;w)) (e.g., Gaussian for regression or Categorical for classification) parameterized by a neural network f(x;w) and training data D = {(xi,yi)}ni=1 with x 2 X and y 2 Y . In Bayesian neural networks (BNNs), we are interested in the posterior distribution of all likely networks given by p(w|D) / Qn i=1 p(yi|f(xi;w)) p(w), where p(w) is the prior distribution over weights. Crucially, when making a prediction on a test point x⇤, in the Bayesian approach we do not only use a single parameter bw to predict y⇤ = f(x⇤; bw), but we marginalize over the whole posterior, thus taking all possible explanations of the data into account:
p(y⇤|x⇤,D) = Z p(y⇤|f(x⇤;w)) p(w|D) dw (1)
While approximating the posterior of Bayesian neural networks (or sampling from it) is a challenging task, performing maximum a posteriori (MAP) estimation, which corresponds to finding the mode of the posterior, is usually simple. Ensembles of neural networks use the non-convexity of the MAP optimization problem to create a collection of K independent—and possibly different— solutions. Considering n weight configurations of a neural network {wi}ni=1 with wi 2 Rd, the dynamics of the ensemble under the gradient of the posterior lead to the following update rule at iteration t:
w t+1 i wti + ✏t (wti)
with (wti) = rwti log p(w t i |D) , (2)
with step size ✏t. Ensemble methods have a long history [e.g., 45, 26, 6] and were recently revisited for neural networks [42] and coined deep ensembles. The predictions of the different members are combined to create a predictive distribution by using the solutions to compute the Bayesian model average (BMA) in Eq. (1). Recent works [65] have shown that deep ensembles can outperform some of the Bayesian approaches for uncertainty estimation. Even more recently, Wilson and Izmailov [82] argued that deep ensembles can be considered a compelling approach to Bayesian model averaging. Despite these ideas, the ability of deep ensembles to efficiently average over multiple hypotheses and to explore the functional landscape of the posterior distribution studied in [18] does not guarantee sampling from the right distribution. Indeed, the additional Langevin noise introduced in [77], which is not considered in deep ensembles, is crucial to ensure samples from the true Bayesian posterior.
From a practical standpoint, since the quality of an ensemble hinges on the diversity of its members, many methods were recently proposed to improve this diversity without compromising the individual accuracy. For instance, Wenzel et al. [80] propose hyper-deep ensembles that combine deep networks with different hyperparameters. Similarly, cyclical learning-rate schedules can explore several local minima for the ensemble members [33]. Alternatively, Rame and Cord [67] proposed an informationtheoretic framework to avoid redundancy in the members and Oswald et al. [63] studied possible interactions between members based on weight sharing. However, the absence of a constraint that prevents particles from converging to the same mode limits the possibility of improvement by introducing more ensemble members. This means that any hopes to converge to different modes must exclusively rely on:
1. the randomness of the initialization 2. the noise in the estimation of the gradients due to minibatching 3. the number of local optima that might be reached during gradient descent.
Moreover, the recent study of Geiger et al. [25] showed how the empirical test error of the ensemble converges to the one of a single trained model when the number of parameters goes to infinity, leading to deterioration of the performance. In other words, the bigger the model, the harder it is to maintain diversity in the ensemble and avoid collapse to the same solution. This is intuitively due to the fact that bigger models are less sensitive to the initialization. Namely, in order for them to get stuck in a local minimum, they must have second derivatives that are positive simultaneously in all directions. As the number of hidden units gets larger, this becomes less likely.

heading: 2.1 Repulsive force in weight space
text: To overcome the aforementioned limitations of standard deep ensembles, we introduce, inspired by SVGD [51], a deep ensemble with members that interact with each other through a repulsive component. Using a kernel function to model this interaction, the single models repel each other based on their position in the weight space, so that two members can never assume the same weights. Considering a stationary kernel k(·, ·) : Rd ⇥ Rd ! R acting in the parameter space of the neural network, a repulsive term R can be parameterized through its gradient:
(wti) = rwti log p(w t i |D) R ✓n rwtik(w t i ,w t j) on j=1 ◆ . (3)
To get an intuition for the behavior of this repulsive term and its gradients, we can consider the RBF kernel k(wi,wj) = exp 1h ||wi wj || 2 with lengthscale h and notice how its gradient
rwtik(w t i ,w t j) =
2 h (wtj wti)k(wti ,wtj)
drives wi away from its neighboring members wj , thus creating a repulsive effect. Naturally, not all the choices of R induce this effect. One of the simplest formulations to obtain it is via a linear combination of the kernel gradients scaled by a positive factor, that is,
Pn j=1rwtik(w t i ,w t j) with
2 R+⇤ . We will see in Section 3 how the choice of can be justified in order to obtain convergence to the Bayesian posterior together with alternative possible formulations of R that preserve this convergence.

heading: 2.2 Repulsive force in function space
text: To overcome the aforementioned overparameterization issue, the update in Eq. (3) can be formulated in function space instead of weight space. Let f : w 7! f(· ;w) be the map that maps a configuration of weights w 2 Rd to the corresponding neural network regression function and denote as fi := f(· ;wi) the function with a certain configuration of weights wi. We can now consider n particles in function space {fi}ni=1 with f 2 F and model their interaction with a general positive definite kernel k(·, ·). We also consider the implicit functional likelihood p(y|x,f), determined by the measure p(y|x,w) in the weight space, as well as the functional prior p(f), which can either be defined separately (e.g., using a GP) or modeled as a push-forward measure of the weight-space prior p(w). Together, they determine the posterior in function space p(f |D). The functional evolution of a particle can then be written as:
f t+1i f ti + ✏t (f ti )
with (f ti ) = rfti log p(f t i |D) R ✓n rfti k(f t i ,f t j ) on j=1 ◆ . (4)
However, computing the update in function space is neither tractable nor practical, which is why two additional considerations are needed. The first one regards the infinite dimensionality of function space, which we circumvent using a canonical projection into a subspace: Definition 1 (Canonical projection). For any A ⇢ X , we define ⇡A : RX ! RA as the canonical projection onto A, that is, ⇡A(f) = {f(a)}a2A.
In other words, the kernel will not be evaluated directly in function space, but on the projection k ⇡B(f),⇡B(f 0) , with B being a subset of the input space given by a batch of training data points. The second consideration is to project this update back into the parameter space and evolve a set of particles there, because ultimately we are interested in representing the functions by parameterized neural networks. For this purpose, we can use the Jacobian of the i-th particle as a projector:
(wti) = ✓ @f ti @w
t i
◆>  rfti log p(f t i |D) R ✓n rfti k(⇡B(f t i ),⇡B(f t j )) on j=1 ◆ . (5)

heading: 2.3 Comparison to Stein variational gradient descent
text: Note that our update is reminiscent of SVGD [51], which in parameter space can be written as:
(wti) = nX
j=1
k(wti ,w t j)rwti log p(w t i |D) +
nX
j=1
rwtjk(w t j ,w t i) . (6)
It is important to notice that here, the gradients are averaged across all the particles using the kernel matrix. Interestingly, SVGD can be asymptotically interpreted as gradient flow of the KL divergence under a new metric induced by the Stein operator [16, 50] (see Appendix D for more details). Moving the inference from parameter to function space [76] leads to the update rule
(wti) = ✓ @f ti @w
t i
◆> ✓ 1
n
nX
j=1
k(f ti ,f t j )rftj log p(f t j |D) +rftj k(f t i ,f t j )
◆ . (7)
This way of averaging gradients using a kernel can be dangerous in high-dimensional settings, where kernel methods often suffer from the curse of dimensionality. Moreover, in Eq. (6), the posterior gradients of the particles are averaged using their similarity in weight space, which can be misleading in multi-modal posteriors. Worse yet, in Eq. (7), the gradients are averaged in function space and
are then projected back using exclusively the i-th Jacobian, which can be harmful given that it is not guaranteed that distances between functions evaluated on a subset of their input space resemble their true distance. Our proposed method, on the other hand, does not employ any averaging of the posterior gradients and thus comes closest to the true particle gradients in deep ensembles.

heading: 3 Repulsive deep ensembles are Bayesian
text: So far, we represented the repulsive force as a general function of the gradients of a kernel. In this section, we show how to determine the explicit form of the repulsive term, such that the resulting update rule is equivalent to the discretization of the gradient flow dynamics of the KL divergence in Wasserstein space. We begin by introducing the concepts of particle approximation and gradient flow.

heading: 3.1 Particle approximation
text: A particle-based approximation of a target measure depends on a set of weighted samples {(xi, wi)}ni=1, for which an empirical measure can be defined as
⇢(x) = nX
i=1
wi (x xi) , (8)
where (·) is the Dirac delta function and the weights wi satisfy wi 2 [0, 1] and Pn
i=1 wi = 1. To approximate a target distribution ⇡(x) using the empirical measure, the particles and their weights need to be selected in a principled manner that minimizes some measure of distance between ⇡(x) and ⇢(x) (e.g., a set of N samples with weights wi = 1/N obtained using an MCMC method).

heading: 3.2 Gradient flow in parameter space
text: Given a smooth function J : Rd ! R in Euclidean space, we can minimize it by creating a path that follows its negative gradient starting from some initial conditions x0. The curve x(t) with starting point x0 described by that path is called gradient flow. The dynamics and evolution in time of a considered point in the space under this minimization problem can be described as the ODE1
dx dt = rJ(x) . (9)
We can extend this concept to the space of probability distributions (Wasserstein gradient flow) [3]. Let us consider the space of probability measures P2(M), that is, the set of probability measures with finite second moments defined on the manifold M:
P2(M) = ⇢ ' : M! [0,1) Z
M d' = 1,
Z
M |x|2'(x)dx < +1 .
Taking ⇧(µ, ⌫) as the set of joint probability measures with marginals µ, ⌫, we can define the Wasserstein metric on the space P2(M) as:
W 2 2 (µ, ⌫) = inf
⇡2⇧(µ,⌫)
Z |x y|2 d⇡(x, y) . (10)
Considering the optimization problem of a functional J : P2(M)! R, such as the KL divergence between the particle approximation in Eq. (8) and the target posterior ⇡(x),
inf ⇢2P2(M) DKL(⇢,⇡) =
Z
M (log ⇢(x) log ⇡(x))⇢(x) dx ,
the evolution in time of the measure ⇢ under the equivalent of the gradient, the Wasserstein gradient flow, is described by the Liouville equation2 [38, 3, 64]:
@⇢(x)
@t = r ·
✓ ⇢(x)r
⇢ DKL(⇢,⇡)
◆
= r · ✓ ⇢(x)r log ⇢(x) log ⇡(x) ◆ ,
(11)
1Together with the initial condition, this is know as the Cauchy problem. 2Also referred to as continuity equation.
where r ⇢DKL(⇢,⇡) =: rW2DKL(⇢,⇡) is the Wasserstein gradient and the operator ⇢ : P2(M) ! R represents the functional derivative or first variation (see Appendix C for more details). In the particular case of the KL functional, we can recover the Fokker-Planck equation,
@⇢(x)
@t = r · ⇢(x)r(log ⇢(x) log ⇡(x))
= r · ⇢(x)r log ⇡(x) +r2⇢(x) ,
that admits as unique stationary distribution the posterior ⇡(x). The deterministic particle dynamics ODE [2] related to Eq. (11), namely mean-field Wasserstein dynamics, is then given by:
dx dt = r log ⇢(x) log ⇡(x) . (12)
Considering a discretization of Eq. (12) for a particle system {x}ni=1 and small stepsize ✏t, we can rewrite Eq. (12) as:
x t+1 i = x t i + ✏t r log ⇡(xti) r log ⇢(xti) . (13)
Unfortunately, we do not have access to the analytical form of the gradient r log ⇢, so an approximation is needed. At this point, it is crucial to observe the similarity between the discretization of the Wasserstein gradient flow in Eq. (13) and the repulsive update in Eq. (3) to notice how, if the kernelized repulsion is an approximation of the gradient of the empirical particle measure, the update rule minimizes the KL divergence between the particle measure and the target posterior. Different sample-based approximations of the gradient that use a kernel function have been recently studied. The simplest one is given by the kernel density estimation (KDE) (details in Appendix E) ⇢̃t(x) = 1 n Pn i=1 k(x, x i t), where k(·, ·) : Rd ⇥ Rd ! R and the gradient of its log density is given by [70]:
r log ⇢(xti) ⇡ Pn j=1rxtik(x t i, x t j)Pn
j=1 k(x t i, x t j)
. (14)
Using this approximation in Eq. (13) we obtain:
x t+1 i = x t i + ✏t ✓ r log ⇡(xti) Pn j=1rxtik(x t i, x t j)Pn
j=1 k(x t i, x t j)
◆ , (15)
where, if we substitute the posterior for ⇡, we obtain an expression for the repulsive force in Eq. (3). This shows that if the repulsive term in Eq. (3) is the normalized sum of the gradients R =⇣Pn
j=1 k(w t i ,w t j) ⌘ 1 Pn j=1rwtik(w t i ,w t j), we do not only encourage diversity of the ensemble members and thus avoid collapse, but surprisingly—in the asymptotic limit of n!1, where the KDE approximation is exact [66]—also converge to the true Bayesian posterior!
Nevertheless, approximating the gradient of the empirical measure with the KDE can lead to suboptimal performance, as already studied by Li and Turner [46]. They instead introduced a new Stein gradient estimator (SGE) that offers better performance, while maintaining the same computational cost. Even more recently, Shi et al. [69] introduced a spectral method for gradient estimation (SSGE), that also allows for a simple estimation on out-of-sample points. These two estimators can be used in Eq. (13), to formulate the following update rules with two alternative repulsive forces. The one using the Stein estimator, that we will call SGE-WGD, is:
x t+1 i = x t i + ✏t ✓ r log ⇡(xti) + nX
j=1
(K + ⌘I) 1ij nX
k=1
rxtkk(x t k, x t j)
◆ , (16)
where K is the kernel Gram matrix, ⌘ a small constant, and I the identity matrix. We can notice an important difference between KDE and SGE, in that the former is only considering the interaction of the i-th particle being updated with all the others, while the latter is simultaneously considering also the interactions between the remaining particles. The spectral method, that we will call SSGE-WGD, leads to the following update rule:
x i t+1 = x i t + ✏t ✓ r log ⇡(xit) + JX
j=1
1
2 j
nX
m=1
nX
k=1
ujkrxmk(xtm, xtk) · nX
l=1
ujlk(x t i, x t l)
◆ (17)
where j is the j-th eigenvalue of the kernel matrix and ujk is the k-th component of the j-th eigenvector. Computationally, both SSGE and SGE have a cost of O(M3 +M2d), with M being the number of points and d their dimensionality. SSGE has an additional cost for predictions of O(M(d+ J)), where J is the number of eigenvalues.

heading: 3.3 Gradient flow in function space
text: To theoretically justify the update rule introduced in function space in Eq. (5), we can rewrite the Liouville equation for the gradient flow in Eq. (11) in function space as
@⇢(f)
@t = r ·
✓ ⇢(f)r
⇢ DKL(⇢,⇡)
◆
= r · ✓ ⇢(f)r log ⇢(f) log ⇡(f) ◆ .
(18)
Following this update, the mean field functional dynamics are
df dt = r log ⇢(f) log ⇡(f) . (19)
Using the same KDE approximation as above, we can obtain a discretized evolution in function space and with it an explicit form for the repulsive force in Eq. (4) as
f it+1 = f i t + ✏t ✓ rf log ⇡(f it ) Pn j=1rfti k(f t i ,f t j )Pn
j=1 k(f t i ,f t j )
◆ . (20)
The update rules using the SGE and SSGE approximations follow as for the parametric case. It is important to notice that this update rule requires the function space prior gradient: rfj log p(fj |x,y) = rfj log p(y|x,fj) +rfj log p(fj). If one wants to use an implicit prior defined in weight space, an additional estimator is needed due to its analytical intractability. We again adopted the SSGE, introduced by Shi et al. [69], which was already used for a similar purpose in Sun et al. [71]. It is also interesting to note that the update rule in Eq. (20) readily allows for the use of alternative priors that have an analytical form, such as Gaussian processes. This is an important feature of our method that allows for an explicit encoding of function space properties that can be useful for example in achieving better out of distribution detection capabilities [14].

heading: 3.4 The choice of the kernel
text: A repulsive effect can always be created in the ensemble by means of the gradient of any kernel function that is measuring the similarity between two members. Nevertheless, it is important to keep in mind that to ensure the asymptotic convergence to the Bayesian posterior, the repulsive component must be a consistent estimator of the gradient in Eq. (13), as shown in Section 3. Therefore, some important constraints over the kernel choice are needed. In particular, the SGE and SSGE need a kernel function belonging to the Stein class (see Shi et al. [69] for more details). On the other hand, for the KDE, any symmetric probability density function can be used. On this subject, the work of Aggarwal et al. [1] has shown how the Manhattan distance metric (L1 norm) might be preferable over the Euclidean distance metric (L2 norm) for high-dimensional settings. We performed some additional experiments using the L1 norm (Laplace kernel) for the KDE but we could not observe any substantial difference compared to using the L2 norm. Further investigations regarding this hypothesis are left for future research.

heading: 4 Experiments
text: In this section, we compare the different proposed WGD methods with deep ensembles and SVGD on synthetic sampling, regression, and classification tasks and real-world image classification tasks. We use an RBF kernel (except where otherwise specified) with the popular median heuristic [51] to choose the kernel bandwidth. In our experiments, an adaptive bandwidth leads to better performance than fixing and tuning a single constant value for the entire evolution of the particles. We also quantitatively assess the uncertainty estimation of the methods in terms of calibration and OOD detection. In our experiments, we report the test accuracy, negative log-likelihood (NLL), and the expected calibration error (ECE) [58]. To assert the robustness on out-of-distribution (OOD) data, we report the ratio between predictive entropy on OOD and test data points (Ho/Ht), and the OOD detection area under the ROC curve AUROC(H) [47]. Moreover, to assess the diversity of the ensemble generated by the different methods in function space, we measure the functional diversity using the model disagreement (MD) (details in Appendix B). In particular, we report the ratio between the average model disagreement on the OOD and test data points (MDo/MDt) and additionally the OOD detection AUROC(MD) computed using this measure instead of the entropy.
Sampling from synthetic distributions As a sanity check, we first assessed the ability of our different approximations for Wasserstein gradient descent (using KDE, SGE, and SSGE) to sample from a two-dimensional Gaussian distribution (Figure 2). We see that our SGE-WGD, SSGE-WGD and the SVGD fit the target almost perfectly. We also tested the different methods in a more complex two-dimensional Funnel distribution [59] and present the results in Figure F.1 in the Appendix. There, SGE-WGD and SVGD also perform best.
BNN 1D regression We then assessed the different methods in fitting a BNN posterior on a synthetically generated one-dimensional regression task. The results are reported in Figure 1, consisting of the mean prediction and ±1, 2, 3 standard deviations of the predictive distribution. We can see that all methods performing inference in the weight space (DE, w-SVGD, WGD) are unable to capture the epistemic uncertainty between the two clusters of training data points. Conversely, the functional methods (f-SVGD, fWGD) are perfectly able to infer the diversity of the hypotheses in this region due to the lack of training evidence. They thereby achieve a predictive posterior that very closely resembles the one obtained with the gold-standard HMC sampling.
BNN 2D classification Next, we investigated the predictive performance and quality of uncertainty estimation of the methods in a two-dimensional synthetic classification setting. The results are displayed in Figure 3. We can clearly observe that the weight-space methods are overconfident and do not capture the uncertainty well. Moreover, all the functions seems to collapse to the optimal classifier. These methods thus only account for uncertainty close to the decision boundaries and to the origin
region, for which the uncertainty is purely aleatoric. In this setting, f-SVGD suffers from similar issues as the weight space methods, being overconfident away from the training data. Conversely, our fWGD methods are confident (low entropy) around the data but not out-of-distribution, thus representing the epistemic uncertainty better. This suggests that the functional diversity captured by this method naturally leads to a distance-aware uncertainty estimation [49, 20], a property that translates into confident predictions only in the proximity of the training data, allowing for a principled OOD detection.
FashionMNIST classification Moving on to real-world data, we used an image classification setting using the FashionMNIST dataset [83] for training and the MNIST dataset [43] as an out-ofdistribution (OOD) task. The results are reported in Table 1 (top). We can see that all our methods improve upon standard deep ensembles and SVGD, highlighting the effectiveness of our proposed repulsion terms when training neural network ensembles. In particular, the sge-WGD offers the best accuracy, whereas the methods in function space all offer a better OOD detection. This is probably due to the fact that these methods achieve a higher entropy ratio and functional diversity measured via the model disagreement when compared to their weight-space counterparts. Interestingly, they also reach the lowest NLL values. We can also notice how the model disagreement (MD) not only serves its purpose as a metric for the functional heterogeneity of the ensemble but also allows for a better OOD detection in comparison to the entropy. To the best of our knowledge, this insight has not been described before, although it has been used in continual learning [31]. Interestingly, using this metric, the hyper-deep ensemble [80] shows OOD detection performance comparable with our repulsive ensemble in function space.
CIFAR classification Finally, we use a ResNet32 architecture [29] on CIFAR-10 [41] with the SVHN dataset [61] as OOD data. The results are reported in Table 1 (bottom). We can see that in this case, the weight-space methods achieve better performance in accuracy and OOD detection using the entropy than the ones in function space. Nevertheless, all our repulsive ensembles improve functional diversity, accuracy, and OOD detection when compared to standard SVGD, whereas the standard deep ensemble achieves the best OOD detection using the entropy.

heading: 5 Related Work
text: The theoretical and empirical properties of SVGD have been well studied [40, 48, 13] and it can also be seen as a Wasserstein gradient flow of the KL divergence in the Stein geometry [16, 50]
(see Appendix D for more details). Interestingly, a gradient flow interpretation is also possible for (stochastic gradient) MCMC-type algorithms [48], which can be unified under a general particle inference framework [10]. Moreover, our Wasserstein gradient descent using the SGE approximation can also be derived using an alternative formulation as a gradient flow with smoothed test functions [48]. A projected version of WGD has been studied in Wang et al. [75], which could also be readily applied in our framework. Besides particle methods, Bayesian neural networks [54, 59] have gained popularity recently [79, 22, 19, 36], using modern MCMC [59, 79, 22, 24, 21] and variational inference techniques [5, 72, 17, 34]. On the other hand, ensemble methods have also been extensively studied [42, 18, 82, 23, 80, 32, 85, 78].Moreover, repulsive interactions between the members have also been studied in Wabartha et al. [74]. Moreover, providing Bayesian interpretations for deep ensembles has been previously attempted through the lenses of stationary SGD distributions [56, 8], ensembles of linear models [57], additional random functions [62, 12, 27], approximate inference [82], Stein variational inference [15], and marginal likelihood lower bounds [53], and ensembles have also been shown to provide good approximations to the true BNN posterior in some settings [36]. Furthermore, variational inference in function space has recently gained attention [71] and the limitations of the KL divergence have been studied in Burt et al. [7].

heading: 6 Conclusion
text: We have presented a simple and principled way to improve upon standard deep ensemble methods. To this end, we have shown that the introduction of a kernelized repulsion between members of the ensemble not only improves the accuracy of the predictions but—even more importantly—can be seen as Wasserstein gradient descent on the KL divergence, thus transforming the MAP inference of deep ensembles into proper Bayesian inference. Moreover, we have shown that incorporating functional repulsion between ensemble members can improve the quality of the estimated uncertainties on simple synthetic examples and OOD detection on real-world data and can approach the true Bayesian posterior more closely.
In future work, it will be interesting to study the impact of the Jacobian in the fWGD update and its implications on the Liouville equation in more detail, also compared to other neural network Jacobian methods, such as neural tangent kernels [37] and generalized Gauss-Newton approximations [35]. Moreover, it would be interesting to derive explicit convergence bounds for our proposed method and compare them to the existing bounds for SVGD [40].

heading: Acknowledgments
text: VF would like to acknowledge financial support from the Strategic Focus Area “Personalized Health and Related Technologies” of the ETH Domain through the grant #2017-110 and from the Swiss Data Science Center through a PhD fellowship. We thank Florian Wenzel, Alexander Immer, Andrew Gordon Wilson, Pavel Izmailov, Christian Henning, and Johannes von Oswald for helpful discussions. We also thank Dr. Sheldon Cooper, Dr. Leonard Hofstadter and Penny Hofstadter for their support and inspiration.

Review summaries:
Maintaining functional diversity between ensemble members are a large, and still not completely solved problem. The authors are introducing a kernelized repulsive term to the update rule of an ensemble. They show empirically that this increases the functional diversity between ensemble members. The authors also demonstrate that their improvement translates into a proper Bayesian inference.

This paper introduces a novel but simple repulsive term to the training objective for deep ensembles which prevents the ensemble members from collapsing to the same point in *function* space. This training procedure is closely related to that of Stein variational gradient descent. Furthermore, the authors show that this repulsive term turns deep ensemble training into approximate Bayesian inference. Finally, the authors demonstrate that this method improves on both SVGD and deep ensembles in a range of benchmark tasks. 

The paper tackles the problem of building ensemble methods to improve prediction accuracy and uncertainty estimates. The authors propose adding a repulsive term to the gradient update rule for each ensemble element to produce diverse ensembles, and they design repulsive terms for which the empirical distribution over ensemble members converges to the true posterior distribution. Specifically, the authors propose three repulsive terms that arise from three alternative approximations of the gradient of the empirical distribution over ensemble members, and they compare the resulting ensembles with standard deep ensembles in synthetic tasks and image classification tasks. Finally, to account for the fact that different weights can parameterize the same function (and therefore a repulsive term in the weight space may not achieve the desired effect), the authors consider updates (and repulsive terms) in both weight space and function space and compare both approaches in each evaluation task.  The claims are that the proposed approach: * enables using deep ensembles to perform proper Bayesian inference, * improves the quality of the uncertainty estimates in synthetic tasks, * improves OOD detection on image datasets.

The paper proposes a particle based inference method that, unlike Deep Ensembles, approximates the Bayesian posterior of a deep neural network. The proposed approaches are variants of Stein Variational Gradient Descent. They employ a repulsive force (in the form of a kernel) that prevents the particles from collapsing into the same mode.  The main contribution of the work is showing that this particle-based training method is equivalent to gradient descent minimizing the KL-divergence of the true posterior and the particle approximation in Wasserstein space.  The paper also presents empirical results. On a 1d synthetic benchmark, it is shown that the proposed methods are able to capture in-between uncertainty. On a 2d synthetic benchmark, it is shown that the uncertainty estimates are akin to that of HMC. On larger, image classification datasets, it is shown that the methods outperform both deep ensembles and svgd.