Amortized Bayesian Meta-Learning

Meta-learning , or learning-to-learn , has proven to be a successful strategy in attacking problems in supervised learning and reinforcement learning that involve small amounts of data . State-of-the-art solutions involve learning an initialization and/or optimization algorithm using a set of training episodes so that the metalearner can generalize to an evaluation episode quickly . These methods perform well but often lack good quantification of uncertainty , which can be vital to realworld applications when data is lacking . We propose a meta-learning method which efficiently amortizes hierarchical variational inference across tasks , learning a prior distribution over neural network weights so that a few steps of Bayes by Backprop will produce a good task-specific approximate posterior . We show that our method produces good uncertainty estimates on contextual bandit and few-shot learning benchmarks . 1 INTRODUCTION . Deep learning has achieved success in domains that involve a large amount of labeled data ( Oord et al. , 2016 ; Huang et al. , 2017 ) or training samples ( Mnih et al. , 2013 ; Silver et al. , 2017 ) . However , a key aspect of human intelligence is our ability to learn new concepts from only a few experiences . It has been hypothesized that this skill arises from accumulating prior knowledge and using it appropriately in new settings ( Lake et al. , 2017 ) . Meta learning attempts to endow machine learning models with the same ability by training a metalearner to perform well on a distribution of training tasks . The meta-learner is then applied to an unseen task , usually assumed to be drawn from a task distribution similar to the one used for training , with the hope that it can learn to solve the new task efficiently . Popular meta-learning methods have advanced the state-of-the-art in many tasks , including the few-shot learning problem , where the model has to learn a new task given a small training set containing as few as one example per class . Though performance on few-shot learning benchmarks has greatly increased in the past few years , it is unclear how well the associated methods would perform in real-world settings , where the relationship between training and evaluation tasks could be tenuous . For success in the wild , in addition to good predictive accuracy , it is also important for meta-learning models to have good predictive uncertainty - to express high confidence when a prediction is likely to be correct but display low confidence when a prediction could be unreliable . This type of guarantee in predictive ability would allow appropriate human intervention when a prediction is known to have high uncertainty . Bayesian methods offer a principled framework to reason about uncertainty , and approximate Bayesian methods have been used to provide deep learning models with accurate predictive uncertainty ( Gal & Ghahramani , 2016 ; Louizos & Welling , 2017 ) . By inferring a posterior distribution over neural network weights , we can produce a posterior predictive distribution that properly indicates the level of confidence on new unseen examples . Accordingly , we consider meta-learning under a Bayesian view in order to transfer the aforementioned benefits to our setting . Specifically , we extend the work of Amit & Meir ( 2018 ) , who considered hierarchical variational inference for meta-learning . The work primarily dealt with PAC-Bayes bounds in meta-learning and the experiments consisted of data with tens of training episodes and small networks . In this paper , we show how the meta-learning framework of Finn et al . ( 2017 ) can be used to efficiently amortize variational inference for the Bayesian model of Amit & Meir ( 2018 ) in order to combine the former ’ s flexibility and scalability with the latter ’ s uncertainty quantification . 2 META-LEARNING VIA HIERARCHICAL VARIATIONAL INFERENCE . We first start by reviewing the hierarchical variational bayes formulation used in Amit & Meir ( 2018 ) for meta-learning . Assume we observe data from M episodes , where the ith episode consists of data Di containing N data items , meaning Di = { ( xi , j , yi , j ) } Nj=1 . We assume a hierarchical model with global latent variable θ and episode-specific variables φi , i = 1 , . . .M ( see Figure 1 ) . Hierarchical variational inference can then be used to lower bound the likelihood of the data : log [ M∏ i=1 p ( Di ) ] = log [ ∫ p ( θ ) [ M∏ i=1 ∫ p ( Di|φi ) p ( φi|θ ) dφi ] dθ ] ≥ Eq ( θ ; ψ ) [ log ( M∏ i=1 ∫ p ( Di|φi ) p ( φi|θ ) dφi ) ] − KL ( q ( θ ; ψ ) ‖p ( θ ) ) = Eq ( θ ; ψ ) [ M∑ i=1 log ( ∫ p ( Di|φi ) p ( φi|θ ) dφi ) ] − KL ( q ( θ ; ψ ) ‖p ( θ ) ) ≥ Eq ( θ ; ψ ) [ M∑ i=1 Eq ( φi ; λi ) [ log p ( Di|φi ) ] − KL ( q ( φi ; λi ) ‖p ( φi|θ ) ) ] − KL ( q ( θ ; ψ ) ‖p ( θ ) ) = L ( ψ , λ1 , . . . , λM ) . Here , ψ and λ1 , . . . , λM are the variational parameters of the approximate posteriors over the global latent variables θ and the local latent variables φ1 , . . . , φM , respectively . Thus , variational inference involves solving the following optimization problem : argmax ψ , λ1 ... , λM Eq ( θ ; ψ ) [ M∑ i=1 Eq ( φi ; λi ) [ log p ( Di|φi ) ] − KL ( q ( φi ; λi ) ‖p ( φi|θ ) ) ] − KL ( q ( θ ; ψ ) ‖p ( θ ) ) ( 1 ) ≡ argmin ψ , λ1 ... , λM Eq ( θ ; ψ ) [ M∑ i=1 −Eq ( φi ; λi ) [ log p ( Di|φi ) ] + KL ( q ( φi ; λi ) ‖p ( φi|θ ) ) ] + KL ( q ( θ ; ψ ) ‖p ( θ ) ) ( 2 ) Amit & Meir ( 2018 ) solve this optimization problem via mini-batch gradient descent on the objective starting from random initialization for all variational parameters . They maintain distinct variational parameters λi for each episode i , each of which indexes a distribution over episode-specific weights q ( φi ; λi ) . While they only consider problems with at most 10 or so training episodes and where each φi is small ( the weights of a 2-layer convolutional network ) , this approach is not scalable to problems with large numbers of episodes - such as few-shot learning , where we can generate millions of episodes by randomizing over classes and examples - and requiring deep networks . 3 AMORTIZED BAYESIAN META-LEARNING . 3.1 SCALING META-LEARNING WITH AMORTIZED VARIATIONAL INFERENCE . Learning local variational parameters λi for a large number of episodes M becomes difficult as M grows due to the costs of storing and computing each λi . These problems are compounded when each φi is the weight of a deep neural network and each λi are variational parameters of the weight distribution ( such as a mean and standard deviation of each weight ) . Instead of maintaining M different variational parameters λi indexing distributions over neural network weights φi , we compute λi on the fly with amortized variational inference ( AVI ) , where a global learned model is used to predict λi from Di . A popular use of AVI is training a variational autoencoder ( Kingma & Welling , 2013 ) , where a trained encoder network produces the variational parameters for each data point . Rather than training an encoder to predict λi given the episode , we show that inference can be amortized by finding a good initialization , a la MAML ( Finn et al. , 2017 ) . We represent the variational parameters for each episode as the output of several steps of gradient descent from a global initialization . Let LDi ( λ , θ ) = −Eq ( φi ; λ ) [ log p ( Di|φi ) ] + KL ( q ( φi ; λ ) ‖p ( φi|θ ) ) be the part of the objective corresponding to data Di . Let the procedure SGDK ( D , λ ( init ) , θ ) represent the variational parameters produced after K steps of gradient descent on the objective LD ( λ , θ ) with respect to λ starting at the initialization λ ( 0 ) = λ ( init ) and where θ is held constant i.e . : 1. λ ( 0 ) = λ ( init ) 2. for k = 0 , . . . , K − 1 , set λ ( k+1 ) = λ ( k ) − α∇λ ( k ) LD ( λ ( k ) , θ ) We represent the variational distribution for each dataset qθ ( φi|Di ) in terms of the local variational parameters λi produced after K steps of gradient descent on the loss for dataset Di , starting from the global initialization θ : qθ ( φi|Di ) = q ( φi ; SGDK ( Di , θ , θ ) ) . Note that θ here serves as both the global initialization of local variational parameters and the parameters of the prior p ( φ | θ ) . We could pick a separate prior and global initialization , but we found tying the prior and initialization did not seem to have a negative affect on performance , while significantly reducing the number of total parameters necessary . With this form of the variational distribution , this turns the optimization problem of ( 2 ) into argmin ψ Eq ( θ ; ψ ) [ M∑ i=1 −Eqθ ( φi|Di ) [ log p ( Di|φi ) ] + KL ( qθ ( φi|Di ) ‖p ( φi|θ ) ) ] + KL ( q ( θ ; ψ ) ‖p ( θ ) ) . ( 3 ) Because each qθ ( φi|Di ) depends on ψ via θ ( the initialization for the variational parameters before performing K steps of gradient descent ) , we can also backpropagate through the computation of q via the gradient descent process to compute updates for ψ . Though this backpropagation step requires computing the Hessian , it can be done efficiently with fast Hessian-vector products , which have been used in past work involving backpropagation through gradient updates ( Maclaurin et al. , 2015 ; Kim et al. , 2018b ) . This corresponds to learning a global initialization of the variational parameters such that a few steps of gradient descent will produce a good local variational distribution for any given dataset . We assume a setting where M > > N , i.e . we have many more episodes than data points within each episode . Accordingly , we are most interested in quantifying uncertainty within a given episode and desire accurate predictive uncertainty in qθ ( φi|Di ) . We assume that uncertainty in the global latent variables θ should be low due to the large number of episodes , and therefore use a point estimate for the global latent variables , letting q ( θ ; ψ ) be a dirac delta function q ( θ ) = 1 { θ = θ∗ } . This removes the need for global variational parameters ψ and simplifies our optimization problem to argmin θ [ M∑ i=1 −Eqθ ( φi|Di ) [ log p ( Di|φi ) ] + KL ( qθ ( φi|Di ) ‖p ( φi|θ ) ) ] + KL ( q ( θ ) ‖p ( θ ) ) , ( 4 ) where θ∗ is the solution to the above optimization problem . Note that KL ( q ( θ ) ‖p ( θ ) ) term can be computed even when q ( θ ) = 1 { θ = θ∗ } , as KL ( q ( θ ) ||p ( θ ) ) = − log p ( θ∗ ) .

Summary: This work proposes an adaptation to MAML-type models that accounts for posterior uncertainty in task specific latent variables. This is achieved via a hierarchical Bayesian view of MAML, employing variational inference for the task-specific parameters. The key intuition of this paper is that one can perform fast and efficient test-time variational inference for the task-specific latent variables by learning a good initialization during meta-training. This is achieved in a very similar fashion to MAML, and allows for an interesting form of amortization of test-time inference.