Source-linked AI summary
Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives
George Tucker, Dieterich Lawson, Shixiang Gu, Chris J. Maddison
TL;DR
IWAE bounds tighten as sample counts rise, but standard inference-network gradients suffer diminishing signal-to-noise ratio, and an earlier lower-variance estimator lacked an unbiasedness proof. This paper develops DReG estimators through a second application of reparameterization and reports unbiased variance reduction and improved performance across tasks.
Problem
Method
DReG applies a second reparameterization to construct an unbiased, computationally efficient gradient estimator for IWAE and related RWS and JVI updates.
Source: §1 · p. 2Results
Takeaways & Limitations
DReG has the same computational cost as the original estimators, and the authors recommend using it instead of typical gradient estimators.
Source: §7 · p. 8RWS inference-network updates are not gradients of a unified objective, which could potentially lead to instability or divergence.
Source: §4.1 · p. 4
Abstract
from arXiv · showhide
Deep latent variable models have become a popular model choice due to the scalable learning algorithms introduced by (Kingma & Welling, 2013; Rezende et al., 2014). These approaches maximize a variational lower bound on the intractable log likelihood of the observed data. Burda et al. (2015) introduced a multi-sample variational bound, IWAE, that is at least as tight as the standard variational lower bound and becomes increasingly tight as the number of samples increases. Counterintuitively, the typical inference network gradient estimator for the IWAE bound performs poorly as the number of samples increases (Rainforth et al., 2018; Le et al., 2018). Roeder et al. (2017) propose an improved gradient estimator, however, are unable to show it is unbiased. We show that it is in fact biased and that the bias can be estimated efficiently with a second application of the reparameterization trick. The doubly reparameterized gradient (DReG) estimator does not suffer as the number of samples increases, resolving the previously raised issues. The same idea can be used to improve many recently introduced training techniques for latent variable models. In particular, we show that this estimator reduces the variance of the IWAE gradient, the reweighted wake-sleep update (RWS) (Bornschein & Bengio, 2014), and the jackknife variational inference (JVI) gradient (Nowozin, 2018). Finally, we show that this computationally efficient, unbiased drop-in gradient estimator translates to improved performance for all three objectives on several modeling tasks.
1 INTRODUCTION
Deep latent-variable models make direct likelihood training intractable, so they are commonly trained with variational bounds. This section motivates and introduces DReG, an unbiased gradient estimator that addresses IWAE gradient problems as sample counts increase.
1 INTRODUCTION
- Intractable latent-variable marginalization makes direct maximum-likelihood training impractical, motivating tractable lower bounds such as the ELBO.A commonly used factorized variational family can limit model expressiveness and lead to overly simplistic learned models.Source: §1 · p. 1
- IWAE becomes tighter as sample count increases, but its standard inference-network gradient estimator performs poorly because its signal-to-noise ratio diminishes.This motivates seeking alternative gradient estimators.Source: §1 · pp. 1–2
- DReG uses a second reparameterization to construct an unbiased, computationally efficient IWAE gradient estimator, and extends to RWS and JVI updates.It is presented as a drop-in replacement that does not suffer as the number of samples increases.Source: §1 · p. 2
- Across MNIST and Omniglot generative modeling and MNIST structured prediction, DReG estimators reduce variance and improve performance over the original estimators.The evaluation covers IWAE, RWS, and JVI estimators.Source: §1 · p. 2
2 BACKGROUND
The IWAE bound tightens with more samples, but this does not guarantee a better inference-network gradient estimator: its signal-to-noise ratio can deteriorate. The section distinguishes this estimator problem from the behavior of the bound itself and notes that STL is biased when K > 1.
2 BACKGROUND
- In amortized inference, qφ(z|x) is a learnable inference network, and the variational-family expressiveness affects the lower bound’s tightness.Direct marginalization is generally intractable, so the approach maximizes a variational lower bound such as the ELBO.Source: §2 · p. 2
- The standard IWAE inference-network gradient estimator’s SNR goes to zero as K increases, motivating lower-variance estimators.The bound can become tighter even as this estimator’s SNR worsens.Source: §2 · p. 2
- The IWAE bound converges to log pθ(x) regardless of qφ as K increases, so φ’s effect on the bound diminishes, but this alone does not imply deteriorating gradient-estimator SNR.SNR depends on how the estimator's variance decreases relative to the gradient magnitude.Source: §2 · p. 3
- For K > 1, dropping the first term of the IWAE gradient as in STL introduces bias; for K = 1, that term vanishes in expectation.Roeder et al. proposed dropping the term to reduce gradient variance.Source: §2 · p. 3
3 DOUBLY REPARAMETERIZED GRADIENT ESTIMATORS (DREGS)
DReG applies a second reparameterization to estimate a term in the IWAE inference-network gradient, yielding an unbiased estimator. Its behavior differs from STL for K > 1, and its SNR improves with K.
3 DOUBLY REPARAMETERIZED GRADIENT ESTIMATORS (DREGS)
- A second reparameterization estimates the first term of the IWAE gradient efficiently, and the resulting expression can be estimated with a single Monte Carlo sample.The derivation uses an equivalence between REINFORCE and reparameterization gradients; reparameterization typically has lower variance because it uses the derivative of f.Source: §3 · p. 3
- The derivation treats zi as constant in a partial derivative with respect to φ, then changes the expectation’s random variable to z1:K.It focuses on one of the K terms and interprets the normalized importance weight as a REINFORCE reward.Source: §3 · p. 3
- When zi is not reparameterizable, the passage proposes using a control variate.Source: §3 · p. 3
- For K = 1, this term vanishes exactly, recovering the ELBO estimator proposed by Roeder et al. (2017); for K > 1, there is no reason to expect analytical vanishing.Source: §3 · p. 4
- When q(z|x) is optimal, the IWAE-DReG inference-network estimator vanishes exactly and has zero variance; its SNR improves with K, unlike the standard IWAE estimator.The analysis reports O(K) SNR scaling for both generation and inference-network gradients.Source: §3 · p. 4
4 ALTERNATIVE TRAINING ALGORITHMS
The section extends doubly reparameterized gradients to alternative training algorithms, deriving versions for RWS and JVI. It also relates RWS-DReG to IWAE-DReG and STL through a weighted family of estimators.
4.1 REWEIGHTED WAKE SLEEP (RWS)
- The RWS wake update does not suffer diminishing SNR as K increases, but its separate p and q updates are not gradients of a unified objective and could lead to instability or divergence.RWS optimizes a separate inference-network objective rather than a joint objective for p and q.Source: §4.1 · p. 4
- RWS-DReG applies reparameterization to the wake update for the inference network and uses a single-sample Monte Carlo estimate.The wake update approximates the KL divergence from pθ(z|x) to qφ(z|x) using self-normalized importance sampling.Source: §4.1 · p. 4, DOUBLY REPARAMETERIZED REWEIGHTED WAKE… · p. 5
- DReG(α) recovers RWS-DReG at α=1, IWAE-DReG at α=0, and STL at α=0.5.The Roeder et al. (2017) inference-network estimator can be viewed as the sum of the IWAE gradient estimator and the wake update, motivating convex combinations of IWAE-DReG and RWS-DReG.Source: DOUBLY REPARAMETERIZED REWEIGHTED WAKE… · p. 5
4.2 JACKKNIFE VARIATIONAL INFERENCE (JVI)
- The JVI estimator combines K- and K−1-sample IWAE estimators, so DReG can be applied to each term.The authors focus on first-order JVI for simplicity and note that the approach extends to higher-order estimators.Source: DOUBLY REPARAMETERIZED JACKKNIFE… · p. 5, §4.2 · p. 5
5 RELATED WORK
Related work frames IWAE and ELBO as Monte Carlo objectives and discusses problems with multi-sample gradient estimators. It also describes RWS and wake-sleep as alternative training approaches.
5 RELATED WORK
- Monte Carlo objectives (MCOs) use the log of an unbiased marginal-likelihood estimator as a lower bound; ELBO and IWAE correspond to single- and K-sample estimators, respectively.The tightness of an MCO is directly related to the variance of its underlying marginal-likelihood estimator.Source: §5 · p. 5
- The standard IWAE gradient estimator’s SNR scales poorly with many samples, while RWS does not share this issue and can outperform IWAE-trained models.Rainforth et al. (2018) identify the SNR problem, and Le et al. (2018) report empirical degradation across multiple tasks and compare RWS favorably.Source: §5 · p. 5
- Wake-sleep originated as a method for training Helmholtz machines and was later extended to multi-sample and sequential settings.The cited extensions are Bornschein & Bengio (2014) for multi-sample training and Gu et al. (2015) for sequential training.Source: §5 · p. 6
6 EXPERIMENTS
Experiments test DReG on a controlled Gaussian example and on MNIST and Omniglot modeling tasks. Across generative modeling and structured prediction, DReG reduces gradient variance and improves performance, with task-dependent comparisons among estimators.
6.1 TOY GAUSSIAN
- As K increases, IWAE-DReG’s SNR rises while the standard IWAE estimator’s SNR approaches 0; the STL estimator is biased, but IWAE-DReG’s observed bias is not significant.The toy Gaussian evaluation compares estimator SNR, squared bias, and variance across increasing K.Source: §6.1 · p. 6, Figure 1 · p. 6
6 EXPERIMENTS
- Across MNIST generative modeling, Omniglot generative modeling, and MNIST structured prediction, DReG reduces variance and improves performance over the original estimators.The structured-prediction comparison also finds that biased estimators underperform unbiased IWAE gradients, with RWS becoming unstable later in training.Source: §6.2 · p. 7, §6.3 · p. 8
6.2 GENERATIVE MODELING
- On MNIST generative modeling, DReG improves performance for IWAE, RWS, and JVI; the best convex-combination weighting varies by task.For this generative-modeling task, combinations weighted heavily toward RWS-DReG perform best; Omniglot experiments also find reduced variance and improved test performance.Source: §6.2 · p. 7, Figure 2 · p. 7, Figure 3 · p. 8
6.3 STRUCTURED PREDICTION ON MNIST
- On MNIST structured prediction, DReG improves all three updates, while biased STL and RWS underperform unbiased IWAE gradients and RWS becomes unstable later in training.The task models the bottom half of a binarized MNIST digit from its top half.Source: §6.3 · p. 8, Figure 4 · p. 8
7 DISCUSSION
The discussion summarizes DReG estimators for IWAE, RWS, and JVI, reporting unbiased variance reduction and improved performance at the same computational cost. It also identifies task-dependent estimator weighting and sequential extensions as directions for further work.
7 DISCUSSION
- DReG estimators for IWAE, RWS, and JVI provide unbiased variance reduction and improved performance across tasks at the same computational cost as the original estimators.The authors recommend using DReG instead of typical gradient estimators.Source: §7 · p. 8
- DReG estimators for variational sequential Monte Carlo and Neural Adaptive Sequential Monte Carlo remain an open development direction.These approaches extend IWAE and RWS, respectively, to sequential latent-variable models.Source: §7 · pp. 8–9
- The best convex combination of IWAE-DReG and RWS-DReG is task dependent, motivating future work on automatically adapting its weighting to data.The paper reports that heavily weighting RWS-DReG performs best on MNIST generative modeling.Source: §7 · p. 9
- The simple form of IWAE-DReG suggests that a more direct derivation may extend to general Monte Carlo objectives.Source: §7 · p. 9
8 APPENDIX
The appendix derives a gradient identity, analyzes how IWAE-DReG’s inference-gradient signal-to-noise ratio scales with sample count, and gives surrogate objectives for shared parameters. Its figures describe estimator and test-performance comparisons across MNIST and Omniglot tasks.
8 APPENDIX
- Figures 5–8 specify estimator, sample-count, and test-performance comparisons across Omniglot and MNIST; Figure 6 varies IWAE-DReG/RWS-DReG mixtures.Figure 5 covers Omniglot generative modeling, Figure 6 Omniglot structured prediction, and Figures 7–8 MNIST tasks.Source: Figure 5 · p. 11, Figure 6 · p. 11, Figure 7 · p. 12
8.1 EQUIVALENCE BETWEEN REINFORCE GRADIENT AND REPARAMETERIZATION TRICK GRADIENT
- For reparameterizable qφ(z), the appendix derives an identity between the REINFORCE gradient and the reparameterization-trick gradient.The derivation evaluates the gradient using each approach and concludes that the identity follows.Source: §8.1 · pp. 12–13
8.2 INFORMAL ASYMPTOTIC ANALYSIS
- For the unbiased IWAE-DReG estimator, an informal delta-method argument gives inference-network gradient standard deviation scaling as K^-3/2; its SNR improves with K, as does the generation-network SNR.The variance scales as K^-3, yielding standard deviation scaling as K^-3/2.Source: §8.2 · p. 13
8.3 UNIFIED SURROGATE OBJECTIVES FOR ESTIMATORS
- For shared p and q parameters, the appendix gives RWS surrogate objectives and defines modifiers that stop gradients through z_i, θ, or all variables.It also notes that DReG(α = 0.5) does not exactly correspond to STL because of scaling between terms.Source: §8.3 · pp. 13–14