Title: On importance-weighted autoencoders

Abstract: The importance weighted autoencoder (IWAE) (Burda et al., 2016 ) is a popular variational-inference method which achieves a tighter evidence bound (and hence a lower bias) than standard variational autoencoders by optimising a multi-sample objective, i.e. an objective that is expressible as an integral over K > 1 Monte Carlo samples. Unfortunately, the IWAE multi-sample objective leads to inference-network gradients which break down as K is increases (Rainforth et al., 2018) . This breakdown can only be circumvented by removing high-variance score-function terms, either by heuristically ignoring them (which yields the 'sticking-the-landing' IWAE (IWAE-STL) gradient from Roeder et al. ( 2017)) or through an identity from Tucker et al. ( 2019) (which yields the 'doubly-reparametrised' IWAE (IWAE-DREG) gradient). In this work, we develop an encompassing framework which directly optimises the proposal distribution in importance sampling as in the reweighted wake-sleep (RWS) algorithm from Bornschein & Bengio (2015) . From this unified framework, most of the previously proposed gradient estimators can be naturally derived. This permits a better understanding of the assumptions and trade-offs that are at play. Importantly, the derived gradient estimators are guaranteed to not degenerate as K → ∞.

1. Introduction:

Let x be some observation and let z be some latent variable taking values in some space Z. These are modeled via the generative model p θ (z, x) = p θ (z)p θ (x|z) which gives rise to the marginal likelihood p θ (x) = Z p θ (z, x) dz of the model parameters θ. In this work, we analyse algorithms for variational inference, i.e. algorithms which aim to 1. learn the generative model, i.e. find a value θ which is approximately equal to the maximum-likelihood estimate (MLE) θ ml := arg max θ p θ (x); 2. construct a tractable variational approximation q φ,x (z) of p θ (z|x) = p θ (z, x)/p θ (x), i.e. find the value φ such that q φ ,x (z) is as close as possible to p θ (z|x) in some suitable sense.

A few comments about this setting are in order. Firstly, as is common in the literature, we restrict our presentation to a single latent representation-observation pair (z, x) to avoid notational clutter -the extension to multiple independent observations is straightforward. Secondly, we assume that no parameters are shared between the generative model p θ (z, x) and the variational approximation q φ,x (z). This is common in neural-network applications but could be relaxed. Thirdly, our setting is general enough to cover amortised inference.

For this reason, we often refer to φ as the parameters of an inference network.

Two main classes of stochastic gradient-ascent algorithms for optimising ψ := (θ, φ) which employ K ≥ 1 Monte Carlo samples ('particles') to reduce errors have been proposed.

• IWAE. The importance weighted autoencoder (IWAE) (Burda et al., 2016) maximizes a joint lower bound L K ψ ≤ p θ (x) whose bias decreases as K → ∞. The gradients of this objective can be unbiasedly approximated via the Monte-Carlo method. Unfortunately, the signal-to-noise ratio of the IWAE φ-gradient vanishes as K grows (Rainforth et al., 2018) . Two modified IWAE φ-gradients avoid this breakdown by removing high-variance 'score-function' terms:

-IWAE-STL. The 'sticking-the-landing' IWAE (IWAE-STL) φ-gradient (Roeder et al., 2017) heuristically drops the problematic score-function terms from the IWAE φ-gradient. This induces bias for the IWAE objective. -IWAE-DREG. The 'doubly-reparametrised' IWAE (IWAE-DREG) φ-gradient (Tucker et al., 2019) unbiasedly removes the problematic score-function terms from the IWAE φ-gradient using a formal identity.

• RWS. The reweighted wake-sleep (RWS) algorithm (Bornschein & Bengio, 2015) optimises two separate objectives for θ and φ. Its gradients are approximated by self-normalised importance sampling with K particles: this induces a bias which vanishes as K → ∞. RWS can be viewed as an adaptive importance-sampling approach which iteratively improves its proposal distribution while simultaneously optimising θ via stochastic approximation. Crucially, the RWS φ-gradients do not degenerate as K → ∞.

Of these two methods, the IWAE is the most popular and Tucker et al. (2019) demonstrated empirically that RWS can break down, conjecturing that this is due to the fact that RWS does not optimise a joint objective (for θ and φ). Meanwhile, the IWAE-STL gradient performed consistently well despite lacking a firm theoretical footing. Yet, IWAE suffers from the above-mentioned φ-gradient breakdown and exhibited inferior empirical performance to RWS (Le et al., 2019) . Thus, it is not clear whether the multi-sample objective approach of IWAE or the adaptive importance-sampling approach of RWS is preferable.

In this work, we show that directly optimising the proposal distribution, e.g. as done by RWS, is preferable to optimising the IWAE multi-sample objective because (a) the multi-sample objective typically relies on reparametrisations and, even if these are available, leads to the φ-gradient breakdown, (b) modifications of the IWAE φ-gradient which avoid this breakdown (i.e. IWAE-STL and IWAE-DREG) can be justified in a more principled manner by taking an RWS-type adaptive importance-sampling view. This conclusion was already reached by Le et al. (2019) based on numerical experiments. They demonstrated that the need for reparametrisations can make IWAE inferrior to RWS e.g. for discrete latent variables. Our work complements theirs by formalising this argument. To this end, we slightly generalise the RWS algorithm to obtain a generic adaptive importance-sampling framework for variational inference which we term adaptive importance sampling for learning (AISLE) for ease of reference. We then show that AISLE admits not only RWS but also the IWAE-DREG and IWAE-STL gradients as special cases.

