PaperlayerView exact revision on arXiv

Source-linked AI summary

Reducing Reparameterization Gradient Variance

Andrew C. Miller, Nicholas J. Foti, Alexander D'Amour, Ryan P. Adams

arXiv:1705.07880v1stat.MLstat.COstat.ME

TL;DR

Noisy reparameterization gradients can make Monte Carlo variational inference slow or unstable, while averaging more samples is costly. The paper constructs a cheap, unbiased control-variate estimator that reduces gradient variance by 20–2,000× in demonstrated models and supports faster, more stable optimization.

  • Problem

    MCVI needs lower-variance reparameterization gradients because excessive noise can slow or destabilize optimization, while averaging estimates is computationally expensive.

  • Method

    The paper uses an analytical linear approximation of the known Monte Carlo generating process as a cheap control variate, preserving unbiasedness while reducing estimator variance.

  • Results

    20–2,000× reduction in gradient-norm variance was observed under some conditions, with the approach effective on non-conjugate hierarchical Bayesian models and a neural network.

  • Takeaways & Limitations

    The approach produces tighter ELBO-gradient estimators that lead to faster and more stable optimization in the demonstrated models.

  • Takeaways & Limitations

    The approach relies on the reparameterization-gradient distribution being well captured by linear structure, which may not hold for all models.

Abstract

from arXiv · show

Optimization with noisy gradients has become ubiquitous in statistics and machine learning. Reparameterization gradients, or gradient estimates computed via the "reparameterization trick," represent a class of noisy gradients often used in Monte Carlo variational inference (MCVI). However, when these gradient estimators are too noisy, the optimization procedure can be slow or fail to converge. One way to reduce noise is to use more samples for the gradient estimate, but this can be computationally expensive. Instead, we view the noisy gradient as a random variable, and form an inexpensive approximation of the generating procedure for the gradient sample. This approximation has high correlation with the noisy gradient by construction, making it a useful control variate for variance reduction. We demonstrate our approach on non-conjugate multi-level hierarchical models and a Bayesian neural net where we observed gradient variance reductions of multiple orders of magnitude (20-2,000x).

1. Introduction

MCVI uses stochastic gradient estimates whose variance can hinder optimization, especially when reducing noise by averaging samples is costly. The paper introduces an unbiased control-variate estimator based on the known Monte Carlo generating process and reports substantial variance reduction and improved optimization behavior.

  • Stochastic gradient methods replace expensive objective-gradient evaluations with cheap, unbiased estimates, enabling optimization for massive datasets and flexible probabilistic models.Examples include mini-batch stochastic gradient descent and gradients estimated by sampling from generative or approximating distributions Mohamed and Lakshminarayanan [2016], Tran et al. [2017], Arjovsky et al. [2017], Goodfellow et al. [2014].
  • MCVI reframes variational inference as optimization and estimates gradients by sampling from an approximate posterior, with data subsampling available as an additional acceleration.The VI objective is often intractable, motivating Monte Carlo gradient estimates Kingma and Welling [2014], Ranganath et al. [2014], Rezende et al. [2014].
  • Reparameterization gradients typically have lower variance than score-function gradients but apply only to continuous random variables, making variance control important for MCVI success.Score-function estimates can handle discrete and continuous variables but often have high variance Ranganath et al. [2014].
  • 20-2,000 × reductions in gradient-norm variance were observed under some conditions, alongside faster convergence and more stable optimization traces.The approach was applied to Gaussian approximating families and evaluated under various conditions.
  • The method constructs a cheap control variate from an analytical linear approximation and combines it with a naïve gradient estimate to retain unbiasedness while reducing variance.Because MCVI gradient randomness is completely determined by a known Monte Carlo generating process, that process can be used to de-noise the estimator.

2. Background

Variational inference approximates an intractable posterior by optimizing the ELBO over a parameterized family, using unbiased reparameterization gradients estimated by Monte Carlo. Gradient noise creates a step-size–convergence trade-off, motivating control variates that reduce variance without bias or excessive computation.

  • Variational Inference: Variational inference selects q(z; λ) from a simpler family Q to approximate the posterior p(z|D) by minimizing a divergence.The posterior is generally intractable, so VI optimizes an alternative objective over variational parameters λ.
  • Variational Inference: The ELBO is optimized because maximizing it minimizes the KL divergence between q(z; λ) and p(z|D).Direct KL evaluation requires computing the posterior, motivating the ELBO objective.
  • Reparameterization Gradients: Reparameterization gradients express sampling as a differentiable map z = T(ϵ; λ), then differentiate a Monte Carlo ELBO estimate with respect to λ.For a diagonal Gaussian, the map is z = m_λ + s_λ ⊙ϵ with ϵ ∼N(0, I_D).
  • Gradient Variance and Convergence: Monte Carlo variational inference estimates unbiased reparameterization gradients, but large gradient noise forces smaller steps and slows convergence.The reparameterization gradient satisfies E[ĝ_λ] = ∇_λL(λ), while its Monte Carlo estimation introduces noise.
  • Gradient Variance and Convergence: In a Bayesian neural network with D = 653 parameters, the single-sample estimator converges extremely slowly, whereas the 50-sample estimator converges quickly at 50 times the cost.Figure 1 compares Adam with fixed learning rates across different sample counts; larger steps are faster but become chaotic with noisy gradients.
  • Gradient Variance and Convergence: The method targets lower-variance gradient estimates by using information from the variational model q(·; λ) to construct a control variate without introducing bias.Control variates inject information into an unbiased estimator to produce another unbiased estimator with lower variance.

