PaperlayerView exact revision on arXiv

Source-linked AI summary

A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning

Iris A. M. Huijben, Wouter Kool, Max B. Paulus, Ruud J. G. van Sloun

arXiv:2110.01515v2cs.LGstat.ML

TL;DR

Discrete stochasticity makes sampling and gradient estimation difficult in deep learning, especially for unnormalized or exponentially large distributions. This survey explains the Gumbel-max trick, organizes its sampling and gradient-estimation extensions and applications, and reviews design choices; it also identifies variance and temperature-selection trade-offs.

  • Problem

    Deep learning models need methods for sampling from discrete distributions and estimating gradients through discrete stochastic nodes, including when domains are unnormalized or exponentially large.

  • Method

    The article surveys the Gumbel-max trick, its sampling and gradient-estimation extensions, machine-learning applications, algorithm-selection guidance, and design choices.

  • Results

    Scaling independent Gumbel variates changes the temperature of the Boltzmann distribution while still yielding exact categorical samples, and Gumbel-based methods cover unstructured and structured sampling settings.

  • Takeaways & Limitations

    The review connects Gumbel-max sampling with Poisson processes and weighted reservoir sampling and organizes applications to support method selection for discrete stochasticity.

  • Takeaways & Limitations

    Gumbel-Softmax temperature selection involves a bias-variance trade-off, and the empirically optimal value may exceed the value suggested by the relaxation limit.

Abstract

from arXiv · show

The Gumbel-max trick is a method to draw a sample from a categorical distribution, given by its unnormalized (log-)probabilities. Over the past years, the machine learning community has proposed several extensions of this trick to facilitate, e.g., drawing multiple samples, sampling from structured domains, or gradient estimation for error backpropagation in neural network optimization. The goal of this survey article is to present background about the Gumbel-max trick, and to provide a structured overview of its extensions to ease algorithm selection. Moreover, it presents a comprehensive outline of (machine learning) literature in which Gumbel-based algorithms have been leveraged, reviews commonly-made design choices, and sketches a future perspective.

1 INTRODUCTION

Discrete stochasticity creates challenges for sampling and gradient estimation in deep learning, motivating this review of Gumbel-based methods, applications, design choices, and future directions.

  • Discrete models are relevant to decisions, event occurrence, social networks, and data compression because many aspects of the world are discrete.
  • Discrete stochasticity is difficult when distributions are unnormalized or exponentially large, and gradients through discrete stochastic nodes hinder backpropagation.These challenges affect both sampling and updating parameters preceding stochastic nodes.
  • The Gumbel-max trick and its extensions address discrete sampling, while Gumbel-Softmax variants provide biased alternatives to high-variance REINFORCE estimation.
  • The article develops intuitive and mathematical understanding of the Gumbel-max trick, reviews its extensions, and provides handles for algorithm selection and design choices.
  • The review covers categorical and Gumbel background, Gumbel-based sampling algorithms, gradient estimators, machine-learning applications, practical choices, and future perspectives.

2 PRELIMINARIES

The preliminaries define categorical and Gumbel variables and explain inverse transform sampling, including its simple normalized-distribution setting and categorical interpretation.

  • 2.1 Categorical distribution: A categorical distribution assigns probabilities to N classes and can be parameterized by normalized probabilities, unnormalized probabilities, or temperature-scaled logits.The partition function normalizes the distribution, while temperature controls its entropy from near-deterministic to uniform.
  • 2.1 Categorical distribution: The categorical sample is represented by the index of the selected class or, equivalently, a one-hot vector.
  • 2.2 Gumbel distribution: The Gumbel distribution is a type-I generalized extreme-value distribution parameterized by location µ and scale β, with standard variables denoted G.
  • 2.2 Gumbel distribution: The Gumbel inverse cumulative density function transforms a Uniform variable into a Gumbel sample through a double negative-logarithmic relation.A single negative logarithm instead produces an Exponential sample, linking these parameterizations.
  • 2.3 Inverse transform sampling: Inverse transform sampling converts a Uniform sample through a distribution’s inverse cumulative density function to obtain a categorical class.It requires one random variate per sample but requires the categorical distribution to be normalized.

3 APPLICATIONS

Gumbel-based gradient estimators support learning with discretized data and discrete model choices across latent variables, attention, generation, compression, and neural architecture search.

  • 3 APPLICATIONS: Gumbel-based methods have been used for discrete action selection, active-learning data sampling, dialog generation, and translation.
  • 3 APPLICATIONS: Gumbel-based estimators enable end-to-end training through discrete stochastic components by relaxing the non-differentiable Gumbel-max sampling operation.
  • 3 APPLICATIONS: Applications span discrete latent-variable models, hard attention, inherently discrete generation, data compression, and discrete model selection.The review categorizes applications into discretized data and model selection from a discrete model space.
  • 3.1.2 Attention: Hard attention is relaxed because discrete accept-or-decline choices block gradients, with applications including interpretability, recurrent models, graph neural networks, and agent decisions.
  • 3.1.4 Data compression: In data compression, estimators support learning quantization levels, codebook selection, and adaptive binarization for discrete representations.
  • 3.2 Model selection: For neural architecture search, Gumbel-based algorithms learn hyperparameters such as model width, depth, or layer-wise bit counts rather than relying only on heuristic selection.

