Original: SUMO: U NBIASED E STIMATION OF L OG M ARGINAL P ROBABILITY FOR L ATENT V ARIABLE M ODELS
Revised: SUMO: U NBIASED E STIMATION OF L OG M ARGINAL P ROBABILITY FOR L ATENT V ARIABLE M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: The standard variational lower bounds used to train latent variable models produce biased estimates of most quantities of interest
Revised: The standard variational lower bounds used to train latent variable models produce biased estimates of most quantities of interest

Original: We introduce an unbiased estimator of the log marginal likelihood and its gradients for latent variable models based on randomized truncation of inﬁnite series
Revised: We introduce an unbiased estimator of the log marginal likelihood and its gradients for latent variable models based on randomized truncation of inﬁnite series

Original: If parameterized by an encoder-decoder architecture, the parameters of the encoder can be optimized to minimize its variance of this estimator.
Revised: If parameterized by an encoder-decoder architecture, the parameters of the encoder can be optimized to minimize its variance of this estimator.

Original: 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.
Revised: 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.

Original: This estimator also allows use of latent variable models for tasks where unbiased estimators, rather than marginal likelihood lower bounds
Revised: This estimator also allows use of latent variable models for tasks where unbiased estimators, rather than marginal likelihood lower bounds

Original: , are preferred such as minimizing reverse
Revised: , are preferred such as minimizing reverse

Original: KL divergences and estimating score functions.
Revised: KL divergences and estimating score functions.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Latent variable models are powerful tools for constructing highly expressive data distributions and for understanding how high-dimensional observations might possess a simpler representation.
Revised: Latent variable models are powerful tools for constructing highly expressive data distributions and for understanding how high-dimensional observations might possess a simpler representation.

Original: Latent variable models are often framed as probabilistic graphical models, allowing these relationships to be expressed in terms of conditional independence.
Revised: Latent variable models are often framed as probabilistic graphical models, allowing these relationships to be expressed in terms of conditional independence.

Original: 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.
Revised: 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.

Original: More recently there has been a surge of interest in probabilistic latent variable models that incorporate ﬂexible nonlinear likelihoods via deep neural networks (Kingma & Welling, 2013).
Revised: More recently there has been a surge of interest in probabilistic latent variable models that incorporate ﬂexible nonlinear likelihoods via deep neural networks (Kingma & Welling, 2013).

Original: These models can blend the advantages of highly structured probabilistic priors with the empirical successes of deep learning (Johnson et al., 2016).
Revised: These models can blend the advantages of highly structured probabilistic priors with the empirical successes of deep learning (Johnson et al., 2016).

Original: 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.
Revised: 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.

Original: It can be challenging, however, to ﬁt the parameters of a ﬂexible latent variable model, since computing the marginal probability of the data requires integrating out the latent variables in order to maximize the likelihood with respect to the model parameters.
Revised: It can be challenging, however, to ﬁt the parameters of a ﬂexible latent variable model, since computing the marginal probability of the data requires integrating out the latent variables in order to maximize the likelihood with respect to the model parameters.

Original: Typical approaches to this problem include the celebrated expectation maximization algorithm
Revised: Typical approaches to this problem include the celebrated expectation maximization algorithm

Original: (Dempster et al., 1977), Markov chain Monte Carlo, and the Laplace approximation.
Revised: (Dempster et al., 1977), Markov chain Monte Carlo, and the Laplace approximation.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, a 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;
Revised: However, a 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;

Original: Mnih et al., 2016; Norouzi et al., 2016) which requires minimizing the log probability of the samples under the model.
Revised: Mnih et al., 2016; Norouzi et al., 2016) which requires minimizing the log probability of the samples under the model.

Original: 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 are high dimensional; it is not as susceptible to underﬂow and it has gradients that are numerically sensible.
Revised: 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 underﬂow and it has gradients that are numerically sensible.

Original: However, the log transformation introduces some challenges: Monte Carlo estimation techniques such as importance sampling do not straightforwardly give unbiased estimates of this quantity.
Revised: However, the log transformation introduces some challenges: Monte Carlo estimation techniques such as importance sampling do not straightforwardly give unbiased estimates of this quantity.

Original: Nevertheless, there has been signiﬁcant 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;
Revised: Nevertheless, there has been signiﬁ- cant 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;

Original: Bamler et al., 2017; Nowozin, 2018).
Revised: Bamler et al., 2017; Nowozin, 2018).

Original: 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 inﬁnity.
Revised: 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 inﬁnity.

Original: In this work, we construct an unbiased estimator of the log marginal likelihood.
Revised: In this work, we construct an unbiased estimator of the log marginal likelihood.

Original: Although there is no theoretical guarantee that this estimator has ﬁnite variance, we ﬁnd that it can work well in practice.
Revised: Although there is no theoretical guarantee that this estimator has ﬁnite variance, we ﬁnd that it can work well in practice.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 efﬁcient to sample from.
Revised: 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 efﬁcient to sample from.

Original:  P RELIMINARIES 2. L ATENT VARIABLE MODELS
Revised:  P RELIMINARIES 2. L ATENT VARIABLE MODELS

Original: Latent variable models (LVMs) describe a distribution over data in terms of a mixture over unobserved quantities.
Revised: Latent variable models (LVMs) describe a distribution over data in terms of a mixture over unobserved quantities.

Original: Let p θ ( x ) be a family of probability density (mass) functions on a data space X , indexed by parameters θ .
Revised: Let p θ ( x ) be a family of probability density (mass) functions on a data space X , indexed by parameters θ .

Original: 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.
Revised: 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.

Original: In a latent variable model, p θ ( x ) is deﬁned 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 ) .
Revised: In a latent variable model, p θ ( x ) is deﬁned 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 ) .

Original: This conditional distribution is sometimes called an “observation model” or a conditional likelihood.
Revised: This conditional distribution is sometimes called an “observation model” or a conditional likelihood.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: 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).
Revised: 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).

Original: The downside of this mixing approach is that the integral may be intractable to compute, making it difﬁcult to evaluate p θ ( x ) — a quantity often referred to in Bayesian statistics and machine learning as the marginal likelihood or evidence .
Revised: The downside of this mixing approach is that the integral may be intractable to compute, making it difﬁcult to evaluate p θ ( x ) — a quantity often referred to in Bayesian statistics and machine learning as the marginal likelihood or evidence .

Original: 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.
Revised: 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.

Original: 2. T RAINING LATENT VARIABLE MODELS
Revised: 2. T RAINING LATENT VARIABLE MODELS

Original: Fitting a parametric distribution to observed data is often framed as the minimization of a difference between the model distribution and the empirical distribution.
Revised: Fitting a parametric distribution to observed data is often framed as the minimization of a difference between the model distribution and the empirical distribution.

Original: The most common difference measure is the forward Kullback-Leibler (KL) divergence; if p data ( x ) is the empirical distribution and p θ ( x ) is a parametric family, then minimizing the KL divergence  with respect to θ is equivalent to maximizing the likelihood:
Revised: The most common difference measure is the forward Kullback-Leibler (KL) divergence; if p data ( x ) is the empirical distribution and p θ ( x ) is a parametric family, then minimizing the KL divergence ( D KL ) with respect to θ is equivalent to maximizing the likelihood:

Original: [Equation]
Revised: [Equation]

Original: Thus the optimization problem of ﬁnding the MLE parameters θ comes down to maximizing the expected log probability of the data:
Revised: Thus the optimization problem of ﬁnding the MLE parameters θ comes down to maximizing the expected log probability of the data:

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: However, this requires unbiased estimates of ∇ θ log p θ ( x ) , which are not available for latent variable models.
Revised: However, this requires unbiased estimates of ∇ θ log p θ ( x ) , which are not available for latent variable models.

Original: Instead, a stochastic lower bound of log p θ ( x ) is often used and then differentiated for optimization.
Revised: Instead, a stochastic lower bound of log p θ ( x ) is often used and then differentiated for optimization.

Original: Though many lower bound estimators  are applicable, we focus on an importance-weighted evidence lower bound (Burda et al., 2016).
Revised: 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).

Original: This lower bound is constructed by introducing a proposal distribution q (x ; z  ) and using it to form an importance sampling estimate of the marginal likelihood:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: If K samples are drawn from q ( z ; x ) then this provides an unbiased estimate of p θ ( x ) and the biased “importance-weighted autoencoder” estimator IWAE K ( x ) of log p θ ( x ) is given by
Revised: If K samples are drawn from q ( z ; x ) then this provides an unbiased estimate of p θ ( x ) and the biased “importance-weighted autoencoder” estimator IWAE K ( x ) of log p θ ( x ) is given by

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: While the
Revised: While the

Original: 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.
Revised: 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.

Original: We discuss tasks for which a lower bound estimator would be ill-suited in Section 3.4.
Revised: We discuss tasks for which a lower bound estimator would be ill-suited in Section 3.4.

Original: There are two properties of IWAE that will allow us to modify it to produce an unbiased estimator:
Revised: There are two properties of IWAE that will allow us to modify it to produce an unbiased estimator:

