Meta-Learning Probabilistic Inference for Prediction

1 INTRODUCTION . Many applications require predictions to be made on myriad small , but related datasets . In such cases , it is natural to desire learners that can rapidly adapt to new datasets at test time . These applications have given rise to vast interest in few-shot learning ( Fei-Fei et al. , 2006 ; Lake et al. , 2011 ) , which emphasizes data efficiency via information sharing across related tasks . Despite recent advances , notably in meta-learning based approaches ( Ravi and Larochelle , 2017 ; Vinyals et al. , 2016 ; Edwards and Storkey , 2017 ; Finn et al. , 2017 ; Lacoste et al. , 2018 ) , there remains a lack of general purpose methods for flexible , data-efficient learning . Due to the ubiquity of recent work , a unifying view is needed to understand and improve these methods . Existing frameworks ( Grant et al. , 2018 ; Finn et al. , 2018 ) are limited to specific families of approaches . In this paper we develop a framework for meta-learning approximate probabilistic inference for prediction ( ML-PIP ) , providing this view in terms of amortizing posterior predictive distributions . In Section 4 , we show that ML-PIP re-frames and extends existing point-estimate probabilistic interpretations of meta-learning ( Grant et al. , 2018 ; Finn et al. , 2018 ) to cover a broader class of methods , including gradient based meta-learning ( Finn et al. , 2017 ; Ravi and Larochelle , 2017 ) , metric based meta-learning ( Snell et al. , 2017 ) , amortized MAP inference ( Qiao et al. , 2017 ) and conditional probability modelling ( Garnelo et al. , 2018a ; b ) . The framework incorporates three key elements . First , we leverage shared statistical structure between tasks via hierarchical probabilistic models developed for multi-task and transfer learning ( Heskes , ∗Authors contributed equally †Work done while at Microsoft Research 2000 ; Bakker and Heskes , 2003 ) . Second , we share information between tasks about how to learn and perform inference using meta-learning ( Naik and Mammone , 1992 ; Thrun and Pratt , 2012 ; Schmidhuber , 1987 ) . Since uncertainty is rife in small datasets , we provide a procedure for metalearning probabilistic inference . Third , we enable fast learning that can flexibly handle a wide range of tasks and learning settings via amortization ( Kingma and Welling , 2014 ; Rezende et al. , 2014 ) . Building on the framework , we propose a new method – VERSA – which substitutes optimization procedures at test time with forward passes through inference networks . This amortizes the cost of inference , resulting in faster test-time performance , and relieves the need for second derivatives during training . VERSA employs a flexible amortization network that takes few-shot learning datasets , and outputs a distribution over task-specific parameters in a single forward pass . The network can handle arbitrary numbers of shots , and for classification , arbitrary numbers of classes at train and test time ( see Section 3 ) . In Section 5 , we evaluate VERSA on ( i ) standard benchmarks where the method sets new state-of-the-art results , ( ii ) settings where test conditions ( shot and way ) differ from training , and ( iii ) a challenging one-shot view reconstruction task . 2 META-LEARNING PROBABILISTIC INFERENCE FOR PREDICTION . We now present the framework that consists of ( i ) a multi-task probabilistic model , and ( ii ) a method for meta-learning probabilistic inference . 2.1 PROBABILISTIC MODEL . Two principles guide the choice of model . First , the use of discriminative models to maximize predictive performance on supervised learning tasks ( Ng and Jordan , 2002 ) . Second , the need to leverage shared statistical structure between tasks ( i.e . multi-task learning ) . These criteria are met by the standard multi-task directed graphical model shown in Fig . 1 that employs shared parameters θ , which are common to all tasks , and task specific parameters { ψ ( t ) } Tt=1 . Inputs are denoted x and outputs y . Training data D ( t ) = { ( x ( t ) n , y ( t ) n ) } Ntn=1 , and test data { ( x̃ ( t ) m , ỹ ( t ) m ) } Mtm=1 are explicitly distinguished for each task t , as this is key for few-shot learning . LetX ( t ) and Y ( t ) denote all the inputs and outputs ( both test and train ) for task t. The joint probability of the outputs and task specific parameters for T tasks , given the inputs and global parameters is : p ( { Y ( t ) , ψ ( t ) } Tt=1| { X ( t ) } Tt=1 , θ ) = T∏ t=1 p ( ψ ( t ) |θ ) Nt∏ n=1 p ( y ( t ) n |x ( t ) n , ψ ( t ) , θ ) Mt∏ m=1 p ( ỹ ( t ) m |x̃ ( t ) m , ψ ( t ) , θ ) . In the next section , the goal is to meta-learn fast and accurate approximations to the posterior predictive distribution p ( ỹ ( t ) |x̃ ( t ) , θ ) = ∫ p ( ỹ ( t ) |x̃ ( t ) , ψ ( t ) , θ ) p ( ψ ( t ) |x̃ ( t ) , D ( t ) , θ ) dψ ( t ) for unseen tasks t . 2.2 PROBABILISTIC INFERENCE . This section provides a framework for meta-learning approximate inference that is a simple reframing and extension of existing approaches ( Finn et al. , 2017 ; Grant et al. , 2018 ) . We will employ point estimates for the shared parameters θ since data across all tasks will pin down their value . Distributional estimates will be used for the task-specific parameters since only a few shots constrain them . Once the shared parameters are learned , the probabilistic solution to few-shot learning in the model above comprises two steps . First , form the posterior distribution over the task-specific parameters p ( ψ ( t ) |x̃ ( t ) , D ( t ) , θ ) . Second , compute the posterior predictive p ( ỹ ( t ) |x̃ ( t ) , θ ) . These steps will require approximation and the emphasis here is on performing this quickly at test time . We will describe the form of the approximation , the optimization problem used to learn it , and how to implement this efficiently below . In what follows we initially suppress dependencies on the inputs x̃ and shared parameters θ to reduce notational clutter , but will reintroduce these at the end of the section . Specification of the approximate posterior predictive distribution . Our framework approximates the posterior predictive distribution by an amortized distribution qφ ( ỹ|D ) . That is , we learn a feed-forward inference network with parameters φ that takes any training dataset D ( t ) and test input x̃ as inputs and returns the predictive distribution over the test output ỹ ( t ) . We construct this by amortizing the approximate posterior qφ ( ψ|D ) and then form the approximate posterior predictive distribution using : qφ ( ỹ|D ) = ∫ p ( ỹ|ψ ) qφ ( ψ|D ) dψ . ( 1 ) This step may require additional approximation e.g . Monte Carlo sampling . The amortization will enable fast predictions at test time . The form of these distributions is identical to those used in amortized variational inference ( Edwards and Storkey , 2017 ; Kingma and Welling , 2014 ) . In this work , we use a factorized Gaussian distribution for qφ ( ψ|D ( t ) ) with means and variances set by the amortization network . However , the training method described next is different . Meta-learning the approximate posterior predictive distribution . The quality of the approximate posterior predictive for a single task will be measured by the KL-divergence between the true and approximate posterior predictive distribution KL [ p ( ỹ|D ) ‖qφ ( ỹ|D ) ] . The goal of learning will be to minimize the expected value of this KL averaged over tasks , φ∗ = arg min φ E p ( D ) [ KL [ p ( ỹ|D ) ‖qφ ( ỹ|D ) ] ] = arg max φ E p ( ỹ , D ) [ log ∫ p ( ỹ|ψ ) qφ ( ψ|D ) dψ ] . ( 2 ) Training will therefore return parameters φ that best approximate the posterior predictive distribution in an average KL sense . So , if the approximate posterior qφ ( ψ|D ) is rich enough , global optimization will recover the true posterior p ( ψ|D ) ( assuming p ( ψ|D ) obeys identifiability conditions ( Casella and Berger , 2002 ) ) .1 Thus , the amortized procedure meta-learns approximate inference that supports accurate prediction . Appendix A provides a generalized derivation of the framework , grounded in Bayesian decision theory ( Jaynes , 2003 ) . The right hand side of Eq . ( 2 ) indicates how training could proceed : ( i ) select a task t at random , ( ii ) sample some training data D ( t ) , ( iii ) form the posterior predictive qφ ( ·|D ( t ) ) and , ( iv ) compute the log-density log qφ ( ỹ ( t ) |D ( t ) ) at test data ỹ ( t ) not included in D ( t ) . Repeating this process many times and averaging the results would provide an unbiased estimate of the objective which can then be optimized . This perspective also makes it clear that the procedure is scoring the approximate inference procedure by simulating approximate Bayesian held-out log-likelihood evaluation . Importantly , while an inference network is used to approximate posterior distributions , the training procedure differs significantly from standard variational inference . In particular , rather than minimizing KL ( qφ ( ψ|D ) ‖p ( ψ|D ) ) , our objective function directly focuses on the posterior predictive distribution and minimizes KL ( p ( ỹ|D ) ‖qφ ( ỹ|D ) ) . End-to-end stochastic training . Armed by the insights above we now layout the full training procedure . We reintroduce inputs and shared parameters θ and the objective becomes : L ( φ ) = − E p ( D , ỹ , x̃ ) [ log qφ ( ỹ|x̃ , θ ) ] = − E p ( D , ỹ , x̃ ) [ log ∫ p ( ỹ|x̃ , ψ , θ ) qφ ( ψ|D , θ ) dψ ] . ( 3 ) We optimize the objective over the shared parameters θ as this will maximize predictive performance ( i.e. , Bayesian held out likelihood ) . An end-to-end stochastic training objective for θ and φ is : L̂ ( θ , φ ) = 1 MT ∑ M , T log 1 L L∑ l=1 p ( ỹ ( t ) m |x̃ ( t ) m , ψ ( t ) l , θ ) , with ψ ( t ) l ∼ qφ ( ψ|D ( t ) , θ ) ( 4 ) 1Note that the true predictive posterior p ( y|D ) is recovered regardless of the identifiability of p ( ψ|D ) . and { ỹ ( t ) m , x̃ ( t ) m , D ( t ) } ∼ p ( ỹ , x̃ , D ) , where p represents the data distribution ( e.g. , sampling tasks and splitting them into disjoint training data D and test data { ( x̃ ( t ) m , ỹ ( t ) m ) } Mtm=1 ) . This type of training therefore uses episodic train / test splits at meta-train time . We have also approximated the integral over ψ using L Monte Carlo samples . The local reparametrization ( Kingma et al. , 2015 ) trick enables optimization . Interestingly , the learning objective does not require an explicit specification of the prior distribution over parameters , p ( ψ ( t ) |θ ) , learning it implicitly through qφ ( ψ|D , θ ) instead . In summary , we have developed an approach for Meta-Learning Probabilistic Inference for Prediction ( ML-PIP ) . A simple investigation of the inference method with synthetic data is provided in Section 5.1 . In Section 4 we will show that this formulation unifies a number of existing approaches , but first we discuss a particular instance of the ML-PIP framework that supports versatile learning .

Summary: This paper proposes both a general meta-learning framework with approximate probabilistic inference, and implements an instance of it for few-shot learning. First, they propose Meta-Learning Probabilistic inference for Prediction (ML-PIP) which trains the meta-learner to minimize the KL-divergence between the approximate predictive distribution generated from it and predictive distribution for each class. Then, they use this framework to implement Versatile Amortized Inference, which they call VERSA. VERSA replaces the optimization for test time with efficient posterior inference, by generating distribution over task-specific parameters in a single forward pass. The authors validate VERSA against amortized and non-amortized variational inference which it outperforms. VERSA is also highly versatile as it can be trained with varying number of classes and shots.