On importance-weighted autoencoders
Anonymous authors Paper under double-blind review
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 (IWAEDREG) 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 uniﬁed framework, most of the previously proposed gradient estimators can be naturally derived.
This permits a better understanding of the assumptions and trade-oﬀs that are at play.
Importantly, the derived gradient estimators are guaranteed to not degenerate as K → ∞ .
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 ) =
R Z p θ ( z, x ) d z 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. ﬁnd 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. ﬁnd 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 modiﬁed 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 ﬁrm theoretical footing.
Yet, IWAE suﬀers 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) modiﬁcations of the IWAE φ -gradient which avoid this breakdown (i.e. IWAE-STL and IWAE-DREG) can be justiﬁed in a more principled manner by taking an RWS-type adaptive importance-sampling view.
This conclusion was already reached by Le et al.
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 → ∞ .
Speciﬁcally, 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.
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 justiﬁed 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 modiﬁed 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.
Le et al.
for an extensive empirical comparisons of all the algorithms discussed in this work.
Notation.
We repeatedly employ the shorthand p ( f ) :
=
R Z f ( z ) p ( z ) d z 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 ) :=
Q Kk =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
[Equation]
where γ θ ( z ) :
=
p θ ( z, x ) and where Z θ :=
p θ ( x ) =
R Z γ θ ( z ) d z .
Background 2. 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 ˆ π θ h φ, z i ( f ) :=
K X k = s w kψ f ( z k ) with s w kψ :=
w ψ ( z k )
[Equation]
The notation h φ, z i stresses the dependence of this estimator on φ and z .
The quantity w ψ ( z k ) are called the k th 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 b Z θ h φ, z i :=
[ w ψ ( z 1 ) + . . .
+ w ψ ( z K )] /K is an unbiased (‘importance-sampling’) estimator of
Z θ .
Remark 2.
The self-normalised estimate ˆ π θ h φ, z i ( 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 ˆ π θ h φ, z i ( f ) → π θ ( f ) almost surely as K → ∞ .
2. Importance weighted autoencoder (IWAE)
Objective.
The importance weighted autoencoder (IWAE) , introduced by Burda et al.
(2016), seeks to ﬁnd 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:
[Equation] [Equation]
where the expectation is w.r.t.
z ∼ q ⊗ K φ .
For any ﬁnite 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
Domke & the IWAE also constitutes another VAE on an extended space based on an auxiliary-variable construction developed
Andrieu &
Standard reparametrisation gradient.
The gradient of the IWAE objective from (1): ∇ ψ L Kψ = E (cid:2) ∇ ψ log b
Z θ h φ, z i + G ψ ( z ) (cid:3) , with G ψ ( z ) :=
log b
Z θ h φ, z i P Kk =1 ∇ ψ log q φ ( z k ).
The intractable quantity
E (cid:2) G ψ ( z ) (cid:3) 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 (cid:15) on some space E and a family of diﬀerentiable mappings h φ : E → Z such that if (cid:15) ∼ q (cid:15) we have that z
= z ( (cid:15) )
=
h φ ( (cid:15) ) ∼ q φ .
Under R1 , with (cid:15) 1 , . . .
, (cid:15)
K iid ∼ q (cid:15) and z k :=
z ( (cid:15) k ) :
= h φ ( (cid:15) k ), the gradient can be expressed as
[Equation]
with (cid:72) ψ ( z ) :
= ∇ φ [log ◦ [Equation] ψ 0 ◦ [Equation]
|
[Equation]
ψ ( h − 1 φ ( z )).
Here, the notation ψ 0 indicates that one does not diﬀerentiate 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.
Under R1 , for suitably integrable f ψ : Z → R , we have
[Equation]
We now exclusively focus on the φ -portion of the IWAE gradient, b ∇ iwae φ h θ, z i
[Equation]
Remark 3 (drawbacks of the IWAE φ -gradient).
The gradient b ∇ iwae φ h θ,
z i 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 E [ b ∇ iwae φ
[Equation]
θ, z i ] /
[Equation]
iwae
[Equation]
Rainforth et al., 2018) This
is because b ∇ iwae φ h θ, z i constitutes a self-normalised importance-sampling approximation
of
π
[Equation]
an identity which directly follows from Lemma 1 with f
[Equation]
• Inability to achieve zero variance.
As pointed out in Roeder et al.
(2017), var[ b ∇ iwae φ h θ, z i ] > 0 even in the ideal scenario where q φ = π θ despite the fact that in this case, w ψ is constant and hence var[log b Z θ h φ, z
i
= 0 .
Two modiﬁcations of b ∇ iwae φ h θ,
z i 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 (cid:72) ψ ≡ 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.
, this introduces an additional bias whenever K >
• IWAE-DREG.
The ‘doubly-reparametrised’ IWAE (IWAE-DREG) gradient pro- posed by Tucker et al.
(2019) removes the score-function terms through Lemma 1,
[Equation]
The quantities b ∇ iwae-dreg φ h θ, z
i and b ∇ iwae φ h φ, z i are equal in expectation.
2. Reweighted wake-sleep (RWS)
The reweighted wake-sleep (RWS) algorithm was proposed in Bornschein & Bengio (2015).
Letting KL( p k q ) :
= R Z log[ p ( z ) /q
( z )] p ( z ) d z is the Kullback–Leibler (KL)-divergence from p to q , the RWS algorithm seeks to optimise ψ = ( θ, φ ) as
[Equation]
The θ - and φ -gradients read
[Equation]
These quantities are usually intractable and therefore approximated by replacing π θ by the self-normalised importance sampling approximation
ˆ π θ h φ, z
i (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 b ∇ rws φ h θ, z i 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 ) ψ (cid:28) s w (2) ψ for k ≥ 3.
The transformed weights {F ( s w kψ ) }
Kk =1 are then mainly supported on the two particles with the largest self-normalised weights.
AISLE: A unified adaptive importance-sampling framework 3. Objective
If θ is ﬁxed, 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 ﬁxed, 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
[Equation]
Here, D ƒ ( p k q ) :=
R Z ƒ( p ( z ) /q
( z ))
q ( z ) d z 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 suﬃces for the remainder of this work; we call the resulting algorithm adaptive importance sampling for learning (AISLE) .
As will become clear below, this uniﬁed framework permits a straightforward and principled derivation of robust φ -gradient estimators that do not degenerate as K → ∞ .
3. θ -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 φ :
[Equation]
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. φ -gradient 3.3. General derivation
In the derivations to follow, integrals of the form π θ ([ F ◦ w ψ ] ∇ φ log q φ ) naturally appear.
These can also be expressed as
[Equation]
[Equation] H ( y ) [Equation] = F
[Equation]
y [Equation] y .
By Lemma 1,
[Equation]
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 ƒ ( π θ k q φ ) =
Z κθ R Z ˜ f
[ w ψ ( z )]
q φ ( z ) d z + C ( θ ) for an exponent κ ∈ R and constant C ( θ ) inde-
pendent 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 0 ( y ) − ˜ f
( y ) /y , simple algebra then directly shows that
[Equation]
[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
b
Z θ h φ, z i 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. Special case ‘inclusive’ KL-divergence: RWS and IWAE-STL
We have KL( π θ k q φ ) =
Z κθ R Z ˜ f
( w ψ ( z ))
q φ ( z ) d z + 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 0 ( y )
= 1.
• AISLE-KL-NOREP/RWS.
Without relying on any reparametrisation, Equation (12) yields the following gradient, which clearly equals b ∇ rws
[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 ) , b ∇ aisle-kl φ
[Equation]
Proposition 1 is notable because it shows that IWAE-STL (which avoids the breakdown highlighted in Rainforth et al.
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 justiﬁed.
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 b ∇ rws-dreg φ h θ, z i given in (8) by ﬁrst replacing the exact (but intractable) φ -gradient by the self-normalised importance-sampling approximation b ∇ rws φ h θ, z i 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 ﬁrst applying Lemma 1 to the exact (RWS) φ -gradient and then approximating the resulting expression.
This can potentially reduce both bias and variance.
3.3. Special case α -divergence: IWAE-DREG
Up to some irrelevant additive constant, the α -divergence between two distributions p and q is given by R Z ( p ( z ) /q
( z ))
α q ( z ) d z for some α > 1.
This can also be expressed as
Z κθ R Z ˜ f
( w ψ ( z ))
q φ ( z ) d z 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 0 ( 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 ( π θ k q φ ) = var z ∼ q φ [ w ψ / Z θ ] is the variance of the importance weights.
• AISLE-
α -NOREP.
Without relying on any reparametrisation, Equation (
) yields
[Equation]
with the following special case which is also proportional to the ‘score gradient’ from
Dieng et al.
(2017, Appendix G
: b ∇ aisle-
χ norep φ
[Equation]
• AISLE-
α .
Using reparametrisation, Equation (12) becomes
[Equation]
again with the special case b ∇ aisle-
χ 2 φ
[Equation]
This demonstrates that IWAE-DREG can be derived (up to the proportionality factor 2 K ) in a principled manner from AISLE, i.e. without the need for a multi-sample objective.
Proposition 2.
For any ( θ, φ, z ) , b ∇ aisle-
χ 2 φ
[Equation]
Note that if the implementation normalises the gradients, e.g. as eﬀectively 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. Special case ‘exclusive’ KL-divergence.
For the ‘exclusive’ KL-divergence, we have KL( q φ k π θ ) =
R ˜ f ( w ψ ( z ))
q φ ( z ) d z + C ( θ ) with ˜ f ( y ) =
log( y ).
In that case, with the notation from Section 3.3.1, we have h 0 ( y ) =
1 /y .
This directly leads to the following approximation,
[Equation]
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( π θ k q φ ).
However, care must be taken because minimising the exclusive divergence does not necessarily lead to well behaved or even well-deﬁned importance weights and thus can negatively aﬀect learning of θ (whose gradient is a self-normalised importance-sampling approximation which makes use of those weights).
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.
To formalise this argument, we have introduced a simple, uniﬁed adaptive-importance-sampling framework termed adaptive importance sampling for learning (AISLE) (which slightly generalises the RWS algorithm) and have proved that AISLE allows us to derive the ‘sticking-the-landing’ IWAE (IWAE-STL) gradient from Roeder et al.
(2017) and the ‘doubly-reparametrised’ IWAE (IWAE-DREG) gradient from Tucker et al.
as special cases.
We hope that this work highlights the potential for further improving variational techniques by drawing upon the vast body of research on (adaptive) importance sampling in the computational statistics literature.
Conversely, the methodological connections established in this work may also serve to emphasise the utility of the reparametrisation trick from
Kingma & Welling (2014); Tucker et al.
to computational statisticians.
In a companion article, we are extending the present work to the variational sequential
Monte Carlo methods from Maddison et al.
; Le et al.
; Naesseth et al.
(2018) and to the tensor Monte Carlo approach from Aitchison (2018).
A On the rôle of the self-normalisation bias within RWS/AISLE A.
The self-normalisation bias
Within the self-normalised importance-sampling approximation, the number of particles, K , interpolates between two extremes:
• As K ↑ ∞ , ˆ π θ h φ, z i ( f ) becomes an increasingly accurate approximation of π θ ( f ).
• For K = 1, however, ˆ π θ h φ, z i ( 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 b ∇ aisle-kl
φ h θ, z i and b ∇ aisle-
χ 2 φ h θ, z i .
• As K ↑ ∞ , these two estimators become increasingly accurate approximations of the ‘inclusive’ -divergence gradients −∇ φ KL( π θ k q φ ) =
π θ ( (cid:72) φ ) and −∇ φ χ 2 ( π θ k q φ ) =
2 π θ ([ w ψ / Z θ ] (cid:72) φ ), respectively.
• For K = 1, however, these two estimators reduce to vanilla Monte Carlo approximations of the ‘exclusive’ -divergence gradients −∇ φ KL( q φ k π θ ) =
q φ ( (cid:72) φ ) and − 2 ∇ φ KL( q φ k π θ ) =
q φ ( (cid:72) φ ), respectively.
This is similar to the standard IWAE φ -gradient which also represents a vanilla Monte Carlo approximation of −∇ φ KL( q φ k π θ ) 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 diﬃcult 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. 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( π θ k q φ ) is small as this implies well-behaved importance weights.
In contrast, a small ‘exclusive’ KL-divergence KL( q φ k π θ ) is not suﬃcient 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.
Suﬃciently expressive Q .
For the moment, assume that the family Q is ﬂexible (‘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( π θ k q φ ) is (at least approximately) equal to arg min φ KL( q φ k π θ ).
Insuﬃciently expressive Q .
In general, the family Q is not ﬂexible 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 ) =
Q Dd =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( π θ k q φ ) as discussed above.
Remark 4.
In Scenario 1 above, i.e. for a suﬃciently ﬂexible 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 ﬁnd (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 RWSDREG φ -gradient is then equal to 0.
B Empirical illustration B. Algorithms
In these supplementary materials, we illustrate the diﬀerent φ -gradient estimators (recall that all algorithms discussed in this work share the same θ -gradient estimator).
Speciﬁcally, 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.
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. 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
[Equation]
We model each latent variable–observation pair ( z, x ) as
[Equation]
where θ := µ = µ 1: D ∈ R D , where Σ :=
( σ d,d 0 ) ( d,d 0 ) ∈{ 1 ,...,D } ∈ R D
× D is assumed to be known and where I denotes the D × D -identity matrix.
For any θ ,
[Equation]
[Equation]
with P :=
( Σ − 1 + I) − 1 and ν θ,x := P ( Σ − 1 µ + x ).
In particular, (18) implies that θ ml
= 1 N P Nn =1 x ( n ) .
Proposal/variational approximation.
We take the proposal distributions as a fullyfactored Gaussian:
[Equation]
where A = ( a d,d 0 ) ( d,d 0 ) ∈{ 1 ,...,D } 2 ∈ R D × D , b = b 1: D ∈ R D and, for c 1: D =: c ∈ R D , C :
= diag(e 2 c 1 , . . .
, e 2 c D ).
The parameters to optimise are thus
[Equation]
where a d := [ a d, 1 , a d, 2 , . . .
, a d,D ] T ∈ R D × 1 denotes the column vector formed by the elements in the d th row of A .
Furthermore, for the reparametrisation trick, we take q ( (cid:15) ) :
=
N( (cid:15) ; 0 , I), where 0 ∈ R D is a vector whose elements are all 0, so that
[Equation]
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 speciﬁed both the generative model and the variational approximation to be isotropic Gaussians.
Speciﬁcally, their setting can be recovered by taking Σ := I and ﬁxing c d = log(2 / 3) / 2 so that C = 23 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 (cid:15)
= (cid:15) 1: D = h − 1 φ,x ( z ) to simplify the notation.
Thus,
[Equation]
Let a d :=
[ a d, 1 , a d, 2 , . . .
[Equation]
a d,D ] T ∈ R D × 1 denote the column vector formed by the elements in the d th row of A .
Then, letting (cid:12) denote elementwise multiplication and using the convention that addition or subtraction of the scalar 1 is to be done elementwise,
[Equation]
Furthermore, write h φ,x =
[ h φ,x, 1 , . . .
, h φ,x,D ] T , i.e.
[Equation]
and let ι ( d )
=
. , 0] T ∈ R D be the vector whose entries are all 0 except for the d th entry which is 1.
Then, for d ∈ { 1 , . .
, D } ,
[Equation]
Again writing (cid:15) =
h − 1 φ,x ( z ) implies that
[Equation]
so that [Equation] letting [ ∇ z log w ψ,x ( z )]
d denote the d th element of the vector ∇ z log w ψ,x ( z ),
[Equation]
From this
[Equation]
since
[Equation]
we have that
[Equation]
Impact of the reparametrisation.
We end this subsection by brieﬂy 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 (cid:72) ψ,x := ∇ φ [log ◦ w ψ 0 ,x ◦ h φ,x ] | ψ 0 =
ψ ◦ 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 ? )
=
( 12 I , 12 µ, 12 I).
By (21) and (22), and with the usual convention (cid:15) = h − 1 φ,x ( z ), we then have [Equation]
[Equation]
Note that the only source of randomness in this expression is the multivariate normal random variable (cid:15) .
Thus, by (23) and (24), for any values of A and b and any K ≥ 1, the variance of the A - and b -gradient portion of AISLE-KL/IWAE-STL and AISLE- χ 2 /IWAE-DREG goes to zero as C → C ?
= 12 I.
In other words, in this model, these ‘score-function free’ φ -gradients achieve (near) zero variance for the parameters governing the proposal mean as soon as the variance-parameters fall within a neighbourhood of their optimal values.
Furthermore, (25) combined with (26) shows that for any K ≥ 1, the variance of the C -gradient portion also goes to zero as ( A, b, C )
→ ( A ? , b ? , C ? ).
A more thorough analysis of the beneﬁts of reparametrisation-trick gradients in Gaussian settings is carried out in Xu et al.
B. Simulations
Setup.
We end this section by empirically comparing the algorithms from Subsection B.1.
We run each of these algorithms for a varying number of particles, K ∈ { 1 , 10 , 100 } , and varying model dimensions, D ∈ { 2 , 5 , 10 } .
Each of these conﬁgurations is repeated independently 100 times.
Each time using a new synthetic data set consisting of N =observations sampled from the generative model after generating a new ‘true’ prior mean vector as µ ∼ N(0 , I).
Since all the algorithms share the same θ -gradient, we focus only on the optimisation of φ and thus simply ﬁx θ := θ ml throughout.
We show results for the following model settings.
• Figure 1.
The generative model is speciﬁed via Σ = I.
In this case, there exists a value φ ?
of φ such that q φ,x ( z ) =
π θ,x ( z ).
Note that this corresponds to Scenario 1 in Subsection A.2.
• Figure 2.
The generative model is speciﬁed via Σ
= (0 .
| d − d 0 | +1 )
( d,d 0 ) ∈{ 1 ,...,D } 2 .
Note that in this case, the fully-factored variational approximation cannot fully mimic the dependence structure of the latent variables under the generative model.
That is, in this case, q φ,x ( z ) 6 =
π θ,x ( z ) for any values of φ .
Note that this corresponds to Scenario 2 in Subsection A.2.
To initialise the gradient-ascent algorithm, we draw each component of the initial values φ 0 of φ IID according to a standard normal distribution.
We use both plain stochastic gradient-ascent with the gradients normalised to have unit L 1 -norm (Figures 1a, 2a) and
ADAM (Kingma & Ba, 2015) with default parameter values (Figures 1b, The total number of iterations is 10 , 000; in each case, the learning-rate parameters at the i th step are i − 1 / 2 .
We also ran the algorithms in each of the above-mentioned scenarios with ﬁxed values of c d , e.g. as in Rainforth et al.
; Tucker et al.
However, we omit the results as this did not signiﬁcantly change the relative performance of the diﬀerent algorithms.
For the same reason, we omit results related to the optimisation of A and C .
[Figure]
[Figure] [Figure]
Summary of results.
Below, we outline what we believe to be the main takeaways from these simulation results for this particular model.
However, further theoretical analysis is required to determine whether these hold in more general scenarios.
The ‘score-function free’ KL-divergence based AISLE algorithms typically performed somewhat better than their χ 2 -divergence based counterparts, i.e. AISLE-KL outperformed AISLE-
χ 2 .
We conjecture that this is due to the fact that the χ 2 -divergence based variants square the (self-normalised) importance weights which increases the variance of the φ -gradients.
The performance of the φ -gradients AISLE-KL-NOREP and AISLE-
χ 2 -NOREP (which do not use any reparametrisation) typically beneﬁted strongly from moderate (relative to the dimension of the latent variables) increases in the number of particles.
In the scenario shown in Figure 2, for larger K , these gradients almost attained the performance of the ‘score-function free’ φ -gradient AISLE-KL/IWAE-STL and outperformed AISLE-
χ 2 /IWAE-DREG.
We conjecture that this is due to the fact that in the scenario shown in Figure 2, the variational family does not include the target distribution, i.e. q φ
= π θ for any φ , and as a result, the main advantage of the ‘score-function free’ gradients – i.e. the fact that they can potentially achieve zero variance – cannot be realised.
The standard IWAE φ -gradient performed worse than the other methods in any of the scenarios considered (except in the trivial case K = 1 in which IWAE reduces to the VAE).
Indeed, as expected, the performance of the standard IWAE φ -gradient consistently worsened with increasing K .
This can be attributed to the issue highlighted in Rainforth et al.
(see Subsection 2.2), i.e. to the fact that the signal-to-noise ratio of this gradient vanishes as O ( K − 1 / 2 ) (as this gradient constitutes a self-normalised importance-sampling approximation of an integral which is equal to zero).
More surprisingly, the ‘score-function free’ φ -gradients AISLE-KL/IWAE-STL, AISLE-
χ 2 /IWAE-DREG did not necessarily improve with increasing K .
Indeed, their performance sometimes became worse as can be seen most clearly in Figure 1.
We note that this cannot be explained by the signal-to-noise ratio decay (which Rainforth et al.
(2018) highlighted for the standard IWAE φ -gradient) because the ‘score-function free’ φ -gradients do not constitute self-normalised importancesampling approximations of integrals which are equal to zero.
Instead, we conjecture that as discussed in Remark 4 in the scenario shown in Figure 1, the O ( K − 1 ) self-normalisation bias of these gradients happens to be beneﬁcial and outweighs the O ( K − 1 / 2 ) standard-deviation decrease obtained from increasing K . 5.
The ‘doubly-reparametrised’ RWS-gradient RWS-DREG from Tucker et al.
and given in (8) performed well for a moderate to large number of particles.
Though it crucially requires K > 1 (note that for K = 1 this gradient is simply a vector of zeros).