Meta-Learning for Variational Inference

1 INTRODUCTION . Bayesian inference provides a powerful tool for probabilistic modelling of data , however , exact Bayesian inference is often intractable . Therefore practical Bayesian inference often resort to approximations , and various approximate inference methods have been developed . Variational Inference ( VI ) ( Jordan et al. , 1999 ; Zhang et al. , 2018 ) approximates the intractable target distribution through optimization of a tractable distribution . Compared to Markov chain Monte Carlo ( MCMC ) , VI is biased but more computationally efficient , making it particularly suitable to large-scale models in deep learning such as Bayesian neural networks and deep generative models . VI approximates the target distribution by minimizing a divergence objective . Different divergence metrics essentially define a different inference algorithm which leads to different properties of the approximation . Therefore the selection of this divergence is one of the crucial factors of making VI successful . The most widely used divergence measure is KL ( q||p ) where p is the target distribution and q is the approximate distribution . However , using this KL divergence for VI has been criticized for under-estimating the uncertainty ( Bishop , 2006 ; Blei et al. , 2017 ; Wang et al. , 2018a ) , which leads to poor model performance when uncertainty estimation is essential . Many alternative divergence measures have been proposed for VI to alleviate this issue ( Minka et al. , 2005 ; Hernández-Lobato et al. , 2016 ; Li & Turner , 2016 ; Csiszár et al. , 2004 ; Bamler et al. , 2017 ; Wang et al. , 2018a ) , which provide better bias and variance trade-offs and lead to better predictive results with more accurate uncertainty estimation . However , as illustrated by Figure 1 , this line of work has also shown that the optimal divergence can vary depending on tasks and an undesired divergence objective can lead to mediocre performance ( Li & Turner , 2016 ; Depeweg et al. , 2016 ) . Unfortunately , choosing a suitable divergence objective for a specific task is challenging as it requires a thorough understanding of ( i ) the shape of the target distribution ; ( ii ) the desirable properties of the approximate distribution ( e.g . mass-covering or mode-seeking ) ; and ( iii ) the bias-variance trade-off between the tightness and the variance of the Monte Carlo estimate of the variational bound . A crucial question remains to be addressed in order to make VI a success for a wide range of applications : can we automatically choose a suitable divergence which are tailored to specific type of tasks ? To answer this question , we propose meta-learning for variational inference ( meta-VI ) which utilizes the advantages of meta-learning to improve approximate Bayesian inference . Meta-learning is to design a learner based on several training tasks that can generalize well to future tasks ( Naik & Mammone , 1992 ; Thrun & Pratt , 2012 ; Hochreiter et al. , 2001 ) . Our meta-VI learns an inference algorithm that is tailored to the problem of interest . Additionally , meta-VI can provide a good initialization of the variational parameters which reduces the training time remarkably . We summarize our contributions as the following : • We developed a general framework for meta-learning variational inference ( Section 3.1 ) , which can choose divergence objective automatically by meta-learning . Specifically , we derive meta-VI first for α-divergence based VI , then extend it to f -divergence family and provided a novel parameterization of f -divergence which is then used in meta-VI . • In addition to meta-learning the divergence objective , we further combine meta-VI with meta-learning parameters for the variational distribution ( Section 3.2 ) . This method provides initialization for fast adaptation and automatically learns the divergence to optimize . • We demonstrate the efficacy of our approach on various experimental settings ( Section 4 ) . On both the Gaussian mixture distribution approximation and Bayesian neural network regression tasks , meta-VI significantly out-performs all baseline methods . On a large-scale , real-world example , meta-VI also improves performance when applied to partial variational auto-encoder based recommender systems . 2 PRELIMINARIES . Considering a dataset D = { xn } Nn=1 and a θ-parameterized model . Bayesian inference requires computing the posterior over θ given the datasetD : p ( θ|D ) = p ( D|θ ) p ( θ ) p ( D ) . The exact posterior p ( θ|D ) is generally intractable , therefore one needs to resort to approximation with a tractable approximate posterior q ( θ ) ≈ p ( θ|D ) . Typically the approximation posterior q ( θ ) is obtained by minimizing a divergence , e.g . Variational Inference ( VI ) applications often minimize KL ( q ( θ ) ||p ( θ|D ) ) . This turns Bayesian inference into an optimization task ( divergence minimization ) . In practice , due to the intractability of p ( D ) , VI alternatively maximizes an equivalent objective called the variational lower bound : LVI = Eq [ log p ( D , θ ) q ( θ ) ] = log p ( D ) − KL ( q||p ) ( 1 ) Renyi ’ s α-divergence α-divergence is a rich family that includes many common divergence as special cases ( Minka , 2001 ; Hernández-Lobato et al. , 2016 ; Li & Turner , 2016 ) . In this paper , we focus on Renyi ’ s definition ( Rényi et al. , 1961 ; Li & Turner , 2016 ) rather than the others ( Amari , 2012 ; Tsallis , 1988 ) due to its unified expression of gradient for all finite α . The benefit of having a unified expression will be seen in Section 3.1 . Renyi ’ s α-divergence is defined by Dα ( p||q ) = 1 α− 1 log ∫ p ( θ ) αq ( θ ) 1−αdθ , α > 0 , α 6= 1 , ( 2 ) whereDα ( p||q ) → KL ( p||q ) when α→ 1 . Similar to maximizing the variational lower bound when the divergence objective is KL ( q||p ) , one can maximize the variational Renyi bound ( VR bound ) derived from Renyi ’ s α-divergence : Lα ( q ; D ) = 1 1− α logEq [ ( p ( θ , D ) q ( θ ) ) 1−α ] = log p ( D ) −Dα ( q||p ) ( 3 ) The reparameterization trick ( Salimans et al. , 2013 ; Kingma & Welling , 2013 ) is commonly used in practice for gradient ascent based optimization of the VR bound Eq . ( 3 ) , where sampling θ ∼ qφ ( θ ) is conducted by first sampling ∼ p ( ) from a simple distribution independent with the variational parameter φ ( e.g . Gaussian ) then parameterizing θ = hφ ( ) . Using the reparameterization trick , the VR bound Eq . ( 3 ) becomes Lα ( qφ ; D ) = 11−α logE [ ( p ( hφ ( ) , D ) qφ ( hφ ( ) ) ) 1−α ] . The expectation is usually computed by Monte Carlo ( MC ) approximation . The gradient of VR bound w.r.t . φ after MC approximation with K particles is ∇φLα ( qφ ; x ) = K∑ k=1 [ wα , k∇φ log p ( hφ ( k ) , x ) q ( hφ ( k ) ) ] ( 4 ) where wα , k = ( p ( hφ ( k ) , x ) q ( hφ ( k ) ) ) 1−α/∑K k=1 [ ( p ( hφ ( k ) , x ) q ( hφ ( k ) ) ) 1−α ] . When α = 0 the weights wα , k = 1/K and the corresponding gradient ( 4 ) becomes an unbiased estimate of the gradient of the variational lower bound ( 1 ) . As shown in Figure 1 , approximate inference with different α-divergences results in different variational distributions . Minka et al . ( 2005 ) and Li & Turner ( 2016 ) also showed the optimal αdivergence may vary for different tasks and datasets , and in practice it is difficult to choose an optimal α-divergence a priori . f -divergence f -divergence defines a more general family of divergences ( Csiszár et al. , 2004 ; Minka et al. , 2005 ) . It can be defined using a twice differentiable convex function f : R+ → R ( Csiszár et al. , 2004 ) : Df ( p||qφ ) = Eθ∼qφ [ f ( p ( θ ) qφ ( θ ) ) − f ( 1 ) ] . ( 5 ) This family includes KL-divergences in both directions which can be seen by taking f ( t ) = − log t for KL ( q||p ) and f ( t ) = t log t for KL ( p||q ) . It also contains α-divergences which takes f ( t ) = tα α ( α−1 ) for α ∈ R\ { 0 , 1 } . Although f -divergence family is very rich due to the usage of arbitrary twice differentiable convex function , it requires significant expertise to design a suitable f function for a specific task . 3 META-VI . The goal of meta-learning a variational inference algorithm is to learn a divergence objective , so that the resulting VI algorithm produces an approximate distribution with desired properties on a certain type of tasks . To achieve this , we first construct a learnable divergence family , then design a meta-loss function that gives guidance for updating the divergence . Assume we have M training tasks T1 , ... , TM sampled from an underlying task distribution p ( T ) . Each task has its own probabilistic model pTi ( θi , DTi ) . Let Dη ( ·||· ) be the learnable divergence parameterized by η , then for each task the approximate posterior qφi ( θi ) is computed by minimizing Dη ( pTi ( θi|DTi ) ||qφi ( θi ) ) . In the rest of the paper we also write Dη ( qφi , Ti ) = Dη ( pTi ( θi|DTi ) ||qφi ( θi ) ) for brevity . During meta-training , we define a meta-loss function J ( qφi , Ti ) which is optimized w.r.t . the divergence parameter η . This meta-loss function is designed to evaluate the desired properties of the approximate distribution , e.g . test log-likelihood . During meta-testing , a new task is sampled from p ( T ) , and the learned divergence Dη is used to optimize the variational distribution . Besides the above setting , we also consider a few-shot learning set-up similar to the model-agnostic meta-learning ( MAML ) framework ( Finn et al. , 2017 ; 2018 ; Kim et al. , 2018 ) . In this case the probabilistic model architecture is shared by all the tasks , and the goal of meta-training is to obtain a divergence as well as an initialization of the variational parameters φ for unseen tasks . After metatraining , we will train the model with the learned divergence objective and the learned initialization of variational parameters φ on new tasks . The above meta-learning settings are practical as demonstrated in many previous work ( Finn et al. , 2017 ; 2018 ; Kim et al. , 2018 ) , including meta-learning for Bayesian inference ( Gong et al. , 2018 ) . Attaining common knowledge based on the previous tasks has been proved to be useful for the future tasks . Algorithm 1 Meta-D Input : M : number of training tasks . β , γ : learning rate hyperparameters . Initialize η , φi , i = 1 , . . . , M ( φi can have different structures ) . loop for Ti , i = 1 , . . . , M do for B times do Update the variational parameters with the current divergence : φi ← φi − β∇φiDη ( qφi , Ti ) . end for end for Update η ← η − γ∇η 1M ∑ i J ( qφi , Ti ) end loop Output : η Algorithm 2 Meta-D & φ Input : p ( T ) : distribution over tasks . β , γ , τ : learning rate hyperparameters . Initialize φ , η loop Sample M tasks Ti ∼ p ( T ) . for all Ti do Update the variational parameters with the current divergence : φi ← φ − β∇φDη ( qφ , Ti ) . end for Update φ← φ−τ∇φ 1M ∑ i J ( qφi , Ti ) ; η ← η − γ∇η 1M ∑ i J ( qφi , Ti ) end loop Output : η , φ

Summary: This paper proposes to use a meta-learning approach to learn the divergence used in variational inference and initial variational parameters. It considers two families of learnable divergences, including the alpha-divergence and f-divergence, and proposes two double loop algorithms. In both algorithms, the inner loop adjusts the variational parameters for a specific variational inference task. The outer loop adjusts the parameters in the divergence (and perhaps an initial variational parameter for all tasks) in terms of a human-designed meta-loss. Proposition 1 (following Wang et al. 2018a) shows that the gradients with respect to the meta-parameters in the f-divergence can be obtained using f''. This makes it convenient to parameterize f'' by neural networks to implicitly model f.