SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models
Authors: Yucen Luo
Alex Beatson
Mohammad Norouzi
Jun Zhu
David Duvenaud
Ryan P. Adams
Ricky T. Q. Chen
Venue:  (2020)

Abstract: Standard variational lower bounds used to train latent variable models produce biased estimates of most quantities of interest. We introduce an unbiased estimator of the log marginal likelihood and its gradients for latent variable models based on randomized truncation of infinite series. If parameterized by an encoder-decoder architecture, the parameters of the encoder can be optimized to minimize its variance of this estimator. We show that models trained using our estimator give better test-set likelihoods than a standard importance-sampling based approach for the same average computational cost. This estimator also allows use of latent variable models for tasks where unbiased estimators, rather than marginal likelihood lower bounds, are preferred, such as minimizing reverse KL divergences and estimating score functions.

Published as a conference paper at ICLR 2020
SUMO: UNBIASED ESTIMATION OF LOG MARGINAL PROBABILITY FOR LATENT VARIABLE MODELS
Yucen Luo'
Alex Beatson
Tsinghua University
Princeton University
1uoyc 1@mai 11s t eeiguaa.ed.... cn
abeatson8cs www.nnccspriineeton..o
Mohammad Norouzi Google Research mmorouzi@google.com com
Jun Zhu Tsinghua University
David Duvenaud
University of Toronto
desz jttsinghua........o.. edu. on
duvenaud8cs toronto.
Ryan Adams Princcton University rpa@princeton edu
Ricky Q. Chen
University of Toronto
rtgichenêcs toronto.cdu
ABSTRACT
Standard variational lower bounds used to train latent variable models produce bi- ased estimates of most quantities of interest. We introduce unbiased estimator of the log marginal likelihood and its gradients for latent variable models based on randomized truncation of infinite series. If parameterized by encoder-decoder architecture, the parameters of the encoder can be optimized to minimize its vari- ance of this estimator: We show that models trained using our estimator give better test-set likelihoods than a standard importance sampiiing based approach for the same average computational cost. This estimator also allows use of latent variable models for tasks where unbiased estimators, rather than marginal likelihood lower bounds, are preferred, such as minimizing reverse KI. divergences and estimating score functions.
INTRODUCTION
Latent variable models are powerful tools for constructing highly expressive data distributions and for understanding how nigh-dimensional observations might possess a simpler representation. La- tent variable models are often framed as probabilistic graphical models, allowing these relation- ships to be expressed in terms of conditional independence. Mixture models, probabilistic principal component analysis (Tipping & Bishop, 1999), hidden Markov models, and latent Dirichlet allo- cation (Blei et al., 2003) are all examples of powerful latent variable models. More recently there has been a surge interest in probabilistic latent variable models that incorporate flexible nonlinear likelihoods via deep neural networks (Kingma& & Welling, 2014). These models can blend the advan- tages of highly structured probabilistic priors with the empirical successes of deep learning (Johnson al., 2016; Luo al., 2018). Moreover, these explicit latent variable models atively interpretable representations, in which simple interpolation in the latent space can lead to sieaanically-mmaninge
often yield rel-
changes high-dimensional observations (e.g., Higgins al. (2017)).
can be challenging, however, to fit the parameters of a flexible latent variable model, since com- puting the marginal likelihood of the data requires integrating out the latent variables. Typical approaches to this problem include the celebrated expectation maximization algorithm (Dempster al., 1977), Markov chain Monte Carlo, and the Laplace approximation. Variational inference gen- eralizes expectation maximization by forming lower bound on the aforementioned (log) marginal likelihood, using a tractable approximation to the unmanageable posterior over latent variables. The maximization of this lower found-rather than the true log marginal likelihood-ii often rela- tively straightforward when using automatic differentiation and Monte Carlo sampling. However,
Equull contribution.

---

