PaperlayerView exact revision on arXiv

Source-linked AI summary

Gradient Estimation Using Stochastic Computation Graphs

John Schulman, Nicolas Heess, Theophane Weber, Pieter Abbeel

arXiv:1506.05254v3cs.LG

TL;DR

Many machine-learning objectives are expectations over random variables, but their losses and gradients are often intractable and prior Monte Carlo derivations were problem-specific. The paper introduces stochastic computation graphs and unbiased gradient estimators computed through surrogate losses and modified backpropagation, providing a general framework that also encompasses prior variance-reduction methods.

  • Problem

    Expected-loss gradients in probabilistic modeling and reinforcement learning are often intractable, while prior Monte Carlo estimators did not address the general case.

  • Method

    The paper represents mixed stochastic and deterministic computations as stochastic computation graphs and derives unbiased estimators computed through surrogate losses or modified backpropagation.

  • Results

    The framework automatically derives unbiased gradient estimators and generalizes estimators and variance-reduction techniques from prior work.

  • Takeaways & Limitations

    Gradient estimation for intricate models combining stochastic and deterministic operations can be implemented using automatic differentiation software.

  • Takeaways & Limitations

    The pathwise derivative estimator requires f(x(z, θ)) to be continuous in θ for every z, so it is invalid for a step function.

Abstract

from arXiv · show

In a variety of problems originating in supervised, unsupervised, and reinforcement learning, the loss function is defined by an expectation over a collection of random variables, which might be part of a probabilistic model or the external world. Estimating the gradient of this loss function, using samples, lies at the core of gradient-based learning algorithms for these problems. We introduce the formalism of stochastic computation graphs---directed acyclic graphs that include both deterministic functions and conditional probability distributions---and describe how to easily and automatically derive an unbiased estimator of the loss function's gradient. The resulting algorithm for computing the gradient estimator is a simple modification of the standard backpropagation algorithm. The generic scheme we propose unifies estimators derived in variety of prior work, along with variance-reduction techniques therein. It could assist researchers in developing intricate models involving a combination of stochastic and deterministic operations, enabling, for example, attention, memory, and control actions.

1 Introduction

The paper addresses expected losses involving random variables, where prior gradient estimators were problem-specific and the general case lacked treatment. It introduces stochastic computation graphs and an unbiased estimator computable through a modified backpropagation procedure.

  • Expected losses arise in latent-variable likelihood maximization, policy gradients, attention, and memory models combining stochastic and deterministic operations.
  • These losses and gradients are often intractable because they require summing over exponentially many latent configurations or evaluating high-dimensional integrals.
  • The framework generalizes variance-reduction techniques and other optimization methods from reinforcement learning and variational inference.
  • The paper introduces stochastic computation graphs and derives unbiased estimators for gradients of expected losses in this general setting.
  • The estimator is computed efficiently as the gradient of a differentiable surrogate loss using automatic differentiation or modified backpropagation.

2 Preliminaries

The preliminaries distinguish score-function and pathwise-derivative estimators, then represent mixed stochastic and deterministic computations as stochastic computation graphs. The graph structure determines which estimator terms apply and supports models with multiple inputs.

  • 2.1 Gradient Estimators for a Single Random Variable: The pathwise-derivative estimator differentiates through a deterministic transformation x(z, θ), but requires continuity in θ for f(x(z, θ)).
  • 2.1 Gradient Estimators for a Single Random Variable: The score-function estimator applies under more permissive conditions, including discontinuous functions or discrete random variables.
  • 2.1 Gradient Estimators for a Single Random Variable: Score-function estimation uses sampled function values, whereas pathwise differentiation requires derivatives of f and is usually preferable for high-dimensional x when applicable.
  • 2.2 Stochastic Computation Graphs: A stochastic computation graph is a directed acyclic graph containing externally set inputs, deterministic functions, and conditionally distributed stochastic nodes.
  • 2.2 Stochastic Computation Graphs: Graph structure determines whether score-function, pathwise-derivative, or combined estimators are used, with circles denoting stochastic nodes and squares deterministic nodes.
  • 2.3 Simple Examples: The formalism accommodates multiple inputs that can influence different subsets of stochastic and cost nodes, including multilayer stochastic neural networks.

