Meta-Learning Bayesian Neural Network Priors Based on PAC-Bayesian Theory

1 INTRODUCTION . Bayesian Neural Networks ( BNNs ) offer a probabilistic interpretation of deep learning by inferring distributions over the model ’ s weights ( Neal , 1996 ) . With the potential of combining the scalability and performance of neural networks ( NNs ) with a framework for uncertainty quantification , BNNs have lately received increased attention ( Blundell et al. , 2015 ; Gal & Ghahramani , 2016 ) . In particular , their ability to express epistemic uncertainty makes them highly relevant for applications such as active learning ( Hernández-Lobato & Adams , 2015 ) and reinforcement learning ( Riquelme et al. , 2018 ) . However , BNNs face two major issues : 1 ) the intractability of posterior inference and 2 ) the difficulty of choosing good Bayesian priors . While the former has been addressed in an extensive body of literature on variational inference ( e.g . Blundell et al. , 2015 ; Blei et al. , 2016 ; Mishkin et al. , 2018 ; Liu & Wang , 2016 ) , the latter has only received limited attention ( Vladimirova et al. , 2019 ; Ghosh & Doshi-Velez , 2017 ) . Choosing an informative prior for BNNs is particularly difficult due to the high-dimensional and hardly interpretable parameter space of NNs . Due to the lack of good alternatives , often a zero-centered , isotropic Gaussian is used , reflecting ( almost ) no a priori knowledge about the problem at hand . This does not only lead to poor generalization when data is scarce , but also renders the Bayesian uncertainty estimates poorly calibrated ( Kuleshov et al. , 2018 ) . Meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) acquires inductive bias in a data-driven way , thus , constituting an alternative route for addressing this issue . In particular , meta-learners attempt to extract shared ( prior ) knowledge from a set of related learning tasks ( i.e. , datasets ) , aiming to learn in the face of a new , related task . Our work develops a principled and scalable algorithm for metalearning BNN priors . We build on the PAC-Bayesian framework ( McAllester , 1999 ) , a methodology from statistical learning theory for deriving generalization bounds . Previous PAC-Bayesian bounds for meta-learners ( Pentina & Lampert , 2014 ; Amit & Meir , 2018 ) require solving a difficult optimization problem , involving the optimization of the prior as well as multiple variational posteriors in a nested manner . Aiming to overcome this issue , we present a PAC-Bayesian bound that does not rely on nested optimization and , unlike ( Rothfuss et al. , 2020 ) , can be tractably optimized for BNNs . This makes the resulting meta-learner , referred to as PACOH-NN , not only much more computationally efficient and scalable than previous approaches for meta-learning BNN priors ( Amit & Meir , 2018 ) , but also agnostic to the choice of approximate posterior inference method which allows us to combine it freely with recent advances in MCMC ( e.g . Chen et al. , 2014 ) or variational inference ( e.g . Wang et al. , 2019 ) . Our experiments demonstrate that the computational advantages of PACOH-NN do not result in degraded predictive performance . In fact , across several regression and classification environments , PACOH-NN achieves a comparable or better predictive accuracy than several popular meta-learning approaches , while improving the quality of the uncertainty estimates . Finally , we showcase how metalearned PACOH-NN priors can be used in a real-world bandit task concerning the development of vaccines , suggesting that many other challenging real-world problems may benefit from our approach . 2 RELATED WORK . Bayesian Neural Networks . The majority of research on BNNs focuses on approximating the intractable posterior distribution ( Graves , 2011 ; Blundell et al. , 2015 ; Liu & Wang , 2016 ; Wang et al. , 2019 ) . In particular , we employ the approximate inference method of Liu & Wang ( 2016 ) . Another crucial question is how to select a good BNN prior ( Vladimirova et al. , 2019 ) . While the majority of work ( e.g . Louizos & Welling , 2016 ; Huang et al. , 2020 ) employs a simple zero-centered , isotropic Gaussian , Ghosh & Doshi-Velez ( 2017 ) and Pearce et al . ( 2020 ) have proposed other prior distributions for BNNs . In contrast , we go the alternative route of choosing priors in a data-driven way . Meta-learning . A range of popular methods in meta-learning attempt to learn the “ learning program ” in form of a recurrent model ( Hochreiter et al. , 2001 ; Andrychowicz et al. , 2016 ; Chen et al. , 2017 ) , learn an embedding space shared across tasks ( Snell et al. , 2017 ; Vinyals et al. , 2016 ) or the initialization of a NN such that it can be quickly adapted to new tasks ( Finn et al. , 2017 ; Nichol et al. , 2018 ; Rothfuss et al. , 2019b ) . A group of recent methods also uses probabilistic modeling to also allow for uncertainty quantification ( Kim et al. , 2018 ; Finn et al. , 2018 ; Garnelo et al. , 2018 ) . Although the mentioned approaches are able to learn complex inference patterns , they rely on settings where metatraining tasks are abundant and fall short of performance guarantees . In contrast , we provide a formal assessment of the generalization properties of our algorithm . Moreover , PACOH-NN allows for principled uncertainty quantification , including separate treatment of epistemic and aleatoric uncertainty . This makes it particularly useful for sequential decision algorithms ( Lattimore & Szepesvari , 2020 ) . PAC-Bayesian theory . Previous work presents generalization bounds for randomized predictors , assuming a prior to be given exogenously ( McAllester , 1999 ; Catoni , 2007 ; Germain et al. , 2016 ; Alquier et al. , 2016 ) . More recent work explores data-dependent priors ( Parrado-Hernandez et al. , 2012 ; Dziugaite & Roy , 2016 ) or extends previous bounds to the scenario where priors are meta-learned ( Pentina & Lampert , 2014 ; Amit & Meir , 2018 ) . However , these meta-generalization bounds are hard to minimize as they leave both the hyper-posterior and posterior unspecified , which leads to nested optimization problems . Our work builds on the results of Rothfuss et al . ( 2020 ) who introduce the methodology to derive the closed-form solution of the PAC-Bayesian meta-learning problem . However , unlike ours , their approach suffers from ( asymptotically ) non-vanishing terms in the bounds and relies on a closedform solution of the marginal log-likelihood . By contributing a numerically stable score estimator for the generalized marginal log-likelihood , we are able to overcome such limitations , making PACBayesian meta-learning both tractable and scalable for a much larger array of models , including BNNs . 3 BACKGROUND : THE PAC-BAYESIAN FRAMEWORK . Bayesian Neural Networks . Consider a supervised learning task with data S = { ( xj , yj ) } mj=1 drawn from unknown distribution D. In that , X = { xj } mj=1 ∈ Xm denotes training inputs and Y = { yj } mj=1 ∈ Ym the targets . For brevity , we also write zj : = ( xj , yj ) ∈ Z . Let hθ : X → Y be a function parametrized by a NN with weights θ ∈ Θ . Using the NN mapping , we define a conditional distribution p ( y|x , θ ) . For regression , we set p ( y|x , θ ) = N ( y|hθ ( x ) , σ2 ) , where σ2 is the observation noise variance . For classification , we choose p ( y|x , θ ) = Categorical ( softmax ( hθ ( x ) ) ) . For Bayesian inference , one presumes a prior distribution p ( θ ) over the model parameters θ which is combined with the training data S into a posterior distribution p ( θ|X , Y ) ∝ p ( θ ) p ( Y|X , θ ) . For unseen test data points x∗ , we form the predictive distribution as p ( y∗|x∗ , X , Y ) = ∫ p ( y∗|x∗ , θ ) p ( θ|X , Y ) dθ . The Bayesian framework presumes partial knowledge of the data-generating process in form of a prior distribution . However , due to the practical difficulties in choosing an appropriate BNN prior , the prior is typically strongly misspecified ( Syring & Martin , 2018 ) . As a result , modulating the influence of the prior relative to the likelihood during inference typically improves the empirical performance of BNNs and is thus a common practice ( Wenzel et al. , 2020 ) . Using such a ” tempered ” posterior pτ ( θ|X , Y ) ∝ p ( θ ) p ( Y|X , θ ) τ with τ > 0 is also referred to as generalized Bayesian learning ( Guedj , 2019 ) . The PAC-Bayesian Framework . In the following , we introduce the most relevant concepts of PAC-Bayesian learning theory . For more details , we refer to Guedj ( 2019 ) . Given a loss function l : θ ×Z → R , we typically want to minimize the generalization error L ( θ , D ) = Ez∗∼D l ( θ , z∗ ) . Since D is unknown , the empirical error L̂ ( θ , S ) = 1m ∑m i=1 l ( θ , zi ) is usually employed , instead . In the PAC-Bayesian framework , we are concerned with randomized predictors , i.e. , probability measures on the parameter space Θ , allowing us to reason about epistemic uncertainty . In particular , we consider two such probability measures , a prior P ∈ M ( Θ ) and a posterior Q ∈ M ( Θ ) . In here , byM ( Θ ) , we denote the set of all probability measures on Θ . While in Bayesian inference , the prior and posterior are tightly coupled through Bayes ’ theorem , the PAC-Bayesian framework only requires the prior to be independent of the data S. Using the definitions above , the so-called Gibbs error for a randomized predictor Q is defined as L ( Q , D ) = Eh∼Q L ( h , D ) . Similarly , we define its empirical counterpart as L̂ ( Q , S ) = Eh∼Q L̂ ( h , S ) . The PAC-Bayesian framework provides upper bounds for the unknown Gibbs error in the following form : Theorem 1 . ( Alquier et al. , 2016 ) Given a data distribution D , a prior P ∈ M ( Θ ) , a confidence level δ ∈ ( 0 , 1 ] , with probability at least 1− δ over samples S ∼ Dm , we have : ∀Q ∈M ( Θ ) : L ( Q , D ) ≤ L̂ ( Q , S ) + 1√ m [ DKL ( Q||P ) + ln 1 δ + Ψ ( √ m ) ] ( 1 ) where Ψ ( √ m ) = lnEθ∼PES∈Dm exp [ √ m ( L ( θ , D ) − L̂ ( θ , S ) ) ] . In that , Ψ ( √ m ) is a log moment generating function that quantifies how strong the empirical error deviates from the Gibbs error . By making additional assumptions about the loss function l , we can bound Ψ ( √ m ) and thereby obtain tractable bounds . For instance , if l ( θ , z ) is bounded in [ a , b ] , we obtain Ψ ( √ m ) ≤ ( ( b− a ) 2 ) /8 by Hoeffding ’ s lemma . For unbounded loss functions , it is common to assume bounded moments . For instance , a loss is considered sub-gamma with variance factor s2 and scale parameter c , under a prior P and data distribution D , if its deviations from the mean can be characterized by random variable V : = L ( h , D ) − l ( h , z ) whose moment generating function is upper bounded by that of a Gamma distribution Γ ( s , c ) ( Boucheron et al. , 2013 ) . In such case , we obtain Ψ ( √ m ) ≤ s2/ ( 2− 2c√ m ) ) . Connecting the PAC-Bayesian framework and generalized Bayesian learning . In PACBayesian learning we aim to find the posterior that minimizes the bound in ( 1 ) which is in general a challenging optimization problem over the space of measuresM ( Θ ) . However , to our benefit , it can be shown that the Gibbs posterior is the probability measure that minimizes ( 1 ) . For details we refer to Lemma 2 in the Appendix or Catoni ( 2007 ) and ( Germain et al. , 2016 ) . In particular , this gives us Q∗ ( θ ) : = arg min Q∈M ( Θ ) √ mL̂ ( Q , S ) +DKL ( Q||P ) = P ( θ ) e− √ mL̂ ( θ , S ) /Z ( S , P ) , ( 2 ) where Z ( S , P ) is a normalization constant . In a probabilistic setting , our loss function is the negative log-likelihood , i.e . l ( θ , zi ) : = − log p ( zi|θ ) . In this case , the optimal Gibbs posterior coincides with the generalized Bayesian posterior Q∗ ( θ ; P , S ) ∝ P ( θ ) p ( S|θ ) 1/ √ m/Z ( S , P ) where Z ( S , P ) = ∫ Θ P ( θ ) p ( S|θ ) 1/ √ m dθ is the generalized marginal likelihood of the data sample S .

Summary: The paper proposes a method for learning BNN priors based on optimising a PAC-Bayes bound for meta-learning, which they call PACOH-NN. This extends previous work on optimising PAC-Bayes bounds (PACOH), and is attractive for being principled in construction and effective in practice. The main claims are that fewer tasks are needed to achieve good performance in terms of both utility and calibration, and that the algorithm comes with performance guarantees.