Original: First, it is consistent in the sense that as the number of samples K increases, the expectation of
Revised: First, it is consistent in the sense that as the number of samples K increases, the expectation of

Original: IWAE K ( x ) converges to log p θ ( x ) .
Revised: IWAE K ( x ) converges to log p θ ( x ) .

Original: Second,  is also monotonically non-decreasing in expectation:log p θ ( x ) =
Revised: Second, it is also monotonically non-decreasing in expectation:

Original: lim K →∞ E

Original: [Equation]
Revised: [Equation]

Original: IWAE K ( x )] and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: These properties are sufﬁcient to create an unbiased estimator using the Russian roulette estimator.
Revised: These properties are sufﬁcient to create an unbiased estimator using the Russian roulette estimator.

Original: 2. R USSIAN ROULETTE ESTIMATOR
Revised: 2. R USSIAN ROULETTE ESTIMATOR

Original: In order to create an unbiased estimator of the log probability function, we employ the Russian roulette estimator (Kahn, 1955).
Revised: In order to create an unbiased estimator of the log probability function, we employ the Russian roulette estimator (Kahn, 1955).

Original: This estimator is used to estimate the sum of inﬁnite series, where evaluating any term in the series almost surely requires only a ﬁnite amount of computation.
Revised: This estimator is used to estimate the sum of inﬁnite series, where evaluating any term in the series almost surely requires only a ﬁnite amount of computation.

Original: Intuitively, the Russian roulette estimator relies on a randomized truncation and upweighting of each term to account for the possibility of not computing these terms.
Revised: Intuitively, the Russian roulette estimator relies on a randomized truncation and upweighting of each term to account for the possibility of not computing these terms.

Original: To illustrate the idea, let ∆ k denote the k -th term of an inﬁnite series.
Revised: To illustrate the idea, let ˜∆ k denote the k -th term of an inﬁnite series.

Original: Assume the partial sum of the series (cid:80) ∞
Revised: Assume the partial sum of the series (cid:80) ∞

Original: k =1 ∆ k converges to some quantity we wish to obtain.
Revised: k =1 ˜∆ k converges to some quantity we wish to obtain.

Original: We can construct a simple estimator by always computing the ﬁrst term then ﬂipping a coin b ∼ Bernoulli ( q ) to determine whether we stop or continue evaluating the remaining terms.
Revised: We can construct a simple estimator by always computing the ﬁrst term then ﬂipping a coin b ∼ Bernoulli ( q ) to determine whether we stop or continue evaluating the remaining terms.

Original: With probability 1 − q , we compute the rest of the series.
Revised: With probability 1 − q , we compute the rest of the series.

Original: By reweighting the remaining future terms by 1 / (1 − q ) , we obtain an unbiased estimator:
Revised: By reweighting the remaining future terms by 1 / (1 − q ) , we obtain an unbiased estimator:

Original: [Equation]
Revised: [Equation]

Original: To obtain the “Russian roulette” (RR) estimator (Forsythe & Leibler, 1950), we repeatedly apply this trick to the remaining terms.
Revised: To obtain the “Russian roulette” (RR) estimator (Forsythe & Leibler, 1950), we repeatedly apply this trick to the remaining terms.

Original: In effect, we make the number of terms a random variable K , taking values in 1 , 2 , . . .
Revised: In effect, we make the number of terms a random variable K , taking values in 1 , 2 , . . .

Original: to use in the summation (i.e., the number of successful coin ﬂips) from some distribution with probability mass function p ( K )
Revised: to use in the summation (i.e., the number of successful coin ﬂips) from some distribution with probability mass function p ( K )

Original: =
Revised: =

Original: P ( K = K ) with support over the positive integers.
Revised: P ( K = K ) with support over the positive integers.

Original: With K drawn from p ( K ) , the estimator takes the form:
Revised: With K drawn from p ( K ) , the estimator takes the form:

Original: [Equation]
Revised: [Equation]

Original: 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., (cid:80) ∞
Revised: 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., (cid:80) ∞

Original: k =1 | ∆ k | < ∞ (Chen et al., 2019).
Revised: k =1 | ˜∆ k | < ∞ (Chen et al. ; Lemma 3).

Original: This condition ensures that the average of multiple samples will converge to the value of the inﬁnite series by the law of large numbers.
Revised: This condition ensures that the average of multiple samples will converge to the value of the inﬁnite series by the law of large numbers.

Original: However, the variance of this estimator depends on the choice of p ( K ) and can potentially be very large or even inﬁnite (McLeish, 2011; Rhee & Glynn, 2015; Beatson & Adams, 2019).
Revised: However, the variance of this estimator depends on the choice of p ( K ) and can potentially be very large or even inﬁnite (McLeish, 2011; Rhee & Glynn, 2015; Beatson & Adams, 2019).

Original:  SUMO: U NBIASED ESTIMATION OF LOG PROBABILITY FOR LVM S 3. R USSIAN ROULETTE TO TIGHTEN LOWER BOUNDS
Revised:  SUMO: U NBIASED ESTIMATION OF LOG PROBABILITY FOR LVM S 3. R USSIAN ROULETTE TO TIGHTEN LOWER BOUNDS

Original: We can turn any absolutely convergent series into a telescoping series and apply the Russian roulette randomization to form an unbiased stochastic estimator.
Revised: We can turn any absolutely convergent series into a telescoping series and apply the Russian roulette randomization to form an unbiased stochastic estimator.

Original: We focus here on the IWAE bound described in Section 2.2.
Revised: We focus here on the IWAE bound described in Section 2.2.

Original: Let ∆ k
Revised: Let ∆ k

Original: ( x ) =
Revised: ( x ) =

Original: IWAE k +1 ( x ) − IWAE k
Revised: IWAE k +1 ( x ) − IWAE k

Original: ( x )  .
Revised: ( x ) and ˜∆ k ( x ) = E q [∆ k ( x )] , where the expectations are with respect to q ( z ; x ) .

Original: Then since IWAE converges absolutely, we apply equation 7 to construct our estimator, which we call SUMO (Stochastically 
Revised: Then since ˜∆ k converges absolutely, we apply equation 7 to construct our estimator, which we call SUMO (Stochastically Unbiased Marginalization

Original: Unbiased Marginalization Objective):
Revised: Objective). The detailed derivation of SUMO is in Appendix A.1. [Equation] K ∼ p ( K )

Original: [Equation]
Revised: [Equation]

Original: 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 ) :
Revised: 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 ) :

Original: [Equation]
Revised: [Equation]

Original: Furthermore, we can prove that thegradients of SUMO also have ﬁnite expectation (Appendix A.3). Algorithm 1 shows how we compute a single sample of SUMO, which has the same expected cost as IWAE K
Revised: where the expectation is taken over p ( K ) and q ( z ; x ) (see Algorithm 1 for our exact sampling procedure). Furthermore, under appropriate conditions, we have E [ ∇ θ SUMO ( x )]

Original: where K =  E 
Revised: = ∇ θ E [ SUMO ( x )]

Original: [ K + 1] .
Revised: = ∇ θ log p θ ( x ) (see Appendix A.4).

Original: 3. G RADIENT VARIANCE AND THE CHOICE OF p ( K )
Revised: 3. O PTIMIZING VARIANCE - COMPUTE PRODUCT BY CHOICE OF p ( K )

Original: To efﬁciently optimize a limit, one should choose an estimator to minimize the product of the second moment of the gradient estimates and the expected compute cost per evaluation.
Revised: To efﬁciently optimize a limit, one should choose an estimator to minimize the product of the second moment of the gradient estimates and the expected compute cost per evaluation.

Original: The choice of p ( K ) effects both the variance and computation cost of our estimator.
Revised: The choice of p ( K ) effects both the variance and computation cost of our estimator.

Original: Denoting ˆ G := ∇ θ ˆ Y and ∆ gk := ∇ θ [ IWAE k +1 ( x ) − IWAE k
Revised: Denoting ˆ G := ∇ θ ˆ Y and ∆ gk := ∇ θ [ IWAE k +1 ( x ) − IWAE k

Original: ( x )] , the Russian roulette estimator is optimal across a broad family of unbiased randomized truncation estimators ifand only if the ∆ gk are statistically independent, in which case it has second moment E
Revised: ( x )] , the Russian roulette estimator is optimal across a broad family of unbiased randomized truncation estimators if the ∆ gk are statistically independent, in which case it has second moment E

Original: ||
Revised: ||

Original: ˆ G
Revised: ˆ G

Original: || 

Revised: ||

Original: = (cid:80) ∞
Revised: = (cid:80) ∞

Original: k =1 E || ∆
Revised: k =1 E || ∆

Original: gk || 22 / P ( K≥ k ) (Beatson & Adams, 2019).
Revised: gk || 22 / P ( K≥ k ) (Beatson & Adams, 2019).

Original: For the following derivation of p ( K ) , we assume independence of ∆ gk and restrict ourselves to the Russian roulette estimator.

