[Paper Review] Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation

Yoshua Bengio, Nicholas Léonard+1 more|arXiv (Cornell University)|Aug 15, 2013

Stochastic Gradient Optimization TechniquesComputer Science11 references2,001 citations

TL;DR

The authors compare four gradient-estimation strategies for stochastic or non-smooth neurons and demonstrate their use in a conditional computation setup to gate parts of a network.

ABSTRACT

Stochastic neurons and hard non-linearities can be useful for a number of reasons in deep learning models, but in many cases they pose a challenging problem: how to estimate the gradient of a loss function with respect to the input of such stochastic or non-smooth neurons? I.e., can we "back-propagate" through these stochastic neurons? We examine this question, existing approaches, and compare four families of solutions, applicable in different settings. One of them is the minimum variance unbiased gradient estimator for stochatic binary neurons (a special case of the REINFORCE algorithm). A second approach, introduced here, decomposes the operation of a binary stochastic neuron into a stochastic binary part and a smooth differentiable part, which approximates the expected effect of the pure stochatic binary neuron to first order. A third approach involves the injection of additive or multiplicative noise in a computational graph that is otherwise differentiable. A fourth approach heuristically copies the gradient with respect to the stochastic output directly as an estimator of the gradient with respect to the sigmoid argument (we call this the straight-through estimator). To explore a context where these estimators are useful, we consider a small-scale version of {\em conditional computation}, where sparse stochastic units form a distributed representation of gaters that can turn off in combinatorially many ways large chunks of the computation performed in the rest of the neural network. In this case, it is important that the gating units produce an actual 0 most of the time. The resulting sparsity can be potentially be exploited to greatly reduce the computational cost of large deep networks for which conditional computation would be useful.

Motivation & Objective

• Motivate gradient estimation through stochastic or non-smooth neurons for conditional computation.
• Review and compare four gradient estimation families: unbiased gradient estimators, stochastic-smooth decompositions, noise-injected differentiable graphs, and straight-through approaches.
• Demonstrate feasibility of training with stochastic gates to selectively activate parts of a large network.
• Assess practical performance of proposed methods on a gating/experts architecture with sparsity constraints.

Proposed method

• Formulates stochastic neurons as h_i = f(a_i, z_i) and derives gradient flow opportunities.
• Introduces four approaches: (i) unbiased gradient estimator for stochastic binary neurons (REINFORCE-like); (ii) decomposing a stochastic binary neuron into a stochastic binary part and a smooth first-order approximation; (iii) injecting noise to create a differentiable graph; (iv) straight-through estimator propagating gradients through a binary/stochastic gate.
• Proposes and analyzes the Noisy Rectifier, STS (Stochastic Times Smooth), ST (Straight-Through), and Unbiased REINFORCE-based estimators.
• Discusses variance reduction via centered estimators and unit-specific baselines for the unbiased gradient.

Experimental results

Research questions

• RQ1 Can we back-propagate gradients through stochastic or non-smooth neurons effectively?
• RQ2 What gradient estimators provide unbiased or low-variance updates for stochastic binary or gated units?
• RQ3 Do stochastic gates enable meaningful conditional computation with potential computational savings?
• RQ4 How do these estimators perform in practice on a gating/experts network using MNIST?

Key findings

• Unbiased gradient estimator for stochastic binary neurons is proven to be unbiased with respect to the gradient of the expected loss.
• The STS unit and Noisy Rectifier exhibit favorable properties and enable gradient flow in stochastic gating.
• Straight-Through estimator performs surprisingly well in practice despite bias, often yielding the best validation/test results in the experiments.
• Conditioning the gater with stochastic gates reduces computation by gating a small fraction of units (about 10%), with modest impact on performance.
• All tested estimators allow training to proceed; noise injection can improve both training objectives and generalization.
• Straight-Through units achieved the best validation and test error in the reported MNIST experiments.