4 SAMPLING ALGORITHMS

This section develops the Gumbel-max trick, its distributional foundations, and extensions linking categorical sampling to Poisson processes, temperature scaling, and alternative parameterizations.

  • 4.1.1 Definition and properties: The Gumbel-max trick samples a categorical variable by adding independent Gumbel noise to unnormalized log-probabilities and selecting the maximum.The perturbed logits are shifted independent Gumbels, and the maximum is independent of the argmax.
  • 4.1.1 Definition and properties: The trick is invariant to common probability scaling, so it operates directly on unnormalized probabilities and on arbitrary subdomains.This supports unconstrained optimization of log θ rather than normalized probability vectors.
  • 4.1.3 Exploration vs exploitation: Changing the Gumbel scale β samples from a categorical distribution with Boltzmann temperature Tβ, linking noise scale to exploration temperature.The sampled index remains analytically categorical under nonstandard Gumbel location and scale parameters.
  • 4.1.2 Link to reservoir sampling and Poisson processes: Uniform/Beta, Exponential, and Gumbel variables are alternative parameterizations connected by negative-log transformations, supporting links to reservoir sampling and Poisson processes.A single negative logarithm maps Uniform to Exponential, while a second maps Exponential to Gumbel.
  • 4.1.2 Link to reservoir sampling and Poisson processes: In the Poisson-process view, the first arrival identifies the sampled category while its arrival time is independent of that category, mirroring Gumbel max-stability.Merged arrivals have rate Z, and the shortest arrival time corresponds to the optimum value.

Top-down sampling

Top-down sampling reverses the usual Gumbel-max direction by conditioning perturbed logits on the maximum and/or argmax, while related extensions support large, structured, and multiple-sample domains.

  • Top-down sampling: Top-down sampling starts from the maximum and/or argmax and conditionally samples perturbed logits, exploiting their independence for parallel or sequential procedures.Bottom-up sampling instead generates Gumbels first and then computes the maximum and argmax.
  • 4.3 Extended sampling algorithms: The Gumbel-max trick requires N Gumbel realizations for one N-class sample, motivating extensions for multiple samples and exponentially large domains.The survey organizes these extensions across unstructured and structured sampling scenarios.
  • 4.3 Extended sampling algorithms: The survey also covers multiple samples from structured distributions, including sets or sequences generated from neural sequence models using perturb-and-MAP approaches.These methods extend Gumbel-based sampling beyond single categorical events.
  • Sampling without replacement: Gumbel-top-k reuses one set of perturbed logits and selects the k largest values, yielding k independent samples without replacement.This avoids repeated explicit renormalization on the shrinking domain.
  • Sampling without replacement: Sampling without replacement can also produce ordered sequences, sets, permutations, or Wallenius’ noncentral hypergeometric distributions, depending on the representation.For k = N, ordered samples follow the Plackett-Luce model; unordered probabilities may require O(N!) computation or approximation.

5 GRADIENT ESTIMATION IN NEURAL NETWORK

The section surveys gradient estimators for discrete stochastic nodes, emphasizing relaxed Gumbel-Softmax gradients and their extensions to structured distributions. It contrasts unbiased but high-variance score-function estimators with biased, lower-variance relaxations.

  • 5.1 Gumbel-Softmax distribution: Gumbel-Softmax relaxes categorical samples into continuous vectors for differentiable optimization, converging to hard samples as λ →0+.The relaxation spreads mass across bins while preserving order, so argmax recovers the corresponding hard sample.
  • 5.2 Gradient estimators: REINFORCE applies to discrete variables with minimal assumptions and is unbiased, but its high variance can make practical optimization ineffective.Control variates and multiple samples are common variance-reduction strategies.
  • 5.2 Gradient estimators: Pathwise estimators can have low variance but require differentiability and continuous reparameterization, which discrete random variables do not provide.Relaxed estimators therefore use continuous approximations during training while retaining discrete variables at test time.
  • 5.2.1 Gradients for sampling unstructured distributions: The Gumbel-Softmax estimator trades bias for reduced variance, with lower temperatures generally reducing bias but increasing variance.It is unbiased for the relaxed objective but biased for the original categorical gradient; straight-through variants use hard forward computations.
  • 5.2.2 Gradients for structured distributions: Gumbel-Softmax extensions address structured variables including fixed-size subsets and permutations through repeated samples, successive softmax applications, and Gumbel-Sinkhorn relaxations.Alternative transformations have also been proposed for structured gradient estimation.