Original: While the ∆ gk are not in fact strictly independent with our sampling procedure (Algorithm 1), and other estimators within the family may perform better, we justify our choice by showing that E ∆ i ∆ j for i (cid:54) =
Revised: While the ∆ gk are not in fact strictly independent with our sampling procedure (Algorithm 1), and other estimators within the family may perform better, we justify our choice by showing that E ∆ i ∆ j for i (cid:54) =

Original: j converges to zero much faster than E ∆ 2 k (Appendices A.
Revised: j converges to zero much faster than E ∆ 2 k (Appendices A.2 & A.3).

Original: A.2).
Revised: In the following, we assume independence of ∆ gk and heuristically motivate an optimal p ( K ) in terms of reducing the product of compute and variance.

Original: We  show that E || ∆ gk || 22 is O ( 1 / k 2 )
Revised: We ﬁrst show that E || ∆ gk || 22 is O ( 1 / k 2 )

Original: (Appendix A.3).
Revised: (Appendix A.5).

Original: Therefore, the optimal compute-variance product (Rhee & Glynn, 2015; Beatson & Adams, 2019) is given by P ( K ≥ k ) ∝ (cid:112) E || ∆ gk || 22 ) .
Revised: This implies the optimal compute-variance product (Rhee & Glynn, 2015; Beatson & Adams, 2019) is given by P ( K ≥ k ) ∝ (cid:112) E || ∆ gk || 22 ) .

Original: In our case, this gives P ( K ≥ k ) = 1 / k , which results in an estimator with inﬁnite expected computation and no ﬁnite bound on variance.
Revised: In our case, this gives P ( K ≥ k ) = 1 / k , which results in an estimator with inﬁnite expected computation and no ﬁnite bound on variance.

Original: In fact, any p ( K ) which gives rise to provably ﬁnite variance requires a heavier tail than P ( K ≥ k )
Revised: In fact, any p ( K ) which gives rise to provably ﬁnite variance requires a heavier tail than P ( K ≥ k )

Original: = 1 / k and will have inﬁnite expected computation.
Revised: = 1 / k and will have inﬁnite expected computation.

Original: There are no theoretical guarantees for optimization with an inﬁnite-variance gradient estimator, but we empirically ﬁnd that convergence is stable.
Revised: Though we could not theoretically show that our estimator and gradients have ﬁnite variance, we empirically ﬁnd that convergence—even in the setting of minimizing log probability—is stable.

Original: We plot || ∆ k || 22 for the toy variational inference task used to assess signal to noise ratio in Tucker et al.
Revised: We plot || ∆ k || 22 for the toy variational inference task used to assess signal to noise ratio in Tucker et al.

Original: (2018) and Rainforth et al.
Revised: (2018) and Rainforth et al.

Original: (2018), and ﬁnd that they converge faster than 1 k 2 in practice (Appendix A.4).
Revised: (2018b), and ﬁnd that they converge faster than 1 k 2 in practice (Appendix A.6).

Original: While this indicates the variance is better than the theoretical bound, an estimator havinginﬁnite expected computation cost will always be an issue as it indicates signiﬁcant probability of sampling arbitrarily large K . We therefore modify the
Revised: While this indicates the variance is better than the theoretical bound, an estimator having

Original: Algorithm 1 Computing SUMO, an unbiased estimator of log p ( x ) .
Revised: Algorithm 1 Computing SUMO, an unbiased estimator of log p ( x ) .

Original: Input:
Revised: 1 iid log p

Original: [Equation]
Revised: k − log q

Original: , m
Revised: k

Original: [Equation]
Revised: inﬁnite expected computation cost will always be an issue as it indicates signiﬁcant probability of sampling arbitrarily large K .

Original:  tail of the sampling distribution such that the estimator has ﬁnite expected computation:
Revised: We therefore modify the tail of the sampling distribution such that the estimator has ﬁnite expected computation:

Original: [Equation]
Revised: [Equation]

Original: We typically choose α = 80 , which gives an expected computation cost of approximately 5 terms.
Revised: We typically choose α = 80 , which gives an expected computation cost of approximately 5 terms.

Original: 3.2. T RADING VARIANCE AND COMPUTE
Revised: 3.2. T RADING VARIANCE AND COMPUTE

Original: One way topotentially improve the RR estimator is to construct it so that some minimum number of terms (denoted here as m ) are always computed.
Revised: One way to improve the RR estimator is to construct it so that some minimum number of terms (denoted here as m ) are always computed.

Original: This puts a lower bound on the computational cost, but can potentially lower variance, providing a design space for trading off estimator quality against computational cost.
Revised: This puts a lower bound on the computational cost, but can potentially lower variance, providing a design space for trading off estimator quality against computational cost.

Original: This corresponds to a choice of RR estimator in which P ( K = K ) =for K ≤ m .
Revised: This corresponds to a choice of RR estimator in which P ( K = K ) = 0 for K ≤ m .

Original: This computes the sum out to m terms (effectively computing IWAE m ) and then estimates the remaining difference with Russian roulette:
Revised: This computes the sum out to m terms (effectively computing IWAE m ) and then estimates the remaining difference with Russian roulette:

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: 3. T RAINING q ( z ; x ) TO REDUCE VARIANCE
Revised: 3. T RAINING q ( z ; x ) TO REDUCE VARIANCE

Original: The SUMO estimator does not require amortized variational inference, but the use of an “encoder” to produce the approximate posterior q ( z ; x ) has been shown to be a highly effective way to perform rapid feedforward inference in neural latent variable models.
Revised: The SUMO estimator does not require amortized variational inference, but the use of an “encoder” to produce the approximate posterior q ( z ; x ) has been shown to be a highly effective way to perform rapid feedforward inference in neural latent variable models.

Original: We use φ to denote the parameters of the encoder q φ ( z ; x ) .
Revised: We use φ to denote the parameters of the encoder q φ ( z ; x ) .

Original: However, the gradients of SUMO with respect to φ are in expectation zero precisely because SUMO is an unbiased estimator of log p θ ( x ) , regardless of our choice of q φ ( z ; x ) .
Revised: However, the gradients of SUMO with respect to φ are in expectation zero precisely because SUMO is an unbiased estimator of log p θ ( x ) , regardless of our choice of q φ ( z ; x ) .

Original: Nevertheless, we would expect the choice of q φ ( z ; x ) signiﬁcantly impacts the variance of our estimator.
Revised: Nevertheless, we would expect the choice of q φ ( z ; x ) signiﬁcantly impacts the variance of our estimator.

Original: As such, we optimize q φ ( z ; x ) to reduce the variance of the SUMO estimator.
Revised: As such, we optimize q φ ( z ; x ) to reduce the variance of the SUMO estimator.

Original: We can obtain unbiased gradients in the following way (Ruiz et al., 2016; Tucker et al., 2017):
Revised: We can obtain unbiased gradients in the following way (Ruiz et al., 2016; Tucker et al., 2017):

Original: [Equation]
Revised: [Equation]

Original: Notably, the expectation of this estimator depends on the variance of SUMO, which we have not been able to bound.
Revised: Notably, the expectation of this estimator depends on the variance of SUMO, which we have not been able to bound.

Original: In practice, we observe gradients which are bounded but sometimes very large.
Revised: In practice, we observe gradients which are bounded but sometimes very large.

Original: We apply gradient clipping to the encoder to clip gradients which are excessively large in magnitude.
Revised: We apply gradient clipping to the encoder to clip gradients which are excessively large in magnitude.

Original: This helps stabilize the training progress but does introduce bias into the encoder gradients.
Revised: This helps stabilize the training progress but does introduce bias into the encoder gradients.

Original: Fortunately, the encoder itself is merely a tool for variance reduction, and biased gradient with respect to the encoder can still signiﬁcantly help optimization.
Revised: Fortunately, the encoder itself is merely a tool for variance reduction, and biased gradient with respect to the encoder can still signiﬁcantly help optimization.

Original: 3. A PPLICATIONS OF UNBIASED LOG PROBABILITY
Revised: 3. A PPLICATIONS OF UNBIASED LOG PROBABILITY

Original: Here we list some applications where an unbiased log probability is useful.
Revised: Here we list some applications where an unbiased log probability is useful.

Original: Using SUMO to replace existing lower bound estimates allows latent variable models to be used for new applications where a lower bound is inappropriate.
Revised: Using SUMO to replace existing lower bound estimates allows latent variable models to be used for new applications where a lower bound is inappropriate.

Original: As latent variable models can be both expressive and efﬁcient to sample from, they are frequently useful in applications where the data are high-dimensional and samples from the model are needed.
Revised: As latent variable models can be both expressive and efﬁcient to sample from, they are frequently useful in applications where the data is high-dimensional and samples from the model are needed.

Original: Minimizing log p θ ( x ) .
Revised: Minimizing log p θ ( x ) .

Original: Some machine learning objectives include terms that seek to increase the entropy of the learned model.
Revised: Some machine learning objectives include terms that seek to increase the entropy of the learned model.

