Greedy Bayesian Posterior Approximation with Deep Ensembles

1 INTRODUCTION . Estimation of predictive uncertainty is one of the most important challenges to solve in Deep Learning ( DL ) . Applications in finance , medicine and self-driving cars are examples where reliable uncertainty estimation may help to avoid substantial financial losses , improve patient outcomes , or prevent fatal accidents ( Gal , 2016 ) . However , to date , despite rapid progress , there is a lack of principled methods that reliably estimate the predictive uncertainty of deep neural networks ( DNNs ) . Numerous attempts have been made to develop Bayesian techniques for uncertainty estimation in DL ( Gal & Ghahramani , 2016 ; Lakshminarayanan et al. , 2017 ; Wilson & Izmailov , 2020 ; Maddox et al. , 2019 ; Izmailov et al. , 2020 ; Van Amersfoort et al. , 2020 ; Wenzel et al. , 2020b ) . One of the most practical and empirically best-performing approaches is based on training a series of independent DNNs ( Lakshminarayanan et al. , 2017 ; Wilson & Izmailov , 2020 ; Ashukha et al. , 2020 ; Lu et al. , 2020 ; Wenzel et al. , 2020b ) . The main method in this category , Deep Ensembles ( DE ) ( Lakshminarayanan et al. , 2017 ) , is used as a reference approach in the context of this paper . Recent studies , e.g . Wilson & Izmailov ( 2020 ) interpret ensembles as an approximation of predictive posterior . While this interpretation is correct from a Bayesian point of view , obtaining individual ensemble members via maximum a posteriori probability ( MAP ) estimation , as e.g . done in DE ( Lakshminarayanan et al. , 2017 ) , may not lead to obtaining good coverage of the full support of the posterior distribution , and has arbitrary bad approximation guarantees . For example , the resulting approximation can be poor in the case when the true posterior distribution is unimodal . In this work , we argue that enforcing coverage of posterior ( i.e . ensemble diversity ) is non-trivial , and needs a specialized principled approach . We highlight this graphically in Figure 1 . Another important line of work in modern Bayesian DL ( BDL ) is a paradigm of performing Bayesian inference in the weight space . While distributions over weights induce distributions over functions ( Wilson & Izmailov , 2020 ) , it is rather unclear what the properties of such functions are , and whether Bayesian posteriors obtained in the weight space yield good quality approximations in the function space . For example , it is known that diverse weights do not necessary yield diverse functions , thus sampling from weight-based posteriors may yield poor quality uncertainty estimation , e.g . in detecting out-of-distribution ( OOD ) data ( Garipov et al. , 2018 ; Hafner et al. , 2020 ) . Recent studies ( Wang et al. , 2019 ; D ’ Angelo & Fortuin , 2021 ) show the promise of particle variational inference ( POVI ) done in the function space , however , the performance of those methods is not state-of-the-art due to the use of BNN priors . In practice , a weight decay prior combined with batch normalization , data augmentation and dropout are often employed due to empirical improvement , and theoretical guarantees are traded for practical performance . Furthermore , particle-based function space VI requires training an ensemble of BNNs simultaneously , which is not only difficult to implement , but also requires extensive resources for parallelization . Summary of the contributions . In this paper , we propose a novel and principled methodology for approximate function space posterior inference for DNNs . Contrary to the mainstream approach in BDL , which is based on defining a posterior distribution over the model parameters ( Maddox et al. , 2019 ) , we take a functional view , which allows us to treat the problem of training ensembles as optimization over sets . Namely , selecting a set of functions to approximate the Bayesian posterior is fundamentally similar to problems such as facility location and set cover , which are frequently solved using submodular optimization ( Krause et al. , 2008 ) . Specifically , our contributions are : 1 . We show that fitting a kernel density estimator to a distribution using an f -divergence is a cardinality-fixed non-monotone submodular maximization problem . 2 . Inspired by the Random Greedy algorithm for submodular maximization ( Buchbinder et al. , 2014 ) , we design a new method for the function space Bayesian posterior approximation via greedy training of ensembles with a justified coverage-promoting diversity term . 3 . We demonstrate the effectiveness and competitiveness of our approach compared to DE in the OOD detection task on MNIST , CIFAR , and SVHN , benchmarks on different architectures and ensemble sizes . Furthermore , we show that our method can use state-of-the-art training technqiues , compared to the existing Bayesian approaches , and yeilds state-of-theart performance . 2 PRELIMINARIES . 2.1 PROBLEM STATEMENT . Consider an ensemble to be parameterized by a set of functions Z = { zm } Mm=1 ⊂ F , where F is a class of continuous functions , zm : Rd → Rc , with d the dimensionality of the input data , and c the dimensionality of the output . When training ensembles , we generally want to solve the following optimization problem : min Z , |Z|=M R ( Z ) − ΩλM ( Z ) , ( 1 ) whereR ( Z ) = 1N ∑N i=1 ` ( 1 M ∑M m=1 zm ( xi ) , yi ) is the empirical risk of the ensemble , ` : Y×Y → R+ is a loss function , D = { xi , yi } Ni=1 is a training dataset of size N , and ΩλM ( Z ) is some diversitypromoting term , with diversity regularization strength λM . Empirical observations in the earlier works on ensembles ( Lakshminarayanan et al. , 2017 ; Wilson & Izmailov , 2020 ; Fort et al. , 2019 ) have shown that one can simply ignore ΩλM ( Z ) , and rely on non-convexity of the loss landscape , optimizing the risks of individual ensemble members . From a variational inference ( VI ) perspective ( Zhang et al. , 2019 ) , this can be seen as a mode-seeking method , and it has been shown experimentally that every ensemble member may discover different modes of the posterior distribution in the function space p ( z|D ) . Approximation of p ( z|D ) is the ultimate aim of this paper , and we argue that randomization-based mode-seeking is insufficient to obtain a good quality approximation of p ( z|D ) , as this procedure does not maximize the coverage of the support of the posterior . By contrast , we aim to find an ΩλM ( Z ) such that the posterior coverage is also maximized . Taking a VI perspective again ( Zhang et al. , 2019 ) , ΩλM ( Z ) needs to enforce that minZ , |Z|=M R ( Z ) − ΩλM ( Z ) has rather mean-seeking behavior , while still discovering the high density modes . 2.2 A COMBINATORIAL VIEW OF ENSEMBLE CONSTRUCTION . Having now defined the main criteria for ( 1 ) , we highlight that the problem of constructing an ensemble can be seen from a combinatorial point of view . We therefore treat ensemble construction as subset selection from some ground set of functions , and introduce the main notions of submodular analysis , a powerful tool that enables the analysis of the optimization of set functions ( Bach , 2013 ; Fujishige , 2005 ) . Definition 1 ( Submodularity ) . A set function g : 2V → R , for the power set of a base set V , is submodular if for all A ⊆ B ⊂ V and x ∈ V \B g ( A ∪ { x } ) − g ( A ) ≥ g ( B ∪ { x } ) − g ( B ) . ( 2 ) Definition 2 ( Supermodularity and modularity ) . A set function is called supermodular if its negative is submodular , and modular if it is both submodular and supermodular . Consider now problem ( 1 ) . Assuming that the loss function ` is convex , we can derive an upper-bound on the riskR ( Z ) using Jensen ’ s inequality , and obtain a method , which generalizes DE min Z , |Z|=M 1 M M∑ m=1 R ( zm ) − ΩλM ( Z ) . ( 3 ) If M is fixed during optimization , 1M ∑M m=1R ( zm ) contributes a positive modular term to the overall objective . Adding a positive modular function to any set function does not change its submodularity or supermodularity , thus we focus on ΩλM ( Z ) . A trivial approach would be to enforce pair-wise diversity by computing a norm of the pairwise differences between functions , i.e . setting ΩλM ( Z ) = λM ∑ i6=j ‖zi − zj‖2∗ . However , this is a cardinality-fixed submodular minimization problem . It is known that it is strongly NP-hard , i.e . there exists no general polynomial time approximation algorithm for it ( Svitkina & Fleischer , 2011 ) . The poor quality approximation of this approach is highlighted in Figure 1 . We therefore conclude that the choice of ΩλM ( Z ) has a direct impact on the approximabilty of the objective . 3 SUBMODULAR ANALYSIS OF f -DIVERGENCES 3.1 f -DIVERGENCES ARE SUPERMODULAR FUNCTIONS Main result . We now consider the problem of approximating a Bayesian posterior via minimization of an f -divergence . Here , we specifically aim our optimization procedure to have both mode and mean-seeking behaviors , i.e . cover the posterior distribution as much as possible , ending up in its mode . We furthermore aim to obtain a polynomial time algorithm that yields a good quality approximation guarantee . In this paper , we leverage classic definitions of approximation algorithm and approximation guarantees . Definition 3 ( Approximation algorithm and guarantees ( Williamson & Shmoys , 2011 ) ) . A γapproximation algorithm for an optimization problem is a polynomial-time algorithm that for all instances of the problem produces a solution whose value is within a factor γ of the optimal solution . γ in this case is called approximation guarantee . Let us now formally introduce f -divergences . Definition 4 ( f -divergence ) . Let f : R+ → R be a convex function such that f ( 1 ) = 0 . The f -divergence between distributions Pz and Qz ( Pz being absolutely continuous with respect to Qz ) , admitting densities p ( z ) and q ( z ) is defined as Df ( Pz||Qz ) = ∫ f ( p ( z ) q ( z ) ) q ( z ) dz . ( 4 ) Consider some density p ( z ) over continuous functions . We define qM ( z ) = 1M ∑M m=1K ( d ( z , zj ) ) , Kj ( z ) : = K ( z , zj ) is a kernel centered at zj used to approximate the modes of p ( z ) . Theorem 1 . Any f -divergence Df ( p||qM ) = ∫ f ( p ( z ) 1 M ∑M j=1Kj ( z ) ) 1 M M∑ m=1 Km ( z ) dz ( 5 ) between a distribution p ( z ) and a normalized mixture of M kernels with equal weights is supermodualar in a cardinality-fixed setting , assuming that ∀zmaxqM Df ( p ( z ) ||qM ( z ) ) < ∞ . Proof . The proof is shown in Appendix A.1 . Minimization of ( 5 ) is equivalent to a cardinality-constrained maximization of a non-monotone submodular function of Z = { z1 , . . . , zM } . Approximation guarantees for problems of this form are given for non-negative submodular functions ( Buchbinder et al. , 2014 ) . One can convert ( 5 ) to a non-negative function by defining : F ( Z ) : = −Df ( p||qM ) + C , ( 6 ) where C = maxZ D ( p||qM ) is a pre-defined constant , which is important for understanding approximation guarantees , but does not need to be computed in practice . After the described transformation , which leads to ( 6 ) , we obtain F ( Z ) , a non-negative non-monotone submodular function . Approximation guarantees for f -divergences . In the context of submodular functions , there exists an inapproximability result , obtained in ( Gharan & Vondrák , 2011 ) , which states that no general polynomial time algorithm with guarantees better than 0.491 exists to solve a submodular maximization problem with constrained cardinality . We , however , that in practice , it might still be possible to obtain a better approximation factor than 0.491 for some specific types of divergences , as these guarantees are defined for all instances of the optimization problem ( Definition 3 ) . Let us consider the approximation guarantees for ( 6 ) , denoting by q∗M the optimal solution , and by q̂M a solution found by some algorithm . For an approximation factor γ , we have −Df ( p||q̂M ) + C ≥ γ ( −Df ( p||q∗M ) + C ) . ( 7 ) Simple algebra shows that this implies Df ( p||q̂M ) ≤ γ min qM∈F Df ( p||qM ) + ( 1− γ ) max qM∈F Df ( p||qM ) , ( 8 ) where F is a family of distributions over the spaces of candidate functions . The derived result indicates that the upper bound on the approximate solution found by minimizing an f -divergence can be substantially dominated by ( 1− γ ) maxqM∈F Df ( p||qM ) if F is chosen poorly . The impact of the choice of function space , determined largely by the neural network architecture , is therefore unavoidable when designing algorithms for approximating the Bayesian posterior . 3.2 GREEDY MINIMIZATION OF f -DIVERGENCES Random greedy algorithm . Although submodular optimization has natural parallel extensions and associated approximation guarantees , due to the simplicity of presentation , we focus in this paper on forward greedy selection and use Algorithm 1 for optimizing submodular functions . This algorithm has approximation guarantee of 1/e in general . The only required step for this greedy algorithm is a computation of the marginal gain on the objective function F ( Z ) , i.e . ∆ ( zk|Z ) = F ( Z ∪ { zk } ) − F ( Z ) . We argue that due to randomized initialization in neural networks , one can use a naive greedy algorithm instead of the sampling procedure in lines 7-8 , and we show the implementation of the resulting training procedure in Appendix B.1 . Algorithm 1 Random Greedy algorithm 1 : Input : V – Ground set 2 : Input : F – Arbitrary submodular function 3 : Input : M – Cardinality of the solution 4 : Z ← ∅ 5 : for m = 1 to M do 6 : R← arg maxT⊂V \Z : |T |=M ∑ z′∈T ∆ ( z ′|Z ) 7 : ui ← Uniform ( R ) 8 : Z ← Z ∪ { ui } 9 : end for 10 : return Z Marginal gain . At each step of a greedy algorithm , a marginal gain ∆ ( zk|Z ) = F ( Z ∪ { zk } ) − F ( Z ) of adding a new element zk to an existing mixture ∑k−1 j=1 Kj ( z ) is maximized . For f - divergences , we thus formulate the following proposition : Proposition 1 . Consider C = maxDf ( p||qM ) , where Df ( p||qM ) is an arbitrary f -divergence between some distribution p ( z ) and a mixture of kernels qM ( z ) = 1M ∑M j=1Kj ( z ) , and Df ( p||qM ) < ∞ . Then , maximization of a marginal gain for set function F ( Z ) = − ∫ f ( p ( z ) 1 M ∑M j=1Kj ( z ) ) 1 M M∑ m=1 Km ( z ) dz + C , ( 9 ) at a step k of a greedy algorithm corresponds to arg max zk ∆ ( zk|Z ) = arg min zk Ez∼Kk ( z ) f ( p ( z ) 1 M ∑k j=1Kj ( z ) ) . ( 10 ) Proof . The proof is shown Appendix A.2 . The case of reverse KL divergence . Having mean-seeking behavior is useful to fit the kernel density estimator using forward divergences . However , if one wants to optimize marginal gains , they need to have a mode-seeking behavior , which is achieved via optimizing reverse divergences ( Zhang et al. , 2019 ) . If we consider the generator for the reverse KL-divergence , f ( x ) = − log x , the minimization simplifies to min zk Ez∼Kk ( z ) − log p ( z ) + log 1 M k∑ j=1 Kj ( z ) . ( 11 ) In the sequel , we leverage the standard reverse KL-divergence generator , which aims to discover the modes of p ( z ) , but also enforces the coverage of the approximation . Finally , we note that it is costly to compute the expectation in practice ; we thus use the following point estimate of ( 11 ) : min zk − log p ( zk ) + log 1 M k−1∑ j=1 Kj ( zk ) . ( 12 )

Summary: The paper proposes a novel and principled method to make Deep Ensemble Bayesian. It minimizes an f-divergence between the true posterior and a kernel density estimator fitting the functions associated with the ensemble candidates. The authors formulate the learning as a submodular problem and propose to use a greedy approach to solve it. The resultant learning objective consists of a novel diversity term. Some out-of-distribution detection results in computer vision demonstrate the efficacy of the proposed method.