Repulsive Deep Ensembles are Bayesian
Authors: Francesco D'Angelo
Vincent Fortuin
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Bayesian deep learning
Deep ensembles
Wasserstein gradient flow

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 with 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.

Repulsive Deep Ensembles are Bayesian
Francesco D' D'Angelo
Vincent Fortuin
ETH Zürich
ETH Zürich
Zürich, Switzerland dngfra@gmail.com
Zürich, Switzerland fortuinšinf ethz..cc
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 members, such as a saturation of the ensemble performance, which converges to the performance of single model. Moreover, this does not only affect the quality of its predictions, but even more so the uncertainty estimates of the ensemble, and dittt---distibutiion data. We hypothesize that this limitation
ensemble
thus its performance can be overcome by discouraging different ensemble members from collapsing to the same function. To this end, introduce 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. 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.
Introduction
There have been many recent advances 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 deeper understanding of approximate inference methods but also inspired the inception of new and more efficient inference strategies. Following this direction, Liu and Wang 51l 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-optiiizatinn variational inference (POVI), introduces a repulsive force through a kemel acting in the parameter space to evolve a set of samples towards high-density regions of the target distribution without collapsing
point estimate.
Another method which has achieved great success recently are ensembles of neural networks (so- called deep ensembles), which work well both in terms of predictive performance 42 80 well as uncertainty estimation (51, and have also been proposed as a way to perform approximate inference in Bayesian neural networks [82] [6). 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 do they provably converge to the Bayesian posterior under any meaningful limit. In this work, we show how the introduction of repulsive term between the members the ensemble, inspired by SVGD, not only naïvely guarantees the diversity among the members, avoiding their
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Deep Ensemble
w-SVGD
kde-WGD
ssge-WGD
sge-WGD
HMC
f-SVGD
kde-AWGD
ssge-fWGD
sge -gGD
Figure 1: BNN ID regression. The function-space methods (SVGD and WGD) approach the HMC posterior more closely, while the standard deep ensembles and weight-space methods fail to properly account for the uncertainty, especially the in-between uncertainty.
collapse in parameter space, but also allows for reformulation of the method gradient flow of the KL divergence in the Wasserstein space of distributions. It thus allows to endow deep ensembles with convergence guarantees 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 false sense diversity in the ensemble. This property, that will refer to non-ide non-identifiability of neural networks (see Appendix A can lead to redundancies in the posterior distribution. 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 poor uncertainty estimation. Following this idea, Wang et 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.
76 introduced new method to
We make the following contributions:
We derive several different repulsion terms that can be added as regularizers to the gradient updates of deep ensembles endow them with Bayesian convergence properties.
We show that these terms approximate Wasserstein gradient flows of the KL divergence and 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 synthetic and real-world deep learning tasks and show that our proposed repulsive ensembles can achieve competitive performance and improved uncertainty estimation.
2 Repulsive Deep Ensembles
In supervised deep learning, typically consider likelihood function p(y)f((;; (w):)) (e.g.. Gaus- sian for regression or Categorical for classification) parameterized by a neural network f(x; w) and training data 'D { y:) y:) } with E X' and y y In Bayesian neural net- works (BNNs), we are interested in the posterior distribution of all likely networks given by pl ((W/D) I1 p(y+ [a.; w)) p(w), where p(w) is the prior distribution
weights. Cru- cially, when making a prediction on a test point x", in the Bayesian approach we do not only use single parameter w to predict y" f(x" W), but we marginalize over the whole posterior, thus taking all possible explanations of the data into account:
p< D) plv'lf(="; p(W) dw
(1)

---

While approximating the posterior of Bayesian neural networks (or sampling from it) is a challenging task, performing maximum posterioni (MAP) estimation. which corresponds 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-aane possibly different- solutions. Considering n weight configurations of neural network {w, with w; E IR', the dynamics of the ensemble under the gradient of the posterior lead to the following update rule at iteration E:
+++ f- W + c(o(w!)
with
V w/ log ID),
with step size q.- Ensemble methods have long history [e.g.. 45
and were recently revisited for neural networks [22 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. ① Recent works 665] 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 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 [T&] 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.
26
6
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 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 [ Alternatively, Rame and Cord (77] proposed an information-
theoretic 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 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 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.
2.1 Repulsive force in weight space
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 kernel function 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:
d(w:) = Vw log - R ((TV/
(3)
To get intuition for the behavior of this repulsive term and its gradients, we can consider the RBF kernel k(w,, w,) = exp kll; ws11?) with lengthscale h and notice how its gradient
V Vw:k(wE w/) = (w) w;) w()
3

---

drives w; away from its neighboring members wj, thus creating repulsive effect. Naturally, not all the choices of R induce this effect. One of the simplest formulations to obtain it is via linear combination of the kernel gradients scaled by positive factor, that is, 3 プド= V w k<?; w³) with E R We will in Section how the choice of 3 be justified in order to obtain convergence to the Bayesian posterior together with alternative possible formulations of R that preserve this convergence.
2.2 Repulsive force in function space
To overcome the aforementioned overparameterizatton issue, the update in Eq. e can be formulated in function space instead of weight space. Let weights w R to the corresponding neural network regression function and denote f: := f( w.)
f(:: w) be the map that maps configuration of
the function with a certain configuration of weights w,- We can now consider n particles function space {f:)= with f F and model their interaction with general positive definite kernel ki(;;) We also consider the implicit functional likelihood p (x, f). determined by the measure p(y xz, w) in the weight space, well 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( p(f(2) The functional evolution of particle can then be written as:
81
ft cf((f
with $(f)) y log logp(s D) R ((V)
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 (Canonical projection). For any A C X. we define n R - R as the canonical projection onto A, that is, TA = (J(a))aEA-
In other words, the kernel will not be evaluated directly in function space, but on the projection & ((T(J). TB (f))., with B being 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 are interested in representing the functions by parameterized neural networks. For this purpose, we can use the Jacobian the i-th particle as projector:
<(w = Of! Ow!
T
TT logp(/1)) -R (((V((MM)))
(5)
2.3 Comparison to Stein variational gradient descent
Note that update is reminiscent of SVGD [51]. which in parameter space can be written
<(w))
Σ k(W). V logp(w((D) + Σ Vwy k(w w!). j=1
important to notice that here, the gradients are averaged across all 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 Dovendixffor more details). Moving the inference from parameter to function space [76 leads to the update rule
as! Ow!
@ ((w)) =
M n2
k(f f) logn({))))) - (o
1=1
This way of averaging gradients using a kernel can be dangerous 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-me multi-modal posteriors. Worse yet, in Eq. 77. 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 subset their imput 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 partiele gradients in deep ensembles.
Repulsive deep ensembles are Bayesian
So far, we represented the repulsive force as general function of the gradients of kernel. In this section, we show how 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.
3.1 Particle approximation
particle-based approximation a target measure depends on a set weighted samples {Z,, w; for which empirical measure can be defined as
p(x) w, 6(x xz),
(8)
where 8(-) is the Dirac delta function and the weights wy satisfy w, [0, and Σح we To approximate a target distribution (z)) using the empirical measure, the particles and their weights need to be selected in a principled manner that minimizes some measure of distance between (3r)
and p(x) (e.g., a set of N samples with weights wi = I/N obtained using an MCMC method).
3.2 Gradient flow in parameter space
Given smooth function J R + IR in Euclidean space, minimize it by creating path that follows its negative gradient starting from some initial conditions zo. The curve ztt) with starting point zo described by that path is called gradient flow. The dynamics and evolution in time of considered point in the space under this minimization problem can described as the ODE
dx dt -vxx).
We can extend this concept to the space of probability distributions (Wasserstein gradient (lo)) 3 Let us consider the space of probability measures (2MM)), that is, the set of probability measures with finite second moments defined on the manifold M:
P2(M) = {0= M + [0, oo) / de = 1, f |x| < }
Taking I[ v) as the set of joint probability measures with marginals we can define the Wasserstein metric on space (2(M) as:
1---y/ y/2 d
W (p,v) inf
(10)
xell(a.))
Considering the optimization problem of a functional J P2(M) - TR such as the KI. divergence between the particle approximation in Eq. (8) and the target posterior T(x).
DKI(P.T) = [ (log (log p(x) - log T (2) (()) dz
inf PEPI(M)
the evolution in time of the measure under the equivalent of the gradient, the Wasserstein gradient flow, described by the Liouville equation
Öp(x)
V (p())V D 6p (KL(A.))
Øt
(11)
(lgg)()) log log
'Together with the initial condition, this know as the Cauchy problem. Also referred to continuity equation.
5

---

where V D0 =: (w,DKL(P.T) is the Wasserstein gradient and the operator 6p 6p (2(M) + R represents the functional derivative first variation (see Appendix C for more details). In the particular case of the KL functional, we can recover the Fokker-Planck equation,
0p(z) V (p(x)V(log((a)) - log (1:))
- (p(x) (2)) p²)(2),
that admits as unique stationary distribution the posterior (I)) The deterministic particle dynamics ODE 2 related to Eq. T. namely mean-field Wasserstein dynamics, is then given by:
dr
log
(12)
Considering a discretization of Eq. 12 for a particle system {r]]_1 and small stepsize et, we can rewrite Eq.
as:
x 62 V log (xx)) V log p(x)))
(13)
Unfortunately, we do not have access to the analytical form of the gradient V log p. approx- imation needed. At this point, is crucial to observe the similarity between the discretization of the Wasserstein gradient flow in Eq. (33) 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-bascd 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 B (x 1
Ei= k(x, zt), where k(,,,) Rd x Rd + R and the gradient of its log density is given Vlogp(at)
by 70
EJ=1 V x 2j=1 k(z 27))
(14)
Using this approximation in Eq.
13 we obtain:
Σ Vack(st,
+
V log (2
(15)
where, if we substitute the posterior for æ, we obtain an expression for the repulsive force in Eq. This shows that if the repulsive term in Eq. is the normalized sum of the gradients R =
[Ep k(w), w()) -1
Eÿ= Vw kww w(), we do not only encourage diversity of the ensemble members and thus avoid collapse, but surprisingly in the asymptotic limit of - co, where the KDE approximation exact [66l converge to the true Bayesian posterior!
Nevertheless, approximating the gradient of the empirical measure with the KDE can lead to subop- timal 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
out-of-sampll points. These two estimators be used in Eq. (13), to formulate the following update rules with two alternative repulsive forces. The one using the Stein estimator, that will call SGE-WGI is:
++1 x +
V log K (xf) + (TK+
+ +T
V katkka
(16)
j=1
k=1
where K is the kernel Gram matrix, a small constant, and T 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 simultaneously considering also the interactions between the remaining particles. The spectral method, that we will call SSGE-WGD, leads to the following update rule:
R
2 = 2 Vlog (x)) +
Ugk V 2m k( 2())
Σ Eank(ataat)
(17)
where Xy is the j-th eigenvalue of the kernel matrix and Uik is the k-th component of the j-th eigenvector. Computationally, both SSGE and SGE have cost of O(M3 + M&d), with M being the number of points and d their dimensionality. SSGE has an additional cost for predictions of (d J)), where J is the number of eigenvalues.

---

BSVD
Figure 2: Single Gaussian. We show samples from SVGD, KDE-WGD, SGE-WGD, and SSGE- WGD (from left to right). The upper and right plots show the empirical mne-dimensional marginal distributions obtained using KDE on the samples (red) and the particles (blue).
3.3 Gradient flow in function space
To theoretically justify the update rule introduced in function space in Eq. ( we rewrite the Liouville equation for the gradient flow in Eq. (TD in function space as
Op(f) V. 0t
(p(FIV - V DKL(em)
(18)
=V (D(S)V log T (f))
Following this update, the mean field functional dynamics are
df V log p(f) - log (f))
(19)
di
Using the same KDE approximation above, obtain discretized evolution in function space and with it an explicit form for the repulsive force in Eq. as
Σ ( f)
= = f: + ce v (F))
(20)
The update rules using the SGE and SSGE approximations follow as for the parametric case. It im- portant notice that this update rule requires the function space prior gradient: V,, log p(ff æ, y)
V log az, fÿ) + Vr, log p(fj). If one wants to use an implicit prior defined in weight space, an additional estimator needed due to its analytical intractability. We again adopted the SSGE, introduced by Shi et al. [69], which was already used for similar purpose in Sun et al. [7I]. It is also interesting to note that the update rule in Eq. 220 readily allows for use of alternative priors that have analytical form, such Gaussian processes. This is an important feature of our method that allows an explicit encoding of function space properties that can be useful for example in achieving better out of distribution detection capabilities [14]
3.4 The choice of the kernel
repulsive effect can always be created in the ensemble by means of the gradient of any kernel function that measuring the similarity between two members. Nevertheless, is important to keep mind that 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 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. [ has shown how the Manhattan distance metric (L1 norm) might be preferable over the Euclidean distance metric (L.2 norm) for high-dimensional settings. We performed some additional ex experiments using the L.1 norm (Laplace kernel) for the KDE but we could not observe any substantial difference compared to using the 1.2 norm. Further investigations regarding this hypothesis are left for future research.
7

---

Deep Ensemble
w-SVGD
kde-WGD
sSqe-WGD
sge-WGD
1.0
0.8
0.6
HMC
f-SVGD
kde-IWGD
ssge-IWGD
sge-WWGD
0.4
0.2
0.0
Figure 3: BNN 2D classification. We show the entropy the predictive posteriors. Again, the function-space methods capture the uncertainty better than the weight-space ones, thus approaching the gold-standard HMC posterior.
Experiments
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 (OOD) data,
expected calibration enror (ECE) [38. To assert the robustness we report the ratio between predictive entropy on OOD and test data points (Ho/ H,), and the OOD 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 [). In partieular,
out-of-distribunt
detection area under the ROC curve AUROC(H) 1
report the ratio between the average model disagreement on the OOD and test data points (M D. MM) 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 two-dimensional Gaussian distribution (Figure 2 We see that 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. in the Appendix. There, SGE-WGD and SVGD also perform best.
BNN regression We then assessed the different methods in fitting BNN posterior on synthetically generated one-dimensional regression task. The results are reported in Figure I consisting the mean prediction and +1, 2, standard deviations 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, functional methods (f-SVGD, IWGD) are perfectly able to infer the diversity of the hypotheses in this region due to the lack 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 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

---

Table 1: BNN image classification. (UROC(H) is the AUROC computed using the entropy whereas AUROC(MD) is computed using the model disagreement. Ho/ He is the ratio of the entropies on OOD and test points respectively and MDDMMDD is the ratio for model disagreement We see that the best accuracy is achieved by our WGD methods, while our IWGD methods yield the best OOD detection and funtional diversity. All our proposed methods improve over standard deep ensembles in terms of accuracy and diversity, highlighting the effect of our repulsion.
AUROC(H)
AUROC(MD)
Accuracy
H,/Hz
MD,/MD,
ECE
NLL
Deep ensemble
0.00810.001
0.97510.001
91.1221100013
6,257.10.005
6.394.10.011
0.012.10.011
0.129_.0.00
SVGD 151
0.960-00.01
0.973=0.001
91.13420.024
6.315:20.019
6.39520.018
0.01410.001
0.127=0.001
F-SVGD 776
0.956=0.001
0.975=0.001
89,884=0.015
5.652=00009
6.531+0.005
0.013+0.001
0.150=0.001
FachionMNNSS
hyper-DE [S]]
0.968200001
0.901
91.160=0.007
6.682.10.065
7.05910.152
0.0141:0001
0.128-0.001
GD (ours)
0.960=0.001
0,970=00000
91.238=0.0119
6.587=0.019
6.379=0.018
0.014=00001
0.128=0.001
sge- WGD (ours)
0.960=00:01
0.970=00001
91.312=0.016
6.562:40.007
6.363+0.000
0.012=00.011
0.128-0.001
sage- BGg (ours)
0.000
0.979:00.00
0.0024
6.522
6.610:00012
0.012:00001
0.130-0.000
kde-IWGD (ours)
0.071100001
0.980-10.011
91.260.10.011
7.07910.001
6.88710.015
0.015110.001
sgr-PWGD (ours)
0.969-00.00
0.978-0.001
91.192+0.003
7.076++0.004
6.900+-0.005
0.015++0.001
0.125 0.001
sgg-IWGD (ours)
0.901
0.980
91.240.10.022
0.029
6.95110.005
0.016.10.001
Deep ensemble +2]
0.843:00.004
0.736=0.000
85.552=#0.076
2.244::0006
1.667=00000
0.049=0.001
0.270000011
SVGD 5D
0.825+0.001
0.710-0.002
85.142-+0.017
2. 2.006-00000
15567-0.000
0.052+0.001
0
ISVGD 76
0.783=00001
0.712=0.001
84.510=00.081
1,96==0.004
1.62400000 003
0.049=0.001
0.292-0.001
CIFARIO
hyper-DE [00]
0.789-00000
0.743-0.001
84.742-00011
1.96
6690+000055
0.046+00001
0.288-0.001
kde-WGD (ours)
0.838100.011
0.7351...004
85.904 0.000
2.205
1.60100.008
0.05311..001
0.276 0,001
sge-WGD (ours)
0.837:20.003
0.725=00.0M4
85.792-20.035
2.214420010
1634410.004
0.05110.001
0.275_0.001
sage- Dgge-WGDoo (ours)
0.832=0.003
0.731=0.005
85.638=00.388
2.182=0.015
1.655=0.001
0.049=0.001
0.276=0.001
kde-IWGD (ours)
0.791:10.002
0.708
84.888-10.030
1.97010.004
1.749:0.00S
0.044:0.001
0.282
sge-IWGD (ours)
0.795=00001
0.754=0.002
04.766=200000
1.984=0.003
1.729=0.002
0.047=0.001
0.288=0.001
sgg-fWGD (ours)
0.792::0.02
0.752=0.002
84.762=00034
1.970=0.006
1.723+0.005
0.046+00001
0.286-0.001
region, for which the uncertainty is purely aleatoric. In this setting, f-SVGD suffers from similar issues 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 principled OOD detection.
FashionMNIST classification Moving on to real-world data, we used an image classification setting using the FashionMNIST dataset [+ for training and the MNIST dataset [ as out-of- distribution (OOD) task. The results are reported in Table (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 BGD 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 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 has been used in continual learning KT Interestingly, using this metric, the hyper-deep ensemble [0 shows OOD detection performance comparable with our repulsive ensemble in function space.
CIFAR classification Finally, we use a ResNet32 architecture 20 on CIFAR-10 41 with the SVHN dataset [6 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.
5 Related Work
The theoretical and empirical properties of SVGD have been well studied 40 also be seen as a Wasserstein gradient flow of the KL divergence in the Stein geometry [6
48 K
and it can
50
9

---

(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 [0] 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
A8 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
39
have gained
7 19 On the other hand, ensemble methods have also been extensively Moreover, repulsive interactions between the members have A Moreover, providing Bayesian interpretations for deep ensembles has been previously attempted through the lenses of stationary SGD distributions [56 &
popularity recently [79
[22
36), using modern MCMC [59
79 2
24!
21 and variational
inference techniques [
72
studied 42
J
X2
P
also been studied in Wabartha et al.
ensembles of linear models [57], additional random functions 62 12 27
approximate inference Stein variational inference [5], and marginal likelihood lower bounds 153] and ensembles
{2
have also been shown to provide good approximations to the true BNN posterior some settings Furthermore, variational inference in function space has recently gained attention M and the
66
limitations of the KL divergence have been studied in Burt et al.
Conclusion
We have presented 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 but-even more importantly 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 real-world data and can approach the true Bayesian posterior more closely.
be seen
In future work, it will be interesting to study the impact of the Jacobian in the IWGD update and its implications on the Liouville equation more detail, also compared other neural network Jacobian methods, such neural tangent kernels 37 and generalized Gauss-Newton approximations [85) Moreover, it would be interesting derive explicit convergence bounds for our proposed method and compare them to the existing bounds for SVGD 400
Acknowledgments
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 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.
References
[1] Charu C Aggarwal, Alexander Hinneburg. and Daniel A Keim. On the surprising behavior of distance metrics in high dimensional space. In International conference on database theory, pages 420-434. Springer, 2001.
[2] Luigi Ambrosio and Gianluca Crippa. Continuity equations and ode flows with non-smooth velocity. Proceedings of the Royal Society of Edinburgh: Section A Mathematics, 144(6): 1191-1244, 2014.
[3] Luigi Ambrosio. Nicola Gigli, and Giuseppe Savaré. Gradient flows: the space of probability measures. Springer Science & Business Media, 2008.
metric spaces and in
[4] Vijay Badrinarayanan, Bamdev Mishra, and Roberto Cipolla. Symmetry-invart optimization in deep networks. arXiv preprint arXiv: 2015.
[5] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. arXiv preprint 1055:11555.0542,,
10

---

[6] Leo Breiman. Bagging predictors. Machine learning. 24(2):123-144) 1996.
[7] David R Burt, Sebastian W Ober, Adrià Garriga-Alonso, and Mark van der Wilk. Understanding variational inference in function-space. arXiv preprint arXiv:2011.0942) 2020.
[8] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10, IEEE, 2018.
[9] An Mei Chen, Haw-minn Lu, and Robert Hecht-Nielsen. On the geometry of feedforward neural network error surfaces. Neural computation, 5(6):910-922) 1993.
[10] Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particle- optimization framework for scalable bayesian sampling. arXiv preprint arXiv:1803 /159,, 2018.
[11] Jiefeng Chen, Xi Wu, Yingyu Liang. Somesh Jha, et al. Robust out-of-distribution detection in neural networks. arXiv preprint Stti:22003.097111 2020.
[12] Kamil Ciosek, Vincent Fortuin, Ryota Tomioka, Katja Hofmann, and Richard Turner. Conserva- tive uncertainty estimation by fitting prior networks. In International Conference on Learning Representations, 2019.
[13] Francesco D' Angelo and Vincent Fortuin. Annealed stein variational gradient descent. arXiv preprint arXiv:2101.09815 2021.
[14] Francesco D' Angelo and Christian Henning. Uncertainty-based -utto--distribution detection requires suitable function space priors. arXiv preprint arXiv:2110.0602 2021.
[15] Francesco D Angelo, Vincent Fortuin, and Florian Wenzel. On stein variational neural network ensembles. arXiv preprint arXiv:2106. /0760, 2021.
[16] A Duncan, Nikolas Nuesken, and Lukasz Szpruch. On the geometry of stein variational gradient descent. arXiv preprint arXiv:1912.00894, 2019.
[17] Michael W Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-an Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 1 factors. arXiv preprint wrXii:2005.07180 2020.
[18] Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape perspective. arXiv preprinl arXiv: 192.0077,, 2019.
[19] Vincent Fortuin. Priors in bayesian deep learning: review. arXiv preprint arXiv:2105.066: 2021.
[20] Vincent Fortuin, Mark Collier, Florian Wenzel, James Allingham, Jeremiah Liu, Dustin Tran, Balaji Lakshminarayanan, Jesse Berent, Rodolphe Jenatton, and Effrosyni Kokiopoulou. Deep classifiers with label noise modeling and distance awareness. arXiv preprint arXiv:2110.02609 2021.
[21] Vincent Fortuin, Adrià Garriga-Alonso, Mark van der Wilk, and Laurence Aitchison. Bnnpriors: library for bayesian neural network inference with different prior distributions. Software Impacts, page 100079, 2021.
[22] Vincent Fortuin, Adrià Garriga-Alonse Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. arXiv preprint 20X 2021.
[23] Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. arXiv preprint arXiv:1802.10026 2018.
[24] Adrià Garriga-Alonso and Vincent Fortuin. Exact langevin dynamics with stochastic gradients. arXiv preprint drXth:2102.01691, 2021.

---

[25] Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d' A'scolii Giulio Biroli, Clément Hongler. and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2020()):033401,, 2020.
[26] L. K. Hansen and Salamon. Neural network ensembles. IEEE Trans. Pattern Anal. Mach. Intell., 1010)):9331001, October 1990. ISSN 0162-8828. doi: 10.1109/34.8871. URL attps://doo
111////...1//....11 58871
[27] Bobby He, Balaji Lakshninarayanan, and Yee Whye Teh. Bayesian deep ensembles via the neural tangent kernel. arXiv preprint rrXiv:2007.05864, 2020.
[28] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- performance on imagenet classification, In Proceedings of the IEEE international conference on computer vision, pages 1026-1033 2015.
[29] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[30] Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pages 129-135. Elsevier, 1990.
[31] Christian Henning, Maria R Cervera, Francesco Angelo, Johannes von Oswald, Regina Traber, Benjamin Ehret, Seijin Kobayashi, João Sacramento, and Benjamin F Grewe. Posterior meta-replay for continual learning. arXiv preprint arXiv:2103.01133. 2021.
[32] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E Hopcroft, and Kilian Q Weinberger. Snapshot ensembles: Train 1, get m for free. arXiv preprint arXivv:704..000 00009, 2017.
[33] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E. Hopcroft, and Kilian Q. Weinberger. Snapshot ensembles: Train 1, get M for free. CoRR, abs/1104.00119, 2017. URL http:
//arxiv.org/abart www.////////////// 00109
[34] Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection in deep learning. arXiv preprint arXiv:2104.0497. 2021.
[35] Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics, pages 703-711, PMLR, 2021.
[36] Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Wilson, What are bayesian neural network posteriors really like? arXiv preprint arXiv:2104 2021.
[37] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 8580-8589, 2018
[38] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
[39] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[40] Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A non-asymptotic analysis for stein variational gradient descent. Advances in Neural Information Processing Systems, 33, 2020.
[41] Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
[42] Balaji Lakshninarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances neural information processing systems, pages 6402-6413, 2017. 2017.
12

---

[43] Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/nmist// 1998.
[44] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:/711.00165. 2017.
[45] E. Levin, N. Tishby, and S. A. Solla. A statistical approach to learning and generalization in layered neural networks. Proc. of the IEEE Special issue Neural Networks, 1990.
[46] Yingzhen Li and Richard E Turner. Gradient estimators for implicit models. arXiv preprint arXiv:1705.07107, 2017.
[47] Shiyu Liang, Yixuan Li, and Rayadurgam Srikant. Enhancing the reliability of out-of- distribution image detection in neural networks. arXiv preprint arXiv: (tf.:::70602690) 2017.
[48] Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning, pages 4082-4092. PMLR, 2019.
[49] Jeremiah Zhe Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax- Bedrax-Weiss,, and Balaji Laksh- minarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. arXiv preprint arXiv:2006. 2020.
[50] Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in neural information processing systems, pages 31 3115-3123, 2017.
[51] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. In Advances in neural information processing systems, pages 2378-2388 2016.
[52] Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-ie tests. International conference on machine learning, pages 276-284. PMLR, 2016.
[53] Clare Lyle, Lisa Schut, Robin Ru, Yarin Gal, and Mark van der Wilk. bayesian perspective on training speed and model selection. Advances in Neural Information Processing Systems, 33, 2020.
[54] David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computation, 4(3):448-477 1992.
[55] Andrey Malinin, Bruno Mlodozenice, and Mark Gales. Ensemble distribution distillation. arXiv preprint arXivv1905.00070 2019.
[56] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. arXiv preprint 1eXiv:1704.04:: 04289, 2017.
[57] Alexander G de G Matthews, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Sample- then-optimize posterior sampling for bayesian linear models. In NeurIPS Workshop on Advances in Approximate Bayesian Inference, 2017.
[58] Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
[59] Radford M Neal. Bayesian learning for neural networks. 1995.
[60] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science Business Media, 2012.
[61] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
[62] Ian Oshand, Benjamin Van Roy, Daniel Russo, Zheng Wen, et Deep exploration via randomized value functions. J. Mach. Learn. Res., -06, 2019.
13

---

[63] Johannes Von Oswald, Seijin Kobayashi, Joao Sacramento, Alexander Meulemans, Christian Henning, and Benjamin F Grewe. Neural networks with late-phase weights. In International Conference on Learning Representations, 2021. URL https::/ watpaarrvviieeeeefruriii. id=CDqJUxxdxFb
[64] Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. 2001.
[65] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, pages 13991-1 2019.
[66] Emanuel Parzen. On estimation of probability density function and mode. The annals of mathematical statistics, 336)),00651107, 1962.
[67] Alexandre Rame and Matthieu Cord. Dice: Diversity in deep ensembles via conditional redundancy adversarial estimation. arXiv preprint arX1v:2101.05544, 2021.
[68] Geoffrey Roeder, Luke Metz, and Diederik P Kingma. On linear identifiability of learned representations. arXiv preprint :XXx:200770088000 2020.
[69] Jiaxin Shi, Shengyang Sun, and Jun Zhu. A spectral approach to gradient estimation for implicit distributions. In International Conference on Machine Learning, pages 4644-4653, PMLR, 2018.
[70] Radhey S Singh. Improvement on some known nonparametric uniformly consistent estimators of derivatives of a density. The Annals of Statistics, pages 594-399, 1977.
[71] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. arXiv preprint arXiv:1903 05779, 2019.
[72] Jakub Swiatkowski, Kevin Roth, Bastiaan S Veeling, Linh Tran, Joshua V Dillon, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. The k-tied normal distribution: A compact parameterization of gaussian mean field posteriors in bayesian neural networks. arXiv preprint arXiv:2002. 02655, 2020.
[73] Johannes von Oswald, Christian Hemning, João Sacramento, and Benjamin F. Grewe. Continual learning with hypernetworks. In International Conference Learning Representations, 2020.
URL. httpthttss///x/iiie http://aaxxvv.../bbbbaa www.g//ab///.....//// 00696
[74] Maxime Wabartha, Audrey Durand, Vincent François-Lavett and Joelle Pineau. Handling black swan events in deep learning with diversely extrapolated neural networks. In IJCAI, pages 2140-2147,22 2020.
[75] Yifei Wang, Peng Chen, and Wuchen Li. Projected wasserstein gradient descent for high- dimensional bayesian inference. arXiv preprint arXiv:2102.000 2021.
[76] Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang, Function space particle optimization for bayesian neural networks. arXiv preprint arXiv: / 1902.0974, 2019.
[77] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pages 681-688, 2011.
[78] Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: alternative approach to efficient ensemble and lifelong learning. arXiv preprint arXiv:2002.06715. 2020.
[79] Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Šwiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the bayes posterior in deep neural networks really? arXiv preprint arXiv:2002.02405, 2020.
[80] Florian Wenzel, Jasper Snoek, Dustin Tran, and Rodolphe Jenatton. Hyperparameter ensembles for robustness and uncertainty quantification. In Advances in Neural Information Processing Systems, 2020.
14

---

[81] Christopher KI Williams. Computation with infinite neural networks. Neural Computation, 10 (5):1203-1216, 1998.
[82] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. arXiv preprint arXiv:2002.0879) 2020.
[83] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 2017.
[84] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03553 2016.
[85] Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. arXiv preprint arXiv:1902.03932. 2019.
15


Reviews:
[{"review_id": "zk2jwu5sN7d", "reviewer": "Reviewer_jr1C", "summary": "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.\n\nThe 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.\n\nThe 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Quality:\n\nThe work is good quality. The theoretical results are convincing and the results are interesting and relevant to researchers working on ensembles. It interprets repulsive ensembles as an approximate particle based inference method by showcasing that the proposed update rules correspond to the gradient flow of the KL-divergence in Wasserstein space. This is very interesting as it allows practitioners to use these particle-based methods with the confidence that they approximate the Bayesian posterior.\n\nThe theoretical results are well presented in an organized way. Related works are clearly explained and contrasted.\n\nA shortcoming is that the experiments section has almost no detail on the implementation. Basic information, such as the number of particles used, is missing. I was able to locate some of these in the supplementary material, but I believe they should be in the main text. It would not be possible to replicate these results based on the level of detail given here.\n\nThe baseline methods slightly underperform. A ResNet 32 should achieve higher than 85% accuracy on Cifar 10.\n\n* From the experiments in Figure 3, it seems that the weight-space methods are all identical to Deep Ensembles. Is this the case because of the high dimensional weight space? That the distances of the particles are too large to be able to affect each other?\n\n### Novelty:\n\nThe paper heavily builds on the results of [1], but formulates the update rules slightly differently. To my knowledge, the theoretical results are also novel, although it is possible that I am unfamiliar with related works.\n\n### Presentation:\n\nThe motivation and the narrative are presented very well, but I found parts of the paper difficult to understand. Here are some parts that are unclear to me:\n\n* What form does $k(f_1, f_2)$ take in practice? The RBF kernel is mentioned as an example, but is this the one used in the experiments? Is it computed over a batch of training points? How are the kernel hyperparameters determined?\n* How is $\\log \\pi(f)$ computed? Given a weight-space prior, the prior density is difficult to estimate due to the non-identifiability problem.\n* What are the advantages and drawbacks of SGE and SSGE? Does one give a more accurate approximation? Are there computational tradeoffs?\n\n[1] Wang, Ziyu, et al. \"Function space particle optimization for Bayesian neural networks.”\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They are addressed adequately.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "ud9WHvBjbP", "reviewer": "Reviewer_evs4", "summary": "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.\n\nThe claims are that the proposed approach:\n* enables using deep ensembles to perform proper Bayesian inference,\n* improves the quality of the uncertainty estimates in synthetic tasks,\n* improves OOD detection on image datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Significance*: Deep ensemble methods enable improving prediction accuracy as well as producing better-calibrated uncertainty estimates. The paper proposes introducing a kernel-based repulsive term that improves ensemble performance as well as improves the theoretical motivation for the ensemble, therefore I think the paper is relevant and will be impactful for the community.\n\n*Experiments*:\n* The synthetic experiments illustrate how the proposed approaches operating in function space can produce better uncertainty estimates than standard ensembles and weight-space methods. Furthermore, the 2D classification experiment also illustrates how the proposed approaches in function space can produce better estimates than f-SVGD in some OOD regions. The synthetic experiments do not generally show significant differences between the three repulsion terms, though SGE-WGD appears to fit the target better in Figure 2.\n* The image experiments show that the proposed approaches generally outperform the standard ensembles, although the ranking is not consistent across evaluation tasks/metrics.\n* Overall, the empirical evaluation is not very thorough but it supports the claim that the repulsive terms are beneficial. I think that the section could be improved by 1) more analysis or discussion on the potential trade-offs in using each repulsive term so as to provide guidelines for future use, 2) highlighting how the proposed approach does not plateau as discussed in the abstract.\n\n*Clarity*:\n* The paper is generally well written. That said, the plots could use a bit more information, e.g. the colors for the densities at the top and right of each plot in Figure 2 are not described, and the captions do not highlight the take-aways.\n* Minor: \\beta appears only in line 196 and is not described, seems to be equal to the inverse of the normalizing term in eq (16).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper proposes a method whose end goal is to improve accuracy and uncertainty prediction across a wide range of machine learning prediction problems. From this perspective, the general benefits and potential negative societal impact of these problems, such as reinforcing certain biases, could be applied here.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4-6"}, "extra_text": {}}, {"review_id": "n78KXNJHaSo", "reviewer": "Reviewer_QY5N", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I like the paper. It addresses an important and still unsolved problem of maintaining functional diversity among ensemble members.  I don't have much to critique. I think this is a worthwhile contribution and that it should get published.\n\nQUESTIONS AND SUGGESTIONS:\n\n1) Do we need to train the full ensemble simultaneously? It seems to me that we do. What are the requirements for how simultaneous exactly does the training have to be in order for this method to work?\n\n2) The models you are looking at on CIFAR-10 in Table 1 are surprisingly week, reaching only mid 80s % of test accuracy. It might be a good idea to show that your technique actually works on SOTA-like models. I can imagine that the effect you observe could be a consequences of your models not being trained to their full potential. A CIFAR-100 verification might also be worth adding, given that CIFAR-10 is a relatively simple and small task.\n\nUPDATE AFTER REBUTTAL:\nI am happy with the authors' response and will therefore increase my score from 6 to 7.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "6FOVFjdNGO", "reviewer": "Reviewer_M3i3", "summary": "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. ", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I found this to be an interesting and thought-provoking paper which I am glad to have had the chance to review. I believe that this is an important area of research and that this paper could lead to interesting future work (or be otherwise of significance). As far as I can tell the proposed method is sound and well-motivated from multiple perspectives. \n\nHowever, while the experimental evaluation is decent, I think there is definitely room for improvement, which I will describe in more detail below. Furthermore, I think that the biggest weakness of the paper is in clarity. I found section 3, and to a lesser extent section 2, difficult to follow. I'll elaborate on the clarity issues below. I also think that the authors should add a discussion of the limitations of their method. For example, it seems to me that one potential issue in this work is that it requires that all of the ensemble members fit in GPU memory in order to do the training efficiently, but perhaps this isn't the case? Similarly, it also seems that the choice of kernel function must matter a lot for the performance of the method, but this isn't discussed. In particular, I am wondering whether a poor choice might lead to worse performance if the repulsive term dominates the training objective (i.e. how much tuning was required to get training to work)?\n\n### Experimental suggestions\n\nThe major weakness of the experimental evaluation was the omission of hyper-deep ensembles as a baseline in at least the MNIST and CIFAR10 experiments. The authors do cite this method, but I think that, as a simple to implement and highly performant baseline which also relies on the increased diversity of ensemble members, this is a must-have comparison. \n\nThe other missing ingredient in the experimental sections were some studies to investigate/confirm why/how the proposed method works. For example, it would be interesting to confirm that flipping the sign of the repulsive term does degrade performance as one would expect. It would also be very interesting to see how the LL and repulsive terms evolve during training for a few different datasets. For example, one might imagine that the repulsive term is larger relative to the LL term as training progresses because the ensemble elements might converge to similar functions. \n\n\n\n#### Minor suggestions \n\n1. In the toy regression plots (fig 1), please show some of the individual functions so that one case see how different the fits are for each ensemble member. \n\n2. I would also like to see the same toy regression experiments but with a larger more flexible network i.e. a width of 100 rather than 50 and a depth of 3 instead of 2. In my experience standard deep ensembles of this size are able to capture \"in-between\" uncertainty, so it would be interesting to see what the repulsive term does in this case. \n\n3. In fig 2, it would be great to have a colour bar to give more context for the uncertainty values. It would also be great to have a larger range so that the behaviour as we ger further away from the training date can be understood better. In particular, it seems like the function space methods will behave nicely but it isn't quite clear given the current range of the plots. \n\nFor the CIFAR10 experiments, it would be interesting to see results using CIFAR100 as the OOD dataset. This is much more challenging than SVHN and might be a case where the functional methods can shine given their potential \"distance-aware\" behaviour.\n\n4. I would also like to see rejection-classification experiments (see \"Uncertainty Estimation Using a Single Deep Deterministic Neural Network\" by van Amersfoort et al. for examples).\n\n5. I think table 1 should be split into two. One which compares the various methods for WGD, and one which compares the best WGD method for each dataset to the baselines. Comparing all of the WGD methods to the baselines is a slightly unfair comparison.\n\n### Clarity Issues\n\nUnfortunately, I can't give as many concrete suggestions for how to improve sections 2 and 3. However, at a high level, I would say that while the technical details provided are important, perhaps some of the detail should be moved to the appendix. The space in the main text could be used to help the reader by providing further explanation, discussion, and intuition building. \n\nHere is a list of specific (but individually minor) difficulties I had when readings these sections, I hope they provide some insight into why I struggled:\n\n1. The symbols $\\mathcal{R}$,  $\\mathcal{F}$, and $\\mathcal{X}$, were never explicitly introduced and their meaning had to be inferred. \n\n2. In section 3.2, there is a switch from Euclidian space to a Riemannian manifold but it isn't explained why this is done. Some intuition would be very useful.\n\n3. The metric tensor in 3.2 needs more explanation. What is this and why has it been introduced. Some intuition would also be great here.\n\n4. In the equation after line 173 it isn't clear if $rho$ is the same as in Eq. 8. \n\n5. The definition of the Wasserstein metric is confusing and ultimately unnecessary since it is the same as one would find from a quick Google search. I think it should be removed (or moved to an appendix) to improve the flow of the section.\n\n6. On line 176 it isn't clear why we have an \"equivalent of the gradient\" and not simply a gradient. \n\n7. Equation 12 seems to come out of nowhere. Some intuition would be great.\n\n8. Comparing equations 14 and 3 is difficult when they are on different pages, so it would be helpful if the text explicitly compared the repulsive terms. Also, the sentence from line 186 to 189 is too long and should be split into two.\n\n9. Eqn 16 would be more clear if the substitution of the posterior was already made.\n\n10. Is $\\beta$ on line 196 simply the denominator of the fraction in Eqn 16?\n\n11. The fact that the true Bayesian posterior is only recovered in the limit as $n \\to \\infty$ should be mentioned much earlier in the text. \n\n\n### Summary\n\nWhile I enjoyed this paper and hope to see it published, I am unable to give a higher score due to the clarity and experimental issues described above. However, if the authors can make section 3 more approachable, and address my two main experimental concerns, I would be happy to increase my score from 5 to 7 or 8. \n\n\n### Post-author response\n\nThe authors have addressed the majority of my concerns, have provided additional experimental results (and committed to adding further results in the camera ready), and have also acknowledge and proposed solutions to the issues of clarity. Thus, I am raising my score to 8.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have not gone into much detail about the limitations of their method. As mentioned above it would be great to have a discussion of:\n1. Memory requirements for training repulsive ensembles.\n2. The importance of the choice of kernel and the general importance of hyperparameter tuning.\n", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "9"}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper the authors proposed the addition of a repulsive term to the training objective of deep ensembles, based on Stein variational inference.  They show that this turns the standard deep ensemble training into an approximation to the Bayesian posterior.  The reviewers found the proposed method compelling and found the paper to be of high writing and technical quality.  They also found the work timely and that it addresses an important open problem.  Thus the recommendation is to accept.

Decision: Accept (Spotlight)