[Paper Review] REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models

George Tucker, Andriy Mnih+3 more|arXiv (Cornell University)|Mar 21, 2017

Generative Adversarial Networks and Image SynthesisComputer Science167 citations

TL;DR

REBAR introduces an unbiased, low-variance gradient estimator for discrete latent variable models by using a control variate based on the gap between REINFORCE and reparameterization gradients of a relaxed Concrete distribution, with online temperature tuning.

ABSTRACT

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.

Motivation & Objective

• Motivate and address the high-variance problem in gradient estimation for discrete latent variables.
• Develop a low-variance, unbiased gradient estimator by combining REINFORCE with a reparameterization-based control variate.
• Introduce a relaxed continuous approximation (Concrete) and a conditional marginalization to improve variance reduction.
• Enable online adaptation of the relaxation temperature to remove an extra hyperparameter.
• Demonstrate state-of-the-art variance reduction on generative modeling and structured prediction tasks.

Proposed method

• Define a control variate based on the difference between REINFORCE and reparameterization gradients for the relaxed model.
• Conditionally marginalize the control variate using p(z|b) to create REBAR, then combine with a corrected reparameterization gradient.
• Couple random variables (u, v) to reduce computation and variance.
• Derive an unbiased gradient estimator that includes a temperature parameter lambda governing the relaxation.
• Optimize the temperature lambda online by minimizing the estimator variance.
• Discuss connections to MuProp in the high-temperature limit and extend to multilayer stochastic networks.

Experimental results

Research questions

• RQ1 Can an unbiased gradient estimator with low variance be constructed for discrete latent variable models without hyperparameter tuning of the relaxation?
• RQ2 Does a control variate based on the relaxation (Concrete distribution) substantially reduce variance compared to existing unbiased estimators?
• RQ3 Can online adaptation of the relaxation temperature further reduce gradient variance without introducing bias?
• RQ4 How does REBAR perform on generative and structured prediction tasks relative to NVIL, MuProp, and Gumbel-Softmax/Concrete?
• RQ5 What are the theoretical and practical relationships between REBAR, MuProp, and SimpleMuProp in various temperature regimes?

Key findings

• REBAR achieves state-of-the-art variance reduction among unbiased gradient estimators across tested tasks.
• In MNIST and Omniglot generative modeling, REBAR with online lambda adaptation consistently reduces gradient variance and improves convergence and final log-likelihood.
• REBAR outperforms the Concrete (Gumbel-Softmax) estimator in several setups, especially for linear models, while matching or exceeding performance in nonlinear settings.
• A modified relaxation reveals a connection between REBAR and MuProp in the high-temperature limit, termed SimpleMuProp.
• Online optimization of the relaxation temperature effectively reduces variance without introducing bias, alleviating the need to pre-specify lambda.