Credit Assignment Techniques in Stochastic Computation Graphs
This work addresses efficiency and scalability issues in gradient-based optimization for AI applications, representing an incremental improvement over existing methods.
The paper tackles the problem of high variance and computational cost in gradient estimation for stochastic computation graphs by introducing value functions, baselines, and critics, enabling lower-variance estimates from partial evaluations.
Stochastic computation graphs (SCGs) provide a formalism to represent structured optimization problems arising in artificial intelligence, including supervised, unsupervised, and reinforcement learning. Previous work has shown that an unbiased estimator of the gradient of the expected loss of SCGs can be derived from a single principle. However, this estimator often has high variance and requires a full model evaluation per data point, making this algorithm costly in large graphs. In this work, we address these problems by generalizing concepts from the reinforcement learning literature. We introduce the concepts of value functions, baselines and critics for arbitrary SCGs, and show how to use them to derive lower-variance gradient estimates from partial model evaluations, paving the way towards general and efficient credit assignment for gradient-based optimization. In doing so, we demonstrate how our results unify recent advances in the probabilistic inference and reinforcement learning literature.