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: The authors introduce a deep ensemble kernel learning approach as a linear-based learning combination, from a deep learning scheme, to approximate kernel functions under a Bayesian (GP) framework. Namely, a universal kernel approximation strategy is proposed from eigen-based decomposition and deep learning-based function composition. Then, a variational inference strategy is used to solve the optimization from kernel-based mappings. Two regularization strategies are studied: optimal prior covariance and isotropic covariance. Results demonstrate the benefits of the proposal.