Control Variates

The control-variate estimator preserves the original gradient estimator’s expectation while changing its variance, with optimal reduction achieved by choosing C according to the estimator–control-variate covariance.

  • The new estimator has the same expectation as the original estimator but a different variance.
  • Optimal variance reduction is attained by appropriately setting C.
  • The optimal C resembles a regression coefficient and is related to the covariance between the control variate and the original estimator.

3. Method: Modeling Reparameterization Gradients

The method approximates the distribution of reparameterization gradient estimates using the variational structure and model curvature, then uses the correlated approximation as a control variate. For diagonal Gaussian families, it yields tractable estimators with several Hessian-operation variants, including scalable Hessian-vector-product approaches.

  • 3. Method: Modeling Reparameterization Gradients: The method models reparameterization-gradient randomness from the shared base noise ϵ, using q_λ, the transformation T, and curvature of ln p(z, D) to construct a correlated control variate.The data term is linearized around z_0, using the model Hessian; the approximation requires the model to be twice differentiable.
  • 3.1. Gaussian Variational Families: For diagonal Gaussian variational families, the approximation ˜g_λ shares ϵ with ˆg_λ, has tractable component distributions and means, and is used with its expectation to reduce RGE variance.The mean-parameter component is Gaussian, while the scale-parameter component is a location-shifted, scaled non-central χ2 distribution.
  • 3.2. Reduced Variance Reparameterization Gradient Estimators: RV-RGE provides Full Hessian, Hessian Diagonal, and HVP+Local estimators, respectively targeting feasible full-Hessian use, cheap diagonal approximations, or efficient Hessian-vector products with local diagonal estimation.Hessian-vector products are nearly as efficient as gradients, while HVP+Local is intended for models where exact Hessian diagonals are unavailable or expensive.
  • 3.3. Related Work: The approach complements Roeder et al. [2017]: this approximation is typically more accurate early in optimization, whereas their estimator has lower variance near convergence.The paper also relates its control-variate construction to prior variance-reduction methods for score-function and stochastic-gradient estimators.

4. Experiments and Analysis

Experiments on a hierarchical Poisson GLM and a Bayesian neural network show that RVRGs substantially reduce gradient variance and improve stochastic optimization. HVP+Local converges faster than Monte Carlo baselines, including a 10-sample HVP+Local estimator outperforming 50-sample MC on the Bayesian neural network.

  • Experimental Setup: The experiments cover a 37-dimensional hierarchical Poisson GLM and a 653-dimensional Bayesian neural network applied to the wine dataset.Both models are non-conjugate and serve as real-data examples for evaluating RVRG variance and stochastic optimization.
  • Quantifying Gradient Variance Reduction: 100–5,000× reductions in entire-gradient variance are observed with HVP+Local, including 80×–3,000× reductions for mean-parameter gradient norms.The Hessian diagonal alone reduces mean-parameter variance only 2–5×, highlighting the importance of the full Hessian-vector product.
  • Quantifying Gradient Variance Reduction: HVP+Local reduces scale-parameter variance by approximately 2,000× early, but later becomes comparable to MC while the full Hessian retains substantial reduction.The later limitation indicates that local Hessian diagonal estimation accounts for most remaining noise; most overall variance is in mean parameters.
  • Optimizer Convergence and Stability: For the frisk model, HVP+Local estimators make early progress and converge quickly, whereas the L = 3 pure MC estimator produces noisy optimization paths across the reported optimization traces.Traces compare L = 2 and L = 10 samples with step sizes .05 and .1; true ELBO values are estimated using 2000 Monte Carlo samples at each iteration.
  • Optimizer Convergence and Stability: A 10-sample HVP+Local estimator outperforms a 50-sample MC estimator in objective value versus wall-clock time on the Bayesian neural network.HVP+Local does more work per iteration but tends to converge faster.

5. Conclusion

The work uses information about the variational family to construct tighter ELBO gradient estimators, enabling faster and more stable optimization. Its effectiveness depends on linear structure capturing the reparameterization-gradient distribution, with extensions proposed for more flexible variational families and inference schemes.

  • Conclusion: The method creates tighter ELBO gradient estimators from variational-family information, producing faster and more stable optimization for non-conjugate hierarchical Bayesian models and a neural network.Its efficacy relies on the reparameterization-gradient distribution being well captured by linear structure, which may not hold for all models.
  • Conclusion: The approach instantiates a broader strategy of using cheap linear structure to remove variation from stochastic gradient estimates, with extensions proposed to flows, hierarchical distributions, and recognition networks.The proposed extensions include flow distributions Rezende and Mohamed [2015], hierarchical distributions Ranganath et al. [2016], and recognition networks Kingma and Welling [2014].

