SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models

1 INTRODUCTION . Latent variable models are powerful tools for constructing highly expressive data distributions and for understanding how high-dimensional observations might possess a simpler representation . Latent variable models are often framed as probabilistic graphical models , allowing these relationships to be expressed in terms of conditional independence . Mixture models , probabilistic principal component analysis ( Tipping & Bishop , 1999 ) , hidden Markov models , and latent Dirichlet allocation ( Blei et al. , 2003 ) are all examples of powerful latent variable models . More recently there has been a surge of interest in probabilistic latent variable models that incorporate flexible nonlinear likelihoods via deep neural networks ( Kingma & Welling , 2014 ) . These models can blend the advantages of highly structured probabilistic priors with the empirical successes of deep learning ( Johnson et al. , 2016 ; Luo et al. , 2018 ) . Moreover , these explicit latent variable models can often yield relatively interpretable representations , in which simple interpolation in the latent space can lead to semantically-meaningful changes in high-dimensional observations ( e.g. , Higgins et al . ( 2017 ) ) . It can be challenging , however , to fit the parameters of a flexible latent variable model , since computing the marginal likelihood of the data requires integrating out the latent variables . Typical approaches to this problem include the celebrated expectation maximization algorithm ( Dempster et al. , 1977 ) , Markov chain Monte Carlo , and the Laplace approximation . Variational inference generalizes expectation maximization by forming a lower bound on the aforementioned ( log ) marginal likelihood , using a tractable approximation to the unmanageable posterior over latent variables . The maximization of this lower bound—rather than the true log marginal likelihood—is often relatively straightforward when using automatic differentiation and Monte Carlo sampling . However , a ∗Equal contribution . lower bound may be ill-suited for tasks such as posterior inference and other situations where there exists an entropy maximization objective ; for example in entropy-regularized reinforcement learning ( Williams & Peng , 1991 ; Mnih et al. , 2016 ; Norouzi et al. , 2016 ) which requires minimizing the log probability of the samples under the model . While there is a long history in Bayesian statistics of estimating the marginal likelihood ( e.g. , Newton & Raftery ( 1994 ) ; Neal ( 2001 ) ) , we often want high-quality estimates of the logarithm of the marginal likelihood , which is better behaved when the data is high dimensional ; it is not as susceptible to underflow and it has gradients that are numerically sensible . However , the log transformation introduces some challenges : Monte Carlo estimation techniques such as importance sampling do not straightforwardly give unbiased estimates of this quantity . Nevertheless , there has been significant work to construct estimators of the log marginal likelihood in which it is possible to explicitly trade off between bias against computational cost ( Burda et al. , 2016 ; Bamler et al. , 2017 ; Nowozin , 2018 ) . Unfortunately , while there are asymptotic regimes where the bias of these estimators approaches zero , it is always possible to optimize the parameters to increase this bias to infinity . In this work , we construct an unbiased estimator of the log marginal likelihood . Although there is no theoretical guarantee that this estimator has finite variance , we find that it can work well in practice . We show that this unbiased estimator can train latent variable models to achieve higher test log-likelihood than lower bound estimators at the same expected compute cost . More importantly , this unbiased estimator allows us to apply latent variable models in situations where these models were previously problematic to optimize with lower bound estimators . Such applications include latent variable modeling for posterior inference and for reinforcement learning in high-dimensional action spaces , where an ideal model is one that is highly expressive yet efficient to sample from . 2 PRELIMINARIES . 2.1 LATENT VARIABLE MODELS . Latent variable models ( LVMs ) describe a distribution over data in terms of a mixture over unobserved quantities . Let pθ ( x ) be a family of probability density ( mass ) functions on a data space X , indexed by parameters θ . We will generally refer to this as a “ density ” for consistency , even when the data should be understood to be discrete ; similarly we will use integrals even when the marginalization is over a discrete set . In a latent variable model , pθ ( x ) is defined via a space of latent variables Z , a family of mixing measures on this latent space with density denoted pθ ( z ) , and a conditional distribution pθ ( x | z ) . This conditional distribution is sometimes called an “ observation model ” or a conditional likelihood . We will take θ to parameterize both pθ ( x | z ) and pθ ( z ) in the service of determining the marginal pθ ( x ) via the mixture integral : pθ ( x ) : = ∫ Z pθ ( x | z ) pθ ( z ) dz = Ez∼pθ ( z ) [ pθ ( x | z ) ] . ( 1 ) This simple formalism allows for a large range of modeling approaches , in which complexity can be baked into the latent variables ( as in traditional graphical models ) , into the conditional likelihood ( as in variational autoencoders ) , or into both ( as in structured VAEs ) . The downside of this mixing approach is that the integral may be intractable to compute , making it difficult to evaluate pθ ( x ) — a quantity often referred to in Bayesian statistics and machine learning as the marginal likelihood or evidence . Various Monte Carlo techniques have been developed to provide consistent and often unbiased estimators of pθ ( x ) , but it is usually preferable to work with log pθ ( x ) and unbiased estimation of this quantity has , to our knowledge , not been previously studied . 2.2 TRAINING LATENT VARIABLE MODELS . Fitting a parametric distribution to observed data is often framed as the minimization of a difference between the model distribution and the empirical distribution . The most common difference measure is the forward Kullback-Leibler ( KL ) divergence ; if pdata ( x ) is the empirical distribution and pθ ( x ) is a parametric family , then minimizing the KL divergence ( DKL ) with respect to θ is equivalent to maximizing the likelihood : DKL ( pdata || pθ ) = ∫ X pdata ( x ) log pdata ( x ) pθ ( x ) dx = −Edata [ log pθ ( x ) ] + const . ( 2 ) Equivalently , the optimization problem of finding the MLE parameters θ comes down to maximizing the expected log probability of the data : θMLE = arg min θ DKL ( pdata || pθ ) = arg max θ Edata [ log pθ ( x ) ] . ( 3 ) Since expectations can be estimated in an unbiased manner using Monte Carlo procedures , simple subsampling of the data enables powerful stochastic optimization techniques , with stochastic gradient descent in particular forming the basis for learning the parameters of many nonlinear models . However , this requires unbiased estimates of ∇θ log pθ ( x ) , which are not available for latent variable models . Instead , a stochastic lower bound of log pθ ( x ) is often used and then differentiated for optimization . Though many lower bound estimators ( Burda et al. , 2016 ; Bamler et al. , 2017 ; Nowozin , 2018 ) are applicable , we focus on an importance-weighted evidence lower bound ( Burda et al. , 2016 ) . This lower bound is constructed by introducing a proposal distribution q ( z ; x ) and using it to form an importance sampling estimate of the marginal likelihood : pθ ( x ) = ∫ Z pθ ( x | z ) pθ ( z ) dz = ∫ Z q ( z ; x ) pθ ( x | z ) pθ ( z ) q ( z ; x ) dz = Ez∼q [ pθ ( x | z ) pθ ( z ) q ( z ; x ) ] . ( 4 ) IfK samples are drawn from q ( z ; x ) then this provides an unbiased estimate of pθ ( x ) and the biased “ importance-weighted autoencoder ” estimator IWAEK ( x ) of log pθ ( x ) is given by IWAEK ( x ) : = log 1 K K∑ k=1 pθ ( x | zk ) pθ ( zk ) q ( zk ; x ) , zk iid∼ q ( z ; x ) . ( 5 ) The special case of K = 1 generates an unbiased estimate of the evidence lower bound ( ELBO ) , which is often used for performing variational inference by stochastic gradient descent . While the IWAE lower bound acts as a useful replacement of log pθ ( x ) in maximum likelihood training , it may not be suitable for other objectives such as those that involve entropy maximization . We discuss tasks for which a lower bound estimator would be ill-suited in Section 3.4 . There are two properties of IWAE that will allow us to modify it to produce an unbiased estimator : First , it is consistent in the sense that as the number of samples K increases , the expectation of IWAEK ( x ) converges to log pθ ( x ) . Second , it is also monotonically non-decreasing in expectation : log pθ ( x ) = lim K→∞ E [ IWAEK ( x ) ] and E [ IWAEK+1 ( x ) ] ≥ E [ IWAEK ( x ) ] . ( 6 ) These properties are sufficient to create an unbiased estimator using the Russian roulette estimator . 2.3 RUSSIAN ROULETTE ESTIMATOR . In order to create an unbiased estimator of the log probability function , we employ the Russian roulette estimator ( Kahn , 1955 ) . This estimator is used to estimate the sum of infinite series , where each sample of the estimator almost surely requires only a finite amount of computation . Intuitively , the Russian roulette estimator relies on a randomized truncation and upweighting of each term to account for the possibility of not computing the remaining terms . To illustrate the idea , let ∆̃k denote the k-th term of an infinite series . Assume the partial sum of the series ∑∞ k=1 ∆̃k converges to some quantity we wish to obtain . We can construct a simple estimator by always computing the first term then flipping a coin b ∼ Bernoulli ( q ) to determine whether we stop or continue evaluating the remaining terms . With probability 1− q , we compute the rest of the series . By reweighting the remaining future terms by 1/ ( 1−q ) , we obtain an unbiased estimator : Ỹ = ∆̃1 + ( ∑∞ k=2 ∆̃k 1− q ) 1b=0 + ( 0 ) 1b=1 E [ Ỹ ] = ∆̃1 + ∑∞ k=2 ∆̃k 1− q ( 1− q ) = ∞∑ k=1 ∆̃k . To obtain the “ Russian roulette ” ( RR ) estimator ( Forsythe & Leibler , 1950 ) , we repeatedly apply this trick to the remaining terms . In effect , we make the number of terms a random variable K , taking values in 1 , 2 , . . . to use in the summation ( i.e. , the number of coin flips ) from some distribution with probability mass function p ( K ) = P ( K = K ) with support over the positive integers . With K drawn from p ( K ) , the estimator takes the form : Ŷ ( K ) = K∑ k=1 ∆̃k P ( K ≥ k ) EK∼p ( K ) [ Ŷ ( K ) ] = ∞∑ k=1 ∆̃k . ( 7 ) The equality on the right hand of equation 7 holds so long as ( i ) P ( K ≥ k ) > 0 , ∀k > 0 , and ( ii ) the series converges absolutely , i.e. , ∑∞ k=1 |∆̃k| < ∞ ( Chen et al . ( 2019 ) ; Lemma 3 ) . This condition ensures that the average of multiple samples will converge to the value of the infinite series by the law of large numbers . However , the variance of this estimator depends on the choice of p ( K ) and can potentially be very large or even infinite ( McLeish , 2011 ; Rhee & Glynn , 2015 ; Beatson & Adams , 2019 ) .

Summary: This paper proposes an unbiased estimator of $\log p_\theta(x)$. Many unbiased estimators of $p_\theta$ exist, but $\log p_\theta$ is needed in many other settings, some of which are not well-served by standard estimators of $p_\theta$. The SUMO estimator is essentially a Russian roulette-based extension of IWAE; it is exactly unbiased, but takes a random and unbounded number of samples.