PaperlayerView exact revision on arXiv

Source-linked AI summary

The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables

Chris J. Maddison, Andriy Mnih, Yee Whye Teh

arXiv:1611.00712v3cs.LGstat.ML

TL;DR

Discrete stochastic nodes lack useful differentiable reparameterizations, limiting low-variance gradient optimization in stochastic computation graphs. The paper introduces Concrete random variables, continuous relaxations with closed-form densities and simple reparameterized sampling, and finds competitive performance with VIMCO across the reported tasks. The relaxation provides low-variance unbiased gradients for a continuous surrogate while remaining biased for the original discrete objective.

  • Problem

    Discrete random variables lack useful reparameterizations because their states are discontinuous, while existing unbiased estimators can have high variance and require special handling.

  • Method

    Concrete random variables relax discrete nodes into continuous simplex-valued variables with closed-form densities and sampling by softmax of logits perturbed by fixed noise.

  • Results

    Concrete relaxations were competitive with VIMCO, occasionally outperforming and occasionally underperforming, while operating in an automatic-differentiation library without special casing.

  • Takeaways & Limitations

    Concrete relaxations provide low-variance unbiased gradients for a continuous surrogate objective while enabling gradient-based training of computation graphs containing discrete stochastic nodes.

  • Takeaways & Limitations

    Temperature selection and relaxation performance remain empirical for each arity, and Concrete gradients are biased with respect to the original discrete objective.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Automatic differentiation handles differentiable computation graphs, but discrete stochastic nodes remain difficult because available unbiased estimators can require special handling and have high variance. The paper introduces Concrete variables as continuously relaxable discrete nodes, enabling automatic differentiation and competitive results against VIMCO.

  • Discrete stochastic nodes challenge automatic differentiation because their states are discontinuous and distributional gradients lack a clear general solution.
  • Existing unbiased gradient estimators for discrete nodes can require special-case implementations, surrogate objectives, and variance-reduction techniques.
  • Concrete variables provide continuous distributions on the simplex with closed-form densities and reparameterized sampling via softmax of logits plus fixed noise.
  • Every discrete random variable corresponds to the zero-temperature limit of a Concrete variable, allowing discrete-node graphs to be optimized through relaxed samples.
  • Concrete relaxations support objectives involving latent-node log-probabilities by using the Concrete density during training and evaluating the discrete graph at test time.
  • Concrete gradients were competitive with VIMCO, sometimes outperforming and sometimes underperforming, while using automatic differentiation without special casing.

2 BACKGROUND

Stochastic computation graphs express learnable stochastic mappings, but optimizing their parameter gradients requires different estimators. Score-function methods apply broadly yet can have high variance, while reparameterization transfers parameter dependence into differentiable transformations; latent-variable training uses variational bounds to address intractable likelihoods.

  • Stochastic computation graphs represent deterministic or sampled node computations in directed acyclic graphs and cover objectives in supervised, unsupervised, and reinforcement learning.
  • For a stochastic node, gradients with respect to downstream parameters can be estimated by Monte Carlo expectations of the objective gradient.
  • Score-function estimators rewrite distribution-parameter gradients as expectations and apply to discrete or continuous variables without requiring the objective to be differentiable in the sampled state.
  • The basic score-function estimator can suffer from high variance, motivating baselines and other variance-reduction techniques.
  • The reparameterization trick samples fixed-distribution noise and transforms it with parameterized functions, enabling chain-rule gradients when the transformation and objective are differentiable.
  • Latent-variable maximum likelihood is typically intractable, so auxiliary variational distributions provide ELBO or importance-weighted objectives that trade computational cost against bound tightness.

3 THE CONCRETE DISTRIBUTION

