PaperlayerView exact revision on arXiv

Source-linked AI summary

Implicit Reparameterization Gradients

Michael Figurnov, Shakir Mohamed, Andriy Mnih

arXiv:1805.08498v4cs.LGstat.ML

TL;DR

Classic reparameterization gradients are low variance but unavailable for important continuous distributions because their required parameterizations or inverse CDFs are unavailable or impractical. The paper uses implicit differentiation of numerically tractable CDFs instead, and reports broader applicability with faster, more accurate estimators. These estimators also support LDA and VAE experiments using non-Normal latent variables.

  • Problem

    Classic reparameterization is not applicable to important continuous distributions such as Gamma, Beta, Dirichlet, and von Mises, while alternatives often have high-variance gradients.

  • Method

    The paper computes reparameterization gradients by implicitly differentiating numerically tractable CDFs rather than inverting standardization functions.

  • Results

    The proposed estimators are faster and more accurate than alternative approaches and perform as well or better than existing methods in the reported LDA and VAE applications.

  • Takeaways & Limitations

    Implicit reparameterization gradients allow stochastic computation with a broader set of continuous distributions, including Gamma, Beta, Dirichlet, and von Mises.

  • Takeaways & Limitations

    The approach requires a numerically tractable CDF or another standardization function whose required gradients can be computed.

Abstract

from arXiv · show

By providing a simple and efficient way of computing low-variance gradients of continuous random variables, the reparameterization trick has become the technique of choice for training a variety of latent variable models. However, it is not applicable to a number of important continuous distributions. We introduce an alternative approach to computing reparameterization gradients based on implicit differentiation and demonstrate its broader applicability by applying it to Gamma, Beta, Dirichlet, and von Mises distributions, which cannot be used with the classic reparameterization trick. Our experiments show that the proposed approach is faster and more accurate than the existing gradient estimators for these distributions.

1 Introduction

Classic reparameterization gradients enable backpropagation through some continuous random variables but exclude important distributions. The paper introduces implicit reparameterization gradients to broaden applicability and reports faster, more accurate estimation.

  • Motivation: The reparameterization trick replaces a distribution with a differentiable transformation of a fixed base distribution, enabling backpropagation through stochastic nodes.It supports models including VAEs, automatic variational inference, Bayesian neural networks, and deep-network regularization.
  • Limitations: Location-scale forms, tractable inverse CDFs, or deterministic transformations make classic reparameterization easy to use, but exclude Gamma, Beta, Dirichlet, von Mises, and other distributions.The paper targets these cases with a general reparameterization tool.
  • Related approaches: Score-function estimators handle more distributions but typically produce high-variance gradients requiring problem-specific variance reduction.Generalized reparameterizations combine estimator classes, while another approach approximates the inverse-CDF derivative.
  • Contribution: The proposed method differentiates the CDF implicitly rather than differentiating its inverse, using automatic differentiation for numerically tractable CDFs.This extends the approach to distributions such as Gamma and von Mises.
  • Contribution: The paper reports unbiased estimators for numerically tractable CDFs, faster and more accurate gradients, improved LDA training, and alternative VAE latent-space topologies.The demonstrations cover Gamma, Beta, Dirichlet, and von Mises distributions.

2 Background

Explicit reparameterization rewrites expectations using parameter-free noise so gradients can pass through stochastic variables. Its practicality is limited when standardization requires difficult or expensive inverse-CDF operations, especially in variational inference.

  • Explicit reparameterization: Explicit reparameterization seeks an invertible, continuously differentiable standardization function that removes distribution-parameter dependence from samples.The resulting expectation is expressed with respect to parameter-free noise.
  • Explicit reparameterization: For Gaussian variables, standardization maps samples to standard Normal noise, allowing the objective to be rewritten as an expectation over that noise.The parameter dependence is transferred into the function being differentiated.
  • Explicit reparameterization: The gradient of the expectation can then be computed as the expectation of gradients.
  • Limitations: Although standardization functions exist for many continuous distributions, inverting a CDF is often complicated and expensive, and differentiating the inverse is harder.This is a practical limitation of explicit reparameterization.
  • Stochastic variational inference: Stochastic variational inference optimizes model and variational-posterior parameters using sampled estimates for otherwise intractable expectation terms.The KL term and its gradients are often analytic, while the remaining term is typically approximated from posterior samples.
  • Stochastic variational inference: Score-function estimators are broadly applicable but relatively high variance, whereas reparameterization is easier and lower variance when applicable.Many important continuous distributions remain outside the classic method’s scope.