Contributions. Novel material is presented in Section 3, where we introduce the AISLEframework. From this, most of the previously proposed gradient estimators can be naturally derived in a principled manner. Importantly, the derived gradient estimators are guaranteed to not degenerate as K → ∞. Specifically, we establish the following connections.

• We prove that the IWAE-STL gradient can be recovered as a special case of AISLE via a principled and novel application of the 'double-reparametrisation' identity from Tucker et al. (2019) . This indicates that the breakdown of RWS observed in Tucker et al. (2019) may not be due to its lack of a joint objective as previously conjectured (since IWAE-STL avoided this breakdown despite having the same idealised objective as RWS). Our work also provides a theoretical foundation for IWAE-STL which was hitherto only heuristically justified as a biased IWAE-gradient.

• We prove that AISLE also admits the IWAE-DREG gradient as a special case. Our derivation also makes it clear that the learning rate should be scaled as O(K) for the IWAE φ-gradient (and its modified version IWAE-DREG) unless the gradients are normalised as implicitly done by popular optimisers such as ADAM (Kingma & Ba, 2015) . In contrast, the learning rate for AISLE need not be scaled up with of K.

• When applied to the family of α-divergences, AISLE leads to a new family of gradient estimators that generalises some previously derived in the literature.

• In the supplementary materials, we provide insights into the impact of the selfnormalisation bias on some of the importance-sampling based gradient approxima-tions (Appendix A) and empirically compare the main algorithms discussed in this work (Appendix B).

We stress that the focus of our work is not necessarily to derive new algorithms nor to establish which of the various special cases of AISLE is preferable. Indeed, while we compare all algorithms discussed in this work empirically on Gaussian models in the supplementary materials, we refer the reader to Tucker et al. (2019) ; Le et al. (2019) for an extensive empirical comparisons of all the algorithms discussed in this work.

Notation. We repeatedly employ the shorthand p(f ) := Z f (z)p(z) dz for the integral of some p-integrable test function f ; thus, p(f

) = E z∼p [f (z)] if p is a probability measure.

Furthermore, q ⊗K (z 1:K ) := K k=1 q(z k ). To keep the notation concise, we hereafter suppress dependence on the observation x, i.e. we write q φ (z) := q φ,x (z) as well as

π θ (z) := p θ (z|x) = p θ (z, x) p θ (x) = γ θ (z) Z θ ,

where γ θ (z) := p θ (z, x) and where Z θ := p θ (x) = Z γ θ (z) dz.

2. Background:



2.1. Importance sampling:

The expectation q φ (f ) of a test function f : Z → R can be unbiasedly estimated by the quantity [f (z 1 ) + . . . + f (z K )]/K using a set of K particles, z := (z 1 , . . . , z K ) ∼ q ⊗K φ , which are independent and identically distributed (IID) according to q φ . Similarly, expectations of the type π θ (f ) can be approximated by the self-normalised importance sampling estimate πθ φ, z (f ) :

= K k=1 s w k ψ f (z k ) with s w k ψ := w ψ (z k ) K l=1 w ψ (z l )

and w ψ (z) := γ θ (z) q φ (z) .

The notation φ, z stresses the dependence of this estimator on φ and z. The quantity w ψ (z k ) are called the kth importance weight and s w k ψ is its self-normalised version. For readability, we have dropped the dependence of s w k ψ on z ∈ Z K from the notation. Remark 1. Since Z θ = q θ (w ψ ), the quantity Z θ φ, z := [w ψ (z 1 ) + . . . + w ψ (z K )]/K is an unbiased ('importance-sampling') estimator of Z θ .

Remark 2. The self-normalised estimate πθ φ, z (f ) is typically not unbiased. Under mild assumptions (e.g. if sup w ψ < ∞), its bias vanishes at rate O(K -1 ), its standard deviation vanishes at Monte-Carlo rate O(K -1/2 ) and πθ φ, z (f ) → π θ (f ) almost surely as K → ∞.

2.2. Importance weighted autoencoder (IWAE):

Objective. The importance weighted autoencoder (IWAE), introduced by Burda et al. (2016) , seeks to find a value θ of the generative-model parameters θ which maximises a lower bound L K ψ on the log-marginal likelihood ('evidence'). This bound depends on the inference-network parameters φ and the number of samples, K ≥ 1:

ψ := (θ , φ ) := arg max ψ L K ψ , L K ψ := E log Z θ φ, z . (1)

where the expectation is w.r.t. z ∼ q ⊗K φ . For any finite K, optimisation of the inferencenetwork parameters φ tightens the evidence bound. Burda et al. (2016) prove that for any φ we have that L K ψ ↑ log Z θ as K → ∞. If K = 1, the IWAE reduces to the variational autoencoder (VAE) from Kingma & Welling (2014) . However, for K > 1, as pointed out in Cremer et al. (2017) ; Domke & Sheldon (2018) , the IWAE also constitutes another VAE on an extended space based on an auxiliary-variable construction developed in Andrieu & Roberts (2009) ; Andrieu et al. (2010) ; Lee (2011) (see, e.g. Finke, 2015, for a review).

