Repulsive Deep Ensembles are Bayesian

1 Introduction . 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 . 2 Repulsive Deep Ensembles . 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 . 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 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 . 2.2 Repulsive force in function space . 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 )

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