Original: The “reverse KL” objective—often used for training models to perform approximate posterior inferences—minimizes E x ∼ p θ ( x ) [log p θ ( x ) − log π ( x )] where π ( x ) is a target density that may only be known up a normalization constant.
Revised: The “reverse KL” objective—often used for training models to perform approximate posterior inferences—minimizes E x ∼ p θ ( x ) [log p θ ( x ) − log π ( x )] where π ( x ) is a target density that may only be known up a normalization constant.

Original: Local updates of this form are the basis of the expectation propagation procedure (Minka, 2001).
Revised: Local updates of this form are the basis of the expectation propagation procedure (Minka, 2001).

Original: This objective has also been used for distilling autoregressive models that are inefﬁcient at sampling (Oord et al., 2018).
Revised: This objective has also been used for distilling autoregressive models that are inefﬁcient at sampling (Oord et al., 2018).

Original: Moreover, reverse KL is connected to the use of entropy-regularized objectives (Williams & Peng, 1991; Ziebart, 2010; Mnih et al., 2016;
Revised: Moreover, reverse KL is connected to the use of entropy-regularized objectives (Williams & Peng, 1991; Ziebart, 2010; Mnih et al., 2016;

Original: Norouzi et al., 2016) in decision-making problems, where the goal is to encourage the decision maker toward exploration and prevent it from settling into a local minimum.
Revised: Norouzi et al., 2016) in decision-making problems, where the goal is to encourage the decision maker toward exploration and prevent it from settling into a local minimum.

Original: Unbiased score function ∇ θ log p θ ( x ) .
Revised: Unbiased score function ∇ θ log p θ ( x ) .

Original: The score function is the gradient of the log-likelihood with respect to the parameters and has uses in estimating the Fisher information matrix and performing stochastic gradient Langevin dynamics (Welling & Teh, 2011), among other applications.
Revised: The score function is the gradient of the log-likelihood with respect to the parameters and has uses in estimating the Fisher information matrix and performing stochastic gradient Langevin dynamics (Welling & Teh, 2011), among other applications.

Original: Of particular note, the REINFORCE gradient estimator—which is generally applicable for optimizing objectives of the form min θ E x ∼ p θ ( x ) [ R ( x )] —is estimated using the score function.
Revised: Of particular note, the REINFORCE gradient estimator (Williams, 1992)—generally applicable for optimizing objectives of the form max θ E x ∼ p θ ( x ) [ R ( x )] —is estimated using the score function.

Original: This can be replaced with the gradient of SUMO which itself is an estimator of the score function ∇ θ log p θ ( x ) .
Revised: This can be replaced with the gradient of SUMO which itself is an estimator of the score function ∇ θ log p θ ( x ) .

Original: [Equation]
Revised: [Equation]

Original: where the inner expectation is over the stochasticity of the SUMO estimator.
Revised: where the inner expectation is over the stochasticity of the SUMO estimator.

Original: Such estimators are often usedin for reward maximization in reinforcement learning where p θ ( x ) is a stochastic policy.
Revised: Such estimators are often used for reward maximization in reinforcement learning where p θ ( x ) is a stochastic policy.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: There is a long history in Bayesian statistics of marginal likelihood estimation in the service of model selection.
Revised: There is a long history in Bayesian statistics of marginal likelihood estimation in the service of model selection.

Original: 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 inﬁnite variance (Murray & Salakhutdinov,  2009).
Revised: 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 inﬁnite variance (Murray & Salakhutdinov, 2009) and exhibits simulation psuedo-bias (Lenk, 2009).

Revised: The

Original: 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).
Revised: Chib estimator (Chib, 1995), the Laplace approximation, and nested sampling (Skilling, 2006) are alternative proposals that can often have better properties (Murray & Salakhutdinov, 2009).

Original: Annealed importance sampling (Neal, 2001) probably represents the gold standard for marginal likelihood estimation.
Revised: Annealed importance sampling (Neal, 2001) probably represents the gold standard for marginal likelihood estimation.

Revised: These, however, turn into consistent estimators at best when estimating the log marginal probability (Rainforth et al., 2018a).

Revised: Bias removal schemes such as jackknife variational inference (Nowozin, 2018) have been proposed to debias log-evidence estimation, IWAE in particular.

Revised: Hierarchical IWAE (Huang et al., 2019) uses a joint proposal to induce negative correlation among samples and connects the convergence of variance of the estimator and the convergence of the lower bound.

Original: Russian roulette also has a long history.
Revised: Russian roulette also has a long history.

Original: It dates back to unpublished work from von Neumann and
Revised: It dates back to unpublished work from von Neumann and

Original: Ulam, who used it to debias Monte Carlo methods for matrix inversion (Forsythe & Leibler, 1950) and particle transport problems (Kahn, 1955).
Revised: Ulam, who used it to debias Monte Carlo methods for matrix inversion (Forsythe & Leibler, 1950) and particle transport problems (Kahn, 1955).

Original: It has gained popularity in statistical physics (Spanier & Gelbard, 1969; Kuti, 1982; Wagner, 1987), for unbiased ray tracing in graphics and rendering (Arvo & Kirk, 1990), and for a number of estimation problems in the statistics community (Wei & Murray, 2017; Girolami et al., 2013; Rychlik, 1990; 1995; Jacob & Thiery, 2015; Jacob et al., 2017).
Revised: It has gained popularity in statistical physics (Spanier & Gelbard, 1969; Kuti, 1982; Wagner, 1987), for unbiased ray tracing in graphics and rendering (Arvo & Kirk, 1990), and for a number of estimation problems in the statistics community (Wei & Murray, 2017; Girolami et al., 2013; Rychlik, 1990; 1995; Jacob & Thiery, 2015; Jacob et al., 2017).

Original: It has also been independently rediscovered many times (Fearnhead et al., 2008; McLeish, 2011; Rhee & Glynn, 2012; Tallec & Ollivier, 2017).
Revised: It has also been independently rediscovered many times (Fearnhead et al., 2008; McLeish, 2011; Rhee & Glynn, 2012; Tallec & Ollivier, 2017).

Original: The use of Russian roulette estimation in deep learning and generative modeling applications has been gaining traction in recent years.
Revised: The use of Russian roulette estimation in deep learning and generative modeling applications has been gaining traction in recent years.

Original: It has been used to solve short-term bias in optimization problems (Tallec & Ollivier, 2017; Beatson & Adams, 2019).
Revised: It has been used to solve short-term bias in optimization problems (Tallec & Ollivier, 2017; Beatson & Adams, 2019).

Original: Wei & Murray (2017) estimates the reciprocal normalization constant of an unnormalized density.
Revised: Wei & Murray (2017) estimates the reciprocal normalization constant of an unnormalized density.

Original: Han et al.
Revised: Han et al.

Original: (2018) uses a similar random truncation approach to estimate the distribution of eigenvalues in a symmetric matrix.
Revised: (2018) uses a similar random truncation approach to estimate the distribution of eigenvalues in a symmetric matrix.

Original: Along similar motivations with our work, Chen et al.
Revised: Along similar motivations with our work, Chen et al.

Original: (2019) uses this estimator to construct an unbiased estimator of the change of variables equation in the context of normalizing ﬂows (Rezende & Mohamed, 2015), and Xu et al.
Revised: (2019) uses this estimator to construct an unbiased estimator of the change of variables equation in the context of normalizing ﬂows (Rezende & Mohamed, 2015), and Xu et al.

Original: (2019) uses it to construct unbiased log probability for a nonparameteric distribution in the context of variational autoencoders (Kingma & Welling, 2013).
Revised: (2019) uses it to construct unbiased log probability for a nonparameteric distribution in the context of variational autoencoders (Kingma & Welling, 2013).

Original: Though we extend latent variable models to applications that require unbiased estimates of log probability and beneﬁt from efﬁcient sampling, an interesting family of models already fulﬁll these requirements.
Revised: Though we extend latent variable models to applications that require unbiased estimates of log probability and beneﬁt from efﬁcient sampling, an interesting family of models already fulﬁll these requirements.

Original: Normalizing ﬂows (Rezende & Mohamed, 2015; Dinh et al., 2017) offer exact log probability and certain models have been proven to be universal density estimators (Huang et al., 2018).
Revised: Normalizing ﬂows (Rezende & Mohamed, 2015; Dinh et al., 2017) offer exact log probability and certain models have been proven to be universal density estimators (Huang et al., 2018).

Original: However, these models often require restrictive architectural choices with no dimensionalityreduction capabilities, and make use of many more parameters to scale upthan alternative generative models (Kingma & Dhariwal, 2018).
Revised: However, these models often require restrictive architectural choices with no dimensionalityreduction capabilities, and make use of many more parameters to scale up (Kingma & Dhariwal, 2018) than alternative generative models.

Original: Discrete variable versions of these models are still in their infancy and must use  biased gradients (Tran et al., 2019; Hoogeboom et al., 2019), whereas latent variable modelsextend naturally  to discrete observations. 
Revised: 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. [Table]

Original:  D ENSITY M ODELING E XPERIMENTS
Revised:  D ENSITY M ODELING E XPERIMENTS

