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: This paper presents a GP model with a deep kernel defined in terms of a (finite) deep ensemble (DE). A variational approximation and a regularization scheme were introduced to optimize the GP. The proposed methods were then demonstrated on several benchmark regression/classification datasets.   Here, the claimed contributions are (1) a DE-parameterized GP posterior that provides a Bayesian justification for DE; and (b) a regularization scheme that improves the generalization of the trained GP.