Partial Rejection Control for Robust Variational Inference in Sequential Latent Variable Models

1 INTRODUCTION . Exact inference in latent variable models is usually intractable . Markov Chain Monte-Carlo ( MCMC ) ( Andrieu et al. , 2003 ) and variational inference ( VI ) methods ( Blei et al. , 2017 ) , are commonly employed in such models to make inference tractable . While MCMC has been the traditional method of choice , often with provable guarantees , optimization-based VI methods have also enjoyed considerable recent interest due to their excellent scalability on large-scale datasets . VI is based on maximizing a lower bound constructed through a marginal likelihood estimator . For latent variable models with sequential structure , sequential Monte-Carlo ( SMC ) ( Doucet & Johansen , 2009 ) returns a much lower variance estimator of the log marginal likelihood than importance sampling ( Bérard et al. , 2014 ; Cérou et al. , 2011 ) . In this work , we focus our attention on designing a low variance , unbiased , and computationally efficient estimator of the marginal likelihood . The performance of SMC based methods is strongly dependent on the choice of the proposal distribution . Inadequate proposal distributions propose values in low probability areas under the target , leading to particle depletion ( Doucet & Johansen , 2009 ) . An effective solution is to use rejection control ( Liu et al. , 1998 ; Peters et al. , 2012 ) which is based on an approximate rejection sampling step within SMC to reject samples with low importance weights . In this work , we leverage the idea of partial rejection control ( PRC ) within the framework of SMC based VI for sequential latent variable models . To this end , we construct a novel lower bound , VSMC-PRC , and propose an efficient optimization strategy for selecting the variational parameters . Compared to other recent SMC based VI approaches ( Naesseth et al. , 2017 ; Maddison et al. , 2017 ; Le et al. , 2017 ) , our approach consists of an inbuilt accept-reject mechanism within SMC to prevent particle depletion . The use of accept-reject within SMC makes the particle weight intractable , therefore , we use a generalization of the Bernoulli factory ( Asmussen et al. , 1992 ) to construct unbiased estimators of the marginal likelihood for SMC-PRC . Although the idea of combining VI with an inbuilt accept-reject mechanism is not new ( Salimans et al. , 2015 ; Ruiz & Titsias , 2019 ; Grover et al. , 2018 ; Gummadi , 2014 ) , a key distinction of our approach is to incorporate an accept-reject mechanism along with a resampling framework . In contrast to standard sampling algorithms that may reject the entire stream of particles , we use a partial accept-reject on the most recent update , increasing the sampling efficiency . Further , the variational framework of SMC-PRC is interesting in itself as it combines accept-reject with particle filter methods . Therefore , our proposed bound VSMC-PRC generalizes several existing approaches for example : Variational Rejection Sampling ( VRS ) ( Grover et al. , 2018 ) , FIVO ( Maddison et al. , 2014 ) , IWAE ( Burda et al. , 2015 ) , and standard variational Bayes ( Blei et al. , 2017 ) . Another key distinction is that , while existing approaches using Bernoulli factory are limited to niche one-dimensional toy examples , our proposed approach is scalable . To the best of our knowledge , there is no prior work that has used Bernoulli factories for such a general case like variational recurrent neural networks ( VRNN ) ; therefore , we believe this aspect to be a significant contribution as well . The rest of the paper is organized as follows : In Section 2 , we provide a brief review on SMC , partial rejection control , and dice enterprise . In Section 3 , we introduce our VSMC-PRC bound and provide new theoretical insights into the Monte-Carlo estimator and design efficient ways to optimize it . Finally , we discuss related work and present experiments on the Gaussian state-space model ( SSM ) and VRNN . 2 BACKGROUND . We denote a sequence of T real-valued observations as x1 : T = ( x1 , x2 , . . . , xT ) , and assume that there is an associated sequence of latent variables z1 : T = ( z1 , z2 , . . . , zT ) . We are interested in inferring the posterior distribution of the latent variables , i.e. , p ( z1 : T |x1 : T ) . The task is , in general , intractable . For the rest of the paper we have used some common notations from SMC and VI literature where zit : i th particle at time t ; Ait−1 : ancestor variable for the i th particle at time t ; θ and φ : model and variational parameters , respectively . 2.1 SEQUENTIAL MONTE CARLO WITH PARTIAL REJECTION CONTROL . An SMC sampler approximates a sequence of densities { pθ ( z1 : t|x1 : t ) } Tt=1 through a set of N weighted samples generated from a proposal distribution . Let the proposal density be qφ ( z1 : T |x1 : T ) = T∏ t=1 qφ ( zt|x1 : t , z1 : t−1 ) . ( 1 ) Consider time t−1 at which we have uniformly weighted samples { N−1 , zi1 : t−1 , Ait−1 } Ni=1 estimating pθ ( z1 : t−1|x1 : t−1 ) . We want to estimate pθ ( z1 : t|x1 : t ) such that particles with a low importance weight are automatically rejected . PRC achieves this by using an approximate rejection sampling step ( Liu et al. , 1998 ; Peters et al. , 2012 ) . The overall procedure is as follows : 1 . Generate zit ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) where i = 1 , 2 , . . . , N . 2 . Accept zit with probability aθ , φ ( z i t|z Ait−1 1 : t−1 , x1 : t ) = 1 + M ( i , t− 1 ) qφ ( zit|x1 : t , zAit−11 : t−1 ) pθ ( xt , zit|x1 : t−1 , z Ait−1 1 : t−1 ) −1 , ( 2 ) where M ( i , t− 1 ) is a hyperparameter controlling the acceptance rate ( see Proposition 3 and Section 3.3 for more details ) . Note that PRC applies accept-reject only on zit , not on the entire trajectory . 3 . If zit is rejected go to step 1 . 4 . The new incremental importance weight of the accepted sample is αt ( z i 1 : t ) = c i tZ ( z Ait−1 1 : t−1 , x1 : t ) , ( 3 ) where cit is cit = pθ ( xt , z i t|x1 : t−1 , z Ait−1 1 : t−1 ) qφ ( zit|x1 : t , z Ait−1 1 : t−1 ) aθ , φ ( z i t|z Ait−1 1 : t−1 , x1 : t ) , ( 4 ) and the intractable normalization constant Z ( . ) ( For simplicity of notation , we have ignored the dependence of Z ( . ) on M ( i , t− 1 ) ) Z ( z Ait−1 1 : t−1 , x1 : t ) = ∫ aθ , φ ( zt|z Ait−1 1 : t−1 , x1 : t ) qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) dzt . ( 5 ) 5 . Compute Monte-Carlo estimator of unnormalized weight w̃it = pθ ( xt , z i t|x1 : t−1 , z Ait−1 1 : t−1 ) 1 K ∑K k=1 aθ , φ ( δ i , k t |z Ait−1 1 : t−1 , x1 : t ) qφ ( zit|x1 : t , z Ait−1 1 : t−1 ) aθ , φ ( z i t|z Ait−1 1 : t−1 , x1 : t ) , ( 6 ) where δi , kt ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) and k = 1 , 2 , . . . , K. Note that w̃ i t would be essential for constructing unbiased estimator for pθ ( x1 : T ) . 6 . Generate ancestor variables Ait through dice-enterprise and set new weights w i t = N −1 for i = 1 , 2 , . . . , N : Ait ∼ Categorical ( αt ( z 1 1 : t ) ∑N j=1 αt ( z j 1 : t ) , αt ( z 2 1 : t ) ∑N j=1 αt ( z j 1 : t ) , . . . , αt ( z N 1 : t ) ∑N j=1 αt ( z j 1 : t ) ) . ( 7 ) 2.2 DICE ENTERPRISE . Simulation of ancestor variables in Eq . 7 is non-trivial due to intractable normalization constants in the incremental importance weight ( see Eq . 3 ) . Vanilla Monte-Carlo estimation of αt ( . ) yields biased samples of ancestor variables from Eq . 7 . To address this issue , we leverage a generalization of Bernoulli factory , called dice-enterprise ( Morina et al. , 2019 ) . Note that multinoulli extensions of Bernoulli factory ( Dughmi et al. , 2017 ) have also been used for resampling within intractable SMC before ( Schmon et al. , 2019 ) , a key distinction of our approach is to design a scalable Bernoulli factory methodology especially useful for VI applications . Suppose we can simulate Bernoulli ( pit ) outcomes where p i t is intractable . Bernoulli factory problem simulates an event with probability f ( pit ) , where f ( . ) is some desired function . In our case , the intractable coin probability pit is the intractable normalization constant , pit = Z ( z Ait−1 1 : t−1 , x1 : t ) = ∫ aθ , φ ( zt|z Ait−1 1 : t−1 , x1 : t ) qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) dzt . ( 8 ) Since pit ∈ [ 0 , 1 ] and we can easily simulate this coin , we obtain the dice-enterprise algorithm below . 1 . Required : Constants { cit } Ni=1 see Eq . 4 . 2 . Sample C ∼ Categorical ( c1t∑N j=1 c j t , c2t∑N j=1 c j t , . . . , cNt∑N j=1 c j t ) 3 . If C = i , generate Ui ∼ U [ 0 , 1 ] and zt ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) • If Ui < aθ , φ ( zt|z Ait−1 1 : t−1 , x1 : t ) output i • Else go to step 2 The dice-enterprise produces unbiased ancestor variables . Note that we can easily control the efficiency of the proposed dice-enterprise through the hyper-parameter M ( similar to Eq . 2 ) in contrast to existing Bernoulli factory algorithms ( Schmon et al. , 2019 ) . For details on efficiency and correctness , please refer to Section 3.1 and Section 3.3 . Our proposed VSMC-PRC bound is constructed through a marginal likelihood estimator obtained by combining the SMC sampler with a PRC step and dice-enterprise . The variance of estimators obtained through SMC-PRC particle filter is usually low ( Peters et al. , 2012 ) . Therefore , we expect VSMC-PRC to be a tighter bound compared to the standard SMC based bounds used in recent works ( Maddison et al. , 2017 ; Naesseth et al. , 2017 ; Le et al. , 2017 ) . Algorithm 1 summarizes the generative process to simulate the VSMC-PRC bound . Please see Figure 1 to visualize the generative process for VSMC-PRC . Algorithm 1 Estimating the VSMC-PRC lower bound 1 : Required : N , K , and M 2 : for t ∈ { 1 , 2 , . . . , T } do 3 : for i ∈ { 1 , 2 , . . . , N } do 4 : zit , c i t , w̃ i t ∼ PRC ( q , p , M ( i , t− 1 ) ) 5 : zi1 : t = ( z Ait−1 1 : t−1 , z i t ) 6 : end for 7 : for i ∈ { 1 , 2 , . . . , N } do 8 : Ait = DICE-ENT ( { cit , zi1 : t } Ni=1 ) 9 : end for 10 : end for 11 : return log ∏T t=1 ( 1 N ∑N i=1 w̃ i t ) 12 : 13 : PRC ( q , p , M ( i , t− 1 ) ) 14 : while sample not accepted do 15 : Generate zit ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) 16 : Accept zit with probability aθ , φ ( z i t|z Ait−1 1 : t−1 , x1 : t ) 17 : end while 18 : Sample { δi , kt } Kk=1 ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) 19 : Calculate w̃it from Eq . 6 20 : Calculate cit from Eq . 4 21 : return ( zit , c i t , w̃ i t ) 22 : 23 : DICE-ENT ( { cit , zi1 : t } Ni=1 ) 24 : Sample C ∼ Multinoulli ( cit∑N j=1 c j t ) N i=125 : if C == i then 26 : Sample Ui ∼ U [ 0 , 1 ] 27 : zit ∼ qφ ( zt|x1 : t , z Ait−1 1 : t−1 ) 28 : end if 29 : if Ui < aθ , φ ( zit|z Ait−1 1 : t−1 , x1 : t ) then 30 : return ( i ) 31 : else 32 : return DICE-ENT ( { cit , zi1 : t } Ni=1 ) 33 : end if

Summary: This paper describes an SMC algorithm to sample the posterior distribution of latent states $p_\theta(z_{1:T}|x_{1:T})$ in a latent variable models $p_\theta(x_{1:T}, z_{1:T})$. The authors consider a completely general setting (the authors assume Eq.(1) but clearly there is nothing to assume here, this the standard Bayes rule). It is well known that the vanilla SMC sampler is a good candidate for ELBO because it provides an unbiased estimator of the likelihood. But the authors prefer here to use a more sophisticated version of the  SMC algorithm, which features a partial rejection algorithm, which amounts to eliminate proposed particles which are "large enough" likelihood.