Original: We ﬁrst compare the performance of SUMO when used as a replacement to IWAE with the same expected cost on density modeling tasks.
Revised: We ﬁrst compare the performance of SUMO when used as a replacement to IWAE with the same expected cost on density modeling tasks.

Original: We make use of two benchmark datasets: dynamically binarized MNIST (LeCun et al., 1998) and binarized OMNIGLOT
Revised: We make use of two benchmark datasets: dynamically binarized MNIST (LeCun et al., 1998) and binarized OMNIGLOT

Original: (Lake et al., 2015).
Revised: (Lake et al., 2015)

Original: [Table]
Revised: [Table]

Revised: We use the same neural network architecture as IWAE (Burda et al., 2016).

Revised: The prior p ( z ) is a 50dimensional standard Gaussian distribution.

Revised: The conditional distributions p ( x i | z ) are independent

Revised: Bernoulli, with the decoder parameterized by two hidden layers, each with 200 tanh units.

Revised: The approximate posterior q ( z ; x ) is also a 50-dimensional Gaussian distribution with diagonal covariance, whose mean and variance are both parameterized by two hidden layers with 200 tanh units.

Original: We reimplemented and tuned IWAE, obtaining strong baseline results which are better than those previously reported.
Revised: We reimplemented and tuned IWAE, obtaining strong baseline results which are better than those previously reported.

Original: We use the same hyperparameters to train with the SUMO estimator.
Revised: We then used the same hyperparameters to train with the SUMO estimator.

Original: We ﬁnd clipping very large gradients can help performance, as large gradients may be infrequently sampled.
Revised: We ﬁnd clipping very large gradients can help performance, as large gradients may be infrequently sampled.

Original: This may introduce a small amount of bias into the gradients while reducing variance, but can nevertheless help achieve faster convergence and should still result in a less-biased estimator.
Revised: This introduces a small amount of bias into the gradients while reducing variance, but can nevertheless help achieve faster convergence and should still result in a less-biased estimator.

Revised: A posthoc study of the effect on ﬁnal test performance as a function of this bias-variance tradeoff mechanism is discussed in Appendix A.7.

Revised: We note that gradient clipping is only done for the density modeling experiments.

Original: The averaged test log-likelihoods and standard deviations over 3 runs are summarized in Table 1.
Revised: The averaged test log-likelihoods and standard deviations over 3 runs are summarized in Table 1.

Original: To be consistent with existing literature, we evaluate our model using IWAE with 5000 samples.
Revised: To be consistent with existing literature, we evaluate our model using IWAE with 5000 samples.

Original: In all the cases, SUMO achieves slightly better performance than IWAE with the same expected cost.
Revised: In all the cases, SUMO achieves slightly better performance than IWAE with the same expected cost.

Original: We also bold the results that are statistically signiﬁcant according to an unpaired t -test with signiﬁcance level 0 .
Revised: We also bold the results that are statistically insigniﬁcant from the best performing model according to an unpaired t -test with signiﬁcance level 0 .

Original: 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 of IWAE.
Revised: 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 of IWAE.

Original:  L ATENT VARIABLES  FOR ENTROPY MAXIMIZATION
Revised:  L ATENT VARIABLES MODELS FOR ENTROPY MAXIMIZATION

Original: We move on to our ﬁrst task for which a lower bound estimate of log probability would not sufﬁce.
Revised: We move on to our ﬁrst task for which a lower bound estimate of log probability would not sufﬁce.

Original: The reverse KL objective is useful when we have access to a (possibly unnormalized) target distribution but no efﬁcient sampling algorithm.
Revised: The reverse KL objective is useful when we have access to a (possibly unnormalized) target distribution but no efﬁcient sampling algorithm.

Original: [Equation]
Revised: [Equation]

Original: A major problem with ﬁtting latent variables models to this objective is the presence of an entropy maximization term, effectively a minimization of log p θ ( x ) .
Revised: A major problem with ﬁtting latent variables models to this objective is the presence of an entropy maximization term, effectively a minimization of log p θ ( x ) .

Original: Estimating this log marginal probability with a lower bound estimator could result inthe optimizing θ to maximize the bias of the estimator instead of the true objective.
Revised: Estimating this log marginal probability with a lower bound estimator could result in optimizing θ to maximize the bias of the estimator instead of the true objective.

Original: Our experiments demonstrate that this causes IWAE to often fail to optimize the objective. Modiﬁed IWAE for entropy maximization.
Revised: Our experiments demonstrate that this causes IWAE to often fail to optimize the objective unless we use a large amount of computation. Modifying IWAE.

Original: The bias of the IWAE estimator can be interpreted as the KL between an importance-weighted approximate posterior q IW ( z ; x ) implicitly deﬁned by the encoder and the  true posterior p ( z | x ) (Domke & Sheldon, 2018).
Revised: The bias of the IWAE estimator can be interpreted as the KL between an importance-weighted approximate posterior q IW ( z ; x ) implicitly deﬁned by the encoder and the [Figure] true posterior p ( z | x ) (Domke & Sheldon, 2018).

Original: Both the encoder and decoder parameters can therefore affect this bias.
Revised: Both the encoder and decoder parameters can therefore affect this bias.

Original: In practice, we ﬁnd that the encoder optimization proceeds at a faster timescale than the decoder optimization: i.e., the encoder can match q  ( z ; x ) to the decoder’s p (x | z  ) more quickly than the latter can match an objective.
Revised: In practice, we ﬁnd that the encoder optimization proceeds at a faster timescale than the decoder optimization: i.e., the encoder can match q IW ( z ; x ) to the decoder’s p ( z | x ) more quickly than the latter can match an objective.

Original: For this reason, we train the encoder to reduce bias and use a minimax training objective
Revised: For this reason, we train the encoder to reduce bias and use a minimax training objective

Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Though this is still a lower bound with unbounded bias, it makes for a stronger baseline than optimizing q ( z ; x ) in the same direction as p ( x, z ) .
Revised: Though this is still a lower bound with unbounded bias, it makes for a stronger baseline than optimizing q ( z ; x ) in the same direction as p ( x, z ) .

Original: We ﬁnd that this approach can work well in practice when k is set sufﬁciently high.
Revised: We ﬁnd that this approach can work well in practice when k is set sufﬁciently high.

Original: [Figure]
Revised: [Figure]

Original: We choose a “funnel” target distribution (Figure 1) similar to the distribution used as a benchmark for inference in Neal et al.
Revised: We choose a “funnel” target distribution (Figure 1) similar to the distribution used as a benchmark for inference in Neal et al.

Original: (2003), where p ∗ has support in R 2 and is deﬁned p ∗ ( x 1 , x 2 ) =
Revised: (2003), where p ∗ has support in R 2 and is deﬁned p ∗ ( x 1 , x 2 ) =

Original: N ( x 1 ; 0 , 1 . 35 2 ) N ( x 2 ; 0 , e 2 x 1 )
Revised: N ( x 1 ; 0 , 1 . 35 2 ) N ( x 2 ; 0 , e 2 x 1 )

Original: We use neural networks with one hidden layer of 200 hidden units and tanh activations for both the encoder and decoder networks.
Revised: We use neural networks with one hidden layer of 200 hidden units and tanh activations for both the encoder and decoder networks.

Original: We use 20 latent variables, with p ( z ) , p θ ( x | z ) , and q φ ( z ; x ) all being Gaussian distributed.
Revised: We use 20 latent variables, with p ( z ) , p θ ( x | z ) , and q φ ( z ; x ) all being Gaussian distributed.

Original: Figure 2 shows the learning curves when using
Revised: Figure 2 shows the learning curves when using

Original: IWAE and SUMO.
Revised: IWAE and SUMO.

Original: Unless k is set very large,
Revised: Unless k is set very large,

Original: IWAE will at some point start optimizing the bias instead of the actual objective.
Revised: IWAE will at some point start optimizing the bias instead of the actual objective.

Original: The reverse
Revised: The reverse

Original: KL is a non-negative quantity, so any estimate signiﬁcantly below zero can be attributed to the unbounded bias.
Revised: KL is a non-negative quantity, so any estimate signiﬁcantly below zero can be attributed to the unbounded bias.

Original: On the other hand, SUMO correctly optimizes for the objective even with a small expected cost.
Revised: On the other hand, SUMO correctly optimizes for the objective even with a small expected cost.

Original: Increasing the expected cost k for SUMO reduces variance.
Revised: Increasing the expected cost k for SUMO reduces variance.

Original: For the same expected cost, SUMO can optimize the true objective but IWAE cannot.
Revised: For the same expected cost, SUMO can optimize the true objective but IWAE cannot.

Original: We also found that if k is set sufﬁciently large, then IWAE can work when we train using the minimax objective in equation 15, suggesting that a sufﬁciently debiased estimator can also work in practice.
Revised: We also found that if k is set sufﬁciently large, then IWAE can work when we train using the minimax objective in equation 15, suggesting that a sufﬁciently debiased estimator can also work in practice.