3 Implicit reparameterization gradients

Implicit reparameterization avoids inverting the standardization function by obtaining parameter gradients through implicit differentiation. CDF-based standardization makes this approach applicable to univariate and multivariate distributions with numerically tractable CDFs.

  • Implicit reparameterization gradients: The proposed gradient method avoids inversion of the standardization function by changing variables in the explicit-gradient formulation.
  • Implicit differentiation: The key insight is to compute the sample’s parameter gradient by implicitly differentiating the equality between standardization and parameter-independent noise.
  • Implicit differentiation: Solving the differentiated standardization equation yields a gradient expression requiring differentiation of the standardization function rather than its inverse.Its value is unchanged by invertible transformations of the standardization function.
  • Example: For a Normal distribution, explicit and implicit reparameterizations produce equivalent gradients, while the implicit version avoids inverting the standardization function.
  • CDF standardization: For univariate distributions, the CDF standardizes samples to Uniform(0, 1), so implicit gradients require only differentiating the CDF.The CDF must be strictly monotonic and continuously differentiable in the sample and parameters.
  • Multivariate extension: In multivariate settings, the distributional transform uses conditional CDFs, producing a diagonal or triangular linear system solvable in O(D) or O(D^2).The diagonal case applies to factorized distributions; otherwise the system is triangular.
  • Implementation: Samples can be obtained with suitable methods such as rejection sampling, while standardization-function gradients can be computed analytically or with automatic differentiation.

4 Applications of implicit reparameterization gradients

Implicit reparameterization is applied by differentiating numerically tractable standardization functions, extending reparameterization to truncated, mixture, Gamma, Beta, Dirichlet, and von Mises distributions. Automatic differentiation provides the strongest reported combination of speed and accuracy among the compared Gamma estimators.

  • General strategy: The method constructs standardization functions, such as CDFs or multivariate distributional transforms, and differentiates them implicitly.Required gradients can be obtained analytically or through automatic differentiation.
  • Applications: Truncated and mixture distributions use CDF-based or distributional-transform standardization, with gradients obtained through automatic differentiation.For mixtures, the multivariate transform uses posterior component weights conditioned on preceding dimensions.
  • General strategy: Implicit reparameterization avoids inverting the standardization function while producing the same result as explicit reparameterization.This makes the approach easier to implement for distributions such as Gamma.
  • Applications: Gamma shape gradients use forward-mode automatic differentiation through a numerical regularized incomplete Gamma-function computation.Gamma rate reparameterization uses the scaling property z/β ∼ Gamma(α, β).
  • Applications: Beta and Dirichlet samples can be constructed from Gamma samples, extending the Gamma-based approach to both distributions.A normalized pair of Gamma samples yields a Beta sample.
  • Applications: Von Mises concentration gradients use forward-mode automatic differentiation through an efficient numerical CDF method.The location parameter is standardized by shifting samples from a zero-location von Mises distribution.
  • Accuracy and speed: 4.5× slower and 3× less accurate than automatic differentiation, the Jankowiak–Obermeyer method underperforms the proposed automatic-differentiation approach.Table 2 reports average error and time per element; automatic differentiation achieves the lowest error and highest speed.

5 Related work

