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: The paper deals with scalable approximate inference for hierarchical models and proposes a framework for hierarchical approximations that exhibit the same factorization as the exact posterior of the hierarchical model. This efficiently enables subsampling the data when evaluating the evidence lower bound and its gradients. The authors further propose a pooling-based strategy for amortized inference to avoid the usual linear scaling in the number of approximation parameters.  The paper is concluded with a set of numerical experiments with both synthetic and real data demonstrating the merits of the proposed framework.