Source-linked AI summary
MuProp: Unbiased Backpropagation for Stochastic Neural Networks
Shixiang Gu, Sergey Levine, Ilya Sutskever, Andriy Mnih
TL;DR
Discrete latent-variable models are difficult to train because backpropagation does not directly handle discrete sampling, while unbiased likelihood-ratio estimators often have high variance. MuProp combines mean-field backpropagation with a likelihood-ratio residual whose control variate is a first-order Taylor expansion, preserving unbiasedness. MuProp outperforms competing estimators in generative-model experiments and converges about 3–4 times faster than NVIL for the categorical model.
Problem
Discrete latent-variable models are difficult to train because backpropagation does not directly handle discrete sampling, while unbiased likelihood-ratio estimators often have high variance.
Method
MuProp combines mean-field backpropagation with a likelihood-ratio residual whose control variate is a first-order Taylor expansion, preserving unbiasedness.
Results
MuProp outperforms competing estimators in generative-model experiments and converges about 3–4 times faster than NVIL for the categorical model.
Takeaways & Limitations
MuProp combines the statistical efficiency of backpropagation with the correctness of likelihood-ratio estimation for stochastic computational graphs.
Takeaways & Limitations
The paper identifies fitted-function baselines and rollout procedures as extensions that could reduce variance or improve deep-network performance, but would add complexity or computational cost.
Abstract
from arXiv · showhide
Deep neural networks are powerful parametric models that can be trained efficiently using the backpropagation algorithm. Stochastic neural networks combine the power of large parametric functions with that of graphical models, which makes it possible to learn very complex distributions. However, as backpropagation is not directly applicable to stochastic networks that include discrete sampling operations within their computational graph, training such networks remains difficult. We present MuProp, an unbiased gradient estimator for stochastic networks, designed to make this task easier. MuProp improves on the likelihood-ratio estimator by reducing its variance using a control variate based on the first-order Taylor expansion of a mean-field network. Crucially, unlike prior attempts at using backpropagation for training stochastic networks, the resulting estimator is unbiased and well behaved. Our experiments on structured output prediction and discrete latent variable modeling demonstrate that MuProp yields consistently good performance across a range of difficult tasks.
1 INTRODUCTION
Stochastic neural networks are difficult to train when discrete variables prevent direct backpropagation, despite their expressive power. MuProp addresses this gap by combining backpropagation efficiency with unbiased-estimator guarantees.
- Motivation: Discrete latent-variable models are harder to train than continuous-variable models because reparameterization-based backpropagation does not directly apply.Likelihood-ratio estimators are unbiased but often less effective than biased alternatives.
- Motivation: An effective training method should exploit backpropagation while retaining the convergence guarantees of an unbiased estimator.
- Contribution: MuProp is an unbiased, backpropagation-based gradient estimator for deep stochastic networks with arbitrary mixtures of stochastic and deterministic nodes.It handles continuous and discrete variables, though the experiments focus on discrete Bernoulli and multinomial models.
- Evaluation: Experiments evaluate MuProp on supervised and unsupervised tasks involving discrete stochastic neural networks that are difficult to train.The introduction frames these settings as cases where biased methods often outperform unbiased ones.
2 RELATED WORK
Prior work provides expressive stochastic models and unbiased likelihood-ratio gradients, but practical training is limited by estimator variance. Control variates and baseline techniques reduce this variance while preserving unbiasedness.
- Related work: Continuous latent-variable models can use reparameterization and backpropagation, whereas discrete latent-variable training remains considerably more difficult.Earlier probabilistic neural models also faced scalability or objective-definition limitations.
- Likelihood-ratio estimation: Likelihood-ratio estimation is broadly applicable and underlies the unbiased estimators discussed in the paper, including MuProp.It requires only that pθ(x) be differentiable with respect to θ.
- Variance reduction: High likelihood-ratio variance can make convergence slow or impossible, especially with minibatch training that typically uses m = 1.Increasing the sample count reduces variance but is computationally expensive.
- Variance reduction: Control variates reduce Monte Carlo variance by subtracting an analytically tractable term and adding back its expectation to preserve an unbiased estimator.
- Variance reduction: The paper considers centering, input-dependent baselines, and variance normalization as practical likelihood-ratio variance-reduction techniques.Variance normalization adaptively divides the learning signal by max(1, √v) and is not itself a baseline.
3 MUPROP
MuProp combines mean-field backpropagation with a likelihood-ratio residual based on a first-order Taylor control variate. This construction extends unbiased gradient estimation to arbitrary stochastic computational graphs while remaining compatible with automatic differentiation.
- MuProp estimator: MuProp combines a deterministic mean-field backpropagation term with a likelihood-ratio term that accounts for residuals.The resulting estimator is unbiased.
- Extensions: Rollout extends the mean-field construction by conditioning later mean-field passes on sampled preceding-layer values, potentially improving deep-network performance at additional O(n^2) partial-pass cost.
- Taylor control variate: Its control variate is the first-order Taylor expansion of f around a fixed value ¯x, subtracting f(¯x) and f′(¯x)(x − ¯x) from the sampled cost.
- Generality and implementation: MuProp extends Taylor-expansion baselines to arbitrary stochastic computational graphs and provides an automatic-differentiation algorithm for integrating the estimator into existing libraries.
- Mean-field network: Because discrete sampling is not differentiable, MuProp computes the Taylor expansion through a deterministic mean-field network that propagates means instead of samples.
4 COMPARISON WITH OTHER GRADIENT ESTIMATORS
The comparison contrasts MuProp with biased estimators for stochastic binary networks, emphasizing that only likelihood-ratio and MuProp have principled deep-network extensions. On MNIST imputation, MuProp is the only unbiased estimator that competes with biased methods such as ST and 1/2.
- Only LR and MuProp have principled extensions to deep networks; other methods require heuristic extensions.
- 4.1 THE STRAIGHT-THROUGH ESTIMATOR: The Straight-Through estimator is biased but low-variance, using backpropagation through binary thresholding as if it were the identity.
- 4.1 THE STRAIGHT-THROUGH ESTIMATOR: MuProp significantly outperforms Straight-Through estimators on certain tasks, suggesting that unbiased estimators are more reliable.
- 4.2 THE 1/2 ESTIMATOR: The 1/2 estimator is a biased method specialized for stochastic binary networks and uses a Taylor-expansion baseline around a fixed point.
- 4.2 THE 1/2 ESTIMATOR: For multinomial variables, the 1/2 estimator has no established sensible expansion point, so the experiments test 1/2, 1/k, and the mean µx(θ).
- MuProp is the only unbiased estimator that can compete with biased estimators such as ST and 1/2 on MNIST imputation.
5 EXPERIMENTS
The experiments compare MuProp with likelihood-ratio and biased estimators across structured prediction and discrete latent-variable tasks. MuProp generally converges faster than likelihood-ratio methods and achieves competitive or superior final performance, while biased estimators can be inconsistent.
- 5 EXPERIMENTS: The experiments evaluate LR, ST, 1/2, and MuProp across diverse architectures, including direct expected-objective optimization and variational training with an inference network.
- 5.1 STRUCTURED OUTPUT PREDICTION: MuProp is slightly worse than ST and 1/2 on the reported structured-prediction task, partly because its single-trunk mean-field pass poorly represents higher-layer conditional distributions.Extensions address this limitation through additional computation or an auxiliary network.
- 5.1 STRUCTURED OUTPUT PREDICTION: MuProp significantly outperforms LR and closely matches ST and 1/2 in convergence speed and final accuracy on structured output prediction.
- 5.2 VARIATIONAL TRAINING OF GENERATIVE MODELS: MuProp outperforms competing estimators on variational training, while variance-reduced LR reaches similar final bounds but converges about 3–4 times slower on the categorical model.
- 5.2 VARIATIONAL TRAINING OF GENERATIVE MODELS: ST and 1/2 have significantly worse final variational lower-bounds than MuProp or LR and vary inconsistently across SBN and categorical architectures.
6 DISCUSSION
The discussion presents MuProp as an unbiased derivative estimator combining backpropagation’s efficiency with likelihood-ratio correctness. It identifies extensions that may improve deep-network performance but require richer baselines, extra computation, or auxiliary training.
- 6 DISCUSSION: MuProp combines the statistical efficiency of backpropagation with the correctness of a likelihood-ratio method.
- 6 DISCUSSION: Learned baselines resembling fitted Q-functions could reduce variance and improve MuProp for very deep graphs, recurrent networks, and reinforcement learning.The proposed extension is motivated by fitted Q-functions that summarize future costs for a policy.
- 6 DISCUSSION: Sample-dependent mean-field rollouts could improve deeper-network performance by reducing divergence between deterministic mean values and fully stochastic samples, at additional computational cost.
7 RECURSIVE DERIVATION OF MUPROP
The recursive derivation defines MuProp for layered discrete stochastic models using mean functions and first-order Taylor expansions around fixed points. The resulting estimator preserves a backpropagation-like recursive structure and extends naturally to directed acyclic graphs.
- 7 RECURSIVE DERIVATION OF MUPROP: The loss is an expectation over layered discrete random variables, with the model distribution factorized across conditional probabilities.
- 7 RECURSIVE DERIVATION OF MUPROP: MuProp assumes differentiable conditional mean functions and recursively expands downstream functions around fixed points that do not depend on the corresponding sampled variables.
- 7 RECURSIVE DERIVATION OF MUPROP: The final estimator resembles classical backpropagation and can be computed efficiently through forward and backward passes.
- 7 RECURSIVE DERIVATION OF MUPROP: The recursive gradient terms combine likelihood-ratio contributions with Taylor-based mean-function corrections across layers.
- 7 RECURSIVE DERIVATION OF MUPROP: Choosing fixed points from a deterministic mean-field forward pass recovers the standard MuProp estimator while avoiding direct differentiation through discrete sampling.
8 MUPROP WITH AUTOMATIC DIFFERENTIATION
MuProp gradient computation uses automatic differentiation across a mean-field network and stochastic passes, with stopped gradients around sampling operations.
- The algorithm computes MuProp gradients using automatic differentiation through a symbolic computational graph.It relies on Gradients and StopGradient functionalities supplied by the automatic differentiation library.
- A deterministic mean-field pass supplies differentiable quantities, while stochastic passes sample node values without propagating gradients through sampling.The default configuration uses one stochastic sample, with runtime equal to one deterministic pass plus m stochastic passes.
- The procedure assumes a graph with n stochastic nodes, a loss function, and stochastic-parent relationships used to construct the estimator.The loss may combine multiple loss functions at different parts of the graph.
- Algorithm 1 is the paper’s explicit procedure for computing the MuProp gradient estimator.