Source-linked AI summary
The Generalized Reparameterization Gradient
Francisco J. R. Ruiz, Michalis K. Titsias, David M. Blei
TL;DR
Standard reparameterization gradients are low-variance but do not easily apply to common beta and gamma variational distributions without approximations. The paper introduces generalized reparameterization gradients using invertible transformations and shows their effectiveness in two nonconjugate models, where a single Monte Carlo sample suffices for an effective low-variance gradient.
Problem
Standard reparameterization gradients do not easily generalize beyond Gaussian variational distributions to common beta and gamma distributions without further approximations.
Method
G-REP uses invertible latent-variable transformations that make transformed variables weakly dependent on variational parameters, combining reparameterization and score function gradients.
Results
A single Monte Carlo sample is enough to obtain an effective low-variance gradient in two nonconjugate probabilistic models.
Takeaways & Limitations
G-REP extends reparameterization gradients to variational families involving gamma and beta distributions while retaining black-box applicability without approximations.
Takeaways & Limitations
Standard reparameterization remains less generic than BBVI and typically requires Gaussian variational distributions, with beta and gamma cases not directly supported.
Abstract
from arXiv · showhide
The reparameterization gradient has become a widely used method to obtain Monte Carlo gradients to optimize the variational objective. However, this technique does not easily apply to commonly used distributions such as beta or gamma without further approximations, and most practical applications of the reparameterization gradient fit Gaussian distributions. In this paper, we introduce the generalized reparameterization gradient, a method that extends the reparameterization gradient to a wider class of variational distributions. Generalized reparameterizations use invertible transformations of the latent variables which lead to transformed distributions that weakly depend on the variational parameters. This results in new Monte Carlo gradients that combine reparameterization gradients and score function gradients. We demonstrate our approach on variational inference for two complex probabilistic models. The generalized reparameterization is effective: even a single sample from the variational distribution is enough to obtain a low-variance gradient.
1 Introduction
Variational inference becomes difficult in nonconjugate models, motivating Monte Carlo gradient methods. Reparameterization gradients are efficient but limited mainly to Gaussian distributions, while G-REP extends them to broader variational families.
- Motivation: Nonconjugate models require effective optimization algorithms that work beyond exponential-family conditional distributions.The paper targets applications including probabilistic neural networks, Bayesian generalized linear models, deep neural networks, and probabilistic programming.
- Monte Carlo gradients: BBVI provides generic score-function gradients but typically suffers from high variance and potentially slow convergence.Variance-reduction methods include Rao-Blackwellization, control variates, local expectations, and importance sampling.
- Monte Carlo gradients: Reparameterization gradients generally have lower variance than BBVI and often need only one Monte Carlo sample.Their main limitation is that they do not easily generalize beyond commonly used Gaussian variational distributions without further approximations.
- Contribution: G-REP uses invertible latent-variable transformations whose transformed distributions depend only weakly on variational parameters.This combines reparameterization and score-function gradients while retaining black-box reuse of variational families and avoiding approximations.
- Contribution: Experiments on a nonconjugate factorization model and a deep exponential family show that one sample can produce an effective low-variance gradient.G-REP outperforms BBVI in speed and ADVI in accuracy in the reported experiments.
2 Background
The background formulates VI through the ELBO and reviews score-function and reparameterization gradients. Reparameterization lowers estimator variance but requires transformations that are unavailable for some important distributions, motivating a weaker dependence condition.
- Variational inference: VI approximates the intractable posterior by optimizing variational parameters to maximize the evidence lower bound.The ELBO is equivalent to minimizing the KL divergence between the variational approximation and the posterior.
- Monte Carlo gradients: When the ELBO expectation is intractable, stochastic optimization constructs Monte Carlo gradient estimates with respect to the variational parameters.The paper reviews score-function and reparameterization approaches for this purpose.
- Score function method: The score-function method samples from q(z; v), but its estimator suffers from high variance and often needs variance-reduction techniques.Control variates and Rao-Blackwellization are cited as variance-reduction methods.
- Reparameterization: Reparameterization writes z as an invertible function of auxiliary variables whose distribution does not depend on variational parameters.This moves the gradient inside the expectation and typically produces much lower-variance estimates; one sample is often sufficient.
- Generalization: The standard reparameterization condition does not hold for gamma and beta distributions under the transformations reviewed.G-REP relaxes the condition by allowing the standardized variable's distribution to depend weakly on variational parameters.
3 The Generalized Reparameterization Gradient
The generalized reparameterization gradient extends reparameterization to distributions such as gamma and beta by transforming latent variables so the transformed density depends weakly on variational parameters. Its gradient combines a reparameterization term with a correction term, with transformations chosen to reduce correction variance.
- Generalized reparameterization: G-REP introduces an invertible transformation z = T(ϵ; v) whose inverse makes ϵ weakly dependent on the variational parameters.The transformed density remains qϵ(ϵ; v), rather than being assumed fully parameter-independent.
- Algorithm and scope: Unlike inverse-CDF and integer-shape composition approaches, the proposed method does not require CDF computation and applies when gamma shape parameters are not integers.The algorithm draws one sample, computes auxiliary functions, estimates both gradient terms, and updates the variational parameters.
- Gradient decomposition: The full gradient combines a standard reparameterization term, g_rep, with a correction term, g_corr, caused by parameter dependence in the transformed density.When qϵ does not depend on v, g_corr becomes zero.
- Gradient decomposition: The framework recovers the score-function gradient under the identity transformation, while interpreting G-REP alternatively as a control variate for that estimator.With the identity transformation, the auxiliary functions vanish and g_rep is zero.
- Choosing transformations: Transformations are selected to make g_corr small by standardizing sufficient statistics so the inverse-transformed variable depends weakly on v.For exponential-family distributions, the paper uses centered sufficient statistics divided by a scale factor.
- Examples: For gamma distributions, standardizing log(z) gives ϵ zero mean and unit variance, making its first two moments independent of α and β.The auxiliary functions then produce the gamma-specific g_rep and g_corr terms.
- Examples: For log-normal variables, standardizing log(z) produces a standard normal ϵ and therefore g_corr = 0.The paper also develops transformations for beta variables based on standardizing logit(z), with a sample-dependent denominator that can set g_corr to zero but introduces bias.
4 Experiments
The experiments evaluate G-REP against BBVI and ADVI on nonconjugate models using beta and gamma variational distributions. G-REP achieves competitive variational bounds and test performance while using substantially fewer samples than BBVI and remaining faster than BBVI.
- Models and datasets: The experiments apply G-REP to sparse gamma deep exponential families and beta-gamma matrix factorization on four image and document datasets.The datasets are Olivetti, NIPS 2011, binarized MNIST, and Omniglot.
- Computational cost: G-REP is 1-4 times slower than ADVI but above one order of magnitude faster than BBVI.Table 1 reports average time per iteration in seconds.
- Experimental setup: G-REP and BBVI use beta and gamma variational distributions, whereas ADVI uses Gaussian distributions corresponding to log-normal or logit-normal distributions.Gradient estimates use 30 Monte Carlo samples for BBVI and 1 sample for both G-REP and ADVI.
- Variational objective: BBVI converges slower than the other methods because each iteration draws multiple samples and evaluates the log-joint for each sample.The reported comparison uses a fixed computational budget of CPU time.
- Variational objective: ADVI and G-REP achieve similar bounds except on MNIST, where G-REP obtains a higher ELBO and a variational approximation closer to the posterior.The passage attributes this difference to the greater fit of sparse gamma and beta variational families.
- Test performance: G-REP provides higher likelihood values than ADVI because ADVI's log-normal and logit-normal families cannot capture sparsity in the considered models.Supplementary visual comparisons report that ADVI samples are more blurry or lack some details.
5 Conclusion
The conclusion presents G-REP as an extension of standard reparameterization to wider variational families and demonstrates it on gamma and beta approximations. The experiments indicate that one Monte Carlo sample can produce a noisy gradient estimate suitable for fast inference.
- G-REP extends the standard reparameterization gradient to a wider class of variational distributions.
- The method applies to differentiable probabilistic models and is demonstrated with gamma and beta variational distributions in two nonconjugate models.
- A single Monte Carlo sample is enough to obtain a noisy gradient estimate, leading to a fast inference procedure.
A Derivation of the Generalized Reparameterization Gradient
The derivation rewrites the gradient under an invertible transformation and separates it into a reparameterization component and a correction component. The correction term arises from the transformed distribution's parameter dependence and Jacobian.
- The generalized gradient is decomposed into the reparameterization term grep and a correction term gcorr.
- The correction term uses qϵ(ϵ; v) = q(T(ϵ; v); v)J(ϵ, v), then applies the chain rule and a change of variables.
- The derivation applies the chain rule to ∇v log q(T(ϵ; v); v) and then changes variables back to the original latent-variable space.
- The auxiliary function u(ϵ; v) is defined as the variational-parameter derivative of the log-Jacobian, u(ϵ; v) ≜ ∇v log J(ϵ, v).
B Particularization for the Gamma Distribution
The paper specializes G-REP transformations and gradient components for gamma and beta distributions. For beta variables, a standardized logit transformation makes the transformed variable mean-zero, while its scale is chosen to reduce the correction term, introducing some bias.
- Gamma distribution: For gamma variables, the transformation has Jacobian J(ϵ, α, β) = T(ϵ, α, β)ψ1(α).
- Gamma distribution: The gamma derivation obtains derivatives of log q with respect to α and uses auxiliary functions h and u to construct gradient components.
- Beta distribution: For beta variables, the method represents z through two gamma-distributed variables and standardizes the logit function directly.
- Beta distribution: The beta transformation is z = 1 / (1 + exp(−ϵσ − ψ(α) + ψ(β))) and has Jacobian T(ϵ; α, β)(1 − T(ϵ; α, β))σ.
- Beta distribution: The beta transformation ensures ϵ has mean zero, while σ is chosen as a function of α and β so gcorr = 0 for the sampled z, although this introduces some bias.
- Beta distribution: The beta gradient uses derivatives of log q with respect to α and β involving log(z) and log(1 − z).
- Beta distribution: The quantity ϵσ can be computed from z as logit(z) − ψ(α) + ψ(β), without knowing σ separately.
- Beta distribution: The derivatives of φ are selected so the correction term vanishes for the sampled value of z, requiring a simple linear equation.
D Particularization for the Dirichlet Distribution
The Dirichlet specialization derives the quantities needed for G-REP from the distribution’s covariance structure and log-density derivatives. It also notes a speed–variance trade-off in simplifying the covariance transformation.
- Dirichlet standardization: The Dirichlet specialization standardizes the distribution using its mean vector and covariance matrix.The covariance matrix has dimension K × K and can be decomposed for computation.
- Covariance computation: The covariance matrix combines diagonal and rank-one structure, enabling faster computation.This structure can be exploited when evaluating the Dirichlet transformation.
- Covariance computation: Because the covariance is positive semidefinite, its square root can be obtained by eigendecomposition.If Σ = VDV⊤, then Σ^1/2 = VD^1/2V⊤.
- Approximation trade-off: Ignoring off-diagonal covariance terms makes the transformation faster but increases estimator variance.The full-covariance transformation is also applicable to the beta distribution.
- Gradient construction: The derivation computes auxiliary functions and their intermediate derivatives before assembling the generalized reparameterization gradient.These quantities are defined from the transformation and derivatives of log q(z; α).
E.1 Using more than 1 sample
The experiment evaluates G-REP with 2, 5, 10, and 20 Monte Carlo samples over the first 200 inference iterations. More samples reduce gradient-estimator variance, especially early in inference.
- Experimental setup: 2, 5, 10, and 20 Monte Carlo samples are compared during the first 200 inference iterations.The experiment tracks both the ELBO and average sample variance of the gradient estimator.
- Effect of sample count: Increasing the number of samples reduces the resulting gradient-estimator variance.The comparison is performed on the Olivetti dataset using G-REP.
E.2 Reconstructed images
The reconstructed-image experiments compare G-REP and ADVI on Olivetti, binarized MNIST, and Omniglot. Across these datasets, G-REP reconstructions retain more detail, while ADVI reconstructions are blurrier or less recognizable.
- Olivetti: For Olivetti, the figure places true observations on the left, G-REP reconstructions in the center, and ADVI reconstructions on the right.G-REP images show more facial detail than ADVI images.
- Olivetti: G-REP reconstructs Olivetti faces with more distinguishable glasses, mustaches, and facial expressions than ADVI.The paper attributes this difference to ADVI’s inability to capture sparse posterior distributions well.
- Cross-dataset comparison: Across MNIST and Omniglot, G-REP samples are visually closer to ground truth, while ADVI outputs can be blurry or unrecognizable.This behavior is reported for both handwritten digits and characters.
- Binarized MNIST: For binarized MNIST, ADVI produces blurrier images than G-REP.The reconstructions are generated by drawing one variational-posterior sample and computing the corresponding mean observations.
- Omniglot: For Omniglot, ADVI likewise produces blurrier images than G-REP.The same one-sample reconstruction procedure is used for the image datasets.