Published as a conference paper at ICLR 2020
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- catropy-regularized reinforcement learn- ing (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 long history in Bayesian statistics of estimating the marginal likelihood (e.g.. New- ton & Raftery (1994); Neal (2001)), we often want high-quality estimates of the logarithm of the marginal likelihood, which is better behaved when the data high dimensional; it is not as suscepti- ble 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 straightforwardyy give unbiased estimates of this quantity. Nevertheless, there has been signifi- cant work to construct estimators of the log marginal likelihood in which
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 ap- proaches 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- bog-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-dimensions action spaces, where an ideal model 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 unob- served quantities. Let (2( be family of probability density (mass) functions a data space X, indexed by parameters 0. We will generally refer to this as a "density" for consistency, even when the data should be understood be discrete; similarly we will use integrals even when the marginal- ization is over a discrete set. In a latent variable model, P=(x) is defined via a space of latent variables Z, family of mixing measures
this latent space with density denoted po()), and conditional distribution po (I z). This conditional distribution is sometimes called an "observation model" or a conditional likelihood. We will take 0 to parameterize both pe 22) and pe (2) in the service of determining the marginal (I via the mixture integral:
po = po(x 2))o(2) dz Exwpo(s) [p][
This simple formalism allows for 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 po (z) quantity often referred to in Bayesian statistics and machine learning the marginal likelihood or evidence. Various Monte Carlo techniques have been developed to provide consistent and of- ten unbiased estimators of px)),, but it is usually preferable work with log (or) and unbiased estimation of this quantity has, to our knowledge. not been previously studied.
2.2 TRAINING LATENT VARIABLE MODELS
Fitting parametric distribution observed data is often framed as the minimization of a difference between the model distribution and the empirical distribution. The most common difference measure the forward Leibbec (KL) divergence; if (a)ta(x) (z) is the empirical distribution and (o(x) is parametric family, then minimizing the KI. divergence (DRL.) with respect to 0 is equivalent to maximizing the likelihood:
Pdata (x) dr -Edata (log po (x:) + const.
DKL(Pean Po) (data(()) log
(2)
(z)
2

---

Published as a conference paper at ICLR 2020
Equivalently, the optimization problem of finding the MLE parameters & comes down to maximizing the expected log probability of the data:
OMLE arg min DKL, ppp)) arg max Edata [log
(3)
9
Since expectations be estimated in an unbiased manner using Monte Carlo procedures, simple subsampling of the data enables powerful stochastic optimization techniques, with stochastic gradi- ent descent in particular forming the basis for learning the parameters of many nonlinear models. However, this requires unbiased estimates of V log po (2),, which are not available for latent vari- able models. Instead, stochastic lower bound of log pa 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--eighted evidence lower bound (Burda et al., 2016). This lower bound is constructed by introducing a proposal distribution ((2; and using form an importance sampling estimate of the marginal likelihood:
(x)) -f6 po( 22) po(2) de 9(2:z) (0 2) po(z)
[o 2) Po(=)
de Ezzy
(4)
((2;x)
((2
If K samples are drawn from ((;; then this provides an unbiased estimate of po (x) and the biased "importance einportance-wiighted autoencoder" estimator IWAEx of log pa (x) is given by
IWAE (x) = log
1 K
K
pa(æ 24) po(2)) ((x;)))
and 24 {(2;)).
(5)
The special case of K = 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 Pa(x) in maximum likelihood training, it may not be suitable for other objectives such those that involve entropy maximization. We discuss tasks for which 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 IWAE (z) converges to log (z)).. Second, it also monotonically non decrescan expectation:
log
lim E[IWAE (=)] K-000
and
E[IWAEK +1 (=)) z E[IWAE (x))
(6)
These properties are sufficient to create an unbiased estimator using the Russian roulette estimator.
2.3 RUSSIAN ROULETTE ESTIMATOR
In order create unbiased estimator of the log probability function, we employ the Russian roulette estimator (Kahn, 1955). This estimator used 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 randomized truncation and upweighting of each term to account for the possibility of not computing the remaining terms.
To illustrate the idea, let Ás denote the k-th term of an infinite series. Assume the partial sum of the series De Åд converges to some quantity wish obtain. We can construct simple estimator by always computing the first term then flipping a coin Bernoulli(g) to determine whether we stop or continue evaluating the remaining terms. With probability q. compute the rest of the series. By reweighting the remaining future terms by 1/(1-)), we obtain an unbiased estimator:
A 1-q 1 (0)1
A - +
B = A
-
(1
4)
AR
q
k=
To obtain the "Russian roulette" (RR) estimator (Forsythe & Leibler, 1950), repeatedly apply this trick to the remaining terms. In effect, we make the number of terms a random variable IC, taking values in ... use in the summation (i.e., the number of coin flips) from some distribution
3

---

Published as a conference paper at ICLR 2020
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:
Dx P(K ==1
K
20
((K)
Σ
K (k) [(K)]
A
k)
The equality on the right hand of equation 7 holds so long as (i) P(K 2 k) > O, Vk and (ii) the series converges absolutely, i.e., Σt= |åx < oo (Chen 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 pl K") and
can potentially be very large or even infinite (McLeish, 2011; Rhee & Glynn, 2015; Beatson & Adams, 2019).
SUMO: UNBIASED ESTIMATION OF LOG PROBABILITY FOR LVMs 3.1
RUSSIAN ROULETTE TO DEBIAS LOWER BOUNDS
We turn any absolutely convergent series into a telescoping series and apply the Russian roulette randomization to form an unbiased stochastic estimator. We focus here on the IWAE bound de- scribed in Section 2.2. Let A& (z) IWAE&+1 (2) (x), then since E, [x]])()))] converges absolutely, we apply equation to construct our estimator, which we call SUMO (Stochastically
Unbiased Marginalization Objective). The detailed derivation of SUMO is in Appendix A.I.
(Z) k=1 PIKZA
SUMO(x)) (x) = IWAE, (x) + +
where K K~ pkk(K)
(8)
The randomized truncation of the series using the Russian roulette estimator means that this is an unbiased estimator of the log marginal likelihood, regardless of the distribution p (K):
[SUMO(2)) = pogpp(x))
where the expectation is taken over p(K) and q(2; v) (see Algorithm 1 for our exact sampling procedure). Furthermore, under some conditions, we have E (>)) V (SUMO(z)) = Va log pa (z) (see Appendix A.4).
3.2 OPTIMIZING VARIANCE- COMPUTTE PRODUCT BY CHOICE OF p(K)
To efficiently optimize limit, one should choose an estimator to minimize the product of the sec- ond moment of the gradient estimates and the expected compute per evaluation. The choice of p(K) effects both the variance and computation cost of our estimator. Denoting G = VeY and A V [I (z) IWAE& [z)]. the Russian roulette estimator optimal across broad family of unbiased randomized truncation estimators if the Ag are statistically independent, in which (Beatson & Adams, 2019). While the Ag are not in fact strictly independent with our sampling procedure (Algorithm 1), and other estima- tors within the family may perform better, we justify our choice by showing that EA,A, for i + converges zero much faster than EA2 (Appendices A.2 & A.3). In the following. we assume independence of A and choose p(K) to minimize the product of compute and variance.
case it has second moment | Exz BI
Algorithm Computing SUMO, an unbiased estimator of log p(z).
Input: x, m 2 1, encoder (22;,)), decoder p(a, 2), p(K), reverse.cdf() P(K 2
1: Sample K p(K)
2: Sample {2+] mm iid
(2;xx)
3: log Wi < loggpppz p =x) log ((
4: ks + [I,...,K
m]
5: cum < log__cmssu_eep(lleeeee
- log(ks(:k++])
6: inv.weights = berpighhs==/reeerseeee 1/reverse.cdf((s)
return cum mn1] cum(inv.weightt & (cum_iwae(m:] cum sum_iwae(mm-:-l
4

---

Published as a conference paper at ICLR 2020
We first show that B 2 is O(1)2) (Appendix A.5). This implies the optimal compute- variance product (Rhee & Glynn, 2015; Beatson & Adams, 2019) is given by P(K 2 k) x V [2). In case, this gives P(K 2 k) 1/k, which results in an estimator with infinite expected computa- tion and no finite bound on variance. In fact, any p( ( which gives rise to provably finite variance requires heavier tail than IP(K 2 k) 1/4 and so will have infinite expected computation.
Though we could not theoretically show that our estimator and gradients have finite variance, we empirically find that gradient descent converges- -even in the setting of minimizing log probability. We plot | 11 for the toy variational inference task used to assess signal to noise ratio in Tucker al. (2018) and Rainforth al. (2018b), and find that they converge faster than E in practice (Appendix A.6). While this indicates the variance is better than the theoretical bound, an estimator having infinite expected computation cost will always be an issue as indicates significant proba- bility of sampling arbitrarily large K. We therefore modify the tail of the sampling distribution such that the estimator has finite expected computation:
i < a 1//- (1 0.1)
2//x
P(K 2 k))
(10)
if N a
We typically choose 80, which gives expected computation cost of approximately 5 terms.
3.2.1 TRADING VARIANCE AND COMPUTE
One way to improve the RR estimator is to construct it so that some minimum number of terms (de- noted here as m) are always computed. This puts lower bound on the computational cost, but can potentially lower variance, providing design space for trading off estimator quality against com- putational cost. This corresponds to a choice of RR estimator in which P(K = K) for K S m. This computes the sum out to terms (effectively computing IWAE.. and then estimates the re- maining difference with Russian roulette:
K
Ax( 2 k)
SUMO (x) IWAB (I)
Σ k=
K K~ p(K(K)
(11)
P(K
In practice, instead of tuning parameters of p(K), we set m to achieve a given expected computation cost per estimator evaluation for fair comparison with IWAE and related estimators.
3.3 TRAINING ((2;x)) TO REDUCE VARIANCE
The SUMO estimator does not require amortized variational inference, but the use of an "encoder" to produce an approximate posterior q(>; x) has been shown to be a highly effective way to perform rapid feedforward inference in neural latent variable models. We use denote the parameters of the encoder q0 (s(2; z). However, the gradients of SUMO with respect to are in expectation zero precisely because SUMO is an unbiased estimator of log po (ar). regardless of our choice of (2; Nevertheless, we would expect the choice of (=(2; x) significantly impacts the variance our esti- mator. As such, optimize 4s (2;x)) reduce the variance of the SUMO estimator. We can obtain unbiased gradients in the following way (Ruiz et al., 2016; Tucker et al., 2017):
V. Var[[UUMO
V V.E[SUMO?] - V [E(SUMO)) = EIV (UMMO)
(12)
Notably, the expectation of this estimator depends on the variance of SUMO, which we have not been able to bound. In practice, we observe gradients which are sometimes very large. We apply gradient clipping to the encoder to clip gradients which are excessively large in magnitude. This helps stabilize the training progress but introduces bias into the encoder gradients. Fortunately, the encoder itself is merely a tool for variance reduction, and biased gradients with respect to the encoder can still significantly help optimization.
3.4 APPLICATIONS OF UNBIASED LOG PROBABILITY
Here we list some applications where unbiased log probability is useful. Using SUMO to replace existing lower bound estimates allows latent variable models to be used for new applications where lower bound is inappropriate. As latent variable models can be both expressive and efficient to sample from, they are frequently useful in applications where the data nigh-diiensional and samples from the model are needed.
5

---

Published as a conference paper at ICLR 2020
Minimizing log Po()). Some machine learning objectives include terms that seek to increase the entropy of the learned model. The "reverse KL" objective- often used for training models to per- form approximate posterior inference-minne target density that may only be known up a normalization constant. Local updates of this form are the basis of the expectation propagation procedure (Minka, 2001). This objective has also been used for distilling autoregressive models that are inefficient at sampling (Oord et al., 2018). Moreover, re- verse KL is connected to the use of entropy-regularizeed objectives (Williams & Peng, 1991; Ziebart, 2010; Mnih et al., 2016; Norouzi et al., 2016) in decision decision-making problems, where the goal is to encourage the decision maker toward exploration and prevent it from settling into local minimum.
(z=po(s) (z)) log ' T22)] where T(x)) is
Unbiased score function Vo log (z)). The score function is the gradient of the log- log-likelihood with respect to the parameters and has uses in estimating the Fisher information matrix and per- forming stochastic gradient Langevin dynamics (Welling & Teh, 2011). among other applications. Of particular note, the REINFORCE gradient estimator (Williams, 1992) 1 applicable for optimizing objectives of the form maxy (x [R(x))]iss (e-= estimated using the score function. This can be replaced with the gradient of SUMO which itself is an estimator of the score func- tion Va log (z)).
V (0))
[p [R(x) (o))
(x) [R ((
(13)
( V
where the inner expectation is over the stochasticity of the SUMO estimator. Such estimators often used for reward maximization in reinforcement learning where po is a stochastic policy.
4 RELATED WORK
There is a long history in Bayesian statistics of marginal likelihood estimation in the service of model selection. The harmonic mean estimator (Newton & Raftery, 1994), for example, has a long (and notorious) history as a consistent estimator of the marginal likelihood that may have infinite variance (Murray & Salakhutdinov, 2009) and exhibits simulation psuedo-bias (Lenk, 2009). The Chib estimator (Chib, 1995), the Laplace approximation, and nested sampling (Skilling, 2006) are alternative proposals that can often have better properties (Murray & Salakhutdinov, 2009). An- nealed importance sampling (Neal, 2001) probably represents the gold standard for marginal like- lihood estimation. These, however, turn into consistent estimators at best when estimating the log marginal probability (Rainforth et al., 2018a). Bias removal schemes such as jackknife variational inference (Nowozin, 2018) have been proposed to debias log-evidence estimation, IWAE in partic- ular. Hierarchical IWAE (Huang et al., 2019) uses joint proposal to induce negative correlation among samples and connects the convergence of variance of the estimator and the convergence of the lower bound.
Russian roulette also has a long history. It dates back to unpublished work from von Neumann and Ulam, who used it debias Monte Carlo methods for matrix inversion (Forsythe & Leibler, 1950) and particle transport problems (Kahn, 1955). It has gained popularity in statistical physics (Spanier & Gelbard, 1969; Kuti, 1982; Wagner, 1987), for unbiased ray tracing in graphics and render- ing (Arvo & Kirk, 1990). and for a number of estimation problems in the statisties community (Wei & Murray, 2017: Lyne al., 2015; Rychlik, 1990; 1995; Jacob & Thiery, 2015; Jacob al., 2017). It has also been independently rediscovered many times (Fearnhead et al., 2008; McLeish, 2011; Rhee & Glynn, 2012; Tallec & Ollivier, 2017).
The use of Russian roulette estimation in deep learning and generative modeling applications has been gaining traction in recent years. has been used to solve short-term bias optimization problems (Tallec & Ollivier, 2017; Beatson & Adams, 2019). Wei & Murray (2017) estimates the reciprocal normalization constant of an unnormalized density. Han et al. (2018) uses a similar ran- dom truncation approach to estimate the distribution of eigenvalues in symmetric matrix. Along similar motivations with our work, Chen et al. (2019) uses this estimator to construct an unbiased estimator of the change of variables equation in the context of normalizing flows (Rezende & Mo- hamed, 2015), and Xu et al. (2019) uses it to construct unbiased log probability for a nonparameteric distribution in the context of variational autoencoders (Kingma & Welling, 2014).

---

Published as a conference paper at ICLR 2020
Though extend latent variable models applications that require unbiased estimates of log probability and benefit from efficient sampling, an interesting family of models already fulfill these requirements. Normalizing flows (Rezende & Mohamed, 2015; Dinh al., 2017) offer act log probability and certain models have been proven to be universal density estimators (e.g. Huang et al. (2018)). However, these models often require restrictive architectural choices with no dimensionality-rccuttiin capabilities, and make use of many more parameters to scale up (Kingma & Dhariwal, 2018) than alternative generative models. Discrete variable versions of these models are still in their infancy and make use of biased gradients (Tran et al., 2019; Hoogeboom et al., 2019), whereas latent variable models naturally extend to discrete observations.
DENSITY MODELING EXPERIMENTS
We first compare the performance of SUMO when used as a replacement to IWAE with the same expected cost on density modeling tasks. We make use of two benchmark datasets: dynamically binarized MNIST (LeCun et al., 1998) and binarized OMNIGLOT (Lake et al., 2015).
We use the same neural network architecture as IWAE (Burda et al., 2016). The prior p(z) is 50- dimensional standard Gaussian distribution. The conditional distributions p(zx|2) are independent Bernoulli, with the decoder parameterized by two hidden layers, each with 200 tanh units. The approximate posterior <(2; is also 50-dimensional Gaussian distribution with diagonal covari- ance, whose mean and variance are both parameterized by two hidden layers with 200 tanh units. We reimplemented and tuned IWAE, obtaining strong baseline results which are better than those previously reported. We then used the same hyperparameters to train with the SUMO estimator. We find clipping very large gradients can help performance, large gradients may be infrequently sampled. This introduces a small amount of bias into the gradients while reducing variance, but nevertheless help achieve faster convergence and should still result in less-biased estimator. A posthoc study of the effect on final test performance as function of this bias variace tradeoff' mechanism is discussed in Appendix A.7. We note that gradient clipping is only done for the density modeling experiments.
The averaged test log- lilkliiheoodds and standard deviations over 3 runs are summarized in Table To be consistent with existing literature, we evaluate our model using IWAE with 5000 samples. In all the cases, SUMO achieves slightly better performance than IWAE with the same expected cost. We also bold the results that are statistically insignificant from the best performing model according to an unpaired f-test with significance level 0.05. However, we do see diminishing returns as we increase k, suggesting that as we increase compute, the variance of our estimator may impact performance more than the bias IWAE.
LATENT VARIABLES MODELS FOR ENTROPY MAXIMIZATION
We move to our first task for which a lower bound estimate of log probability would not suf- fice. The reverse KL objective is useful when we have access to a (possibly unnormalized) target
Table 1: Test negative log-likelihood of the trained model, estimated using IWAE(K=5000), For SUMO, i refers the expected number of computed terms.
MNIST
OMNIGLOT
Training Objective
k=5
k=15
k=50
k=5
k=15
k=550
ELBO (Burda et al., 2016)
86.47
86.35
107.62
107.80
IWAE (Burda et 2016)
85.54
84.78
106.12
104.67
ELBO (Our impl.)
85.97:10.01
85.9910.05
85.88+0.07
106.79+0.08
106.98+0.19
106.844-0.13
IWAF (Our impl.)
85.28+00.01
84.89+0.03
84.50+0.02
104.96+0.04
104.53+0.05
103.99+0.11
JVI (Our impl.)
84.75+0.03
104.08+0.11
SUMO
85.09-10.01
84.71+00.02
84.40+00.03
104.85+0.04
104.29+t0.12
103.79+0.14

---

Published as a conference paper at ICLR 2020
Target log prohability
Training w/ IWAE (k=15) Training w/ SUMO (k=15)
Model samples
Figure I: We trained latent variable models for posterior inference, which requires minimizing log probability under the model. Training with IWAE leads to optimizing for the bias while leaving the true model in an unstable state, whereas training with SUMO EaMO-though noisy-leads leads to convergence.
distribution but no efficient sampling algorithm.
(14) A major problem with fitting latent variables models to this objective is the presence of an entropy maximization term, effectively a minimization of log (z). Estimating this log marginal probability with a lower bound estimator could result in optimizing 0 to maximize the bias of the estimator instead of the true objective. Our experiments demonstrate that this causes IWAE to often fail to optimize the objective unless use a large amount of computation.
min (PL (x)ll"(")"""((( (x)) min (=)) pe))) log p" (x:]
Modifying IWAE. The bias of the IWAE estimator can be interpreted as the KI. between an apportance-weigghed approximate posterior yrw (2;ar) implicitly defined by the encoder and the true posterior p(2(2) (Domke & Sheldon, 2018). Both the encoder and decoder parameters can therefore affect this bias. In practice, we find that the encoder optimization proceeds timescale than the decoder optimization: i.e., the encoder can match ęrw (2; x) to the decoder's s posterior p(z(ar) more quickly than the latter can match an objective. For this reason, we train the encoder to reduce bias and use minimax training objective
faster
min max Bz=p(zr) [IWAEx (x) llgg (x)]
(15)
p(z,2) ((:;))
Though this is still a lower bound with unbounded bias, it makes for stronger baseline than op- timizing g(2; 2) in the same direction as p(x, 2). We find that this approach can indeed work in practice, but requires setting & extremely high.
We choose "funnel" target distribution (Fig- ure 1) similar to the distribution used as benchmark for inference in Neal al. (2003), where p" has support in R² and is defined p" (x1, = (xr 0, 1.352 (a)2; 0, e221
IWAE (K= 15)
IWAE (K = 1500)
Estimate of KL[pallp
75)
SUMO
IWAE (k=
SUMO 150)
We use neural networks with one hidden layer of 200 hidden units and tanh activations for both the encoder and decoder networks. We
-2
use 20 latent variables, with p(z), px(1))) and (=(;)) all being Gaussian distributed.
6000
Training Iteration
Figure shows the learning curves when using IWAE and SUMO. Unless l: is very large, IWAE will at some point start optimizing the bias instead of the actual objective. The reverse KL is non-negative quantity, any estimate significantly below zero can be attributed to the unbounded bias. On the other hand, SUMO correctly optimizes for the objective even with a small expected cost. Increasing the expected cost for SUMO simply reduces variance. We also found that if A: is set sufficiently large, then IWAE can work when we train using the minimax objective in equation 15, suggesting that a sufficiently debi- ased estimator
Figure 2: Training with reverse KL requires min- imizing log p(z). SUMO estimates are unbiased and trains well, but minimizing the lower bound IWAE with small i leads to estimates of -co.
also work in practice. However, this requires much more compute and likely does not scale compared to SUMO. We also visualize the contours of the resulting models in Figure 1. For IWAE, we visualize the model few iterations before it reaches numerical instability.
8

---

Published as a conference paper at ICLR 2020
Expecteed Rewarr
Expection Cenected Retwe
Expect wed Beeto
(SUMO)
Automgreente Proppensert
(SJNDD)
Artamgrensioe troepensesk
LVN (SUNO)
Arromggeeske Intependere
1000
1903
1089
1800
QPBO w/ 100 variables
QPBO w/ 500 variables
QPBO w/ 500 variables & Entropy regularization
Figure 3: Latent variable policies allow faster exploration than autoregressive policy models, while being more expressive than independent policy. SUMO works well with entropy regularization, whereas IWAE is unstable and converges to similar performance as the non-latent variable model.
LATENT VARIABLE POLICIES FOR COMBINATORIAL OPTIMIZATION
Let us now consider the problem of finding the maximum of non-differentiable function, a special case of reinforcement learning without an interacting environment. Variational optimization (Staines & Barber, 2012) can be used reformulate this as the optimization of parametric distribution,
max R(x) 2 max (x)) [R(z)),
(16)
Jr
which is now a differentiable function with respect to the parameters 9, whose gradients can be estimated using a combination of the REINFORCE gradient estimator and the SUMO estima- tor (equation 13). Furthermore, entropy regularized reinforcement learning-where we maximize R(x) + (H( with H(pp) being the entropy of (a)) encourages exploration and inherently related to minimizing reverse KL objective with the target being exponentiated reward (Norouzi et al., 2016).
For concreteness, we focus on the problem of quadratic pseudo-Boolean optimization (QPBO) where the objective is to maximize
R(x) = Σ Ewi(x:) +
Σ way (wn((ss i<f
(17)
where {2,}/=1 E {0.1} are binary variables. Without further assumptions, QPBO NP- hard (Boros & Hammer, 2002). As there exist complex dependencies between the binary variables and optimization of equation 16 requires sampling from the policy distribution po (1). a model that is both expressive and allows efficient sampling would be ideal. For this reason, motivate the use of latent variable models with independent conditional distributions, trained using the SUMO estimator. Our baselines are an autoregressive policy, which captures dependencies but for which sampling must be performed sequentially, and an independent policy, which is easy to sample from but captures no dependencies.
d
(Z) = / /I pa (=:))))))
PAutorog (x) = II p( (x)))
Pindop (x) = II p(z:)
1=1
1=1
We note that Haarnoja et al. (2018) also argued for latent variable policies in favor of learning diverse strategies but ultimately made use of normalizing flows which do not require marginalization.
We constructed one problem instance for each d E {100, 500}. which we note are already intractable for exact optimization. For each instance, we randomly sampled the weights ws and wij uniformly from the interval [-1, 1]. Figure shows the performance of each policy model. In general, the independent policy quick to converge to a local minima and is unable to explore different re- gions, whereas more complex models have better grasp of the "frontier" of reward distributions during optimization. The autoregressive model works well overall but is much slower to train due to its sequential sampling procedure; with d 500, it is 19.2x slower per iteration than SUMO. Surprisingly, we find that estimating the REINFORCE gradient with IWAE results in decent perfor- mance when entropy regularization is present. With entropy regularization, all policies improve

---

Published as a conference paper at ICLR 2020
significantly; however, training with IWAE in this setting results in performance similar to the inde- pendent model, suggesting that it may not be using the latent state. On the other hand, SUMO works with both REINFORCE gradient estimation and entropy regularization, albeit at the cost of slower convergence due to variance.
8 CONCLUSION
We introduced SUMO, a new unbiased estimator of the log probability for latent variable models, and demonstrated tasks for which this estimator performs better than standard lower bounds. Specif- ically, we investigated applications involving entropy maximization where a lower bound performs poorly, but our unbiased estimator can train properly with relatively smaller amount of compute.
In the future, we plan to investigate new families of gradient-based optimizers which can handle heavy-tailed stochastic gradients. It may also be fruitful to investigate the of convex combination of consistent estimators within the SUMO approach, as any convex combination is unbiased, or to apply variance reduction methods to increase stability of training with SUMO.
ACKNOWLEDGEMENTS
This work was partially funded by NSF I1S-1421780. Y.L and J.Z were supported by the NSF China Project (No. 61620106010) Beijing NSF Project (No. L172037), the JP Morgan Faculty Research Program and the NVIDIA NVAIL Program with GPU/DGX Acceleration.
REFERENCES
James Arvo and David Kirk. Particle transport and image synthesis. ACM SIGGRAPH Computer Graphics, 24(4):63-66 6-, 1990.
Robert Bamler, Cheng Zhang, Manfred Opper, and Stephan Mandt. Perturbative black box varia- tional inference. In Advances in Neural Information Processing Systems. 2017.
Alex Beatson and Ryan Adams. Efficient optimization of loops and limits with randomized telescoping sums. In International Conference on Machine Learning, 2019.
Mikoaj Bikowski, Dougal J. Sutherland, Michael Arbel, and Arthur Gretton. Demystifying MMD GANs. In International Conference on Learning Representations, 2018.
David M Blei, Andrew Y Ng, and Michael I Jordan. Latent Dirichlet allocation. Journal of Machine Learning Research, 30ann):93-10022 2003.
Endre Boros and Peter I. Hammer. Pseudo-Booles optimization. Discrete Applied Mathematics, 123(1-3):155-225, 2002.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders, In Inter- national Conference on Learning Representations, 2016.
Ricky TQ Chen, Jens Behrmann, David Duvenaud, and Jörn-Henrik Jacobsen. Residual flows for invertible generative modeling. Advances in Neural Information Processing Systems, 2019.
Siddhartha Chib. Marginal likelihood from the Gibbs output. Journal of the American Statistical
Association, 90(432):1313-132) 1995.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the Royal Statistical Society: Series B (Methodological), 39(1): 1-22, 1977.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
Justin Domke and Daniel R Sheldon. Importance weighting and variational inference. In Advances in Neural Information Processing Systems, pp. 4470 2018.
10

---

Published as a conference paper at ICLR 2020
Paul Fearnhead, Omiros Papaspiliopoulos, and Gareth O Roberts. Particle filters for partially ob- served diffusions. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 70 (4):755-777, 2008.
George E Forsythe and Richard A Leibler. Matrix inversion by a Monte Carlo method. Mathematies of Computation, 4(31):127-129, 1950.
Tuomas Haarmoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hierarchical reinforcement learning. In International Conference on Machine Learning, pp. 1846-1855, 2018.
Insu Han, Haim Avron, and Jinwoo Shin. Stochastic Chebyshev gradient descent for spectral opti- mization. In Advances in Neural Information Processing Systems, pp. 7386-7396, 2018.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerehner. beta- VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Machine Learning, 2017.
Emiel Hoogeboom, Jorn WT Peters, Rianne van den Berg. and Max Welling. Integer discrete flows
and lossless compression. arXiv preprint
2019.
Chin-We Wei Huang. David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Learning, 2018
Chin-Was Huang, Kris Sankaran, Eeshan Dhekane, Alexandre Lacoste, and Aaron Courville. Hier- archical importance weighted autoencoders. In International Conference on Machine Learning, pp. 2869-2878, 2019.
Pierre E Jacob and Alexandre Thiery. On nonnegative unbiased estimators. The Annals of Statis- tics, 43(2):769-784, 2015.
Pierre E Jacob, John O'Leary, and Yves F Atchadé. Unbiased Markov chain Monte Carlo with 2017.
couplings. arXiv preprint arXiv:
Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan Adams, and Sandeep R Datta. Com- posing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems, pp. 2946-2954, 2016.
Herman Kahn. Use of different Monte Carlo sampling techniques. Santa Monica, CA: RAND Corporation, 1955. URL https: /ttp://wwwwww////////////w rand. org/pubs prpubs/papers/eee httpps/ppprspeee// html.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In Internatioona Conference on Learning Representations, 2014.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems, pp. 10215-10224, 2018.
Julius Kuti. Stochastic method for the numerical study of lattice fermions. Physical Review Letters, 49(3):183, 1982.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 132---333) 2015.
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1) 1):2278-2324, 1998.
Peter Lenk. Simulation pseudo-bias correction to the harmonic mean estimator of integrated likeli- hoods. Journal of Computational and Graphical Statistics, 18(4):941960, 2009.
Yucen Luo, Tian Tian, Jiaxin Shi, Jun Zhu, and Bo Zhang. Semi-crowdsoureed clustering with deep generative models. In Advances in Neural Information Processing Systems, pp. 3212-3222. 2018.
Anne-Marie Lyne, Mark Girolami, Yves Atchadé, Heiko Strathmann, Daniel Simpson, et al. On likelihoods. Statistical
Russian roulette estimates for Bayesian inference with science, 30(4):443-467, 2015.
11

---

Published as a conference paper at ICLR 2020
Don McLeish. A general method for debiasing Monte Carlo estimator. Monte Carlo Methods and Applications, 17(4):301315,, 2011.
Thomas P Minka. Expectation propagation for approximate Bayesian inference. In Proceedings of the Seventeenth Conference Uncertainty in Artificial Intelligence, pp. 362-369, Morgan Kaufmann Publishers Inc., 2001.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International Conference Machine Learning, pp. 1928-1937, 2016.
lain Murray and Ruslan Salakhutdinov. Evaluating probabilities under high-dimensional latent vari- able models. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou (eds.), Advances in Neural Information Processing Systems 21, pp. 1137-1144, 2009.
Radford M Neal. Annealed importance sampling. Statistics and Computing, 11(2):225-139 2001.
Radford M Neal et al. Slice sampling. The Annals of Statistics, 31(3):705-767, 2003.
Michael A Newton and Adrian E Raftery. Approximate Bayesian inference with the weighted like- lihood bootstrap. Journal of the Royal Statistical Society: Series (Methodological), 56(1):3-26. 1994.
Mohammad Norouzi, Samy Bengio, Navdeep Jaitly, Mike Schuster, Yonghui Wu, Dale Schuurmans, et al. Reward augmented maximum likelihood for neural structured prediction. In Advances In Neural Information Processing Systems, pp. 1723-1731, 2016.
Sebastian Nowozin. Debiasing evidence approximations: On importance-weigheee autoencoders and jackknife variational inference. International Conference on Learning Representations, 2018.
Aaron van den Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George van den Driessche, Edward Lockhart, Luis C Cobo, Florian Stimberg, et al. Parallel Wavenet: Fast high-fidelity speech synthesis. In International Conference on Machine Learning, 2018.
Tom Rainforth, Robert Comiish, Hongseok Yang, Andrew Warrington, and Frank Wood. On nesting Monte Carlo estimators. In International Conference on Machine Learning, 2018a.
Tom Rainforth, Adam R Kosiorek, Tuan Anh Le, Chris J Maddison, Maximilian Igl. Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In International Con- ference on Machine Learning, 2018b.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of Adam and beyond. In International Conference on Learning Representations, 2018.
Danilo Jimenez. Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference Machine Learning, 2015.
Chang-han Rhee and Peter W Glynn. new approach to unbiased estimation for SDEs. In Pro- ceedings of the Winter Simulation Conference, pp. 17. Winter Simulation Conference, 2012.
Chang-han Rhee and Peter W Glynn. Unbiased estimation with square root convergence for SDE
models. Operations Research, 63(5): 1043, 2015.
Francisco JR Ruiz, Michalis K Titsias, and David M Blei. Overdispersed black-box variational in- ference. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, 2016.
Tomasz Rychlik. Unbiased nonparametric estimation of the derivative of the mean. Statistics & probability letters, 10(4):329-33, 1990.
Tomasz Rychlik. A class of unbiased kernel estimates of a probability density function. Applica- fiones Mathematicae, 22(4):485 1995.
12

---

Published as a conference paper at ICLR 2020
John Skilling. Nested sampling for general Bayesian computation. Bayesian Analysis, 1(4):833- 859, 2006.
Jerome Spanier and Ely M Gelbard. Monte Carlo Principles and Neutron Transport Problems. Addison-Wesley Publishing Company, 1969.
Joe Staines and David Barber. Variational optimization. arXiv preprint arXiv: 12lp.45027 2012.
Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time. arXiv preprint arXiv:1705.00209 2017.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 61(3):61 (0036611-222 1-622, 1999.
Dustin Tran, Keyon Vafa, Kumar Krishna Agrawal, Laurent Dinh, and Ben Poole. Discrete flows: Invertible generative models of discrete data, arXiv preprint arXiv:1905. 2019.
George Tucker, Andriy Mnih, Chris Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems, pp. 2627-2636, 2017.
George Tucker, Dieterich Lawson, Shixiang Gu, and Chris Maddison. Doubly reparameterized gradient estimators for Monte Carlo objectives. In International Conference Learning Repre- sentations, 2018.
Wolfgang Wagner. Unbiased Monte Carlo evaluation of certain functional integrals. Journal of Computational Physics, 71(1):21-33, 1987.
Colin Wei and lain Murray. Markov Chain Truncation for Doubly-Intractable Inference. In Pro- ceedings of the 20th International Conference Artificial Intelligence and Statistics, 2017.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning, pp. 681-688, 2011.
Ronald J Williams. Simple statistical gradient-folloinng algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229 256, 1992.
Ronald J Williams and Jing Peng. Function optimization using connectionistt reinoooormnetriiinnneeeeeeee learning algorithms. Connection Science, 3(3):241-268, 1991.
Kai Xu, Akash Srivastava, and Charles Sutton. Variational Russian roulette for deep Bayesian nonparametrics. In International Conference on Machine Learning, pp. 6963-6972, 2019.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. PhD thesis, figshare, 2010.
13

---

Published as a conference paper at ICLR 2020
A
APPENDIX
A.1
DERIVATION OF SUMO
Let
1 K
pa =k) po po ((2xi=)
Bg (I) = E s (=
log
where 21,.., Fi are sampled independently from ((;;)). And we define the k-th term of the infinite series Åu (x) = E, [[WAE++ - E [IWAE (z)]. Using the properties of IWAE in equation 6,
we have { (I) 2 0, and
[x(+)| = Σ [x(x)= lim IWAE - E. [WAE, (x)]
k=1 k-+00
(o) - E. [WWAE (x)] < oo,
(18)
which means the series converges absolutely. This is sufficient condition for finite expectation of the Russian roulette estimator (Chen et al. (2019); Lemma 3). Applying equation 7 to the series:
Dggpo(>) = (IAA,, + Σ (s(x)
(19)
(.( P(K: k)
IIAE, (I)] + E BK-p(R)
Σ
(20)
Egg [IWAE x1 (x) IWAE& (x)
= -K<p((K
Eg [WAE, (x)l +
(21)
P(K 2 k:)
IWAR ( (x) - IWAF& (x)
= -KK~o()
E
IWAE, (x)
(22)
P(K 2 k)
Let (x(x) (1 IWAF& (x:), Hence our estimator is constructed:
K
(x
K K~ K~)(K), lid g(æ;)z).
SUMO( IWAE, (x) +
(23)
P(K2
And it can be easily seen from equation 22 and equation 23 that SUMO is an unbiased estimator of the log marginal likelihood:
E K<upl K) (K..p(K))).....-ddae) [SUMO())) log Pa(x).
(24)
A.2 CONVERGENCE OF A
We follow the analysis of JVI (Nowozin, 2018), which applied the delta method for moments to show the asymptotic results on the bias and variance of IWAE, both at a rate O We build on this analysis analyze the convergence of Ak.
Let wi p(n) as )((s)) and we define Y3 1 Eh we the sample mean and we have E|y,] k.
((aaa))
B[w] =
IWAE, = log Yx = loog lu + (Yo - 11]
20
(-1)²
(25)
log Σ
(Yr
14)*
tue
We note that we rely on ||Y - < 1 for this power series to converge. This condition was implic- itly assumed, but not explicitly noted, in (Nowozin, 2018). This condition will hold for sufficiently
14

---

Published as a conference paper at ICLR 2020
large k so long the moments of W, exist: could bound the probability ||x=~Ju) 1 by Cheby- shev's inequality or by the Central Limit Theorem. We use the central moments 1/ == [Y n) and Mr El(w 11)²] for t 2. E E(IWAE++1 IWAE2)2
(26)
72
1 (Ve+1 - -413 - -log/ 1 +
(-1) txe
= E
log 1
Σ
(Yx - 11)2)
(27)
tu
E
[(Yx - 11²- - -11++1 - -1)9
(28)
ty
Expanding Eq. 28 to order two gives
1
1
2
EA
E
(Yx- p Yk+1 p))
[Yrx - 11)2 -(k++1- u]
+ to(k
(29)
11
24²
1
(Y + Ya+1 2Y)(Yr Yauti ]²
E
Yu+1 Yaa +
+ o(A=
(30)
112
241
73
E
2(Y - -xx +
(Yn + Yx+1)(X - Ye+1)
+ o(k
(31)
11²
2p
Since we use cumulative sum compute Yx and Ys+1, we obtain Yk+11 Kast1m111
k+l
EA
E
Wk+1
Wk+1 +
26+1 k+1
Ei
Wkk
Wh++1
-
2
+ o(k-2)
(32)
112
k +
2kju
We note that #6+1-1
((
and
Me ETT
Σ
mik
= O1). Therefore As is O and EAA
2kp
0
A.3 CONVERGENCE OF AxA,
Without loss of generality, suppose j 2 k 1,
E E
(-1)²
[(rs 14)* - -(Y++1 -1)
Σ
(-1)* tue
[(Y, - 11)² - -(X++ -1)²
(33)
For clarity, let C+ Yx- yu be the zero-mean random variable. Nowozin (2018) gives the relations
E|C = 72
42
(34)
K
BIC
= 73
H3
(35)
h2
3p2
14 3p3
B|CH| 71 122
(36)
h-3
1-)
EA:A, E
(c
ck+1)
Σ
(C)
0++1)
(37)
tu
tu
Expanding both the sums inside the brackets to order two:
EALA, ss (CK+1 Cx) C2)
(1)
112
E
(C2+1
C2)(CH+1 C2)
(2)
24
1
E (Co+1
C2)(CH 03)
(3)
24
1
+
E
(C++1
(2) C3)
(4)
44
15

---

Published as a conference paper at ICLR 2020
We will proceed by bounding each of the terms (1). (2), (3), (4). First, we decompose Cj. Let
BR. = Ei= - -w).
Cj = h:Ck
j E(us
(w: 11)
k Ck Cx
1
Σ (wi 11) h -C Bk. (38)
3 1=++1
f=++1
We know that Bu.j is independent of Cs and E[Bry] = 0, implying E[Cr Br.sl= = O. Note C3 = k²
jz O2 + Cr B
Now we show that (1) is zero:
1
1
h ++
El (Ck+1 (c)(1 C,)l
E
Ck+1
+ Cst1Bs41At1
j ++1
+ 1
C - -B,a-1R41 - Ck k
j
j+1
k c CCRBAA
C&BJ+1
++1
k Ch+1+ C "ii+ +1)
R
12
j(
k+1
122
+
122
k j(++1) =0
k
n12
j(++1)++1
We now investigate (2): 1
E[ 2/r3 (h - 0B)( - - -
C k
+ CBB - - -C
2/,3
j+1
++1
- C
C2
+ Ch 11, 5
k+1 [
k+1
k
B1
C2 0h-11 + C2 +] j(i
112
P3
Ps
240²
kjG +1)
(k: 1)j(j 1)
2p3
k(k + 1)5(1 + 1)
We now show that (3) is zero:
E
(C++1 -CH = E(C,. C9 C.+1C7111111 - C.C + -
CCCC++ +.CHH1))
2103
24r²
1
Ps
Ps
H3
2
jz
1)²
jz
0
Finally, we investigate (4):
(Ch+1 - -B)(CCH -
Anz
Using the relation in equation 36, we have
22
2k CrBj, + +B,,I
Jz 0
+
(39)
j
122
(j k) k
J274 + 72
(40)
j2
(2k j 3) + 14
(41)
j2b
16

---

Published as a conference paper at ICLR 2020
1 4u
(2k: + j- -3)p2 + p4
+j 2) Mr
E (C2+1
C3)1
52h
(i+ + 1)2h
(2k + 1 1-1113 + 14
(2h j)n pr
j² - 1)
(j (i+ 1++ 1)
(12 5j - -3)1
44
j + 1)3k(k +1 (j + 1)2k(k +)
je [u3 - m44) -
jā(j + +1)2k(k + 1)
=O(j= 3)
In summary, EALA, is 3---))-
A.4 GRADIENT OF SUMO
Assume that V is bounded: it is sufficient that V,IWAE; is bounded and that the sampling
probabilities are chosen such that the partial sums of V. Ag converge, i.e. P(K 2 k) > ckl Vo Il
k)
for some constant Then we have E [V>SUMO(z)] = V EE [SUMO Ve log (x directly by the dominated convergence theorem, as long as SUMÓ is everywhere differentiable, which satisfied by all of our experiments. If ReLU neural networks are to be used, one may be able to show the same property using Theorem of Bikowski al. (2018), assuming finite higher moments and Lipschitz continuity.
A.5 CONVERGENCE OF V [WAAE++1 IWAE,
The IWAE log likelihood estimate is:
1 h T
( 21)
Cx = loog
qx (x)
The gradient of this with respect to A, where À is either 1 or v/, is
dCk
k
d (x(x,
po (2,3.
qo xx)
We abbreviate w. do (2: I and M dus. In both ف and x cases, it suffices to treat the dx
(z.es)
wy and it as i.i.d. random variables with finite variance and expectation. Being a likelihood ratio, wy could be ill behaved when the importance sampling distribution (= (r) is is partieularly mis-
matched from the true posterior p(z.(z)) Po(
pa(2,2; However, the analysis from IWAE (Burda p=(2.a)
al., 2016) requires assuming that the likelihood ratios w; zu (=,l)) are bounded, and adopt this assumption. Reasoning about when this assumption holds, and the behavior of IWAE-like esti- mators when it does not, is an interesting area for future work.
Consider the differences between two gradients: we label A as follows:
dL
dCr
A
dd
dX
We have:
1
A
Uk+l +
Vi
1= Wi
1=1 Wi
Σ4 1
Wi
1=1
k
Wk++1
Σ
2 1=
k-1
Wi
Vk+1 +
1=1 w;)
k+1
Σ 1==
Wi)
1=1
1/
17

---

Published as a conference paper at ICLR 2020
SUMO (A1
inqpote)
SUMO gommerr MAKE
SUMO MIKE
675-
97.3
$7.0
90.2
Endercce ectmaate of
sxiaats
00.0-
85.2-
04
85.5
Figure 4: A comparison of SUMO estimations with different distributions and IWAE estimations of test negative log-likelihood on a trained model with IWAE, objective on MNIST. The expected cost is K + 5 for each evaluation. The results averaged over 100 runs (mean in bold and std shaded).
We again let Y denote the kth sample mean E; wi Then:
4
Vk+1 +
Wh+1
kYk
(+++)
The sample means Y& and jo have finite expectation and variance. The variance vanishes k + (but the expectation does not change).
Vk+1
Wk- + 1
112
2
12
Yx
k+ Y
Let
Mk+1
k
Wk+1
Vk
dx
Yk
Y
1 Var(tbe)
=
E =
12
11 + 1.2
The second term vanishes at a rate strictly faster than 3 the variance of Φe goes to zero as + cc. But the first term does not: d% is biased estimator of do so E&O does change with k but it does not necessarily go to zero:
k
WEv
Ev
Edx = E
Ew
k 1 (Bw)2
Ew
Thus, E| is at most O (.).
A.6 EMPIRICAL CONFIRMATION ON THE CONVERGENCE OF EA, AND E|JA
We measure the A and | 222 toy example to verify the convergence rates empirically. We re-implement the toy Gaussian example from Rainforth et al. (2018b); Tucker et al. (2018). The generative model is po (r , 2) W( 0, (x12,, I). where both æ and in RRD. The encoder is (<(2]z) = N Az + b, 2 J), where 6 = (A, b). The synthetic dataset was generated with = 20 and N 1000 data points using the model parameter Omne from standard Gaussian. Alongside || 1, we plot several reference convergence rates such as Ol)ik), e > 1, and O c < 1, as visual guide. The results are shown in Figure Following the setup Rainforth et al. (2018b), we sample group of model parameters close to the optimal values which are perturbed by Gaussian noise from (0, 0.012). The gradient Ag is taken w.r.t. the model parameter
3
A.7 BIAS- VARIANCE TRADEOFF VIA GRADIENT CLIPPING
While SUMO is unbiased, its variance is extremely high potentially infinite. This property leads to poor performance compared to lower bound estimates such as IWAE when maximizing log- likelihood. In order to obtain models with competitive log-likelihood values, gradient clipping. This allows us to ignore rare gradient samples with extremely large values due to the heavy-tailed nature of its distribution.
make of
18

---

Published as a conference paper at ICLR 2020
10.
126
10.5
10
10
10a
10a4
10"7
19"4
0.3
10"
0.3
10P
(a) Mean of estimated of EAĪ with increasing over ten random trials with 1000 samples per trial. X and axis are on log scalc. Empirically the convergence rate of Ag is between O(1/k²) and ((1/k²).
(b) Mean of estimated E| with increasing over ten trials with 1000 samples per trial. Empiri- cally the convergence is faster than theorctical analy- sis (11/k²)
Figure 5: Empricial validation of the convergence rate of the norms of and A
Gradient clipping introduces bias in favor of reduced variance. Figure 6 shows how the performance changes as a function of the clipping value, and more importantly, the percentage of clipped gradi- ents. As shown, neither full clipping nor no clipping are desirable. We performed this experiment after reporting the results in Table so this grid search was used our experiments. As bias is introduced, we do not use gradient clipping for entropy maximization or policy gradient (REINFORCE).
hyperparameters for
Gradient clip percentage
86.75
86.50
86.50
86.25
86.25
80.75
85.75
85.50
85.50
85.00
85.00
84.75
84.75
10a
10a
107
10a
0
10
20
30
50
70
80
90
100
Gradient dipppig norm
Gradient clip percentage %6
Figure 6: Test negative log- ligelioed against the gradient clipping norm and clipping percentage, when training with SUMO (k=15).
A.8 EXPERIMENTAL SETUP
In density modeling experiments, all the models are trained using a batch size of 100 and the AMS- Grad optimizer (Reddi et al., 2018) with parameters Ir 0.001, 81 0.9, Bg 0.999 and 10 The learning rate reduced by factor 0.8 the validation likelihood does not improve for 50 epochs. We use gradient norm scaling in both the inference and generative networks. We train SUMO using the same architecture and hyperparameters as IWAE except the gradient clipping norm. We set the gradient norm 5000 for encoder and {20, 40. 60) for decoder in SUMO. For IWAE, the gradient norm is fixed to 10 in all the experiments. We report the performance of models with early stopping if improvements have been observed for 300 epochs the validation set.
We add additional plots of the test NLL against the norm and percentage of gradients clipped for the decoder in Figure 6. The plot is based MNIST with expected number of compute 15. Gradient clipping was not used the other experiments except the density modeling ones, where was used as simple tool to obtain better bias-variance trade-off.
19

---

Published as a conference paper at ICLR 2020
A.8.1 REVERSE KL AND COMBINATORIAL OPTIMIZATION
These tasks the same encoder and decoder architecture: one hidden layer with tanh non- linearities and 200 hidden units. We set the latent state to be of size 20. The prior is a standard Gaussian with diagonal covariance, while the encoder distribution is Gaussian with parameterized diagonal covariance. For reverse KL, we used independent Gaussian conditional likelihoods for p(x[2), while for combinatorial optimization we used independent Bernoulli conditional distribu- tions. We found it helps stablize training for both IWAE and SUMO to remove momentum and used RMSprop with learning rate 0.00005 and epsilon 1e-3 for fitting reverse KI.. We used Adam with learning rate 0.001 and epsilon le-3, plus standard hyperparameters for the combinatorial optimiza- tion problems. SUMO used an expected compute of 15 terms.
20


Reviews:
[{"review_id": "rJgxdFWhFS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThis allows marginally better optimization of certain models than IWAE with a much smaller average number of samples, and (more importantly) opens new possibilities such as entropy maximization which are not well-served by lower bounds like IWAE. This is a very nice advantage of this estimator.\n\nOne complaint about this class of estimators in general is: yes, the exact SUMO procedure is technically unbiased. But in practice, if SUMO takes more than, say, a day of compute time – something that will happen with extremely small but nonzero probability – then the user will kill it. And the SUMO estimator conditioned on taking less than a day of compute time is actually biased. This also likely means that, though unbiased, these estimators can be potentially skewed or otherwise \"unpleasant.\" For SUMO in particular, the estimator with $K$ truncated probably has bias bounded based on the bias of IWAE with batch size equal to the truncation point, which is likely quite small. But it would be nice to understand this a little more. (Perhaps it's been studied by some of the recent cited work on these types of estimators.)\n\nRelatedly, you don't prove that this estimator has a finite variance, and in fact it seems plausible theoretically that it might be infinite. Like the \"notorious\" harmonic mean estimator, this is troubling. It seems that things are okay in practice, but how can we tell whether the variance is really finite or not? I don't know if there's a good answer, but one diagnostic might be something like the one suggested at https://stats.stackexchange.com/a/9143 . (Your comments about occasional \"bounded but very large\" gradient estimates are troubling in this respect, depending on what exactly you mean by \"bounded\".) When you do gradient clipping, the estimator of course then has a finite variance, but can we get some sense of how much bias that introduces?\n\nOverall, though, I think this is a very nice new estimator that is both well-founded – despite leaving some questions open – and likely to be practically useful. Given that it is also extremely on-topic for ICLR and novel, I'm rating the paper as \"accept.\"\n\n(For slightly more detail on the \"thoroughness assessment\": I did not really check the proofs in the appendix, but did pay attention to the derivations in the main body.)\n\n\nSmaller notes:\n\n- Top of page 3: you comment that SGD \"requires unbiased estimates of\" gradients of the log-density. In fact, SGD can be shown to work with biased gradient estimators, with suboptimality in the results depending on the bias; see e.g. Chen and Luss, http://arxiv.org/abs/1807.11880 .\n\n- In the definition of $\\tilde{Y}$, above (7): it might make more sense to define $\\tilde{Y}$ with some recursive scheme, rather than as an estimator that either computes one of the $\\Delta$ values or infinitely many of them.\n\n- Start of 3.1: presumably $\\Delta_k$ is what converges absolutely, not IWAE?\n\n- Start of 3.2: as you note, it is clearly not true that the $\\Delta_k^g$ are independent. But you don't really \"assume independence\" – the Russian roulette estimator is still a valid estimator, just perhaps not the optimal among that class. It would be better to say something like that since it seems that the $\\Delta_k^g$ are *nearly* independent (or at least nearly uncorrelated), the Russian roulette estimator is probably at least a reasonable choice.\n\n- Re: the discussion after (9) and in (12), as well as a few other places: I think you show in Appendix A.3 that $\\mathbb E[ \\nabla \\operatorname{SUMO}(x) ]$ exists, but you don't show that it equals $\\nabla \\mathbb E[ \\operatorname{SUMO}(x) ]$. This is likely true, particularly if $q$ and $p$ are each everywhere-differentiable, and it's totally fine if you don't want to prove it out formally, but it would be worth at least a footnote that this is a thing that requires proof. (See e.g. https://arxiv.org/abs/1801.01401 Theorem 5, for a formal result of this type supporting ReLU activations, which you may be able to just use directly.)", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJlZ3eOBYB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents an unbiased estimator of marginal log likelihood given a latent variable model. \nThe method extends the importance-weighted log marginal using the Russian roulette estimator. \nThe marginal log probability estimator is motivated for entropy maximized density estimation and use of REINFORCE (log-derivative) gradient for learning a policy with a latent variable. \n\nThe paper is well-organized and provides a contribution for optimizing latent variable models in certain scenarios. \nThus, I vote for its acceptance. \n\nSome questions follow.\n1) Is it trivial to show the absolute convergence of \\Delta_k(x) series?\nThe absolute convergence is mentioned above equation (8), I am not convinced of this point. \nPerhaps, if its expectation with respect to q(z;x) is applied, this can be shown from equation (6). \nOtherwise we need some assumption on q(z;x) like q(z;x) is reasonably close to p(z) or p(z|x).\n\n2) How was parameter m set for the experiments?\n\n3) I assume the expectation operator is taken over z and K in equations (9, 12). \nIs this correct? An explicit notation should be informative.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "HklgDk-gcB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors consider the unbiased estimation of log marginal likelihood (evidence) after integration of latent variable. On top of the importance-weighted autoencoder (IWAE), which is only guaranteed to be asymptotically unbiased, the authors propose to use Russian Roulette estimator (RRE) to compensate the bias caused by the finite summation.\n\nThe proposed method is interesting and can be applied in many other estimators with similar properties as Eq. (6). Bias compensation using RRE is interesting, but it seems there must be many literatures that took advantage of using RRE to improve estimators. The authors have to be thorough in presenting previous research and explaining the authors’ contribution that is distinguished from those.\n\nThe authors showed synthetic and real application of the estimator, but one concern is the variance. Unbiasedness with finite samples often fails because of the variance, and regularization is often useful rather than correcting bias---If unbiasedness in important, regularization definitely breaks the unbiasedness---. The only discussion about variance is a few lines in page 3 after Eq. (7), but it is unclear how the variance problem is mitigated or why the problem does not suffer high variance.\n\nIn general, this paper is well-written and dealing with important problem with interesting method. Several analysis for understanding the advantages of using the proposed method is insufficient.\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new way to train latent variable models. The standard way of training using the ELBO produces biased estimates for many quantities of interest. The authors introduce an unbiased estimate for the log marginal probability and its derivative to address this. The new estimator is based on the importance weighted autoencoder, correcting the remaining bias using russian roulette sampling. The model is empirically shown to give better test set likelihood, and can be used in tasks where unbiased estimates are needed. 

All reviewers are positive about the paper. Support for the main claims is provided through empirical and theoretical results. The reviewers had some minor comments, especially about the theory, which the authors have addressed with additional clarification, which was appreciated by the reviewers. 

The paper was deemed to be well organized. There were some unclarities about variance issues and bias from gradient clipping, which have been addressed by the authors in additional explanation as well as an additional plot. 

The approach is novel and addresses a very relevant problem for the ICLR community: optimizing latent variable models, especially in situations where unbiased estimates are required. The method results in marginally better optimization compared to IWAE with much smaller average number of samples. The method was deemed by the reviewers to open up new possibilities such as entropy minimization. 
Decision: Accept (Spotlight)