Original: However, this requires much more compute and likely does not scale compared to SUMO.
Revised: However, this requires much more compute and likely does not scale compared to SUMO.

Original: We also visualize the contours of the resulting models in Figure 1.
Revised: We also visualize the contours of the resulting models in Figure 1.

Original: For IWAE, we visualize the model a few iterations before it reaches numerical instability.
Revised: For IWAE, we visualize the model a few iterations before it reaches numerical instability.

Original:  L ATENT VARIABLE POLICIES FOR COMBINATORIAL OPTIMIZATION
Revised:  L ATENT VARIABLE POLICIES FOR COMBINATORIAL OPTIMIZATION

Original: Let us now consider the problem of ﬁnding the maximum of a non-differentiable function, a special case of reinforcement learning without an interacting environment.
Revised: Let us now consider the problem of ﬁnding the maximum of a non-differentiable function, a special case of reinforcement learning without an interacting environment.

Original: Variational optimization (Staines & Barber, 2012) can be used to reformulate this as the optimization of a parametric distribution,
Revised: Variational optimization (Staines & Barber, 2012) can be used to reformulate this as the optimization of a parametric distribution,

Original: [Equation]
Revised: [Equation]

Original: (16) which is now a differentiable function with respect to the parameters θ , whose gradients can be estimated using  the REINFORCE gradient estimator and the SUMO estimator (equation 13).
Revised: [Figure] which is now a differentiable function with respect to the parameters θ , whose gradients can be estimated using a combination of the REINFORCE gradient estimator and the SUMO estimator (equation 13).

Revised: Furthermore, entropy regularized reinforcement learning—where we maximize

Revised: R ( x ) + λ H ( p θ ) with H ( p θ ) being the entropy of p θ ( x ) —encourages exploration and is inherently related to minimizing a reverse KL objective with the target being an exponentiated reward (Norouzi et al., 2016).

Original: For concreteness, we focus on the problem of quadratic pseudo-Boolean optimization (QPBO) where the objective is to maximize
Revised: For concreteness, we focus on the problem of quadratic pseudo-Boolean optimization (QPBO) where the objective is to maximize

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original:  where { x i } di =1 ∈ { 0 , 1 } are binary variables.
Revised: 17) where { x i } di =1 ∈ { 0 , 1 } are binary variables.

Original: Without further assumptions, QPBO is NPhard (Boros & Hammer, 2002).
Revised: Without further assumptions, QPBO is NPhard (Boros & Hammer, 2002).

Original: As there exist complex dependencies between the binary variables and optimization of equation 16 requires sampling from the policy distribution p θ ( x ) , a model that is both expressive and allows efﬁcient sampling would be ideal.
Revised: As there exist complex dependencies between the binary variables and optimization of equation 16 requires sampling from the policy distribution p θ ( x ) , a model that is both expressive and allows efﬁcient sampling would be ideal.

Original: For this reason, we motivate the use of latent variable models with independent conditional distributions, which we trained using the
Revised: For this reason, we motivate the use of latent variable models with independent conditional distributions, which we trained using the

Original: SUMO objective.
Revised: SUMO objective.

Original: Our baselines are an autoregressive policy, which captures dependencies but for[Figure] which sampling must be performed sequentially, and an independent policy, which is easy to sample from but captures no dependencies.
Revised: 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.

Original: [Equation]
Revised: [Equation]

Revised: We note that Haarnoja et al.

Revised: (2018) also argued for latent variable policies in favor of learning diverse strategies but ultimately had to make use of normalizing ﬂows which did not require marginalization.

Original: We construct one problem instance for each d ∈ { 500 , 1000 , 10000 } .
Revised: We constructed one problem instance for each d ∈ { 100 , 500 } , which we note are already intractable for exact optimization.

Original: For each instance, we randomly sampled the weights w i and w ij uniformly from the interval [ − 1 , 1] .
Revised: For each instance, we randomly sampled the weights w i and w ij uniformly from the interval [ − 1 , 1] .

Original: Figure 3 shows the performance of each policy model in solving problems of different sizes, where training time (note the log scale) includes model sampling and updating, but not the computation of the reward function.
Revised: Figure 3 shows the performance of each policy model.

Original: Including the reward computation would simply increase each method by a constant amount of time, but would not allow us to compare across problem instances due to  the quadratic nature of our test problem.
Revised: In general, the independent policy is quick to converge to a local minima and is unable to explore different regions, whereas more complex models have a better grasp of the “frontier” of reward distributions during optimization.

Original: All methods were trained with exactly 10000 iterations. We use the same architecture as in the entropy maximization experiment, except with Bernoulli conditional distributions. Theindependent policy uses the same decoder architecture as the latent variable policy but always uses a ﬁxed random z , whereas the autoregressive policy uses the a masked autoencoder architecture (Germain et al., 2015) with around the same number of parameters.
Revised: The autoregressive works well overall but is much slower to train due to its sequential sampling procedure; with d = 500 , it is 19 .

Original: For the masked autoencoder, evaluation of log probability only requires a single forward pass, but sampling requires d forward passes. We note that even with only d = 500 variables, the search space is intractably large for exact optimization. We ﬁnd that at convergence, the independent policy performs worse than the latent variable and autoregressive policies. While more complex models may have a better grasp of the current “frontier” of reward distributions, the simple independent model must prematurely reduce its entropy in order to obey strong dependencies between variables.
Revised: 2 × slower than training with SUMO.

Original: This reduction in  entropy reduces the ﬁnal reward due to insufﬁcient exploration. This behavior is indicated by the dips in expected reward during optimization.
Revised: Surprisingly, we ﬁnd that estimating the REINFORCE gradient with IWAE results in decent performance when no entropy regularization is present.

Original: On the other extreme, the autoregressive policy is slow to update, with each iteration taking more than 3 seconds on the d = 10000 problem instance and is roughly 100 × slower than the other models.

Original: In contrast, the latent variable and independent policies spend similar amounts of time solving all problem instances, and are parallelizable. While the use of SUMO with k = 15 makes our latent variable policy slower than the independent policy, the performance gap is roughly 0 . 5 × to 2 × the size of the problem instance, suggesting that models that capture dependencies can obtain non-trivial improvement in expected reward.
Revised: With entropy regularization, all policies improve signiﬁcantly; however, training with IWAE in this setting results in performance similar to the independent model. On the other hand, SUMO works with both REINFORCE gradient estimation and entropy regularization, albeit at the cost of slower convergence due to variance.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We introduced SUMO, a new unbiased estimator of the log probability for latent variable models, and demonstrated tasks for which thisunbiased estimator performs better than standard lower bounds.
Revised: 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.

Original: In practice, the high variance of our  estimator necessitates gradient clipping during optimization, which re-introduces a small amount of bias. However, even this clipped estimator provides lower-bias estimates than existing estimators at the same compute cost, and works well when minimizing log probability.
Revised: 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.

Original: We plan to investigate new families of gradient-based optimizers which can handle heavy-tailed stochastic gradients.
Revised: In the future, we plan to investigate new families of gradient-based optimizers which can handle heavy-tailed stochastic gradients.

Original: It may  be fruitful to investigate the use of convex combination of consistent estimators within the SUMO approach, as any convex combination is unbiased.
Revised: It may also be fruitful to investigate the use 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 A A PPENDIX A. D ERIVATION OF SUMO Let [Equation]

Revised: where z 1 , .., z k are sampled independently from q ( z ; x ) .

Original: A A PPENDIX A. C ONVERGENCE OF IWAE k +1{++
 x )]+} −  IWAE k
Revised: And we deﬁne the k -th term of the inﬁnite series ˜∆ k ( x ) : = E q [ IWAE k {++
 x )]+} − E q [ IWAE k

Revised: ( x )] .

Revised: Using the properties of IWAE in equation 6, we have ˜∆ k ( x ) ≥ 0 , and

Revised: [Equation]

Revised: which means the series converges absolutely.

Revised: This is a sufﬁcient condition for ﬁnite expectation of the Russian roulette estimator (Chen et al.

Revised: ; Lemma 3).

Revised: Applying equation 7 to the series:

Revised: [Equation]

Revised: [Equation]

Revised: Let ∆ k

Revised: ( x ) := IWAE k +1 ( x ) − IWAE k ( x ) , Hence our estimator is constructed:

Revised: [Equation]

Revised: And it can be easily seen from equation 22 and equation 23 that SUMO is an unbiased estimator of the log marginal likelihood:

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Revised: A. C ONVERGENCE OF ∆ k

Revised: 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 k both at a rate of O ( 1 k ) .

Revised: We build on this analysis to analyze the convergence of ∆ k .

Original: Let w i = p
Revised: Let w i = p

Original: ( x | z i ) p ( z i ) q
Revised: ( x | z i ) p ( z i ) q

Original: ( z i | x ) and we deﬁne Y k :
Revised: ( z i ; x ) and we deﬁne Y k :

Original: = 1 k (cid:80) ki
Revised: = 1 k (cid:80) ki

Original: =1 w i as the sample mean
Revised: =1 w i as the sample mean

