Uncertainty in Multitask Transfer Learning

1 INTRODUCTION . Recently , significant progress has been made to scale Bayesian neural networks to large tasks and to provide better approximations of the posterior distribution ( Blundell et al. , 2015 ) . Recent works extend fully factorized posterior distributions to more general families ( Louizos and Welling , 2017 ; Krueger et al. , 2017 ; Sun et al. , 2017 ) . It is also possible to sample from the posterior distribution trough mini-batch updates ( Mandt et al. , 2017 ; Zhang et al. , 2017 ) . However , for neural networks , the prior is often chosen for convenience . This may become a problem when the number of observations is insufficient to overcome the choice of the prior . In this regime , the prior must express our current knowledge on the task and , most importantly , our lack of knowledge on it . In addition to that , a good approximation of the posterior under the small sample size regime is required , including the ability to model multiple modes . This is indeed the case for Bayesian optimization ( Snoek et al. , 2012 ) , Bayesian active learning ( Gal et al. , 2017 ) , continual learning ( Kirkpatrick et al. , 2017 ) , safe reinforcement learning ( Berkenkamp et al. , 2017 ) , exploration-exploitation trade-off in reinforcement learning ( Houthooft et al. , 2016 ) . Gaussian processes ( Rasmussen , 2004 ) have historically been used for these applications , but an RBF kernel constitute a prior that is unsuited for many tasks . More recent tools such as deep Gaussian processes ( Damianou and Lawrence , 2013 ) show great potential and yet their scalability whilst learning from multiple tasks needs to be improved . Our contributions are as follow : 1 . We provide a simple and scalable procedure to learn an expressive prior and posterior over models from multiple tasks . 2 . We reach state of the art performances on mini-imagenet . 3 . We propose two new benchmarks , each exposing a failure mode of popular meta learning algorithms . In contrast , our method perform well on these benchmarks . • MAML ( Finn et al. , 2017 ) does not perform well on a collection of sinus tasks when the frequency varies . • Prototypical Network ( Snell et al. , 2017 ) ’ s performance decrease considerably when the diversity of tasks increases . Outline : We first describe the proposed approach in Section 2 . In Section 3 , we extend to three level of hierarchies and obtain a model more suited for classification . Section 4 review related methods and outline the key differences . Finally , In Section 5 , we conduct experiments on three different benchmarks to gain insight in the behavior of our algorithm . 2 LEARNING A DEEP PRIOR . By leveraging the variational Bayesian approach , we show how we can learn a prior over models with neural networks . We start our analysis with the goal of learning a prior p ( w|α ) over the weights w of neural networks across multiple tasks . We then provide a reduction of the Evidence Lower BOund ( ELBO ) showing that it is not necessary to explicitly model a distribution in the very high dimension of the weight space of neural networks . Instead the algorithm learns a subspace suitable for expressing model uncertainty within the distributions of tasks considered in the multi-task environment . This simplification results in a scalable algorithm which we refer to as deep prior . 2.1 HIERARCHICAL BAYES . To learn a probability distribution p ( w|α ) over the weights w of a network parameterized by α , we use a hierarchical Bayes approach across N tasks , with hyper-prior p ( α ) . Each task has its own parameters wj , withW = { wj } Nj=1 . Using all datasets D = { Sj } Nj=1 , we have the following posterior:1 p ( W , α|D ) = p ( α|D ) ∏ j p ( wj |α , Sj ) ∝ p ( D|W ) p ( W|α ) p ( α ) = ∏ j ∏ i p ( yij |xij , wj ) p ( wj |α ) p ( α ) , The term p ( yij |xij , wj ) corresponds to the likelihood of sample i of task j given a model parameterized by wj e.g . the probability of class yij from the softmax of a neural network parameterized by wj with input xij . For the posterior p ( α|D ) , we assume that the large amount of data available across multiple tasks will be enough to overcome a generic prior p ( α ) , such as an isotropic Normal distribution . Hence , we consider a point estimate of the posterior p ( α|D ) using maximum a posteriori2 . We can now focus on the remaining term : p ( wj |α ) . Since wj is potentially high dimensional with intricate correlations among the different dimensions , we can not use a simple Gaussian distribution . Following inspiration from generative models such as GANs ( Goodfellow et al. , 2014 ) and VAE ( Kingma and Welling , 2013 ) , we use an auxiliary variable z ∼ N ( 0 , Idz ) and a deterministic function projecting the noise z to the space of w i.e . w = hα ( z ) . Marginalizing z , we have : p ( w|α ) =∫ z p ( z ) p ( w|z , α ) dz = ∫ z p ( z ) δhα ( z ) −wdz , where δ is the Dirac delta function . Unfortunately , directly marginalizing z is intractable for general hα . To overcome this issue , we add z to the joint inference and marginalize it at inference time . Considering the point estimation of α , the full posterior is factorized as follows : ∏N j=1p ( wj , zj |α , Sj ) ( 1 ) = ∏N j=1 p ( wj |zj , α , Sj ) p ( zj |α , Sj ) ∝ ∏N j=1 p ( wj |zj , α ) p ( zj ) ∏nj i=1 p ( yij |xij , wj ) , where p ( yij |xij , wj ) is the conventional likelihood function of a neural network with weight matrices generated from the function hα i.e . : wj = hα ( zj ) . Similar architecture has been used in Krueger et al . ( 2017 ) and Louizos and Welling ( 2017 ) , but we will soon show that it can be reduced to a simpler architecture in the context of multi-task learning . The other terms are defined as follows : p ( zj ) = N ( 0 , I ) ( 2 ) p ( zj , wj |α ) = p ( zj ) δhα ( zj ) −wj ( 3 ) p ( zj , wj |α , Sj ) = p ( zj |α , Sj ) δhα ( zj ) −wj ( 4 ) The task will consist of jointly learning a function hα common to all tasks and a posterior distribution p ( zj |α , Sj ) for each task . At inference time , predictions are performed by marginalizing z i.e . : p ( y|x , D ) = E zj∼p ( zj |α , Sj ) p ( y|x , hα ( zj ) ) . 1p ( xij ) cancelled with itself from the denominator since it does not depend on wj nor α . This would have been different for a generative approach . 2This can be done through simply minimizing the cross entropy of a neural network with L2 regularization . 2.2 REDUCTION OF THE EVIDENCE LOWER BOUND .. In the previous section , we described the different components for expressing the posterior distribution of Equation 4 . While all these components are tractable , the normalization factor is still intractable . To address this issue , we follow the Variational Bayes approach ( Blundell et al. , 2015 ) . Conditioning on α , we saw in Equation 1 that the posterior factorizes independently for all tasks . This reduces the joint Evidence Lower BOund ( ELBO ) to a sum of individual ELBO for each task . Given a family of distributions qθj ( zj |Sj , α ) , parameterized by { θj } Nj=1 and α , the Evidence Lower Bound for task j is : ln p ( Sj ) ≥ E q ( zj , wj |Sj , α ) nj∑ i=1 ln p ( yij |xij , wj ) −KLj = E qθj ( zj |Sj , α ) nj∑ i=1 ln p ( yij |xij , hα ( zj ) ) −KLj ( 5 ) = ELBOj , where , KLj = KL [ q ( zj , wj |Sj , α ) ‖ p ( zj , wj |α ) ] = E qθj ( zj |Sj , α ) E q ( wj |zj , α ) ln qθj ( zj |Sj , α ) p ( zj |α ) δhα ( zj ) −wj δhα ( zj ) −wj = E qθj ( zj |Sj , α ) ln qθj ( zj |Sj , α ) p ( zj |α ) ( 6 ) = KL [ qθj ( zj |Sj , α ) ∥∥ p ( zj |α ) ] Notice that after simplification3 , KLj is no longer over the space of wj but only over the space zj . Namely , the posterior distribution is factored into two components , one that is task specific and one that is task agnostic and can be shared with the prior . This amounts to finding a low dimensional manifold in the parameter space where the different tasks can be distinguished . Then , the posterior p ( zj |Sj , α ) only has to model which of the possible tasks are likely , given observations Sj instead of modeling the high dimensional p ( wj |Sj , α ) . But , most importantly , any explicit reference to w has now vanished from both Equation 5 and Equation 6 . This simplification has an important positive impact on the scalability of the proposed approach . Since we no longer need to explicitly calculate the KL on the space of w , we can simplify the likelihood function to p ( yij |xij , zj , α ) , which can be a deep network parameterized by α , taking both xij and zj as inputs . This contrasts with the previous formulation , where hα ( zj ) produces all the weights of a network , yielding an extremely high dimensional representation and slow training . 2.3 POSTERIOR DISTRIBUTION . For modeling qθj ( zj |Sj , α ) , we can use N ( µj , σj ) , where µj and σj can be learned individually for each task . This , however limits the posterior family to express a single mode . For more flexibility , we also explore the usage of more expressive posterior , such as Inverse Autoregressive Flow ( IAF ) ( Kingma et al. , 2016 ) or Neural Autoregressive Flow ( Huang et al. , 2018 ) . This gives a flexible tool for learning a rich variety of multivariate distributions . In principle , we can use a different IAF for each task , but for memory and computational reasons , we use a single IAF for all tasks and we condition4 on an additional task specific context cj . Note that with IAF , we can not evaluate qθj ( zj |Sj , α ) for any values of z efficiently , only for these which we just sampled , but this is sufficient for estimating the KL term with a Monte-Carlo approxi- 3We can justify the cancellation of the Dirac delta functions by instead considering a Gaussian with finite variance , . For all > 0 , the cancellation is valid , so letting → 0 , we recover the result . 4We follow the architecture proposed in Kingma et al . ( 2016 ) . mation i.e . : KLj ≈ 1 nmc nmc∑ i=1 ln qθj ( z ( i ) j |Sj , α ) − lnN ( z ( i ) j ∣∣∣0,1 ) , where z ( i ) j ∼ qθj ( zj |Sj , α ) . It is common to approximate KLj with a single sample and let the mini-batch average the noise incurred on the gradient . We experimented with nmc = 10 , but this did not significantly improve the rate of convergence .

Summary: The paper presents a method for training a probabilistic model for Multitask Transfer Learning. The key idea is to introduce a latent variable "z" per task which to capture the commonality in the task instances. Since this leads to an intractable likelihood the authors use the standard ELBO with a Variational Distribution over "z" defined as a Gaussian + Inverse Autoregressive Flow. For classification, the authors also show that they can combine the model with the main idea in Prototypical Networks.