[Paper Review] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables

Chris J. Maddison, Andriy Mnih+1 more|arXiv (Cornell University)|Nov 2, 2016

Geotechnical Engineering and Analysis1,036 citations

TL;DR

The paper introduces Concrete random variables, a continuous relaxation of discrete variables on the simplex with a closed-form density and a reparameterization, enabling low-variance gradient-based optimization of stochastic computation graphs with discrete nodes.

ABSTRACT

The reparameterization trick enables optimizing 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 useful 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-probability of latent stochastic nodes) on the corresponding discrete graph. We demonstrate the effectiveness of Concrete relaxations on density estimation and structured prediction tasks using neural networks.

Motivation & Objective

• Motivate gradient-based optimization for stochastic computation graphs with discrete nodes.
• Introduce a continuous relaxation of discrete variables on the simplex (Concrete distribution).
• Show that Concrete relaxations enable low-variance biased gradients via reparameterization for discrete graphs.
• Demonstrate effectiveness on density estimation and structured prediction tasks using neural networks.

Proposed method

• Define Concrete distribution as a continuous relaxation of one-hot discrete variables on the simplex with parameters alpha and temperature lambda.
• Derive closed-form density p_alpha,lambda on the simplex and establish reparameterization via Gumbel perturbations.
• Explain how to replace discrete nodes with Concrete nodes during training and relax log-probability terms accordingly.
• Compare Concrete relaxations to score-function estimators (e.g., VIMCO/NVIL) on MNIST and Omniglot tasks.
• Discuss temperature selection and its effect on the discreteness and integrality gap of the relaxation.

Experimental results

Research questions

• RQ1 Can a continuous relaxation of discrete variables enable gradient-based optimization of stochastic graphs with discrete nodes?
• RQ2 Does the Concrete distribution provide a practical reparameterization with closed-form density for training neural networks with discrete latent variables?
• RQ3 How do Concrete relaxations perform on density estimation and structured prediction compared to state-of-the-art estimators?
• RQ4 What are the empirical effects of temperature (lambda) on discreteness and model performance?

Key findings

• Concrete relaxations yield low-variance gradients for discrete stochastic nodes in stochastic computation graphs.
• The Concrete distribution has a closed-form density and a straightforward reparameterization using Gumbel perturbations and softmax.
• Empirically, Concrete relaxations perform competitively with VIMCO/NVIL on MNIST and Omniglot density estimation, often outperforming in structured prediction scenarios.
• Results show temperature controls discreteness: lower temperatures push toward discrete behavior, while higher temperatures keep interior solutions, impacting integrality gaps.
• Discrete zero-temperature limits of Concrete distributions correspond to the original discrete distributions, enabling evaluation on the discrete graph at test time.