The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
Chris J. Maddison
and
Andriy Mnih
and
Yee Whye Teh
arXiv e-Print archive - 2016 via Local arXiv
Keywords:
cs.LG, stat.ML
First published: 2016/11/02 (8 years ago) Abstract: The reparameterization trick enables the optimization of large scale
stochastic computation graphs via gradient descent. The essence of the trick is
to refactor each stochastic node into a differentiable function of its
parameters and a random variable with fixed distribution. After refactoring,
the gradients of the loss propagated by the chain rule through the graph are
low variance unbiased estimators of the gradients of the expected loss. While
many continuous random variables have such reparameterizations, discrete random
variables lack continuous reparameterizations due to the discontinuous nature
of discrete states. In this work we introduce concrete random variables --
continuous relaxations of discrete random variables. The concrete distribution
is a new family of distributions with closed form densities and a simple
reparameterization. Whenever a discrete stochastic node of a computation graph
can be refactored into a one-hot bit representation that is treated
continuously, concrete stochastic nodes can be used with automatic
differentiation to produce low-variance biased gradients of objectives
(including objectives that depend on the log-likelihood of latent stochastic
nodes) on the corresponding discrete graph. We demonstrate their effectiveness
on density estimation and structured prediction tasks using neural networks.