PaperlayerView exact revision on arXiv

Source-linked AI summary

Backpropagation through the Void: Optimizing control variates for black-box gradient estimation

Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, David Duvenaud

arXiv:1711.00123v3cs.LG

TL;DR

The paper targets unbiased gradient estimation for objectives involving unknown, nondifferentiable, or discrete random-variable functions. It learns neural-network control variates jointly with model parameters or policies, extending the approach across discrete and continuous settings. The framework supports discrete latent-variable training and an unbiased, action-conditional extension of advantage actor-critic, with reported convergence speedups in discrete environments.

  • Problem

    Backpropagation cannot be applied to black-box objectives or discrete latent-variable models, motivating unbiased gradient estimators for these settings.

  • Method

    The paper learns neural-network surrogates as control variates, combining score-function and differentiable estimators to reduce variance while preserving unbiasedness.

  • Results

    In both discrete reinforcement-learning environments, the estimator achieved greater than a 2-times speedup in convergence over the baseline.

  • Takeaways & Limitations

    LAX and RELAX apply to known or black-box functions of discrete or continuous random variables and require only function queries, sampling, and differentiation of p(b|θ).

  • Takeaways & Limitations

    REBAR requires a known differentiable function and evaluates the discrete loss at continuous inputs; direct dependence of the optimized function on θ is also a scope boundary discussed by the paper.

Abstract

from arXiv · show

Gradient-based optimization is the foundation of deep learning and reinforcement learning. Even when the mechanism being optimized is unknown or not differentiable, optimization using high-variance or biased gradient estimates is still often the best strategy. We introduce a general framework for learning low-variance, unbiased gradient estimators for black-box functions of random variables. Our method uses gradients of a neural network trained jointly with model parameters or policies, and is applicable in both discrete and continuous settings. We demonstrate this framework for training discrete latent-variable models. We also give an unbiased, action-conditional extension of the advantage actor-critic reinforcement learning algorithm.

1 INTRODUCTION

The paper addresses gradient estimation when backpropagation cannot be applied to black-box objectives or discrete latent-variable models. It introduces learned neural-network control variates that produce lower-variance, unbiased estimators across discrete and continuous settings.

  • 1 INTRODUCTION: Black-box reinforcement-learning objectives and discrete latent-variable models prevent direct application of backpropagation.In reinforcement learning the optimized function is unknown, while discrete sampling creates discontinuities and zero parameter gradients.
  • 1 INTRODUCTION: The framework generalizes Tucker et al. (2017) by learning a free-form neural-network control variate.The learned surrogate broadens the range of problems beyond continuous relaxations.
  • 1 INTRODUCTION: The resulting gradient estimator is lower-variance and unbiased, including when no continuous relaxation is available.The paper highlights reinforcement learning and black-box function optimization as such cases.

2 BACKGROUND: GRADIENT ESTIMATORS

The section frames unbiased stochastic gradients as necessary for scalable optimization and contrasts generally applicable but high-variance score-function estimators with low-variance reparameterization methods. Control variates preserve expectation while reducing variance when positively correlated with the original estimator.

  • 2 BACKGROUND: GRADIENT ESTIMATORS: The optimization problem includes choosing distribution parameters to maximize expectations in reinforcement learning and latent-variable models.Reinforcement learning maximizes expected trajectory reward, while latent-variable models optimize marginal likelihood expectations.
  • 2 BACKGROUND: GRADIENT ESTIMATORS: Stochastic optimization is scalable for high-dimensional parameters but is guaranteed to converge to a fixed point only when stochastic gradients are unbiased.Unbiasedness means E[ĝ] = ∂L(θ)/∂θ.
  • 2 BACKGROUND: GRADIENT ESTIMATORS: The score-function estimator, or REINFORCE, is broadly applicable but generally has high variance.It uses the final outcome f(b) without information about how f depends on b.
  • 2 BACKGROUND: GRADIENT ESTIMATORS: The reparameterization trick gives low-variance, unbiased gradients when f and the transformation from fixed noise to latent variables are differentiable.This approach exposes dependence on b through ∂f/∂b and is used for continuous latent-variable models.
  • Control variates: A control variate subtracts a function with known mean from an estimator and adds that mean back.The expectation remains unchanged, while variance decreases when the control variate is positively correlated with the estimator.

3 CONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE

