Nested Variational Inference

1 Introduction . Deep generative models provide a mechanism for incorporating priors into methods for unsupervised representation learning . This is particularly useful in settings where the prior defines an inductive bias that reflects the structure of the underlying domain . Training models with structured priors , however , poses some challenges . A standard strategy for training deep generative models is to maximize a reparameterized evidence lower bound with respect to both the generative and an inference model [ Kingma , Welling , 2013 ; Rezende et al. , 2014 ] . This approach works well in variational autoencoders with isotropic Gaussian priors , but often fails for models with more structured priors or likelihoods . In recent years , a range of strategies for improving upon standard reparameterized variational inference have been put forward . These include wake-sleep style variational methods that minimize the forward KL-divergence [ Bornschein , Bengio , 2015 ; Le et al. , 2019 ] as well as sampling schemes that incorporate annealing [ Huang et al. , 2018 ] , Sequential Monte Carlo [ Le et al. , 2018 ; Naesseth et al. , 2017 ; Maddison et al. , 2017 ] , Gibbs sampling [ Wu et al. , 2019 ; Wang et al. , 2018 ] , and MCMC updates [ Salimans et al. , 2015 ; Hoffman , 2017 ; Li et al. , 2017 ] . While these methods offer flexible inference , typically resulting in better approximations to the posterior compared to traditional variational inference methods , they are either model-specific , requiring specialized sampling schemes and gradient estimators , or can not be easily composed with other techniques . In this paper , we propose nested variational inference , a framework for combining nested importance sampling and variational inference . Nested importance sampling formalizes the construction of proposals by way of calls to other importance samplers [ Naesseth et al. , 2015 ; Naesseth et al. , 2019 ] , Many existing importance samplers are instances of nested samplers , including methods based on annealed importance sampling [ Neal , 2001 ] and sequential Monte Carlo [ Del Moral et al. , 2006 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . NVI learns proposals by optimizing a divergence at each level of nesting . Additionally , we combine nested variational objectives with importance resampling to further improve sampling quality , without the need to undergo extra steps to maintain differentiability due to the local nature of the objective . Resampling also allows to compute gradient estimates based on incremental weights , which depend only on variables that are sampled locally , rather than on all variables in the model . Doing so yields lower variance weights , whilst maintaining a high sample diversity relative to existing methods . 2 Background . 2.1 Stochastic Variational Inference .. Stochastic variational methods approximate a target density π ( z ; θ ) = γ ( z ; θ ) /Z with parameters θ using a variational density q ( z ; φ ) with parameters φ . Two common variational objectives are the forward and reverse Kullback-Leibler ( KL ) divergence , which are both instances of f -divergences Df ( π ∥∥ q ) = E q [ f ( π ( z ; θ ) q ( z ; φ ) ) ] , ( 1 ) with f ( w ) = w logw and f ( w ) = − logw respectively . We are typically interested in the setting where π ( z ; θ ) is the posterior pθ ( z |x ) of a model with latent variables z and observations x . In this case , γ ( z ; θ ) = pθ ( x , z ) is the joint density of the model , and Z = pθ ( x ) is the marginal likelihood . Reverse KL-divergence . When optimizing KL ( q ‖π ) , known as the reverse or exclusive KL divergence , it is common practice to maximize a lower boundL = Eq [ log ( γ/q ) ] = logZ−KL ( q ‖π ) with respect to θ and φ . The gradient of L can be approximated using reparameterized samples [ Kingma , Welling , 2013 ; Rezende et al. , 2014 ] , likelihood-ratio estimators [ Wingate , Weber , 2013 ; Ranganath et al. , 2014 ] , or a combination of the two [ Schulman et al. , 2015 ; Ritchie et al. , 2016a ] . Forward KL-divergence . In the case of the forward divergence KL ( π ‖ q ) , also known as the inclusive KL divergence , stochastic variational methods typically optimize separate objectives for the inference and generative model . A common strategy is to train the generative model by either maximizing the ELBO or likelihood on the data . Optimizing the inference model requires samples from π , which itself requires approximate inference . A common strategy , which was popularized in the context of reweighted wake-sleep ( RWS ) methods [ Bornschein , Bengio , 2015 ; Le et al. , 2019 ] , is to use q as a proposal in an importance sampler . 2.2 Importance Sampling . Self-Normalized Importance Samplers . An expectation Eπ [ g ( z ) ] with respect to π can be rewritten with respect to a proposal density q by introducing an unnormalized importance weight w , E π [ g ( z ) ] = 1 Z E q [ w g ( z ) ] , w = γ ( z ; θ ) q ( z ; φ ) . ( 2 ) Self-normalized estimators use weighted samples ( w ( s ) , z ( s ) ) Ss=1 to both approximate the expectation with respect to q , and to compute an estimate Ẑ of the normalizing constant , E π [ g ( z ) ] ' ĝ = 1 Ẑ 1 S S∑ s=1 w ( s ) g ( z ( s ) ) , Ẑ = 1 S S∑ s=1 w ( s ) , z ( s ) ∼ q ( · ; φ ) . ( 3 ) The resulting estimator is consistent , i.e . ĝ a.s.−→ Eπ [ g ( z ) ] as the number of samples S increases , but it is not unbiased , since it follows from Jensen ’ s inequality that 1/Z = 1/Eq [ Ẑ ] ≤ Eq [ 1/Ẑ ] . If the discrepancy between q and π is large , the importance weights will have large variance , which also has an impact on the bias . When q = π , the importance weight w = Z has zero variance , and the Jensen ’ s inequality is tight . In the context of stochastic variational inference , this means that gradient estimates might initially be strongly biased , since there will typically be a large discrepancy between q and π . However , the variance will typically decrease as the quality of the approximation improves . Nested importance sampling and proper weighting Nested Importance sampling formalizes the construction of proposals by way of calls to other importance samplers . Extending the example above , we now assume q is some unnormalized density itself that we can not directly generates samples from . We can employ another importance sampler with proposal η to simulate samples from q , z ( l ) ∼ η ( · , ψ ) , w ( l ) = q ( z ( l ) ; φ ) η ( z ( l ) ; ψ ) , Ẑq = 1 L L∑ l=1 w ( l ) The resulting weighted samples ( w ( l ) , z ( l ) ) Ll=1 can subsequently be used compute the normalizing constant Zq of q and to generate samples that are approximately distributed w.r.t . q via resampling . This allows us to use q as a proposal in a subsequent importance sampling step , whose weighted samples can again be used to compute consistent estimates w.r.t . the target density of interest π . This can be generalized , by introducing the notion of proper weighting . Definition 2.1 ( Proper weighting ) . Let π be a probability density . For some constant c > 0 , a random pair ( w , z ) ∼ Π is properly weighted ( p.w . ) for an unnormalized probability density γ ≡ Zπ if w ≥ 0 and for all measurable functions g it holds that E w , z∼Π [ w g ( z ) ] = c ∫ dz γ ( z ) g ( z ) = cZ E z∼π [ g ( z ) ] . As long as we can generate properly weighted samples for the target density π from some sampler Π , we can compute consistent estimates 1 S ∑S s=1 w ( s ) g ( z ( s ) ) 1 S ∑S s=1 w ( s ) = cZ Ez∼π [ g ( z ) ] cZ = E z∼π [ g ( z ) ] . ( 4 ) 3 Nested Variational Inference . A widely used strategy in importance sampling is to decompose a difficult sampling problem into a series of easier problems . A common approach is to define a sequence of unnormalized densities { γk } Kk=1 that interpolate between an initial density π1 = γ1/Z1 , for which sampling is easy , and the final target density πK = γK/ZK . At each step , samples from the preceding density serve to construct proposals for the next density , which is typically combined with importance resampling or application of a Markov chain Monte Carlo ( MCMC ) operator to improve the average sample quality . Nested variational objectives . NVI defines objectives for optimizing importance samplers that target a sequence of densities . At every step , it minimizes the discrepancy between a forward density π̂k = γ̂k/Zk−1 , which acts as the proposal , and a reverse density π̌k = γ̌k/Zk , which defines an intermediate target . We define the forward density by combining the preceding target γk−1 with a forward kernel qk , and the reverse density by combining the next target γk with a reverse kernel rk−1 , γ̌k ( zk , zk−1 ) = γk ( zk ) rk−1 ( zk−1|zk ) , γ̂k ( zk , zk−1 ) = qk ( zk|zk−1 ) γk−1 ( zk−1 ) . ( 5 ) Our goal is to learn pairs of densities π̌k and π̂k that are as similar as possible . To do so , we minimize a variational objective D that comprises an f -divergence for each step in the sequence , along with a divergence between the first intermediate target π1 and an initial proposal q1 . Concretely , the NVI objective can be formulated as , DNVI = Df ( π1 ∣∣∣∣ q1 ) + K∑ k=2 Df ( π̌k ∣∣∣∣ π̂k ) . ( 6 ) Since each intermediate density πk occurs in both π̌k and in π̂k+1 , this defines a trade-off between maximizing the similarity to πk+1 and the similarity to πk−1 . To optimize this objective , we need to be able to simulate samples from the intermediate densities . Sampling from intermediate densities . Given a pair ( wk−1 , zk−1 ) that is properly weighted for the previous target density γk−1 , we can use a sequential Monte Carlo sampling [ Del Moral et al. , 2006 ] construction to define a pair ( wk , zk ) that is properly weighted for γk , zk ∼ qk ( · | zk−1 ) , wk = vk wk−1 , vk = γ̌k ( zk , zk−1 ) γ̂k ( zk , zk−1 ) . ( 7 ) We refer to the ratio vk as the incremental weight . In this construction , ( wk , zk−1 : k ) is properly weighted for γ̌k which implies that ( wk , zk ) is also properly weighted for γk , since∫ dzk−1 γ̌k ( zk , zk−1 ) = ∫ dzk−1 γk ( zk ) rk−1 ( zk−1 | zk ) = γk ( zk ) . ( 8 ) Sequential importance sampling can be combined with other operations that preserve proper weighting , including rejection sampling , application of an MCMC transition operator , and importance resampling . This defines a class of samplers that admits many popular methods as special cases , including sequential Monte Carlo ( SMC ) [ Doucet et al. , 2001 ] , annealed importance sampling ( AIS ) [ Neal , 2001 ] , and SMC samplers [ Del Moral et al. , 2006 ] . In this work we only consider variational transition kernels and importance resampling . These samplers vary in the sequences of densities they define . For example , AIS and SMC samplers both define intermediate density over a fixed common domain . Here , a common strategy is to define an annealing path γk ( zk ) = γ1 ( zk ) 1−βkγK ( zk ) βk for 0 = β1 < β2 < . . . < βK = 1 . In contrast , when using SMC for state-space models we can define the intermediate target πk to be the the filtering distribution on the first k states . In this setting , the dimensionality of the support increases at each step and we can omit the reverse kernel rk−1 . We are exploring both of these cases in Section 4 .

Summary: The author proposes a nested variational inference (NVI) framework that combines nested importance sampling and variational inference. This framework allows the training of the important sampling proposal by minimizing KL divergences at each level of nesting. This also provides us a way to learn intermediate densities, which serves as a heuristics for guiding the importance sampler. Empirically, the author evaluates NVI by (1) sampling from a multimodal distribution with a learned annealing path; (2) learn heuristics for future observations in the state-space model; (3) perform amortized inference in the hierarchical deep generative model.