PaperlayerView exact revision on arXiv

Source-linked AI summary

REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models

George Tucker, Andriy Mnih, Chris J. Maddison, Dieterich Lawson, Jascha Sohl-Dickstein

arXiv:1703.07370v4cs.LGstat.ML

TL;DR

Discrete latent-variable models suffer from high-variance gradient estimates, while continuous relaxations trade lower variance for bias. REBAR combines control variates with relaxed-model gradients to produce unbiased, low-variance estimates and achieves state-of-the-art variance reduction across benchmark tasks, with faster convergence and better final log-likelihoods.

  • Problem

    Gradient estimation for discrete latent-variable models remains challenging when variables cannot be marginalized, because sampling-based REINFORCE estimators have high variance.

  • Method

    REBAR combines REINFORCE with Concrete-relaxation gradients through a conditionally marginalized control variate and adapts the relaxation temperature online.

  • Results

    Across generative-modeling tasks, REBAR achieved state-of-the-art variance reduction, faster convergence, and better final log-likelihoods.

  • Takeaways & Limitations

    Online temperature adaptation further reduced variance while preserving REBAR’s superior variance reduction during optimization.

  • Takeaways & Limitations

    Training variational-bound improvements did not directly translate into improved test log-likelihood, and additional inference-network training did not completely resolve overfitting.

Abstract

from arXiv · show

Learning in models with discrete latent variables is challenging due to high variance gradient estimators. Generally, approaches have relied on control variates to reduce the variance of the REINFORCE estimator. Recent work (Jang et al. 2016, Maddison et al. 2016) has taken a different approach, introducing a continuous relaxation of discrete variables to produce low-variance, but biased, gradient estimates. In this work, we combine the two approaches through a novel control variate that produces low-variance, \emph{unbiased} gradient estimates. Then, we introduce a modification to the continuous relaxation and show that the tightness of the relaxation can be adapted online, removing it as a hyperparameter. We show state-of-the-art variance reduction on several benchmark generative modeling tasks, generally leading to faster convergence to a better final log-likelihood.

1 Introduction

The paper targets high-variance gradient estimation for discrete latent-variable models by introducing REBAR, a low-variance unbiased estimator with online temperature optimization. Experiments show state-of-the-art variance reduction, faster convergence, and better final log-likelihoods across generative modeling tasks.

  • 1 Introduction: Discrete latent-variable models are widespread, but REINFORCE-like optimization is difficult when variables cannot be analytically marginalized because of high-variance gradients.Examples include mixture models, reinforcement-learning MDPs, structured-prediction generative models, hard attention, and memory networks.
  • 1 Introduction: REBAR combines REINFORCE gradients with Concrete-relaxation gradients through a conditionally marginalized control variate, substantially improving variance reduction.The control variate is based on the difference between REINFORCE and reparameterization-trick estimators for the relaxed model.
  • 1 Introduction: REBAR remains unbiased for all temperatures, allowing the relaxation temperature to be optimized online and eliminating an additional tuned hyperparameter.A modified Concrete relaxation also connects REBAR to MuProp in the high-temperature limit.
  • 1 Introduction: Across MNIST and Omniglot generative-modeling tasks, REBAR achieves state-of-the-art variance reduction that translates into faster convergence and better final log-likelihoods.Experiments include generative sigmoid belief networks and conditional generative models, alongside a toy demonstration of biased-estimator problems.

2 Background

Discrete latent-variable learning requires estimating gradients whose naïve score-function term can have high variance. Continuous relaxations provide low-variance but biased estimates, creating a bias–variance tradeoff controlled by a temperature parameter.

  • Gradient estimation: A naïve single-sample estimator of the score-function term has high variance, while the direct dependence of f(b, θ) on θ is straightforward to account for.The background therefore focuses on the score-function term under the simplifying assumption f(b, θ) = f(b).
  • Control variates: Control variates reduce variance by subtracting a related quantity with a known expectation and correcting the resulting bias.NVIL learns a control variate independent of b, whereas MuProp uses a linear Taylor expansion; high variance can remain.
  • Continuous relaxation: Replacing the discontinuous hard threshold H(z) with the continuous relaxation σ_λ(z) creates a reparameterizable computational graph and low-variance gradient estimates.The discrete variable is parameterized as b = H(z), with z generated from a Uniform(0, 1) variable through a differentiable transformation.
  • Continuous relaxation: The relaxed estimator is low-variance but biased; as λ →0, the approximation becomes exact while Monte Carlo variance diverges, so λ must balance bias and variance.λ > 0 acts as a temperature controlling relaxation tightness, with low temperatures producing a nearly tight relaxation.