Original: and we have E
Revised: and we have E

Original: [ Y k ] =
Revised: [ Y k ] =

Original: E [ w ] = µ .
Revised: E [ w ] = µ .

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: We follow the notations in Nowozin (2018).
Revised: We note that we rely on || Y k − µ || < 1 for this power series to converge. This condition was implicitly assumed, but not explicitly noted, in (Nowozin, 2018).

Revised: This condition will hold for sufﬁciently large k so long as the moments of w i exist: one could bound the probability

Revised: ||

Revised: Y k − µ || ≥ 1 by Chebyshev’s inequality or by the Central Limit Theorem.

Original: We use the central moments γ t := E [( Y k − µ ) t ] and µ t := 
Revised: We use the central moments γ t := E [( Y k − µ ) t ] and µ t := E [

Original: E [( w − µ ) t ] for t ≥ 2 .
Revised: [Equation] w − µ ) t ] for t ≥ 2 .

Original: [Equation]
Revised: [Equation]

Original: Expanding Eq.
Revised: Expanding Eq.

Original: 21 to order two gives
Revised: 28 to order two gives

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Since we use cumulative sum to compute Y k and Y k +1 , we obtain Y k
Revised: Since we use cumulative sum to compute Y k and Y k +1 , we obtain Y k

Original: +1 =
Revised: +1 =

Original: kY k + w k +1 k
Revised: kY k + w k +1 k

Original: + .
Revised: + .

Original: [Equation]
Revised: [Equation]

Original: We note that w k +1 − 1 k + = O ( [-1 k ) and w k +1 + 2 k +1 k +
-]
Revised: We note that w k +1 − 1 k + = O ( [-1 k ) and w k +1 + 2 k +1 k +
-][Equation]

Original: (cid:80)
Revised: ) and

Original: ki =1 w k 2 kµ
Revised: w k [Equation]

Original: = O (1) .
Revised: = O (1) .

Original: Therefore ∆ k is O ( 1 k ) , and E ∆ 2 k =
Revised: Therefore ∆ k is O ( 1 k ) , and E ∆ 2 k =

Original: O ( 1 k 2 ) .
Revised: O ( 1 k 2 ) .

Original: A. C ONVERGENCE OF ∆ k ∆ j
Revised: A. C ONVERGENCE OF ∆ k ∆ j

Original: From the previous proof, we have

Original: [Equation]

Original: Without loss of generality, suppose j ≥ k + 1 ,
Revised: Without loss of generality, suppose j ≥ k + 1 ,

Original: E ∆ k ∆ j = E (cid:34)(cid:32) ∞ (cid:88) t

Original: = (

Original: [Equation]
Revised: [Equation]

Original: 1) t tµ [Equation] −
Revised: For clarity

Original: [Equation]
Revised: [Equation]

Original: For clarity, let C k
Revised: let C k

Original: = Y k − µ be the zero-mean random variable.
Revised: = Y k − µ be the zero-mean random variable.

Original: Nowozin (2018) gives the relations
Revised: Nowozin (2018) gives the relations

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Expanding both the sums inside the brackets to order two:
Revised: Expanding both the sums inside the brackets to order two:

Original: [Equation]
Revised: [Equation]

Original: We will proceed by bounding each of the terms (1) , (2) , (3) , (4) .
Revised: We will proceed by bounding each of the terms (1) , (2) , (3) , (4) .

Original: First, we decompose C j .
Revised: First, we decompose C j .

Original: Let
Revised: Let

Original: B k,j := [Equation]
Revised: B k,j := 1 j (cid:80) ji = k +1 ( w i − µ ) .

Original: [Equation]
Revised: [Equation]

Original: We know that B k,j is independent of C k and E [ B k,j ] = 0 , implying E [ C k B k,j ] = 0 .
Revised: We know that B k,j is independent of C k and E [ B k,j ] = 0 , implying E [ C k B k,j ] = 0 .

Original: Note C 2 j = k 2 j 2 C 2 k+
Revised: Note C 2 j = k 2 j 2 C 2 k

Original: [Equation]

Original:  C k B k [Equation] j + B 2 k,j .
Revised: + 2 kj C k B k,j + B 2 k,j .

Original: Now we show that (1) is zero:
Revised: Now we show that (1) is zero:

Original: [Equation]
Revised: [Equation]

Original: We now investigate (2) :
Revised: We now investigate (2) :

Original: [Equation]
Revised: [Equation]

Original: We now show that (3) is zero:
Revised: We now show that (3) is zero:

Original: [Equation]
Revised: [Equation]

Original: + 1) 2 ] =

Original: Finally, we investigate (4) :
Revised: Finally, we investigate (4) :

Original: [Equation]
Revised: [Equation]

Original: Using the relation in equation
Revised: Using the relation in equation

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: In summary [Equation] E ∆ k ∆ j is O ( k − 2 j − 2 ) 
Revised: In summary, E ∆ k ∆ j is O ( k − 2 j − 2 ) .

Original: [Equation]
Revised: Assume that ∇ SUMO is bounded: it is sufﬁcient that ∇ IWAE 1 is bounded and that the sampling probabilities are chosen such that the partial sums of ∇ θ ∆ k P ( K≥ k ) converge, i.e. P ( K ≥ k ) >

Original: A. C ONVERGENCE OF ∇ (cid:16) IWAE k +1 − IWAE k (cid:17)
Revised: ck ||∇ θ ∆ k || for some constant c . Then we have E [ ∇ θ SUMO ( x )] = ∇ θ E [ SUMO ( x )]

Original: The IWAE log likelihood estimate is: [Equation]
Revised: = ∇ θ log p θ ( x ) directly by the dominated convergence theorem, as long as SUMO is everywhere differentiable, which is satisﬁed by all of our experiments. If ReLU neural networks are to be used, one may be able to show the same property using Theorem 5 of Bi´nkowski et al. (2018), assuming ﬁnite higher moments and

Original: The gradient of this with respect to λ , where λ is either θ or ψ , is
Revised: Lipschitz constant. A. C ONVERGENCE OF ∇ (cid:16) IWAE k

Original: [Equation]
Revised: [Equation]

Original: θ ( x,z i ) q ψ ( z i | x ) k (cid:88)
Revised: 1 − IWAE k (cid:17)

Original: i = d dλ p
Revised: The IWAE log likelihood estimate is:

Original: [Equation]
Revised: [Equation]

Original: ( x [Equation] z
Revised: The gradient of this with respect to λ , where λ is either θ or ψ , is

Original: [Equation]
Revised: [Equation]

Original: ) q ψ ( z i | x )

Original: We abbreviate w i :=
Revised: We abbreviate w i :=

Original: p θ ( x,z i ) q ψ
Revised: p θ ( x,z i ) q ψ

Original: ( z i | x ) , and ν i
Revised: ( z i | x ) , and ν i

Original: = dw i dλ .
Revised: = dw i dλ .

Original: In both λ = ψ and λ = θ cases, it sufﬁces to treat the w i and ν i as i.i.d.
Revised: In both λ = ψ and λ = θ cases, it sufﬁces to treat the w i and ν i as i.i.d.

Original: random variables with ﬁnite variance and expectation.
Revised: random variables with ﬁnite variance and expectation.

Original: They are i.i.d.
Revised: Being a likelihood ratio, w i could be ill behaved when the importance sampling distribution q ψ ( z i | x ) is is particularly mismatched from the true posterior p ( z i | x ) = [Equation] p [Equation]

Original: due to i.i.d.
Revised: ( x [Equation] z ) . However, the analysis from IWAE (Burda et al., 2016) requires assuming that the likelihood ratios w i

Revised: =

Original: sampling of the z i from q , and the conditions on the moments are ensured for reasonable neural network architectures with bounded weights (enforcing a Lipschitz constraint), and are also required for the basic IWAE estimator to have ﬁnite variance and expectation.
Revised: p θ ( x,z i ) q ψ ( z i | x ) are bounded, and we adopt this assumption. Reasoning about when this assumption holds, and the behavior of IWAE-like estimators when it does not, is an interesting area for future work.

Original: Consider the differences between two gradients: we label ∆ g as follows:
Revised: Consider the differences between two gradients: we label ∆ g as follows:

Original: [Equation]
Revised: [Equation]

Original: We have:
Revised: We have:

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: We again let Y k denote the k th sample mean 1 k (cid:80) i w i . Then: [Equation] The sample means Y k and ¯ µ k have ﬁnite expectation and variance.

Original: [Figure]
Revised: The variance vanishes as k → ∞ (but the expectation does not change). [Equation] Let ν k + Y k + k k

Original: [Figure]
Revised: [Equation] [Equation] 1 ¯ ν k [Equation]

Revised: The second term vanishes at a rate strictly faster than 1 k 2 : the variance of φ k goes to zero as k → ∞ .

Revised: But the ﬁrst term does not: φ k is a biased estimator of φ ∞ so E φ k does change with k , but it does not necessarily go to zero:

Original: [Figure]
Revised: [Equation]

