Deep Ensemble as a Gaussian Process Posterior

1 INTRODUCTION . Bayesian treatment of deep neural networks ( DNNs ) is promised to enjoy principled Bayesian uncertainty while unleashing the capacity of DNNs , with Bayesian neural networks ( BNNs ) as popular examples ( MacKay , 1992 ; Hinton & Van Camp , 1993 ; Neal , 1995 ; Graves , 2011 ) . Nevertheless , despite the surge of advance in BNNs ( Louizos & Welling , 2016 ; Zhang et al. , 2018 ) , many of existing BNNs still face challenges in accurate and scalable inference ( Sun et al. , 2019 ) , and exhibit limitations in uncertainty estimation and out-of-distribution robustness ( Ovadia et al. , 2019 ) . Alternatively , Deep Ensemble ( DE ) ( Lakshminarayanan et al. , 2017 ) trains multiple independent , randomly-initialized DNNs for ensemble , presenting higher flexibility and effectiveness than BNNs . However , it is hard to interpret DE as a Bayesian approach which seeks for the Bayesian posterior of a certain model , and there is no guarantee that the uncertainty estimates given by DE are reliable . To chase a Bayesian justification for DE , recent works like RMS ( Lu & Van Roy , 2017 ; Osband et al. , 2018 ; Pearce et al. , 2020 ) and NTKGP ( He et al. , 2020 ) refine DE to be a “ sample-then-optimize ” approach ( Matthews et al. , 2017 ) . Nonetheless , these works rely on the regression likelihood and often make strong assumptions like linearised and/or infinite-width models . To build a Bayesian refinement of DE without reliance on restrictive assumptions , we propose to use DE to define a Gaussian process ( GP ) approximate posterior , which is dubbed as DE-GP for short . The imposition of a GP form on the approximate posterior is empowered by the equivalence between GPs and BNNs ( Khan et al. , 2019 ; Neal , 1996 ; Lee et al. , 2018 ; Garriga-Alonso et al. , 2018 ; Matthews et al. , 2018 ; Novak et al. , 2018 ; He et al. , 2020 ) , as well as the associated computational benefits when adopting a GP prior without loss of generality . Concretely , we evaluate the empirical mean and covariance of the ensemble members and use them to specify DE-GP . This is inspired by the formula of neural network Gaussian process ( NN-GP ) ( Novak et al. , 2018 ) . Nevertheless , we clarify that we define a parameteric and learnable GP kernel with finitely many and adaptive basis functions , in contrast to the NN-GP kernels built with infinitely many yet fixed basis functions . The flexibility of DE-GP is assured by the expressiveness of DNN ensemble members – by tuning them it is enabled to evolve over a rich variety of kernels . Given these setups , we perform function-space variational inference ( VI ) for learning . We maximize the functional evidence lower bound ( fELBO ) ( Sun et al. , 2019 ; Rudner et al. , 2021 ) to push DE-GP towards the true posterior over functions . We further present a posterior regularization scheme in function space to conveniently incorporate structural priors , and present a concrete example on how to control function complexity . By VI , we are enabled to handle classification problems directly and exactly , without casting them into regression ones ( He et al. , 2020 ) . In practice , DE-GP can be easily implemented on top of a standard DE , while only introducing minimally added cost . The function-space Bayesian inference enables DE-GP to ameliorate the pathologies induced by the over-parameterization nature of DNNs ( Sun et al. , 2019 ; Wang et al. , 2019 ) . Thus DE-GP can offer calibrated uncertainty estimates , especially when adopting large networks ( see Fig . 1 ) . We empirically demonstrate that DE-GP outperforms strong baselines on various regression datasets , and presents superior uncertainty estimates and out-of-distribution robustness without compromising accuracy in challenging classification tasks . DE-GP also shows promise in solving contextual bandit problems , where the uncertainty plays a vital role in guiding exploration . 2 RELATED WORK . Bayesian neural networks ( BNNs ) . Bayesian treatment of DNNs is an emerging topic in deep learning yet with a long history ( Mackay , 1992 ; Hinton & Van Camp , 1993 ; Neal , 1995 ; Graves , 2011 ) . BNNs can be learned by virtue of variational inference ( Blundell et al. , 2015 ; HernándezLobato & Adams , 2015 ; Louizos & Welling , 2016 ; Zhang et al. , 2018 ; Khan et al. , 2018 ; Deng et al. , 2020 ) , Laplace approximation ( Mackay , 1992 ; Ritter et al. , 2018 ) , Markov chain Monte Carlo ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Zhang et al. , 2019 ) , particle-optimization based variational inference ( Liu & Wang , 2016 ) , Monte Carlo dropout ( Gal & Ghahramani , 2016 ) , and other methods ( Maddox et al. , 2019 ; Izmailov et al. , 2020 ) . To avoid the difficulties of posterior inference in weight space , some recent works advocate performing Bayesian reasoning in function space ( Sun et al. , 2019 ; Rudner et al. , 2021 ; Wang et al. , 2019 ) . In function space , BNNs of infinite or even finitely width equal to Gaussian processes ( GPs ) ( Neal , 1996 ; Lee et al. , 2018 ; Novak et al. , 2018 ; He et al. , 2020 ; Khan et al. , 2019 ) , which provides valuable insights for this work . Deep Ensemble ( DE ) ( Lakshminarayanan et al. , 2017 ) is a qualified alternative to BNNs for uncertainty estimation ( Ovadia et al. , 2019 ) , yet lacks a proper Bayesian interpretation . Wilson & Izmailov ( 2020 ) conceptually interpreted DE as a method that approximates the Bayesian posterior predictive distribution , but it is hard to judge whether the approximation is reliable or not in practice . Recently , a suite of works based on the notion of “ sample-then-optimize ” ( Matthews et al. , 2017 ) have been developed to make DE Bayesian . For example , RMS ( Lu & Van Roy , 2017 ; Osband et al. , 2018 ; Pearce et al. , 2020 ) regularizes the ensemble members towards randomised priors to obtain posterior samples after training , while it typically assumes linear data likelihood which is impractical for deep models and classification tasks . He et al . ( 2020 ) proposed to add a randomised function to each ensemble member to realise a function-space Bayesian interpretation , but the method is asymptotically exact in the infinite width limit and is limited to regressions . By contrast , DE-GP works without restrictive assumptions . In parallel , D ’ Angelo & Fortuin ( 2021 ) proposed to add a repulsive term to DE to achieve Bayesian inference , which yet relies on unscalable gradient estimators . 3 BACKGROUND . 3.1 BAYESIAN NEURAL NETWORKS . Consider learning on a dataset D = ( X , Y ) = { ( xi , yi ) } ni=1 , with xi ∈ X as data and yi as C-dimensional targets . Let g ( · , w ) : X → RC denote the function represented by a DNN with weights w. Bayesian neural networks ( BNNs ) treat w as random variables and chase the posterior p ( w|D ) by imposing a prior p ( w ) . Among the family , variational BNNs ( VBNNs ) are particularly appealing as they cast posterior inference as an optimization problem ( Graves , 2011 ; Blundell et al. , 2015 ; Louizos & Welling , 2016 ) , where a variational distribution q ( w ) is introduced to approximate the true posterior by maximizing the evidence lower bound ( ELBO ) : Eq ( w ) [ log p ( D|w ) ] −DKL [ q ( w ) ‖p ( w ) ] . ( 1 ) BNNs marginalize over the posteriors to obtain the posterior predictive for new data x∗ : p ( y|x∗ , D ) = Ep ( w|D ) p ( y|x∗ , w ) ≈ Eq ( w ) p ( y|x∗ , w ) ≈ S∑ s=1 p ( y|x∗ , ws ) , ( 2 ) where ws ∼ q ( w ) , s = 1 , ... , S. Such a procedure is expected to propagate the embedded model uncertainty into the prediction . However in practice , most of the existing BNN approaches face obstacles in precise posterior inference due to non-trivial and convoluted posterior dependencies ( Louizos & Welling , 2016 ; Zhang et al. , 2018 ; Shi et al. , 2018 ; Sun et al. , 2019 ) , and deliver unsatisfactory uncertainty estimation and out-of-distribution ( OOD ) robustness ( Ovadia et al. , 2019 ) . 3.2 DEEP ENSEMBLE . As a workaround of BNNs , Deep Ensemble ( DE ) ( Lakshminarayanan et al. , 2017 ) deploys a set of M DNNs { g ( · , wi ) } Mi=1 for ensemble . { wi } Mi=1 are independently trained to interpret the data from different angles , under maximum likelihood estimation ( MLE ) principle in standard DE , or maximum a posteriori ( MAP ) principle in regularized DE ( rDE ) : max w1 , ... , wM LDE = 1 M M∑ i=1 log p ( D|wi ) , max w1 , ... , wM LrDE = 1 M M∑ i=1 [ log p ( D|wi ) + log p ( wi ) ] . ( 3 ) Due to the randomness in network initialization and stochastic gradient descent ( SGD ) , DE can effectively explore the non-convex , multimodal loss landscape of DNNs ( Fort et al. , 2019 ; Wilson & Izmailov , 2020 ) , thus has shown promise in various uncertainty estimation tasks ( Lakshminarayanan et al. , 2017 ; Ovadia et al. , 2019 ) . However , due to the lack of a Bayesian justification , there is no guarantee that the uncertainty estimates given by DE are reliable . 3.3 RELATION OF THE BAYESIAN POSTERIORS OF DNNS TO GAUSSIAN PROCESSES . It is early shown that shallow BNNs converge to Gaussian processes ( GPs ) in the infinite width limit ( Neal , 1996 ) . Recently , there is ongoing effort to extend the result to multiple layers ( Lee et al. , 2018 ; Matthews et al. , 2018 ) , convolutional architectures ( Garriga-Alonso et al. , 2018 ; Novak et al. , 2018 ) , and beyond ( Yang , 2019 ) . Briefly , an infinitely wide BNN amount to the NN-GP : GP ( 0 , κ ( x , x′ ) ) , with κ ( x , x′ ) : = Ew∼p ( w ) [ g ( x , w ) g ( x′ , w ) > ] , ( 4 ) where p ( w ) is a layerwise isotropic Gaussian and the DNN specifying g has infinite width . Notably , κ is a compositional , matrix-valued kernel , with values in the space of C × C matrices , and can be analytically estimated in some cases ( Novak et al. , 2018 ; 2019 ) . Due to the i.i.d . weights , the kernel evaluation κ ( x∗ , x∗ ) of some data x∗ is a scaled identity matrix . Based on the canonical results of GP regression , we know the posteriors of infinitely wide BNNs are still GPs ( in regression tasks ) , corresponding to the Bayesian training of the last readout layer ( Arora et al. , 2019 ; Lee et al. , 2019 ) . Despite a golden standard for small regression problems , analytical posterior inference for NN-GP faces challenges when handling contemporary DNN architectures ( Novak et al. , 2018 ) , large data ( Shi et al. , 2019 ) and classification problems . Interestingly , finitely wide BNNs with Gaussian posteriors on weights can also be connected to GPs ( Khan et al. , 2019 ; Rudner et al. , 2021 ) . We then arrive at the hypothesis that the Bayesian posteriors of DNNs are closely related to GPs . So , we specify a GP approximate posterior with DE .

Summary: The authors propose using Deep Ensemble (DE) as basis functions to train a Gaussian Process (GP) with the primary motivation of making DEs more Bayesian. The Evidence Lower BOund (ELBO) being maximized is now optimized via variational inference (VI) over the functional space of bases (i.e. fELBO from [3] Sun et al.) defining the mean and covariance of a GP similar to NN-GP [4]. The paper evaluates the proposed method on standard UCI regression and small-scale MNIST / CIFAR10 image classification tasks in terms of accuracy, likelihood, and uncertainty estimation.