Generalized Doubly-Reparameterized Gradient Estimators
1. Introduction
In probabilistic machine learning, we are often interested in optimizing expectations of the form L φ , θ = E q φ ( z )
[ f φ , θ ( z )] w.r.t.
to their parameters, where f φ , θ ( z ) is some objective function, and φ and θ denote the parameters of the sampling distribution q φ ( z ) and other (e.g. model) parameters, respectively.
In the case of widely used variational autoencoders (VAEs, Kingma and Welling (2014), Rezende et al.
(2014)), q φ ( z ) is the variational posterior and θ denotes the model parameters.
In most cases of interest, this expectation is intractable, and we estimate it and its gradients, ∇ φ L and ∇ θ L , using Monte Carlo samples z ∼ q φ ( z ).
In this paper, we address gradient estimation for continuous variables in variational objectives.
A naive implementation of ∇ φ L results in a score function , or REINFORCE, estimator (Williams, 1992), which tends to have high gradient variance; however, if f depends on φ only through z , we can use reparameterization (Kingma and Welling, 2014; Rezende et al., 2014) to obtain an estimator with lower variance by replacing the score function estimator of the gradient with a pathwise estimator .
In variational inference, f φ , θ ( z ) typically depends on φ not only through z but also through the value of the log density log q φ ( z ).
Then, the gradient estimators still involve the score function ∇ φ log q φ ( z ) despite using reparameterization.
Roeder et al.
(2017) propose the sticking the landing ( STL ) estimator, which simply drops these score function terms to reduce variance.
Tucker et al.
(2019) show that STL is biased in general, and introduce the doubly-reparameterized gradient ( DReG s) estimator for IWAE objectives, which again yields unbiased lower-variance gradient estimates.
This is achieved by applying reparameterization for the second time, targeting the score function terms.
The DReG s estimator has, however, two major limitations: 1) it only applies to latent variable models with a single latent layer; 2) it only applies in cases where the score function depends on the same parameters as the sampling distribution
In this work we address both limitations and introduce GDReG s, a generalized doublyreparameterized gradient estimator that applies to general score functions; we also extend the
DReG s estimator to hierarchical models and show that both estimators improve performance on conditional and unconditional image modelling tasks.
2. Background
We are interested in computing gradients of variational objectives L φ , θ =
E z ∼ q φ ( z ) [ f φ , θ ( z )] w.r.t.
the variational parameters φ of the sampling distribution q φ ( z ), and parameters θ of a second distribution p θ ( z ), such as a learnable prior.
Here f φ , θ ( z ) is a general function that can depend on both q φ ( z ) and p θ ( z ) explicitly.
One such objective is the importance weighted autoencoder (IWAE) bound (Burda et al., 2016).
For a VAE with likelihood p λ ( x | z ), (learnable) prior p θ ( z ), and variational posterior (or proposal) q φ ( z | x ), the IWAE objective with K importance weights w k is given by
[Equation]
Gradient estimation.
In practise, we approximate the expectation in L φ , θ by Monte
Carlo sampling, so that our estimates of the expectation and of its gradients become random variables.
We can distinguish between two types of gradient estimators in this setting: (i) score function (SF) estimators and (ii) pathwise estimators .
ScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScoreScore functionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctionsfunctions are gradients of a log probability density w.r.t.
its parameters, for example ∇ φ log q φ ( z ); SF estimators treat the function f φ , θ as a black box and often yield high variance gradients.
In contrast, pathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwisepathwise estimators move the parameter-dependence from the probability density into the argument z of the function f φ , θ ( z ) and diﬀerentiate the computation path to often achieve lower variance gradients by using the knowledge of ∇ z f φ , θ ( z ); see Mohamed et al.
for a recent review.
When computing gradients of the objective L φ , θ , we have to diﬀerentiate both the sampling distribution of the expectation, q φ ( z ), as well as the function f φ , θ ( z ),
[Equation]
) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) ∇ θ f φ , θ ( z ) (cid:3) , (3) and both can give rise to score functions
[Table]
functions (note ∇ φ f φ , θ ( z ) = ∇ log q φ ( z ) f φ , θ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z ) ∇ φ log q φ ( z )).
In the following, we recapitulate how to address the score functions in Eq.
(2) using the reparameterization trick and doubly-reparameterized gradients ( DReG s, Tucker et al.
(2019)), respectively.
In Sec.
3 we introduce GDReG s, a generalization of DReG s, that allows us to eliminate the score function in Eq.
Reparameterization.
We can use the reparameterization trick (Kingma and Welling, 2014; Rezende et al., 2014) to turn the score function, ∇ φ log q φ ( z ), inside the expectation in Eq.
(2) into a pathwise derivative of the function f φ , θ ( z ) as follows: we express the latent variables z ∼ q φ ( z ) through a bijection of new random variables (cid:15) ∼ q ( (cid:15) ), which are independent of φ , z = T q ( (cid:15) ; φ ) ⇔ (cid:15) =
T − 1 q ( z ; φ ).
This allows us to rewrite expectations w.r.t.
q φ ( z ) as E q φ ( z ) [ f φ , θ ( z )] =
E q ( (cid:15) )
[ f φ , θ ( T q ( (cid:15) ; φ ))], which moves the parameter dependence into the argument of f φ , θ ( z ) and gives rise to a pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient pathwise gradient: [Equation] [Table]
[Equation]
[Equation]
Double reparameterization.
Tucker et al.
(2019) reduce gradient variance by replacing the remaining score function in Eq.
(4) with its reparameterized counterpart.
Double reparameterization is based on the identity (Eq.
(5) in Tucker et al.
[Equation]
which follows from the fact that both the score function and the reparameterization estimators are unbiased and thus equal in expectation.
For the IWAE objective Eq.
(1), Tucker et al.
(2019) derived the following doubly-reparametererized gradients ( DReG s) estimator:
[Equation]
While the DReG s estimator reparameterizes the score function in Eq.
(4), the previously proposed STL estimator (Roeder et al., 2017) simply drops it and is usually biased as a result.
Crucially, because DReG s relies on reparameterization, it is limited to score functions of the sampling distribution q φ ( z ), making it inapplicable in the more general setting of arbitrary score functions (Eq.
In the following Sec. 3 we introduce the GDReG s estimator that can be applied to these more general score function terms.
In App.
C we discuss that the seemingly pathwise gradient in Eq.
(4) can actually contain score functions for hierarchical models and explain how to extend DReG s to this case.
Additional score functions arise because the distribution parameters (e.g. the mean and covariance) of one stochastic layer depend on the latent variables of previous layers.
3. Generalized DReG s
Here, we generalize DReG s to score function terms that involve distributions p θ ( z ) diﬀerent from the sampling distribution q φ ( z ), such as E q φ ( z ) [ g φ , θ ( z ) ∇ θ log p θ ( z )].
Such terms appear, for example, when training a VAE with a trainable prior p θ ( z ) with the ELBO or
IWAE objectives.
We cannot use DReG s directly here, as it relies on reparameterization of the sampling distribution q φ ( z ), which means the path would then depend on the parameters φ and double-reparameterization would only apply to its parameters φ , whereas the score function is with w.r.t.
parameters θ of a diﬀerent distribution p θ ( z ).
To make progress we need to make the path depend on the parameters θ while still sampling z ∼ q φ ( z ) during training.
Our solution is to: a temporarily change the path such that it depends on θ ; b perform the reparameterized gradient computation; c change z ( q ) ∼ q φ ( z ) re-express z ( q ) as if sampled from p θ ( z ) (cid:15) z ( q ) ( (cid:15) ; φ ) e (cid:15) ( (cid:15) ; φ , θ ) e (cid:15) z ( p ) ( e (cid:15) ; θ ) T q ( (cid:15) ; φ ) T − 1 p ( z ; θ ) stop_grad T p ( e (cid:15) ; θ ) pathwise gradient through it now depends on θ : ∇ θ T p ( (cid:101) (cid:15) ; θ ) | (cid:101) (cid:15)
[Equation] T [Equation] 1 p ( z , θ ) .
Note that (cid:101) (cid:15) =
T − 1 p ( T q ( (cid:15) ; φ ); θ ) has a diﬀerent, usually more complex, distribution than (cid:15) ∼ q ( (cid:15) ).
the path back so we can use samples z ∼ q φ ( z ) to estimate the expectation.
We change the path by ﬁrst using an importance sampling reweighting to temporarily re-write the expectation, E q φ
[Equation]
) [ ∗ ] = E
[Equation]
z [Equation] ∗ (cid:105) , and then applying reparameterization to the new sampling distribution p θ ( z ): z = T p ( (cid:101) (cid:15) ; θ ) with (cid:101) (cid:15) ∼ q ( (cid:101) (cid:15) ).
We derive the gradient identity in
Eq.
(7) for a general g φ , θ ( z ); we refer to it as the generalized DReG s ( GDReG s) identity
[Equation]
In the derivation we have used the identity x ∇ ∗ log x = ∇ ∗ x repeatedly.
By noting that ∇ TD θ E q φ ( z ) [ g φ , θ ( z )] =
E q φ ( z ) [ ∇ θ g φ , θ ( z )], we can cancel these terms on the left hand side of Eq.
(8) and right hand side of Eq.
By moving − E q φ ( z )
[ g φ , θ ( z ) ∇ θ log p θ ( z )] to the other side we obtain the desired result.
Similar to DReG s (Eq.
, GDReG s allows us to transform score function gradients into pathwise gradients.
But, unlike DReG s, GDReG s applies to general score functions and contains a correction term that vanishes when p θ ( z ) and q φ ( z ) are identical (log q φ ( z ) p θ ( z ) in Eq.
Note that the pathwise derivative ∇ θ T p ( (cid:101) (cid:15) ; θ ) in Eq.
(7) looks like a reparameterization of p θ ( z ) in terms of a noise variable (cid:101) (cid:15) =
T − 1 p ( z ; θ ) with z ∼ q φ ( z ).
We can interpret this sequence of transformations as a normalizing ﬂow (Rezende and Mohamed, 2015)
z
→ (cid:101) (cid:15)
→
z , such that T p ( (cid:101) (cid:15) ; θ )
=
T p ( T − 1 p ( z ; θ ); θ ) =
z .
We can think of this procedure as re-expressing the sample z ∼ q φ ( z ) as if it came from p θ ( z ): Its numerical value z remains unchanged and it is still distributed according to q φ ( z ), yet its pathwise gradient ∇ θ T p ( (cid:101) (cid:15) ; θ ) depends on θ .
We illustrate the corresponding computational ﬂow in Fig.
Note that to derive the GDReG s identity, we only require p θ ( z ) to be reparameterizable ( in Fig. 1).
While q φ ( z ) may be reparameterizable as well ( in Fig. 1), this is not necessary; we only need to be able to evaluate its density in Eq.
We use the GDReG s identity to address general score functions of the form Eq.
(3) and derive the GDReG s estimator for the IWAE objective w.r.t.
the prior parameters θ : (cid:98)
[Equation]
(cid:101) (cid:15) k
=
[Equation]
−
[Equation]
(11) with z 1: K ∼ q φ ( z | x ); full derivation is in App.
D.
The second term in Eq.
(11) looks like the
DReG s estimator for φ in Eq.
(6) except that the samples z k are now re-expressed as if they came from p θ ( z ).
In addition we obtain a term that involves the likelihood p λ ( x | z ) and is linear in (cid:101) w k .
We do not apply GDReG s to the likelihood parameters λ because p λ ( x | z ) is a distribution over x rather than z ; in the following we therefore drop the subscript λ .
Similarly to the DReG s estimator, we can also extend the GDReG s estimator to hierarchical objectives as we explain in App. C. [Figure] P o s t er i o r φ
[Figure] SNR [Figure] We learn all parameters by optimizing the same objective Eq.
(1), but employ diﬀerent gradient estimators for diﬀerent subsets of parameters.
In practice, we implement these estimators using diﬀerent surrogate objectives for the likelihood, proposal, and prior parameters, see App.
F for details.
While separate objectives might seem computationally expensive, most terms are shared between them, and modern frameworks would avoid such duplicate computation.
In practise, we found the runtime increase for training with DReG s and
GDReG s estimators to be less than 10% without any optimization of the implementation.
4. Experiments
Here, we empirically evaluate the proposed GDReG s estimator as well as the proposed hierarchical extensions of DReG s and GDReG s, and compare them to the naive IWAE gradient estimator (labelled as IWAE ) as well as STL (Roeder et al., 2017).
Illustrative example.
We ﬁrst consider an extended version of the illustrative example introduced by Rainforth et al.
(2018) and Tucker et al.
(2019) to show that hierarchical
DReG s and GDReG s increase the gradient signal-to-noise ratio (SNR) and reduce gradient variance compared to the naive IWAE gradient estimator.
We consider a 2-layer linear
VAE with hierarchical prior and variational posterior and ﬁnd that (see Fig.
2): (i) for the parameters φ of the variational posterior, our extended version of the DReG s estimator also resolves the vanishing SNR problem of the naive IWAE estimator (Rainforth et al., 2018) in the hierarchical case by reducing the gradient variance at a faster rate with the number of importance samples; (ii) for the paramters θ of the prior, the GDReG s estimator has smaller gradient variance and better SNR than the naive IWAE estimator but scales at the same rate with the number of importance samples.
Image modelling with VAEs.
We also consider conditional and unconditional image modelling tasks with single layer and hierarchical (multi-layer)
VAEs on several standard benchmark datasets: MNIST (LeCun and Cortes, 2010), Omniglot (Lake et al., 2015), and
FashionMNIST (Xiao et al., 2017).
We use the dynamically binarized versions of the datasets to minimize overﬁtting.
In the hierarchical case, the generative path (prior and likelihood) is top-down whereas the variational posterior is bottom-up; for conditional modelling we predict the bottom half of an image given its top half, as in Tucker et al. ; in this [Figure] M N I
S T ,
[Figure] l ay [Figure] case, both the prior and variational posterior additionally depend on a context variable c (
First, we evaluate the choice of estimator for the parameters of q φ ( z ).
Like Tucker et al.
(2019) for the single layer case, we ﬁnd that our extension of DReG s to hierarchical models leads to a dramatic reduction in gradient variance for the variational posterior parameters φ on all tasks (third column in Fig. 3), which translates to an improved test objective in all cases considered.
DReG s is unbiased and typically outperforms the (biased) STL estimator.
We also observed similar improvements on the training objective.
Second, we consider the estimators for the θ parameters of the prior p θ ( z ).
Using the GDReG s estimator instead of the naive IWAE estimator consistently improves the train and test performance when combined with any estimator for the variational posterior, especially for conditional image modelling with deeper models.
For unconditional image modelling the improvements are marginal, though using GDReG s never hurts.
In terms of gradient variance for the prior parameters θ , GDReG s consistently performs better in the beginning of training, when it always has lower variance.
Later in training this is only consistently true when also using the DReG s estimator for the variational posterior parameters φ .
5. Conclusion
In this paper we generalized the recently proposed doubly-reparameterized gradients ( DReG s,
Tucker et al.
(2019)) estimator for variational objectives in two ways.
First, we showed that for hierarchical models such as VAEs seemingly pathwise gradients can actually contain score functions, and how to consistently and eﬀectively extend DReG s to this case.
Second, we introduced GDReG s, a doubly-reparameterized gradient estimator that applies to general score functions, while DReG s is limited to score functions of the variational distribution.
Finally, we demonstrated that both generalizations can lead to better train and test performance on conditional and unconditional image modelling tasks. Acknowledgments
We thank Chris Maddison for useful comments on the manuscript.
Appendix A. Additional results
[Table]
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
c
[Figure]
nd i t i o n a
[Figure]
Appendix B. Related work
Roeder et al.
(2017) observed that the reparameterization gradient estimator for the ELBO contains a score function term and proposed the STL estimator obtained by dropping this score function to reduce the estimator variance.
They applied the estimator to hierarchical
ELBO models but did not discuss how to treat indirect score functions.
While the STL estimator is unbiased for the ELBO objective, Tucker et al.
(2019) showed that it is biased for more general objectives such as the IWAE bound.
They proposed the DReG s estimator that yields unbiased and low variance gradients for IWAE and resolves the diminishing signal-to-noise issue of the naive IWAE gradients ﬁrst discussed by Rainforth et al.
We extend DReG s to hierarchical models, discuss how to treat the indirect score functions, and generalize it to general score functions by introducing GDReG s.
A number of classic techniques from the variance reduction literature have also been applied to variational inference and reparameterization.
For example, Miller et al.
and Geﬀner and Domke (2020) proposed control variates for reparameterization gradients; while Ruiz et al.
(2016) used importance sampling with a proposal optimized to reduce variance.
Such approaches are orthogonal to methods such as ( G ) DReG s and STL, and can be combined with them for greater variance reduction.
Appendix C. DReG s and GDReG s for hierarchical models
We now consider models with hierarchically structured latent variables and show that even terms that look like pathwise gradients, such as ∇ TD z f φ , θ ( z ) in Eq.
(4), can give rise to score functions in this case.
The score functions appear because the distribution parameters of one stochastic layer depend on the latent variables of another layer .
To illustrate this, consider a hierarchical model with two layers where we ﬁrst sample
z 2 ∼ q φ 2 ( z 2 ) and then z 1 ∼ q φ 1 ( z 1 | z 2 ).
1 Note that the conditioning on z 2 is through the distribution parameters of q φ 1 ( z 1 | z 2 ); to highlight this dependence of z 1 on z 2 , we rewrite q φ 1 ( z 1 | z 2 ) =
q α 1 | 2 ( z 2 , φ 1 ) ( z 1 ), where we explicitly distinguish between the distribution parameters α 1 | 2 , such as the mean and covariance of a Gaussian, and the network parameters φ 1 that parameterize them together with the previously sampled latent z 2 .
A derivative w.r.t.
z 2 that looks like a pathwise gradient actually gives rise to a score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term score function term at the subsequent level (we omit the true pathwise gradients ( . . . )):
[Equation]
(C.1)
[Equation]
Similar additional score functions arise for seemingly pathwise gradients of hierarchical and/or autoregressive priors and variational posteriors.
C.1. Extending DReG s to hierarchical VAEs
Here we show how to extend DReG s to hierarchical VAEs to eﬀectively reduce gradient variance for the variational posterior despite the results in the previous section.
We still consider the IWAE objective (Eq.
but now the latent space z is structured, and both p θ and q φ are hierarchically factorized distributions.
Let us consider a 2-layer VAE zzx and examine the term ∇ TD
φ 2 log q φ 1 , φ 2 ( z 1 , z 2 ) in the total derivative of the IWAE objective as a concrete example.
We have sampled zand z 2 hierarchically using reparameterization: z 2 ( φ 2 ) ≡ T q 2 ( (cid:15) 2 ; α 2 ( φ 2 )) and z 1 ( φ 1 , φ 2 ) ≡
T q 1 ( (cid:15) 1 ; α
[Equation]
|
[Equation]
) , φ 1 )):
[Equation]
[Table]
[Equation]
The total derivative w.r.t.
parameters of the upper layer, φ 2 , gives rise to three types of gradients: (true) pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t. z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 pathwise gradients w.r.t.
z 1 and z 2 , a direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function direct score function because the distribution parameters α 2 ( φ 2 ) directly depend on φ 2 , and an indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function indirect score function because α
| 2 ( z 2 ( φ 2 ) , φ 1 ) indirectly depends on φ 2 through z 2 .
Other terms in the gradient of the objective w.r.t.
φ as well as gradients w.r.t.
the θ parameters decompose similarly.
We have three options to estimate each score function individually: (1) leave it—this naive estimator is unbiased but potentially has high variance; (2) drop it, similar to STL —this estimator is generally biased; (3) doubly-reparameterize it using DReG s—this estimator is unbiased, but can generate further score function terms.
For IWAE objectives we ﬁnd that the indirect score functions come up twice: once when computing pathwise gradients of the initial reparameterization, and a second time (with a diﬀerent prefactor) when computing pathwise gradients for the double-reparameterization of the direct score functions.
The same happens for the (true) pathwise gradients, and it is this double-appearance and the resulting cancellation of prefactors that helps reduce gradient variance for DReG s. Moreover, for most model structures it is impossible to consistently replace all successively arising score functions by doubly-reparameterized gradients.
Thus, to extend DReG s to hierarchical models, we leave the indirect score functions unchanged and only doubly reparameterize the direct score functions.
We provide detailed derivations and a general DReG s estimator for arbitrary hierarchical structures in App.
E, and show how to implement the corresponding surrogate loss functions in App.
F.
Roeder et al.
(2017) apply the STL estimator to hierarchical ELBO objectives but do not discuss indirect score functions.
Their experimental results are consistent with dropping the direct score functions while maintaining the indirect ones, similar to how we extend
DReG s to hierarchical models; the STL estimator is biased for IWAE objectives (Tucker et al., 2019).
C.2. Extending GDReG s to hierarchical VAEs
[Equation]
[Table]
[Equation]
[Table]
[Equation]
2 ( e (cid:15) 1 ; β 1 | 2 ( z 2 , θ 1 ))
T p 2 ( e (cid:15) 2 ; β 2 ( θ 2 ))
Their numerical values and distribution remain unchanged but the gradient ﬂow through them changes.
Note that (cid:101) (cid:15) i follows a diﬀerent, usually more complex, distribution from (cid:15) i .
α i and β i denote the distribution parameters of the variatonal posterior and the prior, respectively.
When extending GDReG s to hierarchical models, we again encounter direct and indirect score functions (see App. C), and we apply GDReG s to the direct score functions but leave the indirect score functions.
See App.
E for derivations of a general GDReG s estimator for arbitrary hierarchical structures and App.
F for surrogate losses to implement them.
To apply GDReG s we need to re-express samples from q φ ( z ) as if they came from p θ ( z ).
We do this for the entire hierarchy jointly.
In Fig.
C.
we illustrate the necessary computational ﬂow for the example of a 2-layer VAE with the variational posterior factorized in the opposite direction from the generative process; see App.
E for the general case.
We draw samples z 1 , z 2 ∼ q φ ( z 1 , z 2 ) =
q φ 1 ( z 1 ) q
φ
( z 2 | z 1 ) (by transforming independent variables (cid:15) i ) and then re-express them as if they were sampled from the prior p θ 2 ( z 2 ) p θ 1 ( z 1 | z 2 ), which factorizes in the opposite direction.
While the numerical values of z 1 and z 2 remain unchanged, z 1 is now dependent on z 2 and both depend on the respective θ parameters when computing gradients; we can view ( z 1 , z 2 ) as samples that were obtained by transforming independent variables ( (cid:101) (cid:15) 1 , (cid:101) (cid:15) 2 ) that follow a more complicated distribution than ( (cid:15) 1 , (cid:15) 2 ).
As in the single-layer case, only p θ ( z ) needs to be reparameterizable.
Appendix D. Derivation of the GDReG s estimator for the IWAE objective
In this section we apply the GDReG s identity derived above to derive the GDReG s estimator for the IWAE objective, Eq.
(11) in the main paper.
D.1. Preliminaries on the IWAE objective
The importance weighted autoencoder (IWAE) objective is given by
[Equation]
where w k are the importance weights (Burda et al., 2016).
Due to the structure of the IWAE objective, any gradient w.r.t.
any of its parameters can be written as
[Equation]
using the chain rule and ∇ ∗ w k = w k ∇ ∗ log w k .
(cid:101) w k are the normalized importance weights, and we have reparameterized z k as T q ( (cid:15) k ; φ ).
Typically, the derivatives we are interested in are w.r.t.
the parameters φ and θ .
We also note the following identity that we use in the derivation of the doubly reparameterized estimators,
[Equation]
(D.3) which follows from applying the chain-rule and using ∇ ∗ w k = w k ∇ ∗ log w k .
Tucker et al.
(2019) derive the DReG s identity (Eq.
and use it to derive the following doubly-reparameterized gradient estimator ( DReG s) w.r.t.
the approximate posterior parameters φ as:
[Equation]
D.2. Derivation of the GDReG s estimator
Similarly, we can derive a generalized doubly-reparameterized gradient ( GDReG s) estimator w.r.t.
the prior parameters θ .
We use the GDReG s identity (Eq.
(7)) derived above with g φ , θ ( z )
= (cid:101) w k and note that the reweighting term log q φ ( z ) p θ ( z ) looks like a log importance weight except for the missing likelihood:
[Equation]
Thus, the GDReG s estimator is given by:
[Equation]
Note that the z k are sampled from q φ ( z k | x ) but re-rexpressed as if they came from p θ ( z ).
We can rewrite the importance weights as
[Equation]
Thus, if the variational posterior q φ ( z k | x ) is equal to the true posterior p θ ( z k | x ), all weights w k become equal to p θ ( x ) and thus constant w.r.t.
z k .
In that case the second term in the
GDReG s estimator Eq.
(11) vanishes and the overall expression simpliﬁes to (cid:98)
[Equation]
In contrast, the usual IWAE gradient involves the score function for p θ ( z k ):
[Equation]
Appendix E. Derivation of the DReG s and GDReG s estimator for IWAE objectives of hierarchical VAEs
In this section we derivations of and further details on the extension of DReG s and GDReG s to hierarchical VAEs with the IWAE objective.
E.1. Preliminaries and notation for the hierarchical IWAE objective
For a hierarchically structured model with L stochastic layers the IWAE objective is still given by Eq.
(D.1) but with importance weights w k given by
[Equation]
Here, z kl denotes the k th importance sample ( k ∈ { 1 , . . .
, K } )
for the l th layer ( l ∈ { 1 , . . .
, L } ).
Both the variational posterior and the prior distribution factorize according to their respective hierarchical structure.
While the prior factorizes top-down in most cases, the variational posterior can have many diﬀerent structures.
In order for the distributions to be valid in the context of a VAE, we require the individual dependency graphs for the prior (generative path) and the variational posterior (inference path) to be directed acyclic graphs.
Cycles would mean that a latent variable conditionally dependent on itself.
To keep the dependency structure general, we write the factorization of the variational posterior and prior as follows:
[Equation]
Here, α l ( · ; φ l ) and β l ( · ; θ l ) are the distribution parameters of the variational posterior and prior distribution in the l th layer, respectively, and we have made the dependencies of the conditional distributions explicit; pa α ( l ) denotes the “parents” of the latent variable z kl according to the dependency graph of the inference path (the factorization of the posterior); similarly, pa β ( l ) denotes the latent variables that z kl directly depends on according to the factorization of the prior p θ .
Typically, the prior is assumed to factorize top-down, such that pa β ( l ) =
z k ( l +1) for all but the top-most layer.
The samples z kl are drawn from the variational posterior and can be expressed through reparameterization as z kl
= T q l ( (cid:15) kl ; α l (pa α ( l ) , φ l )), where (cid:15) kl is an independent noise variable per importance sample and layer.
We note that it is these dependencies of the distribution parameters α l and β l on pa α ( l ) and pa β ( l ), respectively, that give rise to the indirect score functions as discussed in App.
C.
E.2. Derivation of the hierarchical DReG s estimator for IWAE
With notation fully set up we consider the reparameterized gradients of the IWAE objective w.r
[Equation]
t. the variational parameters in a particular stochastic layer φ l :
[Equation]
where we have used the chain-rule to arrive at Eq.
(E.5)
; the ﬁrst term contains both the (true) pathwise gradients as well as the indirect score functions; the second term only contains a direct score function as we only take the partial derivative w.r
[Equation]
t. φ l .
We can rewrite this direct score function gradient because only one term in the (log- )importance weight directly depends on φ l ,
[Equation]
Applying the DReG s identity to this term and using Eq.
(D.3) yields:
[Equation]
which agrees with the ﬁrst term in Eq.
(E.5)
up to the prefactor.
Thus, both the true pathwise gradients as well as the indirect score functions appear twice and the prefactors partly cancel to give rise to the DReG s estimator for hierarchical IWAE objectives:
DReG
[Equation]
estimator for hierarchical IWAE objectives
[Equation]
where z kl = T q l ( (cid:15) kl ; α l (pa α ( l ) , φ l )) , ∀ l ∈ { 1 , . . .
, L } , ∀ k ∈ { 1 , . . .
, K } through reparameterization.
We emphasize that the total derivative w.r.t.
z kl contains pathwise gradients as well as indirect score functions for both the variational posterior as well as for the prior.
The hierarhical DReG s estimator otherwise looks very similar to the DReG s estimator in the single layer case (Tucker et al., 2019).
In App.
F.1 we explain how to implement this estimator eﬀectively and in a structureagnostic way.
That is, we do not have to derive a new estimator for each new dependency graph of the variational posterior or the prior.
E.3. Derivation of the hierarchical GDReG s estimator for IWAE
Next, we derive the expression for the GDReG s estimator for hierarchical VAEs with IWAE objective.
Applying the GDReG s identity entails re-expressing the samples z kl from the variational posterior as if they were sampled from the prior.
Starting form a sample ( z k 1 , . .
. , z kL ) ∼ q φ ( z 1 , . . .
, z L | x ), we use the inverse ﬂow of p θ to obtain new noise variables for each layer, ( (cid:101) (cid:15) k 1 , . . .
, (cid:101) (cid:15) kL ).
We then use the forward ﬂow of p θ to obtain back ( z k 1 , . .
. , z kL ) but with the gradient path now depending on θ as discussed in Sec. 3.
More precisely, we ﬁnd that
[Equation]
original sampling of (
[Equation]
inverse prior ﬂow to obtain new “noise” variables
[Equation]
E
[Equation] [Equation]
forward prior ﬂow to re-express the z kl (E.11) where (cid:15) kl ∼ q ( (cid:15) ) follows a simple distribution that is diﬀerent from the more complicated distribution of (cid:101) (cid:15) kl .
Note how the initial reparameterization of a sample z kl depends on the dependency structure of the variational posterior (through pa α ( · )), while the other transformations depend on the dependency structure of the prior (pa β ( · )).
As for DReG s, we note that only one term in the log importance weight directly depends on the variable θ l ,
[Equation]
With these prerequesits
[Equation]
we can compute the GDReG s estimator for parameters θ l of the l th stochastic layer
[Equation]
GDReG s estimator for hierarchical IWAE objectives
[Equation]
with z 1: K ∼ q φ ( z k | x ), and where we suppressed dependencies on pa ∗ ( l ) where they are not necessary to simplify notation.
The estimator looks very similar to the GDReG s estimator for a single layer IWAE model Eq.
Note that just like above for hierarchical DReG s, the total gradients w.r.t.
z kl give rise to both (true) pathwise gradients as well as indirect score functions through the hierarchical dependencies of the variational posterior and prior.
In App.
F.2 we show how to implement the hierarchical GDReG s estimator Eq.
(E.16) eﬀectively and in a way that is agnostic to the structure of the model.
That is, we do not have to derive a separate estimator for every dependency graph of the variational posterior and prior.
E.4. Double reparameterization and indirect score functions
In principle, we could apply double-reparameterization to the indirect score functions as well.
However, as we explain now, we often cannot doubly-reparameterize all indirect score functions; moreover, even in cases where this is possible, it is still impractical, as the corresponding estimator depends on the exact model structure and would require adaptation to each dependency graph of the prior and variational posterior.
Double reparameterization of indirect score functions works in the same way as for the direct score functions except that g φ , θ ( z ) is given by (cid:101) w 2 k instead of (cid:101) w k in this case.
The derivatives of (cid:101)
[Equation]
k have a similar reproducing property as we observed in Eq.
(D.3):
[Equation]
Thus, double reparameterization of the indirect score functions similarly gives rise to further indirect score functions.
We note that these indirect score functions only appear for the “children” of the current stochastic layer, that is, stochastic variables in those layers that depend on the current layer.
In this context, “children” refers to all children w.r.t.
the dependency structure of both, the variational posterior and the prior.
For a particular layer l we obtain indirect score functions from double reparameterization of all of its (direct or indirect) parent nodes.
Following the dependency structure, we could collect all of these terms and reparameterize them to obtain pathwise gradients only.
However, a problem arises, because we need to account for dependencies of both the variational posterior and the prior.
Reparameterization of a score function gives rise to indirect score functions in all its “children” layers for both the variational posterior and the prior.
For general hierarchical structures, this leads to cycles, in that some of the children of one dependency tree (the variational posterior) are the parents in the other (the prior) and/or vice versa.
In this case we are never able to collect all the terms and fully reparameterize all the score functions.
Moreover, even if the joint dependency graph of the variational posterior and the prior were acyclic, this derivation would be structure-speciﬁc and would need to be repeated for each hierarchical structure.
We therefore do not doubly reparameterize the indirect score functions.
Appendix F. Surrogate losses to implement the DReG s and GDReG s estimators for IWAE objectives
As we discussed in Sec. 3 and similar to Tucker et al.
(2019), we use surrogate loss functions to compute the gradients w.r.t.
the likelihood, proposal, and prior parameters.
That is, we use diﬀerent losses, such that backpropagation results in the respective gradient estimator.
While Tucker et al.
(2019) use a single surrogate loss to compute the gradient estimators for all parts of the objective, we choose to use separate surrogate losses for each of the three parameter groups (likelihood, variational posterior, prior).
In principle, we could combine them into a single loss, but in order to keep presentation simple we keep them separate.
Computationally this does not make a diﬀerence as modern deep learning frameworks avoid duplicate computation.
For the likelihood parameters, we use the regular (negative) IWAE objective Eq.
(1) as a loss.
That is, the gradient estimator for the likelihood parameters is given by the gradient of the negative IWAE objective.
To construct the other surrogate losses we need to stop the gradients at various points in the computation graph.
In the following, we use the shorthand notation (cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) to indicate that we stop gradients into the underlined parts of an expression.
Where it might be ambiguous, or to highlight where we do not stop gradients, we use the shorthand (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52) to indicate that gradients ﬂow.
For example, f ( (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) φφ φφφφφφφφφφφφφ φ φ , (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) θθ θθθθθθθθθθθθθ θ θ ) means that we backpropagate gradients into φ but not into θ .
F.1. DReG s for variational posterior parameters φ F.1.1. Single stochastic layer
Here we reproduce part of the surrogate loss for the variational parameters φ by Tucker et al.
for the single stochastic layer case:
[Equation]
(cid:88)
[Equation]
(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:16) log p λ ( x | (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Table]
[Equation]
(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Equation]
That is, we sample z k ∼ q φ ( z k | x ) as usual (by reparameterizing independent noise variables (cid:15) k ) but stop the gradients of the parameters that parameterize the distributions when evaluating their densities, log q α ( (cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) φφφφφφφφφφφφφφφφφ )
( (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
z
k z k z k z k z k z k z k z k z k z k z k z k z k z k z k z k z k ).
In addition we stop the gradients around the normalized importance weights (cid:101) w k .
Diﬀerentiating L DReG s w.r.t.
the proposal parameters φ yields the DReG s estimator Eq.
Note that we do not explicitly stop gradients into λ or θ because we use separate surrogate losses for those parameter groups.
If we were to use a combined loss, we would potentially have to stop gradients into these parameters as well, depending on the estimator used.
To practically implement this surrogate loss, we use two copies of the variational posterior distribution.
An unaltered one (no stopped gradients) to sample z and one with gradients into the proposal parameters stopped to evaluate the log densities.
The stopped gradient makes sure that we do not obtain a direct score function as we have doubly-reparameterized it.
Note that for single-stochastic-layer models we could also stop the gradients of the distribution parameters α instead as they only depend on φ .
We emphasize that this is not possible for hierarchical models as this would eliminate the indirect score functions and thus produce potentially biased gradients.
F.1.2. Multiple stochastic layers
For multiple layers, the surrogate loss for the DReG s estimator Eq.
(E.8) is given by:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Again
[Table]
we do not explicitly stop gradients into λ or θ l as we only take gradients w.r
[Table]
t. φ l .
The indirect score functions arise due to the indirect dependence of the distribution parameters α l (
[Table]
(cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52)
[Table]
(cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) pa α ( l ) and (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) (cid:52) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ), respectively.
Note how the former depends on the hierarchical structure of the variational posterior, whereas the latter depends on the hierarchical structure of the prior.
To implement this surrogate loss eﬀectively, we again use two copies of the variational posterior distribution.
One un-altered one (without stopped gradiends) from which we sample the individual reparameterized z kl and through which gradients can ﬂow; we use these samples to evaluate densities at and to parameterize the distribution parameters at subsequent layers.
Derivatives w.r.t.
φ l will then give rise to pathwise gradients and indirect score functions.
We use the second copy of the variational posterior, where we have stopped the parameters φ l , to evaluate the density at for the log importance weights in the last summand of Eq.
(F.2).
F.2. GDReG s for prior parameter θ F.2.1. Single stochastic layer
[Equation]
[Table]
[Equation]
[Table]
(cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:16) log p λ ( x | (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
z k z k z k z k z k z k z k z k z k z k z k z k z k z k z k z k z k )
+ log p β ( (cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) θθθθθθθθθθθθθθθθθ ) ( (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
z
k z k z k z k z k
[Table]
k (cid:101)
[Table]
k (cid:101)
[Table]
k (cid:101)
[Table]
k (cid:101)
[Table]
k (cid:101)
[Table]
k (cid:101)
[Table]
θθθθθθθθθθθθθθθθθ )
[Table]
(cid:101) (cid:15)
[Table]
(cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101) (cid:15) k (cid:101)
[Table]
k (cid:101)
[Table]
k = (cid:56)
[Table]
(cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) (cid:56) T − 1 p
( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) T − 1 p ( T q ( (cid:15) k ; φ ) ; θ ) (cid:15) k ∼ q ( (cid:15) k ) (F.3)
Taking the derivative of Eq.
(F.3) w.r.t.
θ gives rise to the GDReG s estimator for the single stochastic layer IWAE objective.
As explained in Sec. 3, we need to re-express z k such that its path depends on θ .
In eﬀect, we ﬁrst sample z k = T q ( (cid:15) k ; φ ), then compute the new noise variable (cid:101) (cid:15) k
= T − 1 p ( z ; θ ), and re-compute z k = T p ( (cid:101) (cid:15) k ; θ ).
Note that we have to stop gradients into the noise variables (cid:101) (cid:15)
k to obtain the correct gradient estimator.
This explains the stop_grad in Fig.
As above, we do not explicitly stop gradients into λ and φ as we use separate losses for these parameter groups and only compute gradients of Eq.
(F.3)
w.r.t.
θ .
To eﬀectively implement this loss, we use two copies of the prior distribution.
One that we implement as a normalizing ﬂow and a second one with stopped gradients into the parameters.
We then proceed as follows:
• Compute the new noise variables (cid:101) (cid:15) k by using the inverse ﬂow T − 1 p on the samples z k from the variational posterior.
• Stop the gradients into (cid:101) (cid:15) k .
• Use the forward ﬂow T p ( (cid:101) (cid:15) k ; θ ) to re-compute z k but with path dependent on θ .
These samples when derived w.r.t.
θ will give rise to the pathwise gradients.
• Use the second copy of the prior (with stopped gradients into its parameters) to evaluate the log density at the samples z k .
The stopped gradients make sure that we do not obtain the direct score function.
F.2.2. Multiple stochastic layers
For multiple stochastic layers the surrogate loss that gives rise to the GDReG s estimator
Eq.
(E.16) is given by:
[Equation]
(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k (cid:101) w k log p λ
x | (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Table]
(cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k (cid:101) w 2 k log w k
[Table]
[Table]
(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Table]
(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52) .
[Table]
[Table] [Table]
, (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Table]
(cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl ; β l ( (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) pa β ( l ) , (cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)(cid:52)
[Table]
(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl (cid:101) (cid:15) kl (cid:101)
(cid:15) kl = (cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56)(cid:56) T
[Table]
(cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17)
T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17) T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17)
T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17) T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17)
T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17) T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17)
T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17) T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17)
T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l (cid:1)(cid:17) T − 1 p l (cid:16) z ( q ) kl ; β l (cid:0) pa β ( l ) , θ l
[Equation]
As for the single layer case, we need to re-express variational posterior samples z kl as if they were sampled from the prior.
To obtain the correct gradients, we again have to stop gradients into the new noise variables (cid:101) (cid:15) kl , also see Fig.
C.1.
As for hierarchical DReG s, the indirect score functions stem from the second and third term of log w k and arise because the distribution parameters α l and β l depend on the “parent” stochastic layers.
As before we use two copies of the prior distribution, one with regular gradients that is set up as a ﬂow, and a second with stopped gradients into the parameters.
This allows us to implement the GDReG s estimator regardless of the model structure.