Deep Ensemble Kernel Learning

1 INTRODUCTION . In recent years , there has been a growing interest in Bayesian deep learning ( DL ) , where the point predictions of traditional deep neural network ( DNN ) models are replaced with full predictive distributions using Bayes ’ Rule ( Neal , 2012 ; Wilson , 2020 ) . The advantages of Bayesian DL over traditional DL are numerous and include greater robustness to overfitting and better calibrated uncertainty quantification ( Guo et al. , 2017 ; Kendall & Gal , 2017 ) . Furthermore , the success of traditional DL already rests on a number of probabilistic elements such as stochastic gradient descent ( SGD ) , dropout , and weight initialization– all of which have been given Bayesian interpretations ( Smith & Le , 2018 ; Gal & Ghahramani , 2016 ; Kingma et al. , 2015 ; Schoenholz et al. , 2016 ; Jacot et al. , 2018 ) , so that insights into Bayesian DL may help to advance DL as a whole . Gaussian processes ( GPs ) are nonparametric Bayesian models with appealing properties , as they admit exact inference for regression and allow for a natural functional perspective suitable for predictive modeling ( Rasmussen & Williams , 2005 ) . While at first glance GPs appear unrelated to DL models , a number of interesting connections between GPs and DNNs exist in the literature , suggesting that GPs can constitute a valid approach to Bayesian DL ( Neal , 1996 ; Lee et al. , 2018 ; de Matthews et al. , 2018 ; Jacot et al. , 2018 ; Damianou & Lawrence , 2013 ; Salimbeni & Deisenroth , 2017 ; Agrawal et al. , 2020 ) . A GP prior is typically characterized by its covariance function or “ kernel ” , which determines the class of functions that the GP can model , as well as its generalization properties outside training data . Kernel selection is the primary problem in GP modeling , and unfortunately traditional kernels such as the radial basis function ( RBF ) kernel are not sufficiently expressive for complex problems where more flexible models such as DNNs generally perform well . This is the key motivation for kernel learning , which refers to the selection of an optimal kernel out of a family of kernels in a data-driven way . A number of approaches to kernel learning exist in the literature , including some that parameterize kernels using DNNs ( Zhou et al. , 2019 ; Li et al. , 2019 ; Bullins et al. , 2018 ; Sinha & Duchi , 2016 ) . As these approaches involve learning feature representations , they are fundamentally different from random-feature methods for efficient kernel representation ( Rahimi & Recht , 2007 ; 2008 ) . However , these approaches are not specific to GPs and do not take advantage of a robust Bayesian framework . In contrast , the deep kernel learning ( DKL ) paradigm does exactly this ; In DKL , a DNN is used as a feature extractor that maps data inputs into a latent feature space , where GP inference with some “ base kernel ” is then performed ( Wilson et al. , 2016b ; a ; Jean et al. , 2016 ; Al-Shedivat et al. , 2017 ; Bradshaw et al. , 2017 ; Izmailov et al. , 2018 ; Xuan et al. , 2018 ) . The resulting model is then trained end-to-end using standard gradient-based optimization , usually in a variational framework . We note that the DKL model is just a GP with a highly flexible kernel parameterized by a DNN . By optimizing all hyperparameters ( including the DNN weights ) with type II maximum likelihood estimation , the DKL model is able to learn an optimal kernel in a manner directly informed by the data , while also taking advantage of the robustness granted by the Bayesian framework . A special case of DKL that is worthy of note was considered in Dasgupta et al . ( 2018 ) , who use a linear base kernel and impose a soft orthogonality constraint to learn the eigenfunctions of a kernel . Although similar in spirit to the approach in this paper , their method does not make use of an efficient variational method , nor is distributed training made possible since all of the basis functions are derived from the same feature network . In this work , we introduce the “ deep ensemble kernel learning ” ( DEKL ) model– a simpler and more efficient special case of DKL with two specifications– the base kernel is linear , and the feature network is partitioned into an “ ensemble ” of “ learners ” with common network architecture . In contrast to nonlinear kernels , the linear kernel allows us to derive an efficient training and inference method for DEKL that circumvents the inducing points approximation commonly used in traditional DKL . The hyperparameters of the linear kernel can also be optimized in closed form , allowing us to simplify the loss function considerably . Convenience aside , we show that DEKL remains highly expressive , proving that it is universal in the sense that it can approximate any continuous kernel so long as its feature network is arbitrarily wide . In other words , we may keep the base kernel simple if we are willing to let the feature network be more complex . The second specification of DEKL lets us handle the complexity of the feature network ; because the feature network is partitioned , it admits easy model parallelism , where the learners in the ensemble are distributed . Moreover , our universality result only requires the number of learners to be arbitrarily large ; the learners themselves need not grow ( meaning fixed-capacity learners are sufficient ) , avoiding additional model parallelism . From a different perspective , DEKL may be regarded as an extension of traditional ensembling methods for DNNs and in particular the deep ensemble ( DE ) model of Lakshminarayanan et al . ( 2017 ) , which is also highly parallelizable . In a DE , each DNN learner parameterizes a distribution over the variates ( e.g. , the mean and variance of a Gaussian in regression , or the logits of a softmax vector in classification ) . Each learner is trained independently with maximum likelihood estimation , and the final predictive distribution of the DE is then defined to be a uniform mixture of the individual learner predictive distributions . Although not Bayesian itself , the DE model boasts impressive predictive performance and was shown to outperform Bayesian methods such as probabilistic back propagation ( Hernández-Lobato & Adams , 2015 ) and MC-dropout ( Gal & Ghahramani , 2016 ) . In contrast , in DEKL , the learners are trained jointly via a shared linear GP layer . We surmise that this may help to promote diversity ( i.e. , low correlation ) among the learners by facilitating coordination , which we verify experimentally . Unlike non-Bayesian joint ensemble training methods such as that of Webb et al . ( 2019 ) , we hypothesize that the DEKL learners might learn to diversify in order to better approximate the posterior covariance– an inherently Bayesian feature . We therefore expect DEKL to be more efficient than DKL and more robust than DE , by drawing on the strengths of both ( see Fig . 1 for a comparison of model architectures ) . 2 DEEP ENSEMBLE KERNEL LEARNING . A DKL model is a GP whose kernel encapsulates a DNN for feature extraction ( Wilson et al. , 2016b ; a ) . A deep kernel is defined as Kdeep ( x1 , x2 ; θ , γ ) = Kbase ( ϕ ( x1 ; θ ) , ϕ ( x2 ; θ ) ; γ ) , ( 1 ) where ϕ ( · ; θ ) is a DNN with weight parameters θ and Kbase ( · , · ; γ ) is any chosen kernel—called the “ base kernel ” —with hyperparameters γ . Note that the kernel hyperparameters of Kdeep include all hyperparameters γ of the base kernel Kbase as well as the DNN weight parameters θ ; Given the expressive power of DNNs , the deep kernel is also highly expressive and may be viewed as a method to automatically select a GP model . Our proposed method , DEKL , is a special case of DKL . Whereas RBF and Matern kernels are typically used as base kernels in DKL , in DEKL we take the base kernel to be the linear kernel : Klin ( x1 , x2 ; V ) = x > 1 V x2 , ( 2 ) where V is a symmetric positive-semidefinite matrix . In order to enable parallel computation , in DEKL , we use a feature network ϕ ( · ; θ ) that is partitioned into an “ ensemble ” of subnetworks ϕi ( · ; θi ) called “ learners ” , having identical network architectures ; i.e. , ϕ ( · ; θ ) is the concatenation of the outputs of the ϕi ( · ; θi ) . DEKL offers a number of advantages over general DKL : 1 . Unlike the hyperparameters of the RBF kernel , in DEKL , we can optimize the hyperparameters of the linear base kernel in closed form . 2 . The linear base kernel allows us to think of a DEKL model not just as a GP but as a finitedimensional Bayesian linear model ( BLM ) , conditional on the feature networks ; this lets us derive an efficient inference method that is much simpler than the inducing points method used in general DKL . 3 . Finally , the partitioned architecture of the feature network makes it much more amenable to model parallelism . Note that DEKL is fundamentally different from random-feature methods such as that of Rahimi & Recht ( 2008 ) , where the learners ϕ ( · ; θi ) are random features that are not optimized during training . A potential drawback to partitioning the feature network as we do in DEKL is that , compared to general DKL , the network is less expressive . However , the following universal approximation theorem for DKL implies that this effect can be compensated by adding parallel learners : Theorem 1 ( Universal kernel approximation theorem ) . Let X ⊂ RD be some compact Euclidean domain , and let σ : R → R be a non-polynomial activation function ( Pinkus , 1999 ) . Then , given a continuous , symmetric , positive-definite kernel κ : X × X 7→ R and any > 0 , there exist a finite number H of affine functions βi : X aff.−−−→ R and a symmetric positive semi-definite matrix V ∈ RH×H such that for any x1 , x2 ∈ X , ∣∣∣∣∣∣ H∑ i , j=1 vijσ ( βi ( x1 ) ) σ ( βj ( x2 ) ) − κ ( x1 , x2 ) ∣∣∣∣∣∣ < . ( 3 ) The proof , found in Appendix A , contains a straightforward combination of Mercer ’ s Theorem ( Mercer , 1909 ) and the Universal Approximation Theorem , attributed to Cybenko , Hornik , Leshno , and Pinkus ( Cybenko , 1989 ; Hornik et al. , 1989 ; Leshno et al. , 1993 ; Pinkus , 1999 ) . Note that Thm . 1 lets us represent non-stationary continuous kernels , in contrast to methods such as random Fourier feature expansion ( Rahimi & Recht , 2007 ) . The approximation in Thm . 1 requires a possibly large number H of affine functions βi . However , in DEKL we replace the functions x → σ ( βi ( x ) ) with an ensemble of strictly more flexible DNN learners ϕ ( · ; θi ) , which can help to reduce the numberH of learners required ; in the proof of Thm . 1 , we approximate each eigenfunction of the target kernel κ with a linear combination of the learners ; if the learners are sufficiently expressive , then it may take only one learner per eigenfunction to approximate the target kernel . We also allow each learner ϕ ( · ; θi ) to have multiple outputs M . In the case of the simple learners x → σ ( βi ( x ) ) , a learner with M outputs is simply a concatenation of M single-output learners , suggesting that multi-output learners may help to further reduce the number H of required learners . The summation in Eq . 3 should be understood as a deep kernel as in Eq . 1 , where the DNN ϕ in Eq . 1 is the concatenation of all learners and the base kernel is given by the linear kernel in Eq . 2 . Theorem 1 is thus a statement about the universality of DEKL with a linear base kernel . Given that other kernels such as the RBF are more popular choices of base kernel in the DKL literature , it is natural to wonder if DEKL remains a universal kernel approximator if we change the base kernel . It turns out that not all choices of base kernel give universality , as is implied by the following remark . Remark 2 . For a deep kernel ( Eq . 1 ) with base kernel Kbase : RH × RH 7→ R to be a universal kernel approximator , the base kernel must be unbounded both above and below . We give more details in Appendix B , but intuitively , since the base kernel is the outermost function in the deep kernel , any bound on its range will prevent the deep kernel from approximating kernels with unbounded range , such as the dot product kernel . The class of base kernels with bounded ( or half-bounded ) range , and thus the base kernels that do not give universality , is large and includes many popular kernels such as the RBF kernel and periodic kernel . The linear base kernel in Eq . 2 is therefore special , as it is not only convenient but also grants us universality . We note the converse of Remark 2 is not true ; an unbounded base kernel does not guarantee that a DEKL model is a universal kernel approximator . For example , restricting the matrix V in the linear base kernel ( Eq . 2 ) to a diagonal matrix breaks universality ; this is because Thm . 1 must hold for even the simplest learners x→ σ ( βi ( x ) ) , which fails to happen when we restrict V ( see Appendix B for details ) . Classifying all base kernels for which DEKL is universal remains an important open problem for future work .

Summary: This paper proposes a variant of the Deep Kernel Learning model (DKL) [1] where multiple independent networks are trained for the features instead of a single network. In addition, the paper proposes to use a linear kernel as a base kernel which allows for universal approximation of any arbitrary kernel, as well as allowing for exact inference of the kernel hyperparameters. The use of stochastic variational inference is proposed for inferring the neural networks weights. The model is compared against Deep Ensemble (DE) and DKL on a synthetic dataset and the UCI dataset.