Appendix A: Control Variates

Control variates reduce estimator variance without introducing bias by combining an unbiased estimator with a zero- or known-mean quantity, using correlation to determine how informative the adjustment is.

  • Scalar control variates: Control variates construct an unbiased estimator with lower variance by combining an unbiased estimate with a quantity whose expectation is known.The adjusted estimator retains the original expectation while changing its variance.
  • Scalar control variates: The scalar control-variate coefficient c is chosen to minimize the adjusted estimator’s variance, although the required covariance must generally be estimated.Under the optimal coefficient, the resulting variance depends on the correlation between the estimator and control variate.
  • Vector control variates: For length-D vector estimators, the method replaces the scalar coefficient with a matrix-valued parameter C and minimizes the sum of marginal variances.The optimal matrix is expressed using the control variate covariance and cross-covariance matrices.
  • Interpretation: The control variate is informative when it is correlated with the original estimator: perfect correlation makes estimation unnecessary, whereas zero correlation recovers the original quantity.The method injects linear structure into the estimator, with informativeness decreasing as correlation vanishes.

A.1. Control Variates and Approximate Functions · Appendix B: Algorithm Details · Appendix C: Model Definitions

The appendix develops Taylor-based control variates for functions of Gaussian noise and specifies RV-RGE algorithm variants, including Hessian approximations and computational optimizations.

  • A.1. Control Variates and Approximate Functions: RV-RGE approximates f(ϵ), with ϵ ∼ N(0, I), using a first-order Taylor expansion around 0 in the univariate case.
  • A.1. Control Variates and Approximate Functions: Because f(ϵ) is general, the covariance between f(ϵ) and its control variate is analyzed through a second-order expansion.
  • A.1. Control Variates and Approximate Functions: The derivation uses the Gaussian moment identity E[ϵ^3] = 0 when characterizing the optimal control variate.
  • Appendix B: Algorithm Details: Algorithm 1 summarizes RV-RGE variants based on different forms of H(·, ·) and diag(H).
  • Appendix B: Algorithm Details: The full-Hessian variant computes the required Hessian terms exactly, whereas the diagonal-Hessian variant uses the Hessian-vector product with the Hessian diagonal.
  • Appendix B: Algorithm Details: The HVP+Local variant computes the Hessian-vector product exactly but estimates the scale-approximation mean from other samples.
  • Appendix B: Algorithm Details: The additional Hessian-vector-product computation can be optimized by caching the forward-pass Hessian at the shared mλ and repeating only the backward pass per sample.

C.1. Multi-level Poisson GLM

This section evaluates the method on a 37-dimensional hierarchical Poisson GLM posterior modeling stop-and-frisk rates across ethnicities and precincts.

  • C.1. Multi-level Poisson GLM: The test model is a 37-dimensional posterior from a hierarchical Poisson GLM measuring relative stop-and-frisk rates across ethnicities and precincts.The model is an illustrative multi-level modeling example [Gelman and Hill, 2006, Chapter 15, Section 1] and is associated with Gelman et al. [2007].
  • C.1. Multi-level Poisson GLM: The log rate combines an intercept, ethnicity effects, precinct effects, and the log number of arrests: ln λ_ep = µ + α_e + β_p + ln N_ep.Y_ep counts stop-and-frisk events for ethnicity group e and precinct p, while N_ep counts arrests over the same fixed period.

C.2. Bayesian Neural Network

This section applies the approach to a Bayesian neural network trained on a 100-row Wine Quality dataset subsample, using hierarchical priors over weights and variances.

  • The experiment uses a 50-unit hidden-layer neural network with ReLU activations and normally distributed weights sharing an inverse-Gamma variance prior.An inverse-Gamma prior is also placed on the observation variance.
  • The model approximates the posterior over weights, weight-scale parameters, and observation-variance parameters given the training data.The parameter set is θ ≜ (w, α, τ), where the multilayer perceptron maps inputs to approximate outputs.
  • Evaluation uses a 100-row subsample of the UCI Wine Quality dataset.

Appendix D: Variance Reduction

Appendix D reports additional variance-reduction measurements for the frisk model across different values of L, the samples drawn per iteration. It includes variance-comparison tables for 3-sample and 50-sample estimators.

  • The appendix evaluates frisk-model variance reduction for different L values, where L denotes samples drawn per iteration.
  • Table 2 presents the frisk model’s variance comparison for L = 3-sample estimators.
  • Table 3 presents the frisk model’s variance comparison for L = 50-sample estimators.
Loading 1705.07880v1…