Source-linked AI summary
Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference
Geoffrey Roeder, Yuhuai Wu, David Duvenaud
TL;DR
The paper addresses nonzero variance in reparameterized ELBO gradients even when the variational posterior is exact. It removes the score-function component through a computation-graph modification, yielding an unbiased estimator whose variance approaches zero near the exact posterior and extending the approach to richer bounds and distributions.
Problem
Reparameterized ELBO gradients can retain nonzero variance even when qφ(z|x) exactly matches p(z|x).
Method
The method removes the score-function component from the reparameterized gradient using stop_gradient and generalizes the estimator to mixture and importance-weighted lower bounds.
Results
The estimator is unbiased and has zero variance when the variational approximation is exact, with variance approaching zero as qφ(z|x) approaches p(z|x).
Takeaways & Limitations
Variance reduction gains grow with approximate-posterior complexity, complementing the development of non-Gaussian posterior families.
Takeaways & Limitations
The proposed trick cannot be straightforwardly applied to flow-based posteriors because intermediate gradient information is needed for correct log-Jacobian terms.
Abstract
from arXiv · showhide
We propose a simple and general variant of the standard reparameterized gradient estimator for the variational evidence lower bound. Specifically, we remove a part of the total derivative with respect to the variational parameters that corresponds to the score function. Removing this term produces an unbiased gradient estimator whose variance approaches zero as the approximate posterior approaches the exact posterior. We analyze the behavior of this gradient estimator theoretically and empirically, and generalize it to more complex variational distributions such as mixtures and importance-weighted posteriors.
1 Introduction
Variational inference uses reparameterized gradients to make approximate inference practical, but these gradients still contain a removable score-function component. The proposed estimator removes that component while retaining unbiasedness and lower variance near the exact posterior.
- Reparameterization exposes how sampled latent variables depend on variational parameters, generally yielding lower-variance gradients than REINFORCE.REINFORCE instead relies on the relationship between log qφ(z|x, φ) and its parameters.
- The proposed estimator removes the score-function term from the reparameterized gradient, producing lower variance in many circumstances.The score function is a special case of the REINFORCE estimator.
- As the variational posterior approaches the true posterior, the proposed estimator's variance approaches zero faster, helping optimization stick to the true variational parameters.Figure 1 illustrates fitting a 100-dimensional variational posterior with standard and path derivative gradients.
- The paper presents an unbiased ELBO estimator, implements it through a single automatic-differentiation computation-graph change, and extends it to mixture and importance-weighted bounds.Experiments evaluate variational and importance-weighted autoencoders on MNIST and Omniglot.
2 Estimators of the variational lower bound
The section compares ELBO estimators and identifies a path-derivative estimator that removes the score-function term while preserving unbiasedness and achieving vanishing variance near the exact posterior.
- Estimators of the variational lower bound: The ELBO can be estimated in equivalent forms, including a fully Monte Carlo form and forms that analytically integrate available terms.For multivariate Gaussian p(z) and qφ(z|x), analytic integration is appealing, but the fully Monte Carlo form can have lower variance when the approximation is exact.
- Estimators of the variational lower bound: The fully Monte Carlo ELBO gradient can retain nonzero variance even when qφ(z|x) exactly matches p(z|x).At the exact posterior, the path derivative vanishes, but the score-function component need not vanish for finite samples.
- Estimators of the variational lower bound: The reparameterized gradient decomposes into a path derivative through sampled z and a score-function component acting directly through log qφ.This decomposition exposes the score function as a source of variance in the total derivative estimator.
- Estimators of the variational lower bound: Removing the zero-expectation score-function term yields an unbiased path-derivative estimator whose variance approaches zero as qφ(z|x) approaches p(z|x).The estimator is implemented by changing the automatic-differentiation computation graph, including a stop-gradient operation, and is extended to mixture and importance-weighted bounds.
- Estimators of the variational lower bound: The path-derivative estimator does not have lower variance in every case because the score function can act as a positively correlated control variate.Estimating an optimal scale can introduce additional variance and requires independently estimating the score function, which may be difficult for flexible posteriors.
- Estimators of the variational lower bound: Using a fixed optimal-scale approximation of 1 improved performance on MNIST and Omniglot density-estimation benchmarks and is easy to implement with automatic differentiation.The authors recommend an annealing schedule toward 1 only when independent score-function estimation is feasible and Monte Carlo variance may reduce overall variance away from the optimum.
3 Implementation Details
The implementation replaces the standard reparameterized gradient’s score-function pathway with a stop_gradient-based computation-graph modification, including extensions for mixture and importance-weighted ELBO gradients.
- Standard and path-derivative gradients: Algorithm 1 computes the standard reparameterized ELBO gradient using q_sample, log p(x, z), and log q(z|x, φ) with automatic differentiation.The loss estimate is differentiated at the current variational parameters.
- Standard and path-derivative gradients: Algorithm 2 differs by applying stop_gradient to variational parameters inside the ELBO loss.This is the implementation change that produces the path derivative gradient.
- Extensions and software support: Algorithms 3 and 4 extend the implementation to mixture and IWAE ELBO gradients, while Table 1 lists software-specific stop_gradient functions.The paper discusses implementations in several automatic-differentiation packages.
- Implementation mechanism: The implementation uses a copied parameter variable treated as constant in automatic differentiation while evaluating log qφ at z.This modification is made through a single computation-graph operation.
- Standard and path-derivative gradients: Blocking gradients through variational parameters in log qφ removes the score-function term, leaving updates that rely on the sampled z pathway.The same variational parameters generate z and evaluate its density, so blocking only the density pathway removes the unwanted component.
4 Extensions to Richer Variational Families
The estimator is extended to mixture and importance-weighted variational families, with empirical results reporting marked improvements, while flow-based extensions remain future work because intermediate gradients are unavailable in the listed software.
- Mixture distributions: Mixture variational families introduce a non-differentiable categorical component choice, which the method handles by integrating out the discrete choice from the ELBO and mixture distribution.The resulting estimator is stated to be unbiased.
- Mixture distributions: A mixture of K base distributions combines component densities with mixture weights, and the mixture ELBO marginalizes over component choices while evaluating each component density.The path-derivative estimator is applied to each component density in the inner marginalization.
- Mixture distributions: Algorithm 3 applies stop_gradient to mixture variational parameters, eliminating score-function gradients from the mixture estimator.It otherwise differs from the vanilla gradient estimator only through this operation.
- Importance-weighted bounds: The IWAE extension applies the same trick because the gradient of log importance weights matches the Monte Carlo ELBO gradient.The paper notes that unbiasedness of the modified IWAE estimator is not immediately obvious.
- Importance-weighted bounds: Each importance-weighted gradient sample uses the same zi to evaluate both its normalized weight and partial-derivative term, preventing a direct linearity-of-expectation argument.The paper instead applies variance reduction to each importance-weighted gradient sample.
- Empirical validation: Empirical experiments report markedly improved results and strong improvement in many cases for the proposed estimator, including mixture and multi-sample ELBO settings.These observations support the authors’ conjecture that the estimator is unbiased in those cases.
- Flow distributions: Flow-based extensions are not implemented because intermediate transformed variables contribute path derivatives, while log-Jacobian evaluation requires their gradient information.The paper leaves the necessary software engineering for future work.
5 Related Work
The paper positions its method as a general, automatic-differentiation-based variance reduction technique that extends prior control-variate and gradient-estimation ideas to continuous variational families.
- 5 Related Work: The modification can be interpreted as adding a control variate, connecting the approach to prior score-function variance-reduction work.Earlier work studied score functions as control variates in non-reparameterized inference.
- 5 Related Work: The method generalizes earlier lower-variance unbiased estimators from specialized Gaussian settings to any continuous variational family.This avoids case-by-case gradient analysis and supports implementation in existing automatic-differentiation software.
- 5 Related Work: The technique is a simple drop-in modification to the reparameterized ELBO gradient and can be applied alongside generalized reparameterization gradients.The authors distinguish it from GRG while reporting that it can improve existing algorithms.
6 Experiments
Experiments on MNIST and Omniglot compare the path derivative estimator with vanilla total-derivative training for VAE and IWAE models, generally finding improved NLL results with some exceptions.
- 6 Experiments: The experiments use 28 × 28 binarized MNIST and Omniglot images, Adam optimization, and 20 observations per minibatch.MNIST has 60,000 training and 10,000 test examples; Omniglot has 24,345 training and 8,070 test examples.
- 6 Experiments: The path derivative estimator improves NLL over the original estimator in all but two reported cases.The comparison covers both variational and importance-weighted autoencoders.
- 6 Experiments: On Omniglot, two-layer VAE improvements over prior results are 2.36, 1.44, and 0.6 nats for k={1, 5, 50}.For one-layer VAEs, the corresponding improvements are 0.72, 0.22, and 0.38 nats.
- 6 Experiments: On Omniglot, two-layer IWAE improvements are 2.34, 1.2, and 0.52 nats for k={1, 5, 50}.One-layer IWAE improvements are 0.72, 0.7, and 0.51 nats.
- 6 Experiments: On MNIST, two-layer VAE improvements are 0.56, 0.33, and 0.45 for k={1, 5, 50}, while one-layer results improve by 0.36 and 0.14 for k={1, 5}.The experiments also report comparable-magnitude improvements for IWAE, with isolated underperformance at k=50.
7 Conclusions and Future Work
The paper concludes that removing the score-function component enables further variance reduction beyond reparameterization, with larger gains for more complex approximate posteriors, while noting limits to the current method’s scope.
- 7 Conclusions and Future Work: Further gradient-variance reduction remains possible even when the reparameterization trick applies.The method is expressed as a modification to the computation graph used by automatic differentiation.
- 7 Conclusions and Future Work: The method’s gain grows with approximate-posterior complexity, complementing the development of non-Gaussian posterior families.
- 7 Conclusions and Future Work: The proposed method is specific to variational inference, although the authors suspect analogous unbiased high-variance terms may occur in reinforcement learning and gradient-based MCMC.