3 Main Results on Stochastic Computation Graphs

The paper formalizes stochastic computation graphs and derives an unbiased gradient estimator combining score-function and pathwise terms. A surrogate objective lets automatic differentiation and backpropagation compute this estimator efficiently under stated differentiability conditions.

  • Gradient estimator: The gradient estimator combines a score-function term for distributional influence with a pathwise derivative term for differentiable dependencies.The score-function component weights log-probability gradients by influenced downstream costs, while the pathwise term propagates dependence through differentiable functions.
  • Graph formalism: Stochastic computation graphs represent computations containing input, deterministic, stochastic, and scalar cost nodes.The framework also defines influence, deterministic influence, dependencies, and sampled node values.
  • Assumptions and scope: The framework permits some graph functions or likelihoods to be nondifferentiable when stochastic nodes block the relevant path.This scope is particularly relevant to reinforcement learning with discontinuous dynamics or reward functions.
  • Gradient estimator: Theorem 1 establishes equivalent gradient-estimator expressions for expected sums of costs under the differentiability condition.The condition requires appropriate Jacobians for deterministic nodes and probability-mass-function derivatives for stochastic nodes along relevant influenced paths.
  • Surrogate loss functions: The surrogate objective is constructed so that differentiating it produces an unbiased gradient estimate.Alternative surrogate definitions can yield the same gradient, and one construction becomes an upper bound when all costs are negative and not deterministically influenced by parameters.
  • Surrogate loss functions: Applying standard automatic differentiation to the surrogate objective converts the stochastic graph into a deterministic graph for backpropagation.This provides a practical implementation route using automatic differentiation software.
  • Extensions: The gradient estimator itself forms a stochastic computation graph, allowing gradient-vector products for Hessian-vector estimation.Dense Hessian computation is generally inefficient, whereas Hessian-vector products are usually only modestly more expensive than gradients.

4 Variance Reduction

Variance reduction is incorporated by subtracting baselines from score-function estimands without changing the expectation. In the general graph setting, baselines can be attached to stochastic nodes using variables they do not influence.

  • Baselines: Subtracting a constant baseline from the score-function integrand leaves the expectation unchanged.This preserves unbiasedness while enabling variance reduction.
  • Baselines: Choosing b = Ex[f(x)] generally produces substantial variance reduction.The baseline is often called a baseline in the paper’s terminology.
  • Graph-wide variance reduction: In a stochastic computation graph, each stochastic node can receive a baseline depending on nodes it does not influence.The paper denotes this set as NONINFLUENCED(v).

5 Algorithms

The algorithm computes the stochastic-graph gradient estimator through a reverse traversal, using a modified backpropagation procedure with additional signals at stochastic nodes.

  • Backward computation: The practical algorithm differentiates a surrogate objective during a backward pass through the stochastic computation graph.Automatic differentiation software can facilitate this implementation.
  • Backward computation: Algorithm 1 recursively computes gv for each deterministic and input node during reverse traversal.The traversal processes non-input nodes in reverse topological order.

6 Related Work

Prior work developed score-function, pathwise, policy-gradient, reparameterization, and stochastic-unit estimators for particular settings. This paper positions its stochastic computation graph framework as a general treatment spanning those structures.

  • Gradient estimation: Gradient estimation has established roots in simulation optimization, Monte Carlo estimation, computation graphs, and automatic differentiation.The paper cites reviews and textbooks covering these areas.
  • Prior estimators: Earlier work studied stochastic units and continuous latent-variable reparameterization, but focused on particular network structures or latent-variable settings.The paper distinguishes these works from its treatment of arbitrarily structured models and loss functions.
  • Prior estimators: The score-function method was applied to variational inference and generative models with discrete stochastic units, including amortized inference networks.These applications commonly used specific approximation or model structures.

7 Conclusion