Standard reparametrisation gradient.:

The gradient of the IWAE objective from (1):

∇ ψ L K ψ = E ∇ ψ log Z θ φ, z + G ψ (z) , with G ψ (z) := log Z θ φ, z K k=1 ∇ ψ log q φ (z k

). The intractable quantity E G ψ (z) can be approximated unbiasedly via a vanilla Monte Carlo approach using a single (K-dimensional) sample point z = (z 1 , . . . , z K ) ∼ q ⊗K φ . Unfortunately, this approximation typically has such a large variance that it becomes impracticably noisy (Paisley et al., 2012) . To remove this high-variance term, the well known reparametrisation trick (Kingma & Welling, 2014) is usually employed. It requires the following assumption.

(R1) There exists a distribution q on some space E and a family of differentiable mappings

h φ : E → Z such that if ∼ q we have that z = z( ) = h φ ( ) ∼ q φ .

Under R1, with 1 , . . . , K iid ∼ q and z k := z( k ) := h φ ( k ), the gradient can be expressed as

EQUATION

)

with ψ (z) := ∇ φ [log • w ψ • h φ ]| ψ =ψ (h -1 φ (z)).

Here, the notation ψ indicates that one does not differentiate w ψ w.r.t. ψ. The IWAE then uses a vanilla Monte Carlo estimate of (2),

EQUATION

)

Before proceeding, we state the following lemma, proved in Tucker et al. (2019, Section 8.1) , which generalises of the well-known identity q φ (∇ φ log q φ ) = 0. Lemma 1 (Tucker et al. ( 2019)). Under R1, for suitably integrable f ψ : Z → R, we have

q φ (f ψ ∇ φ log q φ ) = q (∇ φ [f ψ • h φ ]| ψ =ψ ) = q φ (∇ φ [f ψ • h φ ]| ψ =ψ • h -1 φ ).

We now exclusively focus on the φ-portion of the IWAE gradient, ∇ iwae φ θ, z . Remark 3 (drawbacks of the IWAE φ-gradient). The gradient ∇ iwae φ θ, z has three drawbacks. The last two of these are attributable to the 'score-function' terms ∇ φ log q φ (z) in the φ-gradient portion of (3).

• Reliance on reparametrisations. A reparametrisation à la R1 is necessary to remove the high-variance term G ψ (z). For, e.g. discrete, models that violate R1, control-variate approaches (Mnih & Rezende, 2016) or continuous relaxations have been proposed but these incur additional implementation, tuning and computation costs whilst not necessarily reducing the variance (Le et al., 2019) .

• Vanishing signal-to-noise ratio. The φ-gradient breaks down in the sense that its signal-to-noise ratio vanishes as et al., 2018) . This is because ∇ iwae φ θ, z constitutes a self-normalised importance-sampling approximation of π θ ( ψ -∇ φ log q φ ) = 0, an identity which directly follows from Lemma 1 with f ψ = w ψ .

E[ ∇ iwae φ θ, z ]/ var[ ∇ iwae φ θ, z ] 1/2 = O(K -1/2 ) (Rainforth

• Inability to achieve zero variance. As pointed out in Roeder et al. (2017) , var[ ∇ iwae φ θ, z ] > 0 even in the ideal scenario where q φ = π θ despite the fact that in this case, w ψ is constant and hence var[log Z θ φ, z ] = 0.

Two modifications of ∇ iwae φ θ, z have been proposed which (under R1) avoid the scorefunction terms in (3) and hence (a) exhibit a stable signal-to-noise ratio as K → ∞ and (b) can achieve zero variance if q φ = π θ (because then ψ ≡ 0 since w ψ is constant).

• IWAE-STL. The 'sticking-the-landing' IWAE (IWAE-STL) gradient proposed by Roeder et al. (2017) heuristically ignores the score function terms,

EQUATION

)

As shown in Tucker et al. (2019) ), this introduces an additional bias whenever K > 1.

• IWAE-DREG. The 'doubly-reparametrised' IWAE (IWAE-DREG) gradient proposed by Tucker et al. (2019) removes the score-function terms through Lemma 1,

∇ iwae-dreg φ θ, z := K k=1 s w k ψ 2 ψ (z k ).

(5)

The quantities ∇ iwae-dreg φ θ, z and ∇ iwae φ φ, z are equal in expectation.

2.3. Reweighted wake-sleep (RWS):

The reweighted wake-sleep (RWS) algorithm was proposed in Bornschein & Bengio (2015) .1 Letting KL(p q) := Z log[p(z)/q(z)]p(z) dz is the Kullback-Leibler (KL)-divergence from p to q, the RWS algorithm seeks to optimise ψ = (θ, φ) as

θ := arg max θ log Z θ , φ := arg min φ KL(π θ q φ ).

The θ-and φ-gradients read

EQUATION

)

These quantities are usually intractable and therefore approximated by replacing π θ by the self-normalised importance sampling approximation πθ φ, z (this does not require R1):

EQUATION

)

Since ( 7) relies on self-normalised importance sampling, Remark 2 shows that its bias relative to ( 6) is of order O(1/K). Appendix A discusses the impact of this bias on the φ-gradient in more detail. The optimisation of both θ and φ is carried out simultaneously, allowing both gradients to share the same particles and weights. Nonetheless, the lack of a joint objective (for both θ and φ) is often viewed as the main drawback of RWS.

