Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives

1 INTRODUCTION . Following the influential work by ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , deep generative models with latent variables have been widely used to model data such as natural images ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ; Chen et al. , 2016 ; Gulrajani et al. , 2016 ) , speech and music time-series ( Chung et al. , 2015 ; Fraccaro et al. , 2016 ; Krishnan et al. , 2015 ) , and video ( Babaeizadeh et al. , 2017 ; Ha & Schmidhuber , 2018 ; Denton & Fergus , 2018 ) . The power of these models lies in combining learned nonlinear function approximators with a principled probabilistic approach , resulting in expressive models that can capture complex distributions . Unfortunately , the nonlinearities that empower these model also make marginalizing the latent variables intractable , rendering direct maximum likelihood training inapplicable . Instead of directly maximizing the marginal likelihood , a common approach is to maximize a tractable lower bound on the likelihood such as the variational evidence lower bound ( ELBO ) ( Jordan et al. , 1999 ; Blei et al. , 2017 ) . The tightness of the bound is determined by the expressiveness of the variational family . For tractability , a factorized variational family is commonly used , which can cause the learned model to be overly simplistic . Burda et al . ( 2015 ) introduced a multi-sample bound , IWAE , that is at least as tight as the ELBO and becomes increasingly tight as the number of samples increases . Counterintuitively , although Implementation of DReG estimators and code to reproduce experiments : sites.google.com/view/ dregs . the bound is tighter , Rainforth et al . ( 2018 ) theoretically and empirically showed that the standard inference network gradient estimator for the IWAE bound performs poorly as the number of samples increases due to a diminishing signal-to-noise ratio ( SNR ) . This motivates the search for novel gradient estimators . Roeder et al . ( 2017 ) proposed a lower-variance estimator of the gradient of the IWAE bound . They speculated that their estimator was unbiased , however , were unable to prove the claim . We show that it is in fact biased , but that it is possible to construct an unbiased estimator with a second application of the reparameterization trick which we call the IWAE doubly reparameterized gradient ( DReG ) estimator . Our estimator is an unbiased , computationally efficient drop-in replacement , and does not suffer as the number of samples increases , resolving the counterintuitive behavior from previous work ( Rainforth et al. , 2018 ) . Furthermore , our insight is applicable to alternative multisample training techniques for latent variable models : reweighted wake-sleep ( RWS ) ( Bornschein & Bengio , 2014 ) and jackknife variational inference ( JVI ) ( Nowozin , 2018 ) . In this work , we derive DReG estimators for IWAE , RWS , and JVI and demonstrate improved scaling with the number of samples on a simple example . Then , we evaluate DReG estimators on MNIST generative modeling , Omniglot generative modeling , and MNIST structured prediction tasks . In all cases , we demonstrate substantial unbiased variance reduction , which translates to improved performance over the original estimators . 2 BACKGROUND . Our goal is to learn a latent variable generative model pθ ( x , z ) = pθ ( z ) pθ ( x|z ) where x are observed data and z are continuous latent variables . The marginal likelihood of the observed data , pθ ( x ) =∫ pθ ( x , z ) dz , is generally intractable . Instead , we maximize a variational lower bound on log pθ ( x ) such as the ELBO log pθ ( x ) = logEpθ ( z ) [ pθ ( x|z ) ] ≥ Eq ( z|x ) [ log pθ ( x , z ) q ( z|x ) ] , ( 1 ) where q ( z|x ) is a variational distribution . Following the influential work by ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , we consider the amortized inference setting where qφ ( z|x ) , referred to as the inference network , is a learnable function parameterized by φ that maps from x to a distribution over z . The tightness of the bound is coupled to the expressiveness of the variational family ( i.e. , { qφ } φ ) . As a result , limited expressivity of { qφ } φ , can negatively affect the learned model . Burda et al . ( 2015 ) introduced the importance weighted autoencoder ( IWAE ) bound which alleviates this coupling Ez1 : K [ log ( 1 K K∑ i=1 pθ ( x , zi ) qφ ( zi|x ) ) ] ≤ log pθ ( x ) , ( 2 ) with z1 : K ∼ ∏ i qφ ( zi|x ) . The IWAE bound reduces to the ELBO when K = 1 , is non-decreasing as K increases , and converges to log pθ ( x ) as K → ∞ under mild conditions ( Burda et al. , 2015 ) . When qφ is reparameterizable1 , the standard gradient estimator of the IWAE bound is ∇θ , φEz1 : K [ log ( 1 K K∑ i=1 wi ) ] = ∇θ , φE 1 : K [ log ( 1 K K∑ i=1 wi ) ] = E 1 : K [ K∑ i=1 wi∑ j wj ∇θ , φ logwi ] where wi = pθ ( x , zi ) /qφ ( zi|x ) . A single sample estimator of this expectation is typically used as the gradient estimator . As K increases , the bound becomes increasingly tight , however , Rainforth et al . ( 2018 ) show that the signal-to-noise ratio ( SNR ) of the inference network gradient estimator goes to 0 . This does not 1Meaning that we can express zi as z ( i , φ ) , where z is a deterministic , differentiable function and p ( i ) does not depend on θ or φ . This allows gradients to be estimated using the reparameterization trick ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . Notably , the recently introduced implicit reparameterization trick has expanded the class of distributions for which we can compute reparameterization gradients to include Gaussian mixtures among many others ( Graves , 2016 ; Jankowiak & Obermeyer , 2018 ; Jankowiak & Karaletsos , 2018 ; Figurnov et al. , 2018 ) . happen for the model parameters ( θ ) . Following up on this work , Le et al . ( 2018 ) demonstrate that this deteriorates the performance of learned models on practical problems . Because the IWAE bound converges to log pθ ( x ) ( as K → ∞ ) regardless of qφ , φ ’ s affect on the bound must diminish as K increases . It may be tempting to conclude that the SNR of the inference network gradient estimator must also decrease as K →∞ . However , low SNR is a limitation of the gradient estimator , not necessarily of the bound . Although the magnitude of the gradient converges to 0 , if the variance of the gradient estimator decreases more quickly , then the SNR of the gradient estimator need not degrade . This motivates the search for lower variance inference network gradient estimators . To derive improved gradient estimators for φ , it is informative to expand the total derivative2 of the IWAE bound with respect to φ E 1 : K [ K∑ i=1 wi∑K j=1 wj ( − ∂ ∂φ log qφ ( zi|x ) + ∂ logwi ∂zi dzi dφ ) ] . ( 3 ) Previously , Roeder et al . ( 2017 ) found that the first term within the parentheses of Eq . 3 can contribute significant variance to the gradient estimator . When K = 1 , this term analytically vanishes in expectation , so when K > 1 they suggested dropping it . Below , we abbreviate this estimator as STL . As we show in Section 6.1 , the STL estimator introduces bias when K > 1 . 3 DOUBLY REPARAMETERIZED GRADIENT ESTIMATORS ( DREGS ) . Our insight is that we can estimate the first term within the parentheses of Eq . 3 efficiently with a second application of the reparameterization trick . To see this , first note that E 1 : K [ K∑ i=1 wi∑K j=1 wj ∂ ∂φ log q ( zi|x ) ] = K∑ i=1 E 1 : K [ wi∑K j=1 wj ∂ ∂φ log q ( zi|x ) ] , so it suffices to focus on one of the K terms . Because the derivative is a partial derivative ∂∂φ , it treats zi = z ( i , φ ) as a constant , so we can freely change the random variable that the expectation is over to z1 : K . Now , Ez1 : K [ wi∑ j wj ∂ ∂φ log qφ ( zi|x ) ] = Ez−iEzi [ wi∑ j wj ∂ ∂φ log qφ ( zi|x ) ] , ( 4 ) where z−i = z1 : i−1 , i+1 : K is the set of z1 : K without zi . The inner expectation resembles a REINFORCE gradient term ( Williams , 1992 ) , where we interpret wi∑ j wj as the “ reward ” . Now , we can use the following well-known equivalence between the REINFORCE gradient and the reparameterization trick gradient ( See Appendix 8.1 for a derivation ) Eqφ ( z|x ) [ f ( z ) ∂ ∂φ log qφ ( z|x ) ] = E [ ∂f ( z ) ∂z ∂z ( , φ ) ∂φ ] . ( 5 ) This holds even when f depends on φ . Typically , the reparameterization gradient estimator has lower variance than the REINFORCE gradient estimator because it directly takes advantage of the derivative of f . Applying the identity from Eq . 5 to the right hand side of Eq . 4 gives Ezi [ wi∑ j wj ∂ ∂φ log qφ ( zi|x ) ] = E i [ ∂ ∂zi ( wi∑ j wj ) ∂zi ∂φ ] = E i [ ( 1∑ j wj − wi ( ∑ j wj ) 2 ) ∂wi ∂zi ∂zi ∂φ ] = E i [ ( wi∑ j wj − w 2 i ( ∑ j wj ) 2 ) ∂ logwi ∂zi ∂zi ∂φ ] . ( 6 ) This last expression can be efficiently estimated with a single Monte Carlo sample . When zi is not reparameterizable ( e.g. , the models in ( Mnih & Rezende , 2016 ) ) , we can use a control variate ( e.g. , 2logwi depends on φ in two ways : φ and zi = z ( i , φ ) . The total derivative accounts for both sources of dependence and the partial derivative ∂ ∂φ considers zi as a constant . 1 K ∂ ∂φ log qφ ( zi|x ) ) . In both cases , when K = 1 , this term vanishes exactly and we recover the estimator proposed in ( Roeder et al. , 2017 ) for the ELBO . However , whenK > 1 , there is no reason to believe this term will analytically vanish . Substituting Eq . 6 into Eq . 3 , we obtain a simplification due to cancellation of terms ∇φEz1 : K [ log ( 1 K K∑ i=1 wi ) ] = E 1 : K K∑ i=1 ( wi∑ j wj ) 2 ∂ logwi ∂zi ∂zi ∂φ . ( 7 ) We call the algorithm that uses the single sample Monte Carlo estimator of this expression for the inference network gradient the IWAE doubly reparameterized gradient estimator ( IWAE-DReG ) . This estimator has the property that when q ( z|x ) is optimal ( i.e. , q ( z|x ) = p ( z|x ) ) , the estimator vanishes exactly and has zero variance , whereas this does not hold for the standard IWAE gradient estimator . We provide an asymptotic analysis of the IWAE-DReG estimator in Appendix 8.2 . The conclusion of that analysis is that , in contrast to the standard IWAE gradient estimator , the SNR of the IWAE-DReG estimator exhibits the same scaling behaviour of O ( √ K ) for both the generation and inference network gradients ( i.e. , improving in K ) .

Summary: This paper applies a reparameterization trick to estimate the gradients objectives encountered in variational autoencoder based frameworks with continuous latent variables.  Especially the authors use this double reparameterization trick on Importance Weighted Auto-Encoder (IWAE) and Reweighted Wake-Sleep (RWS)  methods. Compared to IWAE, the developed method's SNR does not go to zero with increasing the number of particles.