The Concrete distribution replaces discrete argmax states with differentiable simplex-valued samples while retaining a closed-form density and discrete behavior in the zero-temperature limit. Its practical use is training relaxed stochastic graphs and evaluating the original discrete objective at test time, subject to temperature constraints.

  • 3.1 DISCRETE RANDOM VARIABLES AND THE GUMBEL-MAX TRICK: The Gumbel-Max trick samples one-hot discrete states by adding fixed Gumbel noise to log αk and applying argmax.This refactors discrete sampling into a deterministic function of parameters and fixed-distribution noise.
  • 3.2 CONCRETE RANDOM VARIABLES: Because argmax has zero derivative except at state-change boundaries, it cannot provide a useful automatic-differentiation reparameterization.The Concrete construction replaces this discontinuous computation with a continuous relaxation.
  • 3.2 CONCRETE RANDOM VARIABLES: A Concrete sample is a probability vector in the simplex produced by applying softmax to Gumbel-perturbed logits, with smaller temperatures yielding more discrete outputs.As λ approaches 0, the softmax approaches argmax while preserving the relative order of the perturbed logits.
  • 3.2 CONCRETE RANDOM VARIABLES: The Concrete distribution has a closed-form simplex density, and rounding a sample recovers the discrete distribution determined by the logits log αk.Its zero-temperature limit is equivalent to rounding, while Proposition 1 also gives a convexity property for sufficiently low temperatures.
  • 3.3 CONCRETE RELAXATIONS: During training, discrete nodes and their log-probability terms are replaced with Concrete counterparts; at test time, the original discrete loss is evaluated.This relaxation is especially intended for large stochastic computation graphs optimized from samples, while exact integration remains preferable when computationally feasible.

4 RELATED WORK

Related work includes simplex distributions, concurrent Gumbel-Softmax relaxations, and straight-through estimators. The paper distinguishes Concrete relaxations by using densities for relaxed objectives and by grounding the gradients in an explicit loss formulation.

  • Simplex distributions: The Dirichlet and Logistic Normal are established distributions over the simplex with different structural and simulation properties.The Logistic Normal can be sampled by applying softmax to Gaussian variables and a fixed final logit.
  • Gumbel-Softmax: Gumbel-Softmax concurrently used softmax-transformed Gumbels as a discrete relaxation but did not use the relaxed density in its objective.For variational inference, that relaxed objective is not a lower bound on marginal likelihood.
  • Straight-through estimators: Straight-through estimators pass discrete values forward while differentiating through an expected value, achieving empirical performance without being shown to estimate any loss function.This contrasts with approaches whose gradients correspond to an explicitly defined objective.

5 EXPERIMENTS

The experiments evaluate Concrete gradients on density estimation and structured prediction with neural networks containing many discrete stochastic units. Concrete generally outperformed VIMCO on nonlinear density models and more uniformly on MNIST structured prediction.

  • The experiments evaluate Concrete relaxations on density estimation and structured output prediction involving neural networks with hundreds of latent discrete nodes.Performance is compared with VIMCO, with NVIL used when m = 1.
  • Density estimation: Concrete relaxations outperformed VIMCO for nonlinear density-estimation models, whereas VIMCO performed better for linear models.
  • Density estimation: The best 4-ary and 8-ary density models achieved test/train NLLs of 86.7/83.3 and 87.4/84.6 with Concrete relaxations.
  • Structured output prediction: Concrete relaxations more uniformly outperformed VIMCO on MNIST structured prediction.
  • Structured output prediction: 4-ary and 8-ary structured-prediction models achieved test/train NLLs of 55.4/46.0 and 54.7/44.8, respectively.Increasing arity uniformly improved these models.

6 CONCLUSION

The paper introduces Concrete distributions as continuous relaxations for discrete stochastic nodes and evaluates them for optimizing stochastic computation graphs. Their gradients are biased for the discrete objective but low-variance and unbiased for a continuous surrogate.

  • Concrete is a continuous distribution on the simplex with closed-form density, positive location parameters, and positive temperature.
  • Every Concrete distribution approaches a discrete distribution as temperature approaches zero, and every discrete distribution can be obtained by discretizing a Concrete distribution.
  • Concrete gradients are biased for the original discrete objective but low-variance unbiased estimators of a continuous surrogate objective.
  • The experiments showed that Concrete stochastic nodes can effectively optimize stochastic computation graphs with discrete stochastic nodes.
  • The experiments did not find annealing or automatic temperature tuning important, although these remain possible future work.

A PROOF OF PROPOSITION 1

This appendix proves properties of the Concrete distribution using an invertible transformation and its Jacobian, then derives the density through change of variables and convexity arguments.

  • The proof begins with an invertible transformation between simplex coordinates and auxiliary variables.
  • The inverse transformation expresses the auxiliary coordinates using logarithms of simplex variables and a normalization term.
  • Row operations reduce the Jacobian determinant to a simpler form used in the change-of-variables density derivation.
  • The density is obtained by combining transformed exponential terms with the Jacobian determinant.
  • For λ ≤ (n − 1)^−1, convexity follows because the component terms and their composition are convex.

B THE BINARY SPECIAL CASE