RWS-DREG.:

Under R1, Tucker et al. (2019) proposed the following 'doublyreparametrised' RWS (RWS-DREG) gradient which is equal to ∇ rws φ θ, z in expectation and is derived by applying Lemma 1 to the latter. It reads

EQUATION

where the function F(w) := w(1 -w) is used to transform the self-normalised importance weights s w k ψ . In high-dimensional settings, it is typically the case that the ordered selfnormalised importance weights s w

(K) ψ < . . . < s w (1)

ψ < 1 are such that s w

(1)

ψ ≈ 1 -s w (2) ψ and s w (k) ψ s w (2) ψ for k ≥ 3. The transformed weights {F( s w k ψ )} K k=1

are then mainly supported on the two particles with the largest self-normalised weights.

3 AISLE: A unified adaptive importance-sampling framework

3.1. Objective:

If θ is fixed, the RWS algorithm reduces to an adaptive importance-sampling scheme which optimises the proposal distribution by minimising the 'inclusive' KL-divergence from the target distribution π θ to the proposal q φ (see, e.g., Douc et al., 2007; Cappé et al., 2008) . If instead φ is fixed, the RWS algorithm reduces to a stochastic-approximation algorithm for estimating the MLE of the generative-model parameters θ. The advantage of optimising θ and φ simultaneously is that (a) Monte Carlo samples used to approximate the θ-gradient can be re-used to approximate the φ-gradient and (b) optimising φ typically reduces the error (both in terms of bias and variance) of the θ-gradient approximation.

However, adapting the proposal distribution q φ in importance-sampling schemes need not necessarily be based on minimising the (inclusive) KL-divergence. Numerous other techniques exist in the literature (e.g. Geweke, 1989; Evans, 1991; Oh & Berger, 1992; Richard & Zhang, 2007; Cornebise et al., 2008) and may sometimes be preferable. Indeed, another popular approach with strong theoretical support is based on minimising the χ 2 -divergence (see, e.g., Deniz Akyildiz & Míguez, 2019) . Based on this insight, we slightly generalise the RWS-objective as θ := arg max θ log Z θ , φ := arg min φ Dƒ(πθ q φ ). ( 9)

Here, Dƒ(p q) := Z ƒ(p(z)/q(z))q(z) dz is some ƒ-divergence from p to q. We reiterate that alternative approaches for optimising φ (which do not minimise ƒ-divergences) could be used. However, we state ( 9) for concreteness as it suffices for the remainder of this work; we call the resulting algorithm adaptive importance sampling for learning (AISLE). As will become clear below, this unified framework permits a straightforward and principled derivation of robust φ-gradient estimators that do not degenerate as K → ∞.

3.2. θ-gradient:

Optimisation is again performed via a stochastic gradient-ascent. The intractable θ-gradient

∇ θ log Z θ = π θ (∇ θ log γ θ ) is approximated as in RWS, i.e. for z ∼ q ⊗K φ : ∇ aisle θ φ, z := ∇ rws θ φ, z = ∇ iwae θ φ, z = K k=1 s w k ψ ∇ θ log γ θ (z k ).

The θ-gradient is thus the same for all algorithms discussed in this work although the IWAEparadigm views it as an unbiased gradient of a (biased) lower-bound to the evidence, while AISLE (and RWS) interpret it as a self-normalised importance-sampling (and consequently biased) approximation of the gradient ∇ θ log Z θ for the 'exact' objective.

3.3. φ-gradient:



3.3.1. General derivation:

In the derivations to follow, integrals of the form π θ ([F • w ψ ]∇ φ log q φ ) naturally appear. These can also be expressed as

Z -1 θ q φ ([H • w ψ ]∇ φ log q φ ) with H(y) := F (y)y. By Lemma 1, π θ ([F • w ψ ]∇ φ log q φ ) = 1 Z θ E z∼q φ [w ψ H (w ψ (z)) ψ (z)] .

Approximating the expectation as well as the normalising constant Z θ on the r.h.s. with the vanilla Monte Carlo method with K samples z 1 , . . . , z K iid ∼ q ⊗K φ yields the estimator

EQUATION

)

Remark 2 shows that this approximation has a bias of order O(K -1 ) and a standarddeviation of order O(K -1/2 ). Now, most of the ƒ-divergences used for variational inference in intractable models are such that there exists a function f : R → R satisfying Dƒ(πθ q φ ) = Z κ θ Z f [w ψ (z)]q φ (z) dz + C(θ) for an exponent κ ∈ R and constant C(θ) independent of φ. In other words, for a given value of θ, the optimization of the ƒ-divergence as a function of φ can be carried out without relying on the knowledge of Z θ . Writing g(y) := f (y) -f (y)/y, simple algebra then directly shows that

EQUATION

)

Since the integral in ( 11) is an expectation with respect to π θ , it can be approximated with selfimportance sampling, possibly multiplied an additional importance-sampling approximation Z θ φ, z of Z θ raised to some power. This leads to,

EQUATION

)

Indeed, Equation (10) applies to (11), leading to the reparametrised estimator

EQUATION

where h(y) = g(y)y and g : R → R given immediately above (11). We now describe several particular cases.