3 REBAR

REBAR combines a Concrete-inspired relaxation with a control variate to obtain low-variance, unbiased gradient estimates for discrete latent-variable models. Its relaxation tightness can be optimized online, while extensions cover multilayer stochastic units and learned Q-functions.

  • Control variate construction: REBAR constructs a control variate from the difference between the relaxed-model REINFORCE estimator and the reparameterization-trick estimator.The relaxed-model estimator is strongly correlated with the original and can reduce variance; common random numbers couple the underlying samples.
  • Adaptive relaxation: The REBAR estimator remains unbiased for every λ > 0, allowing λ to minimize estimator variance without changing its expectation.The scaling η is also estimated by minimizing Monte Carlo variance with SGD.
  • Adaptive relaxation: Optimizing λ jointly with model parameters adapts relaxation tightness online and removes the need to choose λ beforehand.The expectation can be estimated with a single-sample Monte Carlo estimator.
  • Multilayer extension: For multiple stochastic layers, REBAR replaces hard thresholds with continuous relaxations and requires n network passes, one per layer.The multilayer control variate exploits factorization and debiases reparameterizable terms, while alternatives avoiding multiple passes are deferred to Appendix F.
  • Applications: Because its derivation is independent of f, the REBAR control variate can replace f with a learned differentiable Q-function for reinforcement learning.This yields a pseudo-action-dependent baseline based on the relaxation of a policy network’s discrete output.

4 Related work

Related work frames discrete-expectation optimization primarily through the unbiased REINFORCE estimator and its baseline-based variance-reduction extensions. In latent-variable models, these methods support both factorized and structured variational inference with learned baselines.

  • REINFORCE estimators: REINFORCE, also called the likelihood-ratio or score-function estimator, optimizes expectations over discrete distributions using sampled gradients.The estimator uses a sample b from the discrete distribution and a baseline c independent of b.
  • REINFORCE estimators: Although REINFORCE estimators are unbiased, their variance can be too high without carefully chosen baselines.This has motivated substantial work on finding effective baselines.
  • Latent-variable inference: For latent-variable models, REINFORCE-like methods implement sampling-based variational inference with fully factorized or structured variational distributions.These methods use learned baselines ranging from simple scalar forms to nonlinear input-dependent baselines.

5 Experiments

Experiments show that REBAR reduces gradient-estimator variance and generally improves optimization relative to unbiased and biased baselines. However, biased estimators can converge to incorrect solutions, and training-bound improvements do not necessarily improve test log-likelihood.

  • Toy problem: On a Bernoulli toy problem, only unbiased estimators converged to the correct deterministic solution, whereas Concrete converged to a stochastic one.The objective was E_p(b)[(b − t)^2] with t ∈ (0, 1), and the illustrative experiment used t = 0.45.
  • Generative modeling: Across MNIST and Omniglot generative-modeling tasks, REBAR produced the lowest variance for both linear and nonlinear models, with especially large reductions for linear models.Variance comparisons used the same random numbers along a single optimization trajectory; the paper reports corresponding MNIST plots and an Omniglot appendix figure.
  • Test performance: Improved training variational lower bounds did not directly translate into improved test log-likelihood, motivating separate test-set evaluation with a 100-sample lower bound.The paper notes that inference-network regularization alone had previously prevented overfitting, suggesting the observed overfitting was primarily related to the inference network.
  • Structured prediction: On structured prediction, REBAR significantly reduced variance and training-objective performance generally mirrored that reduction, although MuProp excelled for the single-layer linear model.The task modeled the bottom half of an MNIST digit conditional on its top half; MuProp’s advantage was attributed to the accuracy of its first-order expansion in that configuration.

