Hindsight Network Credit Assignment: Efficient Credit Assignment in Networks of Discrete Stochastic Units

TL;DR

This paper introduces Hindsight Network Credit Assignment (HNCA), a novel gradient estimation method for discrete stochastic neural networks that reduces variance and improves training efficiency compared to traditional methods like REINFORCE.

Contribution

The paper presents HNCA, a new unbiased gradient estimator for networks of discrete stochastic units with lower variance and similar computational cost to backpropagation.

Findings

HNCA outperforms REINFORCE in contextual bandit tasks.

HNCA effectively trains discrete variational auto-encoders.

Theoretical analysis confirms variance reduction of HNCA.

Abstract

Training neural networks with discrete stochastic variables presents a unique challenge. Backpropagation is not directly applicable, nor are the reparameterization tricks used in networks with continuous stochastic variables. To address this challenge, we present Hindsight Network Credit Assignment (HNCA), a novel gradient estimation algorithm for networks of discrete stochastic units. HNCA works by assigning credit to each unit based on the degree to which its output influences its immediate children in the network. We prove that HNCA produces unbiased gradient estimates with reduced variance compared to the REINFORCE estimator, while the computational cost is similar to that of backpropagation. We first apply HNCA in a contextual bandit setting to optimize a reward function that is unknown to the agent. In this setting, we empirically demonstrate that HNCA significantly outperforms REINFORCE, indicating that the variance reduction implied by our theoretical analysis is significant and impactful. We then show how HNCA can be extended to optimize a more general function of the outputs of a network of stochastic units, where the function is known to the agent. We apply this extended version of HNCA to train a discrete variational auto-encoder and empirically show it compares favourably to other strong methods. We believe that the ideas underlying HNCA can help stimulate new ways of thinking about efficient credit assignment in stochastic compute graphs.