3.3.2. Special case 'inclusive' KL-divergence: RWS and IWAE-STL:

We have KL(π θ q φ ) = Z κ θ Z f (w ψ (z))q φ (z) dz + C(θ) with κ = -1 and f (y) = y log(y). In that case, with the notations of Section 3.3.1, we have g(y) = 1 and h (y) = 1.

• AISLE-KL-NOREP/RWS. Without relying on any reparametrisation, Equation ( 12) yields the following gradient, which clearly equals ∇ rws φ θ, z :

EQUATION

• AISLE-KL. Using reparametrisation, Equation ( 13) yields the gradient:

EQUATION

We thus arrive at the following result which demonstrates that IWAE-STL can be derived in a principled manner from AISLE, i.e. without the need for a multi-sample objective. Proposition 1. For any (θ, φ, z), ∇ aisle-kl φ θ, z = ∇ iwae-stl φ θ, z . Proposition 1 is notable because it shows that IWAE-STL (which avoids the breakdown highlighted in Rainforth et al. (2018) and which can also achieve zero variance) can be derived in a principled manner from AISLE, i.e. without relying on a multi-sample objective. Proposition 1 thus provides a theoretical basis for IWAE-STL which was previously viewed as an alternative gradient for IWAE for which it is biased and only heuristically justified. Furthermore, the fact that IWAE-STL exhibited good empirical performance in Tucker et al. (2019) even in an example in which RWS broke down, suggests that this breakdown may not be due to RWS' lack of optimising a joint objective as previously conjectured.

Finally, recall that Tucker et al. ( 2019) obtained an alternative 'doubly-reparametrised' RWS φ-gradient ∇ rws-dreg φ θ, z given in (8) by first replacing the exact (but intractable) φ-gradient by the self-normalised importance-sampling approximation ∇ rws φ θ, z and then applying the identity from Lemma 1. Note that this may result in a variance reduction but does not change the bias of the gradient estimator. In contrast, AISLE-KL is derived by first applying Lemma 1 to the exact (RWS) φ-gradient and then approximating the resulting expression. This can potentially reduce both bias and variance.

3.3.3. Special case α-divergence: IWAE-DREG:

Up to some irrelevant additive constant, the α-divergence between two distributions p and q is given by Z (p(z)/q(z)) α q(z) dz for some α > 1. This can also be expressed as Z κ θ Z f (w ψ (z))q φ (z) dz with κ = -α and f (y) = y α . In this case, with the notation from Section 3.3.1, we have g(y) = (α -1)y α-1 and h (y) = α(α -1) y α-1 . Note that the case α = 2 is equivalent, up to an irrelevant additive constant, to a standard χ 2 -divergence. Minimising this divergence is natural in importance sampling since χ 2 (π θ q φ ) = var z∼q φ [w ψ /Z θ ] is the variance of the importance weights.

• AISLE-α-NOREP. Without relying on any reparametrisation, Equation ( 13) yields

EQUATION

with the following special case which is also proportional to the 'score gradient' from Dieng et al. (2017, Appendix G) :

EQUATION

)

again with the special case

∇ aisle-χ 2 φ θ, z := 2K K k=1 ( s w k ψ ) 2 ψ (z k ).

This demonstrates that IWAE-DREG can be derived (up to the proportionality factor 2K) in a principled manner from AISLE, i.e. without the need for a multi-sample objective. Proposition 2. For any (θ, φ, z), ∇ aisle-χ 2 φ θ, z = 2K ∇ iwae-dreg φ θ, z . Note that if the implementation normalises the gradients, e.g. as effectively done by ADAM (Kingma & Ba, 2015) , the constant factor cancels out and AISLE-χ 2 becomes equivalent to IWAE-DREG. Otherwise (e.g. in plain stochastic gradient-ascent) this shows that the learning rate needs to be scaled as O(K) for the IWAE or IWAE-DREG φ-gradients.

3.3.4. Special case 'exclusive' KL-divergence.:

For the 'exclusive' KL-divergence, we have KL(q φ π θ ) = f (w ψ (z))q φ (z) dz + C(θ) with f (y) = log(y). In that case, with the notation from Section 3.3.1, we have h (y) = 1/y. This directly leads to the following approximation,

-∇ φ Dƒ(πθ q φ ) ≈ ∇ aisle-rev-kl φ θ, z := 1 K K k=1 ψ (z k ).

This can be recognised as a simple average over K independent replicates of the 'stickingthe-landing' estimator for VAEs proposed in Roeder et al. (2017, Equation 8 ). As we discuss in Appendix A, optimising this 'exclusive' KL-divergence can sometimes lead to faster convergence of φ than optimising the 'inclusive' KL-divergence KL(π θ q φ ). However, care must be taken because minimising the exclusive divergence does not necessarily lead to well behaved or even well-defined importance weights and thus can negatively affect learning of θ (whose gradient is a self-normalised importance-sampling approximation which makes use of those weights).

4. Conclusion:

We have shown that the adaptive-importance sampling paradigm of the reweighted wake-sleep (RWS) (Bornschein & Bengio, 2015) is preferable to the multi-sample objective paradigm of importance weighted autoencoders (IWAEs) (Burda et al., 2016) because the former achieves all the goals of the latter whilst avoiding its drawbacks. A On the rôle of the self-normalisation bias within RWS/AISLE