6 Discussion · Appendix

The discussion highlights REBAR’s variance reduction, its relation to MuProp, and online temperature adaptation, while noting optimization costs and future extensions. The appendix documents additional variance and training-curve evaluations across MNIST, structured prediction, and Omniglot tasks.

  • 6 Discussion: REBAR is introduced as a Concrete-inspired control variate for REINFORCE that greatly reduces gradient-estimator variance.The discussion identifies variance reduction as a central contribution.
  • 6 Discussion: With a modified relaxation, REBAR and MuProp are closely related in the high-temperature limit.The relationship is stated specifically for the high-temperature regime.
  • 6 Discussion: REBAR’s differences of reparameterization gradients implicitly implement Roeder et al.’s recommendation concerning a score-function term that can increase gradient variance.The discussion connects REBAR’s construction to the recommendation described by Roeder et al. (2017).
  • 6 Discussion: Optimizing the relaxation temperature requires differentiating the parameter gradient with respect to λ, although temperature changes empirically more slowly than parameters.The authors suggest amortizing this cost across several parameter updates and leave that exploration to future work.
  • 6 Discussion: Future work includes multi-sample extensions such as VIMCO, Q-functions for layered models, and applications to reinforcement learning.These are proposed as natural directions for extending the approach.
  • Appendix: The appendix reports log gradient-estimator variance for the single-layer linear model on MNIST and structured prediction tasks.Figure App.1 covers both tasks.
  • Appendix: The appendix shows training variational lower bounds for the single-layer linear model on MNIST and structured prediction, using five random-initialization trials with the median highlighted.Figure App.2 summarizes training behavior across the two tasks.
  • Appendix: The appendix reports log gradient-estimator variance for single-layer linear, two-layer linear, and single-layer nonlinear models on Omniglot.Figure App.3 compares three model architectures.

A Control Variates

Control variates reduce Monte Carlo estimator variance when correlated with the target function, while preserving unbiasedness through a zero-mean adjustment. The method can use a low-variance estimator when the control variate’s expectation is unavailable.

  • A Control Variates: The control-variate estimator remains unbiased for any value of η, allowing η to be chosen to minimize variance.The optimal η follows from minimizing the estimator variance.
  • A Control Variates: A control variate reduces estimator variance by a factor of 1 − ρ(f, g)^2 when it is correlated with the target function.The correlation ρ(f, g) determines the variance reduction.
  • A Control Variates: When E[g] cannot be computed, a low-variance estimator ĝ can replace it in the control-variate construction.Defining g̃ = g − ĝ yields a zero-mean control variate but obscures its interpretability.

B Conditional marginalization for the control variate … F Multilayer stochastic network

The appendix explains REBAR’s conditional-marginalization rationale, categorical and binary reparameterizations, and alternative low-variance decompositions. It also revises the relaxation and extends the estimator to multilayer stochastic networks, including computation-saving alternatives and their limitations.

  • B Conditional marginalization for the control variate: Conditional marginalization explains why the left-hand-side Monte Carlo gradient estimator has much lower variance than the right-hand-side estimator.The left-hand side analytically performs a conditional marginalization.
  • B Conditional marginalization for the control variate: The conditional identities use b = H(z) and p(z) = p(z|b)p(b), with analogous manipulation for the control variate.These identities follow when z is sampled from p(z|b).
  • C Reparameterizations for REBAR: For categorical b, the Gumbel-Max construction represents b as the one-hot argmax H(z), with z_i distributed as Gumbel variables having location log p_i.The conditional distribution uses truncated Gumbel variables and yields a reparameterization function for p(z|b).
  • D An alternative view of REBAR: The REBAR objective decomposes into reparameterizable terms and a residual, allowing the first and third terms to be estimated with low variance.The same construction applies to the binary Bernoulli special case through a Uniform-variable reparameterization.
  • E Rethinking the relaxation and a connection to MuProp: The modified relaxation converges to the mean θ as λ →∞ while becoming exact as λ →0, making the relaxation tightness adaptable through a λ-dependent parameter transformation.The appendix characterizes this construction as σ_λ with transformed parameters.
  • F Multilayer stochastic network: For multilayer stochastic networks, the estimator’s baseline form requires n network passes, where n is the number of layers.The appendix introduces two alternatives that avoid additional passes.
  • F Multilayer stochastic network: A coupled multilayer distribution preserves the discrete and relaxed marginals, factorizes conditionally, and makes the individual factors reparameterizable.The discrete path uses H(z), while the relaxed distribution replaces H(z) with σ_λ(z).
  • F Multilayer stochastic network: The two-evaluation estimator converges to multilayer SimpleMuProp as λ →∞, but increasing depth can cause the continuous path to diverge from the discrete path.Additional transition passes can reduce variance at increased computation, while a learned Q-function avoids extra computation but is left for future work.

