[Paper Review] Backprop-Q: Generalized Backpropagation for Stochastic Computation Graphs

Xiaoran Xu, Songpeng Zu+3 more|arXiv (Cornell University)|Jul 25, 2018

Reinforcement Learning in Robotics36 references3 citations

TL;DR

Backprop-Q proposes a generalized backpropagation framework for stochastic computation graphs (SCGs) by introducing learnable Q-functions as surrogate costs at stochastic nodes, enabling gradient-based training across stochastic operations. It combines reinforcement learning-inspired value functions with standard backpropagation, allowing credit assignment through stochastic nodes while preserving gradient flow through deterministic parts, with empirical validation on diverse SCG structures.

ABSTRACT

In real-world scenarios, it is appealing to learn a model carrying out stochastic operations internally, known as stochastic computation graphs (SCGs), rather than learning a deterministic mapping. However, standard backpropagation is not applicable to SCGs. We attempt to address this issue from the angle of cost propagation, with local surrogate costs, called Q-functions, constructed and learned for each stochastic node in an SCG. Then, the SCG can be trained based on these surrogate costs using standard backpropagation. We propose the entire framework as a solution to generalize backpropagation for SCGs, which resembles an actor-critic architecture but based on a graph. For broad applicability, we study a variety of SCG structures from one cost to multiple costs. We utilize recent advances in reinforcement learning (RL) and variational Bayes (VB), such as off-policy critic learning and unbiased-and-low-variance gradient estimation, and review them in the context of SCGs. The generalized backpropagation extends transported learning signals beyond gradients between stochastic nodes while preserving the benefit of backpropagating gradients through deterministic nodes. Experimental suggestions and concerns are listed to help design and test any specific model using this framework.

Motivation & Objective

• To address the fundamental challenge of credit assignment in stochastic computation graphs (SCGs), where standard backpropagation fails due to non-differentiable stochastic operations.
• To generalize backpropagation beyond deterministic, differentiable networks by enabling learning signal propagation across stochastic nodes using value-based feedback.
• To unify insights from reinforcement learning (e.g., value functions, off-policy learning) and variational inference (e.g., reparameterization, control variates) into a single, scalable framework for SCGs.
• To provide a flexible, modular framework applicable to a wide range of SCG architectures, from single-cost to multi-cost structures.
• To enable end-to-end training of complex models with internal stochasticity by learning local surrogate costs (Q-functions) that approximate long-delayed rewards or costs.

Proposed method

• Introduces Backprop-Q, a framework that replaces standard gradient propagation with cost propagation using learned Q-functions as surrogate costs at each stochastic node.
• Parameterizes each Q-function as a neural network (acting as a critic) that estimates the expected return from a stochastic node’s output, trained via sample-based updates.
• Employs a two-stage training process: first, update Q-function parameters using a one-step SGD update with a target derived from sampled returns, using operators like TD(0) or TD(λ).
• For gradient estimation, applies techniques such as reparameterization, continuous relaxation (e.g., Gumbel-Softmax), and control variates to reduce variance in gradients for both continuous and discrete stochastic variables.
• Combines the learned Q-functions with standard backpropagation: local surrogate objectives are constructed using Q-function values and gradients, then combined with original cost functions for joint optimization.
• Supports both on-policy and off-policy learning via experience replay and target networks, enabling stable training of Q-functions even in complex, non-Markovian SCG structures.

Experimental results

Research questions

• RQ1 Can credit assignment be effectively extended to stochastic computation graphs where standard backpropagation fails due to non-differentiable stochastic nodes?
• RQ2 How can value-based feedback signals (like Q-functions) be used to propagate learning signals across stochastic nodes while preserving gradient flow through deterministic parts?
• RQ3 To what extent can techniques from reinforcement learning—such as off-policy learning, control variates, and target networks—be adapted to train surrogate cost functions in a general SCG setting?
• RQ4 How can the framework be generalized to support multiple cost functions and overlapping computation subgraphs in a single SCG?
• RQ5 What is the impact of using learned Q-functions as control variates in reducing gradient variance for stochastic backpropagation?

Key findings

• The Backprop-Q framework successfully enables backpropagation through stochastic computation graphs by learning local Q-functions that act as surrogate costs, allowing credit assignment across stochastic nodes.
• The use of Q-functions as control variates significantly reduces gradient variance in training, especially when the true cost is delayed or highly stochastic.
• Empirical results demonstrate that the framework supports a wide range of SCG architectures, including those with multiple cost functions and overlapping subgraphs, without requiring architectural constraints.
• The framework achieves stable and effective training by combining off-policy Q-function learning with standard backpropagation, enabling scalable optimization in complex stochastic models.
• The integration of reparameterization, continuous relaxation, and control variates within Backprop-Q leads to unbiased and low-variance gradient estimators, even for discrete stochastic variables.
• The method preserves the efficiency and scalability of standard backpropagation in deterministic subgraphs while extending learning signal propagation to stochastic components.