A.1 The self-normalisation bias Within the self-normalised importance-sampling approximation, the number of particles, K, interpolates between two extremes:

• As K ↑ ∞, πθ φ, z (f ) becomes an increasingly accurate approximation of π θ (f ).

• For K = 1, however, πθ φ, z (f ) = f (z 1 ) reduces to a vanilla Monte Carlo approximation of q φ (f ) (because the single self-normalised importance weight is always equal to 1).

This leads to the following insight about the estimators ∇ aisle-kl φ θ, z and ∇ aisle-χ 2 φ θ, z .

• As K ↑ ∞, these two estimators become increasingly accurate approximations of the 'inclusive'-divergence gradients -∇ φ KL(π θ q φ ) = π θ ( φ ) and

-∇ φ χ 2 (π θ q φ ) = 2π θ ([w ψ /Z θ ] φ )

, respectively. • For K = 1, however, these two estimators reduce to vanilla Monte Carlo approximations of the 'exclusive'-divergence gradients -∇ φ KL(q φ π θ ) = q φ ( φ ) and -2∇ φ KL(q φ π θ ) = 2q φ ( φ ), respectively. This is similar to the standard IWAE φ-gradient which also represents a vanilla Monte Carlo approximation of -∇ φ KL(q φ π θ ) if K = 1 as IWAE reduces to a VAE in this case.

Characterising the small-K self-normalisation bias of the reparametrisation-free AISLE φ gradients, AISLE-KL-NOREP and AISLE-χ 2 -NOREP, is more difficult because if K = 1, they constitute vanilla Monte Carlo approximations of q φ (∇ φ log q φ ) = 0. Nonetheless, Le et al.

(2019, Figure 5 ) lends some support to the hypothesis that the small-K self-normalisation bias of these gradients also favours a minimisation of the exclusive KL-divergence.

A.2 Inclusive vs exclusive KL-divergence minimisation:

Recall that the main motivation for use of IWAEs (instead of VAEs) was the idea that we could use self-normalised importance-sampling approximations with K > 1 particles to reduce the bias of the θ-gradient relative to ∇ θ log Z θ . The error of such (self-normalised) importance-sampling approximations can be controlled by ensuring that q φ is close to π θ (in some suitable sense) in any part of the space Z in which π θ has positive probability mass. For instance, it is well known that the error will be small if the 'inclusive' KL-divergence KL(π θ q φ ) is small as this implies well-behaved importance weights. In contrast, a small 'exclusive' KL-divergence KL(q φ π θ ) is not sufficient for well-behaved importance weights because the latter only ensures that q φ is close to π θ in those parts of the space Z in which q φ has positive probability mass.

Let Q := {q φ } (which is indexed by φ) be the family of proposal distributions/the variational family. Then we can distinguish two scenarios.

1. Sufficiently expressive Q. For the moment, assume that the family Q is flexible ('expressive') enough in the sense that it contains a distribution q φ which is (at least approximately) equal to π θ and that our optimiser can reach the value φ of φ. In this case, minimising the exclusive KL-divergence can still yield well-behaved importance weights because in this case, φ := arg min φ KL(π θ q φ ) is (at least approximately) equal to arg min φ KL(q φ π θ ).

2. Insufficiently expressive Q. In general, the family Q is not flexible enough in the sense that all of its members are 'far away' from π θ , e.g. if the D components z 1 , . . . , z D of z = z 1:D are highly correlated under π θ whilst q φ (z) = D d=1 q φ,d (z d ) is fully factorised. In this case, minimising the exclusive KL-divergence could lead to poorly-behaved importance weights and we should optimise φ := arg min φ KL(π θ q φ ) as discussed above.

Remark 4. In Scenario 1 above, i.e. for a sufficiently flexible Q, using a gradient-descent algorithm which seeks to minimise the exclusive divergence can sometimes be preferable to a gradient-descent algorithm which seeks to minimise the inclusive divergence. This is because both find (approximately) the same optimum but the latter may exhibit faster convergence in some applications. In such scenarios, the discussion in Subsection A.1 indicates that a smaller number of particles, K, could then be preferable for some of the φ-gradients because (a) the O(K -1 ) self-normalisation bias outweighs the O(K -1/2 ) standard deviation and (b) the direction of this bias may favour faster convergence.

Unfortunately, simply setting K = 1 for the approximation of the φ-gradients 2 is not necessarily optimal because • even in the somewhat idealised scenario 1 above and even if the direction of the self-normalisation bias encourages faster convergence, increasing K is still desirable to reduce the variance of the gradient approximations and furthermore, even in this scenario, seeking to optimise the exclusive KL-divergence could lead to poorly behaved importance-sampling approximations of the θ-gradient whenever φ is still far away from optimal;

• not using the information contained in all K particles and weights (which have already been sampled/calculated to approximate the θ-gradient) seems wasteful;

• if K = 1, the reparametrisation-free AISLE φ-gradients, AISLE-KL-NOREP and AISLE-χ 2 -NOREP are simply vanilla Monte Carlo estimates of 0 and the RWS-DREG φ-gradient is then equal to 0.

B Empirical illustration B.1 Algorithms:

In these supplementary materials, we illustrate the different φ-gradient estimators (recall that all algorithms discussed in this work share the same θ-gradient estimator). Specifically, we compare the following approximations.

