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 paper describes a framework for learning and inference in probabilistic models with structured latent spaces. It operates in the tradition of IWAE and VSMC-like methods that use importance sampling to tighten a variational bound (or equivalently, uses a variational bound as an objective for learning importance sampler proposals). Unlike methods that optimize a single global bound, NVI is defined in terms of local pairwise divergences between a sequence of target distributions, which acts to decouple the stochastic optimization, reducing gradient variance and potentially offering a strategy for parallel inference.