Revised: Thus, E || ∆ gk || 2 2 is at most O ( 1 k 2 ) .

Revised: [Figure]

Revised: [Figure]

Original: E v i d e n c e e s t
Revised: E v i d e n c e e s t

Original: i m a t e o f l o g p ( x )
Revised: i m a t e o f l o g p ( x )

Original: SUMOIWAE 10101010 3 K 85.
Revised: SUMOIWAE 10101010 3 K 85.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: E v i d e n c e e s t
Revised: E v i d e n c e e s t

Original: i m a t e o f l o g p ( x )
Revised: i m a t e o f l o g p ( x )

Original: SUMO_geometricIWAE 10101010 3 K
Revised: SUMO_geometricIWAE 10101010 3 K

Original: E v i d e n c e e s t
Revised: E v i d e n c e e s t

Original: i m a t e o f l o g p ( x )
Revised: i m a t e o f l o g p ( x )

Original: SUMO poisson IWAE
Revised: SUMO poisson IWAE

Original: We again let Y k denote the k th sample mean 1 k (cid:80) i w i .

Original: Then:

Original: [Equation]

Original: The sample means Y k and ¯ µ k have ﬁnite expectation and variance.

Original: The variance vanishes as k → ∞ (but the expectation does not change).

Original: [Equation]

Original: Let ν [Equation] The second term vanishes at a rate strictly faster than 1 k 2 : the variance of φ k goes to zero as k → ∞ . But the ﬁrst term does not: φ k is a biased estimator of φ ∞ so E φ k does change with k , but it does not necessarily go to zero:
Revised: A.

Original: [Equation]

Original: Thus, E || ∆ gk || 22 is at most O ( 1 k 2 ) .

Original: A. E MPIRICAL C ONFIRMATION ON THE C ONVERGENCE OF E ∆ 2 k AND E || ∆ gk ||
Revised: E MPIRICAL C ONFIRMATION ON THE C ONVERGENCE OF E ∆ 2 k AND E || ∆ gk ||

Original: We measure the ∆ 2 k and || ∆ gk || 2 2 on a toy example to verify the convergence rates empirically.
Revised: We measure the ∆ 2 k and || ∆ gk || 22 on a toy example to verify the convergence rates empirically.

Original: We reimplement the toy Gaussian example from Rainforth et al. 
Revised: We re-implement the toy Gaussian example from Rainforth et al. (2018b)

Original: ; Tucker et al.
Revised: ; Tucker et al.

Original: The generative model is p θ ( x, z ) =
Revised: The generative model is p θ ( x, z ) =

Original: N ( z | θ, I ) N ( x | z, I )  .
Revised: N ( z | θ, I ) N ( x | z, I ) , where both x and z are in R D .

Original: The encoder is q φ ( z | x ) =
Revised: The encoder is q φ ( z | x ) =

Original: N ( z | Ax + b, 23 I )  .
Revised: N ( z | Ax + b, 23 I ) , where φ = ( A, b ) .

Revised: The synthetic dataset was generated with D =and N = 1000 data points using the true model parameter θ true from a standard Gaussian.

Original: Alongside || ∆ k
Revised: Alongside || ∆ k

Original: || 22 , we plot several reference convergence rates such as O ( 1 / k c ) , c > 1 , and O ( c k ) , c < 1 , as a visual guide.
Revised: || 22 , we plot several reference convergence rates such as O ( 1 / k c ) , c > 1 , and O ( c k ) , c < 1 , as a visual guide.

Original: The results are shown in Figure 5.
Revised: The results are shown in Figure 5.

Original: Following the setup in Rainforth et al.
Revised: Following the setup in Rainforth et al.

Original: (2018), we sample a group of model parameters close to the optimal values which are perturbed by Gaussian noise from N (0 , 0 . 
Revised: (2018b), we sample a group of model parameters close to the optimal values which are perturbed by Gaussian noise from N (0 , 0 . 01 2 ) .

Original: The gradient ∆ gk is taken w.r.t.
Revised: The gradient ∆ gk is taken w.r

Revised: [Figure]

Original:  the model parameter θ .
Revised: t. the model parameter θ .

Original:  A. E XPERIMENTAL S ETUP
Revised: [Figure] [Figure] [Figure] A. B IAS -V ARIANCE T RADEOFF VIA G RADIENT C LIPPING

Revised: While SUMO is unbiased, its variance is extremely high or potentially inﬁnite.

Revised: This property leads to poor performance compared to lower bound estimates such as IWAE when maximizing loglikelihood.

Revised: In order to obtain models with competitive log-likelihood values, we can make use of gradient clipping.

Revised: This allows us to ignore rare gradient samples with extremely large values due to the heavy-tailed nature of its distribution.

Revised: Gradient clipping introduces bias in favor of reduced variance.

Revised: Figure 6 shows how the performance changes as a function of the clipping value, and more importantly, the percentage of clipped gradients.

Revised: As shown, neither full clipping and no clipping are desirable.

Revised: We performed this experiment after reporting the results in Table 1, so this grid search was not used to tune hyperparameter for our experiments.

Revised: As bias is introduced, we do not use gradient clipping for entropy maximization or policy gradient (REINFORCE).

Revised: [Figure]

Revised: A. E XPERIMENTAL S ETUP

Original: All the models are trained using a batch size of 100 and an Amsgrad optimizer (Reddi et al., 2018) with parameters lr = 0 .
Revised: In density modeling experiments, all the models are trained using a batch size of 100 and an Amsgrad optimizer (Reddi et al., 2018) with parameters lr = 0 .

Original: 001 , β 1 = 0 .
Revised: 001 , β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: 999 and (cid:15) =
Revised: 999 and (cid:15) =

Original: 10 − 4 .
Revised: 10 − 4 .

Original: The learning rate is reduced by factor 0 .
Revised: The learning rate is reduced by factor 0 .

Original: 8 with a patience of 50 epochs.
Revised: 8 with a patience of 50 epochs.

Original: We use gradient norm scaling in both the inference and generative networks.
Revised: We use gradient norm scaling in both the inference and generative networks.

Original: We train SUMO using the same architecture and hyperparameters as IWAE except the gradient clipping norm.
Revised: We train SUMO using the same architecture and hyperparameters as IWAE except the gradient clipping norm.

Original: We set the gradient norm to 5000 for encoder and { 20, 40, 60 } for decoder in SUMO.
Revised: We set the gradient norm tofor encoder and { 20, 40, 60 } for decoder in SUMO.

Original: For IWAE, the gradient norm is ﬁxed to 10 in all the experiments.
Revised: For IWAE, the gradient norm is ﬁxed to 10 in all the experiments.

Original: We report the performance of models with early stopping if no improvements have been observed for 300 epochs on the validation set.
Revised: We report the performance of models with early stopping if no improvements have been observed for 300 epochs on the validation set.

Original: [Figure] [Figure]
Revised: 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 on MNIST with expected number of compute k = 15 . Gradient clipping was not used in the other experiments except the density modeling ones, where it can be used as a tool to obtain a better bias-variance trade-off.

Original: A. 
Revised: A.

Original:  R EVERSE KL AND C OMBINATORIAL O PTIMIZATION
Revised:  R EVERSE KL AND C OMBINATORIAL O PTIMIZATION

Original: These two tasks use the same encoder and decoder architecture: one hidden layer with tanh nonlinearities and 200 hidden units.
Revised: These two tasks use the same encoder and decoder architecture: one hidden layer with tanh nonlinearities and 200 hidden units.

Original: We set the latent state to be of size 20.
Revised: We set the latent state to be of size 20.

Original: The prior is a standard
Revised: The prior is a standard

Original: Gaussian with diagonal covariance, while the encoder distribution is a Gaussian with parameterized diagonal covariance.
Revised: Gaussian with diagonal covariance, while the encoder distribution is a Gaussian with parameterized diagonal covariance.

Original: For reverse KL, we used independent Gaussian conditional likelihoods for p ( x | z ) , while for combinatorial optimization we used independent Bernoulli conditional distributions.
Revised: For reverse KL, we used independent Gaussian conditional likelihoods for p ( x | z ) , while for combinatorial optimization we used independent Bernoulli conditional distributions.

Original: We found it helps stablize training for both IWAE and SUMO to remove momentum and used
Revised: We found it helps stablize training for both IWAE and SUMO to remove momentum and used

Original: RMSprop with learning rate 0.00005 and epsilon 1e-3 for ﬁtting reverse KL.
Revised: RMSprop with learning rate 0.00005 and epsilon 1e-3 for ﬁtting reverse KL.

Original: We used Adam with learning rate 0.001 and epsilon 1e-3, plus standard hyperparameters for the combinatorial optimization problems.
Revised: We used Adam with learning rate 0.001 and epsilon 1e-3, plus standard hyperparameters for the combinatorial optimization problems.

Original: SUMO used an expected compute of 15 terms, with m = 5 and the tail-modiﬁed telescoping Zeta distribution.
Revised: SUMO used an expected compute of 15 terms, with m = 5 and the tail-modiﬁed telescoping Zeta distribution.