The framework represents computations with stochastic and deterministic operations and automatically produces gradient estimators under appropriate differentiability conditions. These estimators can be computed efficiently by backward traversal and recover prior estimators in reinforcement learning and probabilistic modeling.

  • Stochastic computation graphs describe computations containing both stochastic and deterministic operations.
  • The framework automatically obtains a gradient estimator when node functions satisfy the required differentiability conditions.
  • Backward traversal computes the gradient using either a surrogate loss with standard backpropagation or a roughly equivalent modified procedure.
  • The framework automatically reproduces gradient estimators previously derived in reinforcement learning and probabilistic modeling.

A Proofs

The proofs establish the gradient identities under continuous-variable and differentiability assumptions. They derive the result for individual costs and extend it to the full expected sum through summation and rearrangement.

  • The proof treats continuous random variables, expressing expectations as integrals; discrete variables require replacing integrals with sums.
  • The gradient derivation first differentiates the expectation of a single cost term, then sums over cost nodes to obtain the full result.
  • The proof requires differentiable integrands, ensured when the probability densities and cost functions are differentiable.
  • The argument isolates a node whose associated expectation vanishes and orders uninfluenced nodes before that node in the graph.

B Surrogate as an Upper Bound, and MM Algorithms

The surrogate objective has a role beyond gradient estimation: under stated restrictions and conditions, it upper-bounds the true objective up to a constant. This supports majorization-minimization optimization such as generalized EM.

  • Under certain conditions, the surrogate objective is an upper bound on the true objective plus a constant.
  • The upper-bound analysis restricts costs to be negative and requires that costs are not deterministically influenced by the parameters.
  • The resulting relation permits majorization-minimization algorithms, including the EM algorithm, to optimize with respect to parameters.

C Examples

The stochastic computation graph formalism applies to generalized EM for latent-variable probabilistic models and to reinforcement learning in Markov Decision Processes. In both settings, expected objectives become compositions of stochastic and deterministic steps.

  • The examples cover generalized EM for maximum likelihood estimation in probabilistic models with latent variables and reinforcement learning in Markov Decision Processes.
  • In both settings, the objective is an expectation represented as a composition of stochastic and deterministic operations.

C.1 Generalized EM Algorithm and Variational Inference.

The paper applies stochastic-computation-graph gradient estimators to variational inference and reinforcement learning, including latent-variable models, MDPs, and POMDPs. These applications use sampled variables, baselines, and reparameterization to obtain unbiased gradients without requiring analytic environment dynamics.

  • Generalized EM and variational inference: Generalized EM maximizes a variational lower bound defined as an expectation over latent variables for each observed sample.The model uses observed x, latent z, and parameter θ, while the inference network provides a parameterized approximation to the posterior.
  • Generalized EM and variational inference: The inference model samples hierarchical latent variables h1, h2, and h3 conditionally, forming a stochastic computation graph for multilayer variational inference.The sampled chain is h1 ∼ qφ1(· | x), h2 ∼ qφ2(· | h1), and h3 ∼ qφ3(· | h2).
  • Generalized EM and variational inference: Given a sample h ∼ qφ, Theorem 2 provides an unbiased gradient estimate for the variational objective.The estimator uses cumulative downstream cost terms Q̂1 = r1 + r2 + r3, Q̂2 = r2 + r3, and Q̂3 = r3, together with baseline functions.
  • Generalized EM and variational inference: For conditionally Gaussian qφ, reparameterization expresses h as μφ(x) + ϵσφ(x), while the entropy term can be computed analytically.This provides a continuous-latent-variable route to gradient estimation and can yield a different estimator for the same observed-variable distribution.
  • Reinforcement learning: In reinforcement learning, stochastic computation graphs recover policy-gradient estimators for MDPs and POMDPs using sampled trajectories and observation histories.For MDPs, trajectories depend on environment dynamics and policy probabilities; for POMDPs, policies depend on histories of past observations.
  • Reinforcement learning: Baselines can reduce variance in policy-gradient estimators, with bt(st) often chosen to approximate the state-value function Vt(st).In model-free control, only sampled trajectories are required, and action nodes block differentiation through unknown environment dynamics.
Loading 1506.05254v3…