The binary special case connects discrete Bernoulli sampling to logistic noise and replaces its step function with a sigmoid-based Binary Concrete relaxation. The appendix also explains how this relaxation enters variational objectives.

  • The binary discrete case is represented by a two-state one-hot variable whose state probability is determined by the ratio of positive parameters.
  • The difference of two Gumbel variables is Logistic, yielding a Bernoulli sample through a unit step function applied to noisy logits.
  • Binary Concrete replaces the step function with a sigmoid applied to Logistic noise, producing a continuous sample in the unit interval.
  • The relaxation replaces the discrete expectation over D with an expectation over Z sampled from Concrete(α(x), λ1), enabling derivatives through stochastic computations.
  • Three objective modifications are considered: Concrete densities, relaxed discrete log masses, or the analytic discrete KL.
  • Only the Concrete-density objective is guaranteed to remain a lower bound; the alternatives can have uninterpretable values and may overfit a spurious KL component.

C.3 WHICH RANDOM VARIABLE TO TREAT AS THE STOCHASTIC NODE

The section chooses the stochastic node in relaxed computation graphs to preserve tractable log-probabilities and reparameterizable objectives. It uses ExpConcrete variables, with exp treated as downstream computation, and a logistic-noise construction for the binary case.

  • Using Gumbel variables directly for log-likelihood terms can cause underflow, whereas treating relaxed states as stochastic nodes gives the tighter variational bound.
  • Working with Concrete random variables in log-space provides a numerically effective implementation strategy.
  • ExpConcrete variables are defined so that exponentiation yields Concrete variables, while invertibility preserves KL divergences between the transformed and original distributions.
  • The proposed implementation uses ExpConcrete variables as stochastic nodes and treats exp as downstream computation, producing a fully reparameterizable objective.
  • For binary Concrete variables, the logit plus Logistic noise is the stochastic node, followed by the sigmoid as downstream computation.

C.4 CHOOSING THE TEMPERATURE

Temperature controls whether Concrete relaxations concentrate near discrete vertices or represent interior simplex values. The paper recommends application-specific tuning, including distinct temperatures for prior and posterior nodes.

  • High temperatures can create interior modes, so temperature selection is important for keeping relaxed nodes near the discrete boundary.
  • For n > 1, λ ≤ (n −1)−1 is not necessarily a good guideline because the suitable temperature depends on n and the application.
  • λ = 2/3 was the best single-performing temperature across the considered n ∈ {2, 4, 8} cases, although λ = 1 was best for n = 4 and λ = 2/3 for n = 8.
  • Distinct temperatures for posterior and prior Concrete nodes improved results quite dramatically when the loss included a KL divergence.
  • Both λ1 for the posterior and λ2 for the prior are tunable hyperparameters.

D EXPERIMENTAL DETAILS

The experiments use neural networks with n-ary discrete stochastic layers, replacing discrete nodes with Concrete relaxations and evaluating implementation choices across model architectures and training procedures.

  • The model architectures mirror prior architectures while replacing Gaussian variables with Concrete or discrete random variables.
  • Conditioning functions are either linear or nonlinear, with nonlinear functions consisting of two tanh layers matching the preceding stochastic layer’s size.
  • The networks use n-ary stochastic units whose log2(n)-dimensional states occupy the corners of a hypercube.
  • For an n-ary node, the discrete model maps one-hot samples through a corner matrix C, while the Concrete relaxation maps continuous samples X through the same matrix as Ẏ = CX.
  • In the binary case, sampling uses uniform noise transformed by the logistic expression, with the corresponding Concrete relaxation replacing the discrete step with a continuous relaxation.
  • Training uses Glorot initialization, Adam, minibatches of 64, 10^7 optimization steps, and grid-searched learning-rate and weight-decay hyperparameters.
  • Density-estimation temperatures were tuned separately for prior and posterior distributions, without temperature annealing.
  • Score-function estimators received early stopping because VIMCO and NVIL were more prone to overfitting.

E EXTRA RESULTS

The extra-results material points readers to a density-estimation table and a random-variable cheat sheet. The cheat sheet summarizes parameterizations and sampling procedures used throughout the work.

  • Table 2 reports density-estimation results for Concrete relaxations with different layer arities.
  • The random-variable cheat sheet assumes access to uniform samples and defines later variables by reusing earlier constructions.
  • The cheat sheet highlights atypical Bernoulli and Logistic parameterizations and treats indexed variables such as Gk as independent random-generator calls.
Loading 1611.00712v3…