6 PRACTICAL CONSIDERATIONS

The section reviews practical choices for parameterizing logits, implementing Gumbel-based methods, regularizing distributions, and selecting temperatures. These choices affect model conditioning, entropy, relaxation bias, variance, and inference behavior.

  • 6.1 Parameterizing the logits: Logits may be directly optimized or predicted conditionally by neural networks, while conditional sampling can also infer logits from already drawn samples.The reviewed applications include neural architecture search, attention, active data acquisition, and discrete GANs.
  • 6.2 Initialization and regularization: Regularization is tailored to the application, promoting sparsity or low entropy in some settings and high entropy in discrete GANs to preserve diversity and prevent mode collapse.Other strategies use KL divergence to a prior or categorical cross-entropy from labels.
  • 6.3 Setting the Gumbel-Softmax temperature λ: Temperature λ lacks a universally established setting: higher values increase bias and reduce variance, whereas lower values reduce bias but increase variance and can cause vanishing gradients.Although λ ≤(N −1)−1 ensures log-convexity, empirical optima were often higher; fixed and annealed schedules are both used.
  • 6.4 Relation between temperatures: The two temperatures interact: relaxing a tempered categorical distribution yields a Gumbel-Softmax temperature of λ/T.This couples the Boltzmann diversity-quality trade-off with the relaxation used for gradient estimation.
  • 6.5 Inference: At inference, trained models can either sample from the categorical distribution or select its highest-probability class, with noise-free inference reported more often.The choice is not always documented, and sampled inference does not require adjusted computation in the same way as training relaxation.

7 SUMMARY & FUTURE PERSPECTIVES

The review surveys the exact Gumbel-max trick, its sampling and gradient-estimation extensions, machine-learning applications, practical choices, and future directions.

  • The review presents the Gumbel-max trick as an exact sampler for unnormalized categorical distributions and connects it to Poisson processes and weighted reservoir sampling.
  • Extensions support sampling from structured models, while Gumbel-Softmax relaxations enable error backpropagation through discrete stochasticity in deep neural networks.
  • The review categorizes Gumbel-based machine-learning applications and discusses practical considerations and commonly made algorithmic choices.
  • Future applications include discrete stochasticity in expressive neural networks for areas such as drug interactions and drug discovery.
  • Future theory may develop more efficient single- and multiple-sample algorithms for unnormalized structured models and relaxed Gumbel-based gradient estimators.

APPENDIX A PROOF GUMBEL-MAX TRICK AND MAX-STABILITY

The appendix proves that Gumbel-max sampling exactly produces a categorical draw and that the maximum of shifted independent Gumbels remains Gumbel-distributed.

  • The Gumbel-max trick draws an exact categorical sample from unnormalized probabilities by taking the argmax of log-probabilities perturbed with i.i.d. standard Gumbels.
  • Independence allows the probability of all competing variables falling below the selected variable to factorize, leading to the categorical probability.
  • The proof identifies the event I = ω with the shifted Gumbel for class ω exceeding all competing shifted Gumbels.
  • The maximum of shifted Gumbels has a Gumbel distribution with location log Z, where Z is the sum of the unnormalized probabilities.

BEL SAMPLES

The appendix uses translation and scale invariance of argmax and max to characterize tempered Gumbel sampling and the resulting maximum distribution.

  • Translation and scale invariance reduce the perturbed argmax to a Gumbel-max draw with logits scaled by 1/(Tβ).
  • The resulting sampled index follows a categorical distribution with probabilities proportional to exp(a/(Tβ)).
  • The maximum remains Gumbel-distributed with location µ + β log Z′ and scale β after scaling and shifting.

APPENDIX C SAMPLING EXPERIMENTS

The sampling experiments illustrate how Gumbel noise scale and Gumbel-Softmax temperature affect categorical samples, entropy, and convergence.

  • For β < 1 the effective categorical distribution has lower entropy, whereas β > 1 makes it more uniform.
  • Multiple hard Gumbel-max samples produce a one-hot histogram matching the underlying categorical distribution.
  • Gumbel-Softmax samples spread mass across classes, with their behavior varying as the temperature λ changes.
  • The expectation of multiple Gumbel-Softmax samples converges to the categorical distribution as λ →0+, but becomes more uniform at higher λ.

APPENDIX D TOP-DOWN SAMPLING

Algorithm 1 presents a top-down construction procedure in discrete sampling space, using the paper’s notation to recursively generate samples and associated categorical variables.

  • The algorithm rewrites the top-down construction method from [107] for discrete sampling space and the paper’s notation.
  • It initializes from a sampling space D ∈ {1, . . . , N} and unnormalized probabilities, then samples an initial Gumbel value and categorical variable.
  • During its queue-based loop, the procedure forms child sampling domains, draws truncated Gumbel values, samples categorical variables, and yields pairs (m_k, ω_k).
Loading 2110.01515v2…