Related approaches broaden reparameterization through surrogate distributions, analytic marginalization, generalized reparameterizations, transport-equation views, and rejection-sampling estimators. These alternatives trade distributional fidelity, computation, or estimator construction against broader applicability.

  • Surrogate distributions: Surrogate distributions replace difficult posteriors with reparameterizable alternatives, but may lose properties such as the sparsity of Gamma and Dirichlet distributions.Examples include replacing Beta with Kumaraswamy, Gamma with Weibull, and Dirichlet with Logistic Normal.
  • Integrating out nuisance variables: Analytically integrating nuisance mixture-component variables causes a K-fold computation increase and can become exponentially expensive across mixture chains.The approach trades computation for simpler reparameterization.
  • Implicit reparameterization gradients: Earlier implicit-gradient work treated explicit and implicit pathwise gradients mainly for univariate distributions and simple graphs without backpropagation.The present line of work extends that setting within machine learning computation graphs.
  • Differential-equation view: The transport-equation view characterizes reparameterization gradients as differential-equation solutions, with multiple solutions for non-factorial multivariate distributions.Choosing among those solutions can reduce estimator variance.
  • Generalized reparameterizations: Generalized reparameterization gradients allow the base distribution to depend weakly on parameters, expanding applicability beyond standard reparameterization.The dependence may occur through higher moments.
  • Rejection-sampling variational inference: RSVI combines proposal reparameterization gradients with a score-function term accounting for the rejection step.For Gamma, RSVI gradients can have lower variance than gradients computed using generalized reparameterization.

6 Experiments

The experiments apply implicit reparameterization to Gamma and von Mises distributions across cross-entropy estimation, LDA, and VAE training. The method generally provides lower-variance or faster optimization than alternatives while retaining distribution-specific properties such as sparsity and periodic latent structure.

  • Experimental setup: Implicit reparameterization is evaluated for Gamma and von Mises distributions in cross-entropy estimation, LDA topic modeling, and VAE training.The experiments use RSVI as the main baseline and include datasets such as 20 Newsgroups, RCV1, and dynamically binarized MNIST.
  • Gradient of the cross-entropy: Implicit gradients are faster and have lower variance than RSVI in the cross-entropy optimization problem.For Dirichlet distributions, RSVI can approach the implicit-gradient variance by increasing shape augmentation B, but this adds computational cost and requires tuning.
  • Latent Dirichlet Allocation: In LDA, implicit gradients perform as well or better than LN-LDA, while learned Dirichlet prior weights can become sparse and identify fewer active topics.Lower gradient variance also leads to faster training-objective convergence, and amortized inference can achieve better perplexity than SVI.
  • Variational Autoencoders: For VAE training on MNIST, a uniform prior and cyclic von Mises latent space are advantageous for low-dimensional latent spaces.Factorized von Mises also outperforms multivariate von Mises-Fisher under a uniform prior in low dimensions, while bell-shaped priors give results similar to Normal priors and posteriors.
  • Variational Autoencoders: The learned two-dimensional latent spaces differ in topology: Normal latents concentrate toward the center, whereas Beta and von Mises latents tile the available space.These differences reflect the distinct support and geometry of the alternative latent distributions.

7 Conclusion

The paper extends reparameterization gradients to additional distributions through a practical implicit-differentiation approach, achieving faster estimation with lower gradient variance. It broadens modeling choices while requiring a numerically tractable CDF or another standardization function.

  • 7 Conclusion: Implicit reparameterization gradients extend practical reparameterization to truncated, Gamma, and von Mises distributions.The method offers a route beyond computationally convenient model choices for distributions previously excluded from classic reparameterization.
  • 7 Conclusion: The proposed estimators are faster than existing methods while providing lower gradient variance.The conclusion states both properties as simultaneous characteristics of the new estimators.
  • 7 Conclusion: Applying the method requires a numerically tractable CDF or another standardization function.The paper suggests approximate standardization with a score-function correction as a possible future direction when such a function is unavailable.

A Testing of implicit gradient implementation

The implementation is tested by selecting functions whose expected-value gradients are known, then checking whether Monte Carlo averages recover those quantities for Gamma and von Mises distributions.

  • A Testing of implicit gradient implementation: A correctness check compares Monte Carlo averages of stochastic gradients with the correct gradient quantity for a suitable function f(z).The test uses an appropriate function and verifies that averaging approximates the analytically correct result.
  • A Testing of implicit gradient implementation: For Gamma(α, 1), choosing f(z) = z gives a target gradient of 1.The average stochastic gradient should therefore equal 1.
  • A Testing of implicit gradient implementation: For vonMises(0, κ), the test uses f(z) = cos z.This function is identified as appropriate for checking the corresponding implicit gradient implementation.