The paper constructs LAX and RELAX by combining score-function gradients, reparameterization, and learned neural-network control variates, then jointly optimizes model parameters and control-variate parameters to reduce variance. The framework extends to discrete variables and reinforcement learning, where it yields an unbiased action-dependent baseline under stated assumptions.

  • 3 CONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE: LAX replaces direct differentiation of a nondifferentiable function with gradients of a neural-network surrogate while preserving unbiasedness.The estimator subtracts the surrogate’s score-function estimator and adds its reparameterization estimator.
  • 3 CONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE: LAX is unbiased for any surrogate, and when cφ = f it reduces to the reparameterization estimator, so its variance can be no higher.The surrogate is trained jointly with the model or policy to find a low-variance estimator.
  • 3 CONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE: The surrogate is optimized by estimating the gradient of the estimator’s variance, including with an unbiased single-sample estimator, rather than minimizing only squared prediction error.The paper contrasts direct variance minimization with control-variate training based on (f(b) − cφ(b))^2.
  • DISCRETE VARIABLES: For discrete variables, RELAX introduces a continuous reparameterizable relaxation and evaluates the control variate at both an unconditional and a variable-conditioned relaxed input.The conditional distribution p(z|b, θ) must also be reparameterizable.
  • SURROGATE ARCHITECTURE: The surrogate architecture can exploit known structure: with a differentiable discrete objective, a concrete relaxation recovers REBAR, while unknown functions can use a generic neural network.A learned component may also be added to the concrete relaxation.
  • REINFORCEMENT LEARNING: In reinforcement learning, LAX uses a differentiable surrogate of both action and state, producing an action-dependent baseline that remains unbiased under Markov dynamics.The policy must be reparameterizable, and the variance-control parameters can again be optimized with a single-sample variance-gradient estimator.

4 SCOPE AND LIMITATIONS

LAX and RELAX extend control-variate gradient estimation to black-box functions, requiring only function queries and access to sampling and differentiation of p(b|θ). Their unbiased-estimator framework assumes no direct parameter dependence in f, but known differentiable dependence can be corrected explicitly.

  • Scope: LAX and RELAX can optimize black-box functions because they construct a neural-network surrogate from scratch.This applies when the objective can be queried, even if it is unknown or nondifferentiable.
  • Scope: The black-box setting requires querying the optimized function and sampling from and differentiating p(b|θ).
  • Limitations: The basic formulation assumes f has no direct dependence on θ, although this dependence can arise in probabilistic models or with regularization.
  • Limitations: When direct dependence on θ is known and differentiable, adding ∂f(b, θ)/∂θ to the estimator restores unbiasedness.

5 RELATED WORK

Related work includes model-based control variates, estimators specialized to discrete latent-variable objectives, sampling-distribution methods, and action-dependent policy-gradient baselines.

  • Control variates: Miller et al. reduce reparameterization-gradient variance by approximating the gradient-generating procedure and using the approximation as a control variate.
  • Discrete latent variables: NVIL and VIMCO provide reduced-variance estimators for discrete latent-variable models, with VIMCO targeting Monte Carlo objectives.
  • Sampling and integration: Staines and Barber, Wierstra et al., and Oates et al. study sampling distributions or non-parametric control variates for black-box optimization and Monte Carlo integration.
  • Reinforcement learning: Recent reinforcement-learning work develops action-dependent baselines, including methods using off-policy data and methods related to LAX for continuous control.

6 APPLICATIONS

The applications evaluate RELAX on a toy loss, discrete variational autoencoders, and reinforcement-learning environments. Across these settings, learned control variates improve optimization, especially convergence speed, policy-gradient variance, and sample efficiency.

  • 6.1 TOY EXPERIMENT: RELAX outperforms REBAR on a toy loss because its free-form surrogate can match the function more closely and provide correctly directed reparameterization gradients.REBAR is constrained to a warped quadratic, whereas RELAX learns a free-form relaxation.
  • 6.2 DISCRETE VARIATIONAL AUTOENCODER: RELAX improves training performance over REBAR for discrete VAEs, with faster convergence and better validation performance in both linear models.The experiments use one- or two-layer Bernoulli latent-variable models on MNIST and Omniglot.
  • 6.2 DISCRETE VARIATIONAL AUTOENCODER: RELAX’s most notable VAE improvement is its rate of convergence, despite only modest gains in training and validation scores.Training curves compare the epochs required to match REBAR’s best validation score.
  • 6.3 REINFORCEMENT LEARNING: In all tested reinforcement-learning environments, RELAX or LAX improves performance and sample efficiency relative to the A2C baseline.The evaluation covers discrete Cart Pole and Lunar Lander tasks and the continuous-action Inverted Pendulum environment.
  • 6.3 REINFORCEMENT LEARNING: 2-times speedup in convergence is exceeded by the estimator in both discrete reinforcement-learning environments compared with the baseline.The estimator also produces drastically reduced policy-gradient variance, enabling larger learning rates while maintaining stable training.

7 CONCLUSIONS AND FUTURE WORK