• AISLE-KL-NOREP. The gradient for AISLE based on the KL-divergence without any further reparametrisation from ( 14) i.e. this coincides with the standard RWSgradient from (7). This gradient does not require R1 but does not achieve zero variance even if q φ = π θ .

• AISLE-KL. The gradient for AISLE based on the KL-divergence after reparametrising and exploiting the identity from Lemma 1; it is given by ( 15) and coincides with the IWAE-STL-gradient from (4).

• AISLE-χ 2 -NOREP. The gradient for AISLE based on the χ 2 -divergence without any reparametrisation given in ( 16). This gradient again does not require R1 but does not achieve zero variance even if q φ = π θ .

• AISLE-χ 2 . The gradient for AISLE based on the χ 2 -divergence after reparametrising and exploiting the identity from Lemma 1; it is given by ( 17) and is alsow proportional to IWAE-DREG from Tucker et al. ( 2019) which was stated in (5).

When normalising the gradients (as, e.g. implicitly done by optimisers such as ADAM Kingma & Ba, 2015) the proportionality constant cancels out so that both these gradient approximations lead to computationally the same algorithm.

• IWAE. The gradient for IWAE employing the reparametrisation trick from Kingma & Welling (2014). Its sampling approximation is given in (3). Recall that this is the φ-gradient whose signal-to-noise ratio degenerates with K as pointed out in Rainforth et al. (2018) (and which also cannot achieve zero variance even if q φ = π θ ).

• IWAE-DREG. The 'doubly-reparametrised' IWAE gradient from (5) which was proposed in Tucker et al. (2019) . It is proportional to AISLE-χ 2 .

• RWS-DREG. The 'doubly-reparametrised' RWS φ-gradient from (8) which was proposed in Tucker et al. (2019) who derived it by applying the identity from Lemma 1 to the RWS φ-gradient.

B.2 Model:

Generative model. We have N D-dimensional observations x (1) , . . . , x (N ) ∈ R D and N D-dimensional latent variables z (1) , . . . , z (N ) ∈ R D . Unless otherwise stated, any vector y ∈ R D is to be viewed as a D × 1 column vector.

Hereafter, wherever necessary, we add an additional subscript to make the dependence on the observations explicit. The joint law (the 'generative model'), parametrised by θ, of the observations and latent variables then factorises as

N n=1 p θ (z (n) )p θ (x (n) |z (n) ) = N n=1 γ θ,x (n) (z (n) ).

We model each latent variable-observation pair (z, x) as 

EQUATION

with P := (Σ -1 + I) -1 and ν θ,x := P (Σ -1 µ + x).

In particular, (18) implies that θ ml = 1 N N n=1 x (n) .

Proposal/variational approximation. We take the proposal distributions as a fullyfactored Gaussian: , a d, D ] T ∈ R D×1 denotes the column vector formed by the elements in the dth row of A. Furthermore, for the reparametrisation trick, we take q( ) := N( ; 0, I), where 0 ∈ R D is a vector whose elements are all 0, so that

EQUATION

h φ,x ( ) := Ax + b + C 1/2 , which means that h -1 φ,x (z) = C -1/2 (z -Ax -b).

Note that the mean of the proposal in (20) coincides with the mean of the posterior in (19) if A = P and b = P Σ -1 µ.

This model is similar to the one used as a benchmark in Rainforth et al. (2018, Section 4) and also in Tucker et al. (2019, Section 6.1) who specified both the generative model and the variational approximation to be isotropic Gaussians. Specifically, their setting can be recovered by taking Σ := I and fixing c d = log(2/3)/2 so that C = 2 3 I throughout. Here, in order to investigate a slightly more realistic scenario, we also allow for the components of the latent vectors z to be correlated/dependent under the generative model. However, as the variational approximation remains restricted to being fully factored, it may fail to fully capture the uncertainty about the latent variables.

Gradient calculations. We end this subsection by stating the expressions needed to calculate the gradients in the Gaussian example presented above. Throughout, we use the denominator-layout notation for vector and matrix calculus and sometimes write = 1:D = h -1 φ,x (z) to simplify the notation. Thus,

EQUATION

) (22)

Let a d := [a d,1 , a d,2 , . . . , a d,D ] T ∈ R D×1 denote the column vector formed by the elements in the dth row of A. Then, letting denote elementwise multiplication and using the convention that addition or subtraction of the scalar 1 is to be done elementwise, 

∇ a d log q φ,x (z) = exp(-2c d )(z d -a T d x -b d )x = exp(-c d ) d x ∈ R D , d ∈ {1, . . . , D}, ∇ b log q φ,x (z) = C -1 (z -Ax -b) = C -1/2 ∈ R D , ∇ c log q φ,x (z) = C -1/2 (z -Ax -b) C -1/2 (z -Ax -b) -1 = -1 ∈ R D ,

∇ a d [log • w ψ ,x • h φ,x ]| ψ =ψ ( ) = [∇ z log w ψ,x (z)] d x, ∇ b [log • w ψ ,x • h φ,x ]| ψ =ψ ( ) = ∇ z log w ψ,x (z), ∇ c [log • w ψ ,x • h φ,x ]| ψ =ψ ( ) = C 1/2 ∇ z log w ψ,x (z).

From this, since