G Implementation details

The implementation used Adam with task-dependent constant learning-rate grids, fixed optimization settings, centered inference inputs, and binary variables in {−1, 1}. Preliminary evaluations found λ = 0.1 performed well across tasks and configurations.

  • G Implementation details: Models used Adam with constant learning rates selected from separate grids for linear and nonlinear models.Linear models used {3 × 10−5, 1 × 10−4, 3 × 10−4, 10−3, 3 × 10−3}; nonlinear models used {3 × 10−5, 10−4}. Higher learning rates caused divergence.
  • G Implementation details: Training used minibatches of 24 elements for 2 million optimization steps with β1 = 0.9 and β2 = 0.999996.The input to the inference network was centered using training-data statistics.
  • G Implementation details: All binary variables took values in {−1, 1}, following Maddison et al. (2016).
  • G Implementation details: Preliminary evaluations found λ = 0.1 performed well across tasks and configurations.The evaluation covered both REBAR and Concrete estimators over a range of λ values.

G.1 MuProp

MuProp’s linear baseline term was detrimental in later layers, so the authors learned an additional scaling factor to modulate it. This modification reduced the variance of the MuProp learning signal beyond the method of Gu et al. (2015).

  • G.1 MuProp: The linear term in the MuProp baseline was detrimental for later layers.
  • G.1 MuProp: The authors learned an additional scaling factor to modulate the linear terms.
  • G.1 MuProp: The modification reduced the variance of the MuProp learning signal beyond the algorithm described in Gu et al. (2015).

G.2 REBAR · G.3 Concrete

The REBAR implementation uses parameter-group-specific control-variate scalings and common random numbers, while online temperature estimation requires a large Adam β2 because smaller values can bias updates. The Concrete comparison accounts for its objective not being a stochastic lower bound on marginal log-likelihood and contrasts it with REBAR’s proper stochastic lower bound.

  • G.2 REBAR: REBAR learns separate control-variate scalings η for each parameter group, including first-layer weights and biases, and uses common random numbers when sampling b, z, and z|b.The estimator samples from shared randomness across its discrete and conditional-relaxation variables.
  • G.2 REBAR: REBAR’s variance analysis ignores the variance contribution from reparameterizable terms, which the authors leave for future work.This is an explicit limitation of the reported analysis rather than a modification to the estimator.
  • G.2 REBAR: A large β2 is crucial for online temperature estimation because Adam reuses gradients for both update numerator and denominator, producing greater bias at smaller β2.The bias is especially severe when gradient magnitudes are asymmetric and heavy-tailed, as in this setting.
  • G.2 REBAR: Dependent sampling of u and v can preserve REBAR’s expected value because its estimator terms depend separately on u or on v and b.The construction samples u uniformly, generates z and b, and can then sample from z|b using the corresponding inverse-transformation point.
  • G.3 Concrete: The Concrete objective used by Maddison et al. (2016) and Jang et al. (2016) is not a stochastic lower bound on marginal log-likelihood.The authors therefore compare REBAR most directly with Jang et al. (2016), whose reported results were similar to Maddison et al. (2016).
  • G.3 Concrete: Unlike the Concrete objective, the authors emphasize that their evaluation used a proper stochastic lower bound.This distinction motivates how the REBAR results are characterized relative to the Concrete comparison.
Loading 1703.07370v4…