The paper presents a generic low-overhead estimator for expectations involving known or black-box functions of discrete or continuous random variables, and extends it to reinforcement learning. It identifies future applications in hard attention, memory indexing, nondifferentiable continuous latent-variable models, and broader variance-reduction or optimization methods.

  • 7 CONCLUSIONS AND FUTURE WORK: The proposed generic gradient estimator handles known or black-box functions of discrete or continuous random variables with little computational overhead.The paper also derives a reinforcement-learning extension for both discrete- and continuous-action domains.
  • 7 CONCLUSIONS AND FUTURE WORK: Future applications include models with hard attention or memory indexing and continuous latent-variable models with nondifferentiable likelihoods.The paper gives a 3D rendering engine as an example of the latter setting.
  • 7 CONCLUSIONS AND FUTURE WORK: In reinforcement learning, the method could be combined with generalized advantage estimation, KFAC, or off-policy control-variate training as in Q-prop.These are proposed extensions rather than evaluated results in the supplied passage.

APPENDICES

The appendices establish unbiasedness of RELAX, provide its optimization algorithm, and describe reparameterized sampling for discrete distributions. The method requires unconditional and conditional relaxed samples linked to the discrete variable.

  • Proof: The proof decomposes RELAX into an unbiased score-function term and additional terms whose expectations cancel to zero.The cancellation follows by rewriting terms and pairing opposite-sign contributions.
  • Algorithm 2: Algorithm 2 samples unconditional and conditional relaxed inputs, estimates the model gradient and variance gradient, then updates model and control-variate parameters.The procedure uses reparameterized samplers for both relaxed inputs.
  • Discrete sampling: For discrete variables, RELAX requires a reparameterized distribution p(z|θ), a deterministic mapping H(z) producing b, and reparameterized samples from p(z|b,θ).The construction defines a joint model p(b,z|θ)=p(b|z)p(z|θ).
  • Discrete sampling: For Bernoulli variables, H(z) is an indicator threshold and z is sampled using a transformed uniform random variable.The appendix separately describes categorical sampling with H(z)=argmax(z).
  • Discrete sampling: For categorical variables, H(z)=argmax(z), and conditional sampling fixes the largest component before sampling the remaining components under truncation.The largest conditional relaxed variable is sampled independently of θ.

C DERIVATIONS OF ESTIMATORS USED IN REINFORCEMENT LEARNING

This section derives the LAX and RELAX estimators for reinforcement learning, including the discrete setting’s reparameterized sampling construction and unbiasedness claim.

  • The section introduces the derivation of the LAX estimator for continuous reinforcement-learning tasks.
  • The derivation proceeds by applying the score-function estimator and adding and subtracting the same term.
  • The discrete-control policy uses a soft-max distribution to sample actions, with z_t sampled conditionally on the state and a relaxed variable sampled conditionally on the action.
  • The RELAX estimator uses both unconditional and action-conditional relaxed samples, z_t ∼ p(z_t|s_t) and z̃_t ∼ p(z_t|a_t,s_t).
  • The passage states that the RELAX estimator formed from these reparameterized samples is unbiased.

D FURTHER RESULTS ON DISCRETE VARIATIONAL AUTOENCODERS

This section presents additional discrete variational-autoencoder results through tables comparing REBAR and RELAX validation performance and convergence speed.

  • The experiments compare the REBAR and RELAX models across datasets and model configurations.
  • Table 3 reports the highest obtained validation ELBO.
  • Table 4 reports the epochs needed to achieve REBAR’s best validation score, with “-” marking nonlinear RELAX models that scored lower than REBAR.

E EXPERIMENTAL DETAILS

The experiments specify training, architecture, and hyperparameter procedures for discrete VAE and reinforcement-learning evaluations, including a noted bias in the discrete-control sampler and a corrected continuous-RL implementation.

  • E.1 DISCRETE VAE: RELAX adds hyperparameters for control-variate depth, weight decay, and learning-rate scaling, with two-layer and four-layer 200-unit ReLU networks evaluated.
  • E.1 DISCRETE VAE: The RELAX search was restricted to the best REBAR learning rate and the next-largest tested rate, and RELAX sometimes converged at rates where REBAR diverged.
  • E.1 DISCRETE VAE: The VAE experiments include one- and two-layer linear models, a one-layer nonlinear model, and validation-ELBO selection for reported results.
  • E.2 DISCRETE RL / E.3 CONTINUOUS RL: Policy-gradient experiments used single Monte Carlo episodes, RMSProp, a .99 discount factor, entropy regularization, grid-searched hyperparameters, and task-specific frame budgets.
  • E.2 DISCRETE RL / E.3 CONTINUOUS RL: The reported CartPole and LunarLander results used a slightly biased sampler, while continuous-RL results were generated with corrected publicly available code after an initially released-code bug.
Loading 1711.00123v3…