∇ φ [log • w ψ,x • h φ,x ]( ) = ∇ φ [log • w ψ ,x • h φ,x ]| ψ =ψ ( ) -∇ φ log q φ,x (z),

we have that

∇ a d [log • w ψ,x • h φ,x ]( ) = ([∇ z log w ψ,x (z)] d -C -1/2 d )x, ∇ b [log • w ψ,x • h φ,x ]( ) = ∇ z log w ψ,x (z) -C -1/2 , ∇ c [log • w ψ,x • h φ,x ]( ) = C 1/2 ∇ z log w ψ,x (z) - + 1.

Impact of the reparametrisation. We end this subsection by briefly illustrating the impact of the reparametrisation trick combined with the identity from Tucker et al. (2019) which was given in Lemma 1. Recall that this approach yields φ-gradients that are expressible as integrals of path-derivative functions ψ,x := ∇ φ [log • w ψ ,x • h φ,x ]| ψ =ψ • h -1 φ,x , Thus, if there exists a value φ such that q φ,x = π θ,x then w ψ,x ∝ π θ,x /q φ,x ≡ 1 is constant so that we obtain zero-variance φ-gradients (see, e.g., Roeder et al., 2017 , for a discussion on this).

For simplicity, assume that Σ = I and recall that we then have q φ ,x = π θ,x if the values (A, b, C) implied by φ are (A , b , C ) = ( 1 2 I, 1 2 µ, 1 2 I).

:

Following Tucker et al. (2019) (based on empirical results in Le et al.

2019), we only use the 'wake-phase' φ-updates for RWS.

Within the IWAE-paradigm, using different numbers of particles for the θ and φ-gradients was recently proposed inRainforth et al. (2018);Le et al. (2018) who termed this approach 'alternating evidence lower bounds', albeit their aim was to circumvent the signal-to-noise ratio breakdown of the IWAE φ-gradient which is distinct from the phenomenon discussed here.

Figures and Tables:

p θ (z) := N(z; µ, Σ), p θ (x|z) := N(x; z; I), where θ := µ = µ 1:D ∈ R D , where Σ := (σ d,d ) (d,d )∈{1,...,D} ∈ R D×D is assumed to be known and where I denotes the D × D-identity matrix. For any θ,

) where A = (a d,d ) (d,d )∈{1,...,D} 2 ∈ R D×D , b = b 1:D ∈ R D and, for c 1:D =: c ∈ R D , C := diag(e 2c1 , . . . , e 2c D ). The parameters to optimise are thus φ := (a T 1 , . . . , a T D , b T , c T ), where a d := [a d,1 , a d,2 , . . .

Furthermore, write h φ,x = [h φ,x,1 , . . . , h φ,x,D ] T , i.e.h φ,x,d ( ) = z d = a T d x + b d + exp(c d ) d ,and let ι (d) = [0, . . . , 0, 1, 0, . . . , 0] T ∈ R D be the vector whose entries are all 0 except for the dth entry which is 1. Then, for d ∈ {1, . . . , D},[∇ a d h φ,x,d ]( ) = 1{d = d }x ∈ R D , d ∈ {1, . . . , D}, [∇ b h φ,x,d ]( ) = ι (d) ∈ R D , [∇ c h φ,x,d ]( ) = exp(c d ) d ι (d) ∈ R D . h -1 φ,x (z) implies that ∇ φ [log • w ψ ,x • h φ,x ]| ψ =ψ ( ) = [∇ φ h φ,x,1 , . . . , ∇ φ h φ,x,D ]( )∇ z log w ψ,x (z),so that, letting [∇ z log w ψ,x (z)] d denote the dth element of the vector ∇ z log w ψ,x (z),

Peer review (iclr2020):
Summary: This paper presents a unifying framework through which much of the recent work on maximum likelihood learning in latent variable models (variational autoencoders) via multisample variational approaches and importance weighted approaches can be understood. It is a relatively clean framework that shows how many of the popular approaches can be described a distinct gradient-based approaches for a single underlying framework with two separate objectives for the generative model weights and the variational posterior weights. Strengths: - The framework is elegant and the derivations are simple. - It clarifies the connection between distinct algorithms and shows the consistency of ones that were previously poorly understood (IWAE-STL). - It has the potential of generating new interesting algorithms. Weaknesses: - It is not clear to me why the first bullet of Remark 1 is so crucial. IWAE can be applied in discrete settings (see, e.g., Mnih & Rezende, 2016) and, as you show, reparameterizations can be applied in the adaptive importance sampling type algorithms to potential improve the variance. - The paper currently struggles with some organizational issues. It reads as if the paper was written as a 15 page paper and split in half to satisfy the length requirements. In particular, I would recommend shortening the introduction, cutting out as much of Sec 2 as possible, and moving experiments into the main draft. Derivations are fine left in the Appendix. - I appreciate that it wasn*t the primary aim of the paper to introduce new algorithms, but I think it could strengthen the contribution to consider at least a few. Are there any other interesting divergences to consider for the proposal distribution objective? - The experiments are quite lacking. In tandem with the above point (consider novel algorithms), the ICLR community might rightfully expect some experiments on large scale models. I appreciate that there might not be much consistency in terms of which methods outperform others, but large scale experiments would at least present evidence of this point. Citations: Mnih & Rezende, 2016. https://arxiv.org/abs/1602.06725