Amortized Variational Inference for Simple Hierarchical Models

1 Introduction . Hierarchical Bayesian models are a general framework where parameters of “ groups ” are drawn from some shared distribution , and then observed data is drawn from a distribution specified by each group ’ s parameters . After data is observed , the inference problem is to infer both the parameters for each group and the shared parameters . These models have proven useful in various domains [ 13 ] including hierarchical regression amd classification [ 12 ] , topic models [ 4 , 22 , 3 ] , polling [ 11 , 24 ] , epidemiology [ 23 ] , ecology [ 8 ] , psychology [ 37 ] , matrix-factorization [ 35 ] , and collaborative filtering [ 26 , 33 ] . A proven technique for scaling variational inference ( VI ) to large datasets is subsampling . The idea is that if the target model has the form p ( z , y ) = p ( z ) ∏ i p ( yi|z ) then an unbiased gradient can be estimated while only evaluating p ( z ) and p ( yi|z ) at a few i [ 29 , 16 , 20 , 31 , 30 , 36 , 15 ] . This paper addresses hierarchical models of the form p ( θ , z , y ) = p ( θ ) ∏ i p ( zi , yi|θ ) , where only y is observed . There are two challenges . First , the number of local latent variables zi increases with the dataset , meaning the posterior distribution increases in dimensionality . Second , there is often a dependence between zi and θ which must be captured to get strong results [ 15 , 18 ] . The aim of this paper is to develop a black-box variational inference scheme that can scale to large hierarchical models without losing benefits of a joint approximation . Our solution takes three steps . First , in the true posterior , the different latent variables zi are conditionally independent given θ , which suggests using a variational family of the same form . We confirm this intuition by showing that for any joint variational family q ( θ , z ) , one can define a corresponding `` branch '' family q ( θ ) ∏ i q ( zi|θ ) such that inference will be equally accurate ( theorem 2 ) . We call inference using such a family the `` branch '' approach . Second , we observe that if using the branch approach , the optimal local variational parameters can be computed only from θ and local data ( eq . ( 12 ) ) . Thus , we propose to amortize the computation of the local variational parameters by learning a network to approximately solve that optimization . We 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . show that when the target distribution is symmetric over latent variables , this will be as accurate as the original joint family , assuming a sufficiently capable amortization network ( claim 5 ) . Third , we note that in many real hierarchical models , there are many i.i.d . data generated from each local latent variable . This presents a challenge for learning an amortization network , since the full network should deal with different numbers of data points and naturally reflect the symmetry between the inputs ( that is , without having to relearn the symmetry . ) We propose an approach where a preliminary `` feature '' network processes each datum , after which they are combined with a pooling operation which forms the input for a standard network ( section 6 ) . This is closely related to the `` deep sets '' [ 39 ] strategy for permutation invariance . We validate these methods on a synthetic model where exact inference is possible , and on a userpreference model for the MovieLens dataset with 162K users who make 25M ratings of different movies . At small scale ( 2.5K ratings ) , we show similar accuracy using a dense joint Gaussian , a branch distribution , or our amortized approach . At moderate scale ( 180K ratings ) , joint inference is intractable . Branch distributions gives a meaningful answer , and the amortized approach is comparable or better . At large scale ( 18M ratings ) the amortized approach is thousands of nats better on test-likelihoods even after branch distributions were trained for almost ten times as long as the amortized approach took to converge ( fig . 6 ) . 2 Hierarchical Branched Distributions We focus on two-level hierarchical distributions . A generic model of this type is given by p ( θ , z , y|x ) = p ( θ ) N∏ i=1 p ( zi|θ ) p ( yi|θ , zi , xi ) , ( 1 ) where θ and z = { zi } Ni=1 are latent variables , y = { yi } Ni=1 are observations , and x = { xi } Ni=1 are covariates . As the visual representations of these models resemble branches , we refer them as hierarchical branch distributions ( HBDs ) . Symmetric . We call an HBD symmetric if the conditionals are symmetric , i.e. , if zi = zj , xi = xj , and yi = yj , it implies that p ( zi|θ ) = p ( zj |θ ) , and p ( yi|θ , zi , xi ) = p ( yj |θ , zj , xj ) . ( 2 ) Locally i.i.d . Often local observations yi ( and xi ) are a collection of conditionally i.i.d observations . Then , an HBD takes the form of p ( θ , z , y|x ) = p ( θ ) N∏ i=1 p ( zi|θ ) ni∏ j=1 p ( yij |θ , zi , xij ) , ( 3 ) where yi = { yij } nij=1 and xi = { xij } ni j=1 are collections of conditionally i.i.d observations and covariates ; ni ≥ 1 is the number of observations for branch i . No local covariates . Some applications do not involve the covariates xi . In such cases , HBDs have a simplified form of p ( θ , z , y ) = p ( θ ) N∏ i=1 p ( zi|θ ) p ( yi|θ , zi ) . ( 4 ) In this paper , we will be using eq . ( 1 ) and eq . ( 3 ) to refer HBDs—the results extend easily to case where there are no local covariates . ( For instance , in section 5 , we amortize using ( xi , yi ) as inputs . When there are no covariates , we can amortize with just yi . ) 2.1 Related Work . Bayesian inference in hierarchical models is an old problem . The most common solutions are Markov chain Monte Carlo ( MCMC ) and VI . A key advantage of VI is that gradients can sometimes be estimated using only a subsample of data [ 29 ] . Hoffman et al . [ 16 ] observe that inference in hierarchical models is still slow at large scale , since the number of parameters scales with the dataset . Instead , they assume that θ and zi from eq . ( 1 ) are in conjugate exponential families , and observe that for a mean-field variational distribution q ( θ ) ∏ i q ( zi ) , the optimal q ( zi ) can be calculated in closed form for fixed q ( θ ) . This is highly scalable , though it is limited to factorized approximations and requires a conditionally conjugate target model . A structured variational approximation like q ( θ ) ∏ i q ( zi|θ ) can be used which reflects the dependence of zi on θ [ 15 , 34 , 2 , 18 ] . However , this still has scalability problems in general since the number of parameters grows in the size of the data ( section 7 ) . To the best of our knowledge , the only approach that avoids this is the framework of structured stochastic VI [ 15 , 18 ] , which assumes the target is conditionally conjugate , and that for a fixed θ an optimal `` local '' distribution q ( zi|θ ) can be calculated from local data . Hoffman and Blei [ 15 ] address matrix factorization models and latent Dirichlet allocation , using Gibbs sampling to compute the local distributions . Johnson et al . [ 18 ] use amortization for conjugate models but do not consider the setting where local observations are a collection of i.i.d observations . Our approach is not strictly an instance of either of these frameworks , as we do not assume conjugacy or that amortization can exactly recover optimal local distributions [ 15 , Eq . 7 ] . Still the spirit is the same , and our approach should be seen as part of this line of research . Amortized variational approximations have been used to learn models with local variables [ 20 , 10 , 17 , 5 ] . A particularly related instance of model learning is the Neural Statistician approach [ 10 ] , where a “ statistics network ” learns representations of closely related datasets–the construction of this network is similar to our “ pooling network ” in section 6 . However , in the Neural Statistician model there are no global variables θ , and it is not obvious how to generalize their approach for HBDs . In contrast , our “ pooling network ” results from the analysis in section 5 , reducing the architectural space when θ is present while retaining accuracy . Moreover , learning a model is strictly different from our “ black-box '' setting where we want to approximate the posterior of a given model , and the inference only has access to log p or ∇θ , z log p ( or their parts ) [ 30 , 21 , 2 , 1 ] . 3 Joint approximations for HBD . When dealing with HBDs , a non-structured distribution does not scale . To see this , consider the naive VI objective—Evidence Lower Bound ( ELBO , L ) . Let qJointϕ be a joint variational approximation over θ and z ; we use sans-serif font for random variables . Then , L ( qJointϕ ∥p ) = EqJointϕ ( θ , z ) [ log p ( θ , z , y|x ) qJointϕ ( θ , z ) ] , ( 5 ) where ϕ are variational parameters . Usually the above expectation is not tractable and one uses a Monte-Carlo estimator . A single sample estimator is given by L̂ = log p ( θ , z , y|x ) qJointϕ ( θ , z ) , ( 6 ) where ( θ , z ) ∼ qJointϕ and L̂ is unbiased . Since qJointϕ need not factorize , even a single estimate requires sampling all the latent variables at each step . This is problematic when there are a large number of local latent variables . One encounters the same scaling problem when taking the gradient of the ELBO . We can estimate the gradient using any of the several available estimators [ 30 , 20 , 31 , 32 ] ; however , none of them scale if qJointϕ does not factorize [ 16 ] . 4 Branch approximations for HBD . It is easy to see that the posterior distribution of the HBD eq . ( 1 ) takes the form p ( θ , z|y , x ) = p ( θ|y , x ) N∏ i=1 p ( zi|θ , yi , xi ) . As such , it is natural to consider variational distributions that factorize the same way ( see fig . 2 . ) In this section , we confirm this intuition—we start with any joint variational family which can have any dependence between θ and z1 , · · · , zN . Then , we define a corresponding “ branch ” family where z1 , · · · , zN are conditionally independent given θ . We show that inference using the branch family will be at least as accurate as using the joint family . We formalize the idea of branch distribution in the next definition . Definition 1 . Let qJointϕ be any variational family with parameters ϕ . We define qBranchv , w to be a corresponding branch family if , for all ϕ , there exists ( v , { wi } Ni=1 ) , such that , qBranchv , w ( θ , z ) = qv ( θ ) N∏ i=1 qwi ( zi|θ ) = qJointϕ ( θ ) N∏ i=1 qJointϕ ( zi|θ ) . ( 7 ) θ z1 z2 zN . . . Figure 2 : qBranchv , w as in eq . ( 7 ) . Given a joint distribution , a branch distribution can always be defined by choosing w and vi as the components of ϕ that influence qJointϕ ( θ ) and qJointϕ ( zi|θ ) , respectively , and choosing qw ( θ ) and qwi ( zi|θ ) correspondingly . However , the choice is not unique ( for instance , the parameterization can require transformations—different transformations can create different variants . ) The idea to use qBranchv , w is natural [ 15 , 2 ] . However , one might question if the branch variational family is as good as the original qJointϕ ; in theorem 2 , we establish that this is indeed true . Theorem 2 . Let p be a HBD , and qJointϕ ( θ , z ) be a joint approximation family parameterized by ϕ . Choose a corresponding branch variational family qBranchv , w ( θ , z ) as in definition 1 . Then , min v , w KL ( qBranchv , w ∥p ) ≤ min ϕ KL ( qJointϕ ∥p ) . We stress that qBranchv , w is a new variational family derived from but not identical to q Joint ϕ . Theorem 2 implies that we can optimize a branched variational family qBranchv , w without compromising the quality of approximation ( see appendix D for proof . ) In the following corollary , we apply theorem 2 to a joint Gaussian to show that using a branch Gaussian will be equally accurate . Corrolary 3 . Let p be a HBD , and let qJointϕ ( θ , z ) = N ( ( θ , z ) |µ , Σ ) be a joint Gaussian approximation ( with ϕ = ( µ , Σ ) ) . Choose a variational family qBranchv , w ( θ , z ) = N ( θ|µ0 , Σ0 ) N∏ i=1 N ( zi|µi +Aiθ , Σi ) with v = ( µ0 , Σ0 ) and wi = ( µi , Σi , Ai ) . Then , min v , w KL ( qBranchv , w ∥p ) ≤ min ϕ KL ( qJointϕ ∥p ) . In the above corollary , the structured family qBranchv , w is chosen such that it can represent any branched Gaussian distribution . Notice , the mean of the conditional distribution is an affine function of θ . This affine relationship appears naturally when you factorize the joint Gaussian over ( θ , z ) . For more details see appendix F .

Summary: An approach vor variational inference in large-scale hierarchical models is presented that is more efficient than the base lines. This is achieved using shared parameters (amortization) and a feature pooling network. Experiments confirm the efficiency of the approach on synthetic data as well as the MovieLens dataset.