B Implementation details for the reparameterization gradient

The implementation differentiates numerically tractable CDF computations with automatic differentiation, using distribution-specific numerical methods for Gamma and von Mises and finite differences for comparison.

  • B Implementation details for the reparameterization gradient: The implementation uses an equivalent, more numerically stable expression for Eqn. (8).This is the stated implementation choice for the reparameterization gradient.
  • B Implementation details for the reparameterization gradient: Gamma CDF derivatives are computed by forward-mode differentiation through differentiable continued-fraction or series expansions.The method selects continued fractions for some argument ranges and series expansions for others.
  • B Implementation details for the reparameterization gradient: 30% faster and an order of magnitude more accurate results are obtained by multiplying the Gamma derivative by exp(−log Gamma(z|α, 1)).The implementation stops when the derivative changes by less than a small value, with maximum iterations of 200 for float32 and 500 for float64.
  • B Implementation details for the reparameterization gradient: For von Mises, the CDF uses a truncated series for κ < 50 and a Normal approximation for larger κ, with forward-mode differentiation through the computation.The implementation uses SciPy's numerical routines.
  • B Implementation details for the reparameterization gradient: The Gamma CDF derivative is evaluated for comparison using a hypergeometric 2F2 expression at arbitrary precision.The comparison implementation uses mpmath with default settings resulting in float64 precision.
  • B Implementation details for the reparameterization gradient: The von Mises comparison derivative differentiates its series and truncates it at the 100th term using SciPy Bessel functions.This provides the numerical expression used for comparison.
  • B Implementation details for the reparameterization gradient: The experiment samples 1000 random variables at each parameter-grid value and reports grid-searched finite-difference step sizes δ.Gamma uses α ∈ {1 × 10−2, 1 × 10−1, 1, 1 × 10^1, 1 × 10^2, 1 × 10^3}; von Mises uses κ ∈ {1 × 10−2, 1 × 10−1, 1, 1 × 10^1}.

D Experimental details

The experiments compare implicit gradients with RSVI and evaluate them in gradient estimation, VAE, and LDA settings using specified sampling, variance, optimization, and dataset procedures.

  • D Experimental details: RSVI uses explicitly reparameterizable proposals: standard Normal for Gamma(α, 1) and Uniform for vonMises(0, κ).The Gamma proposal follows Marsaglia and Tsang, while von Mises uses a wrapped Cauchy proposal.
  • D Experimental details: Cross-entropy gradient variance uses 1000 samples, while timings are measured on a single Intel Xeon CPU core.The analytical gradient of the cross-entropy is used as the reference quantity.
  • D Experimental details: Gradient variance is estimated from Adam's first- and second-moment moving averages and averaged across model parameters.The estimate uses v − m^2 for the moment estimates m and v.
  • D Experimental details: Positive distribution parameters use softplus transformations, with most additionally clipped to [10−3, 10^3] for float32 KL stability.The clipping excludes the Normal scale parameter.
  • D Experimental details: For von Mises latents, learning uses the circular representation (cos z, sin z) and computes μ with atan2(x, y).The paper reports that this representation makes learning much easier than using angles directly.
  • D Experimental details: LDA experiments use 50 topics, Adam with batch size 32, dataset-specific training epochs, and jointly trained prior parameters after burn-in.The setup includes 20 Newsgroups and RCV1 experiments with an amortized inference network.
  • D Experimental details: 60 points higher perplexity results from using a fixed prior, while sampling the KL increases perplexity by 80 points on 20 Newsgroups.The latter result emphasizes analytical KL estimation for sparse distributions with density asymptotes.
  • D Experimental details: The VAE uses fully connected ReLU encoder and decoder networks, Adam optimization, minibatches of 64, and KL annealing over the first 10^5 minibatches.Models are trained for 2 million minibatches without early stopping.
Loading 1805.08498v4…