# 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  © 2022 IEEE. Personal use of this
material is permitted. Permission from IEEE must be obtained for all other
uses, in any current or future media, including reprinting/republishing this
material for advertising or promotional purposes, creating new collective
works, for resale or redistribution to servers or lists, or reuse of any
copyrighted component of this work in other works.I.A.M Huijben is affiliated
with the department of Electrical Engineering, Eindhoven University of
Technology, The Netherlands, and Onera Health, Eindhoven, The Netherlands.
E-mail<EMAIL_ADDRESS>R.J.G. van Sloun is affiliated with the
department of Electrical Engineering, Eindhoven University of Technology, The
Netherlands, and Philips Research, Eindhoven, The Netherlands. Email:
<EMAIL_ADDRESS>W. Kool is affilitated with the Amsterdam Machine
Learning Lab (AMLab), University of Amsterdam, The Netherlands, and ORTEC R&D,
Zoetermeer, The Netherlands. E-mail<EMAIL_ADDRESS>M.B. Paulus is
affiliated with the department of Computer Science, ETH, Zurich, Switzerland.
E-mail<EMAIL_ADDRESS>
###### Abstract

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.

###### Index Terms:

Gumbel-max trick, Sampling, Gradient estimation, Gumbel-Softmax, Categorical
distribution, Structured models

##  1 Introduction

The world around us is discrete in many aspects. Think about decision making,
e.g. in traffic (Should I decelerate, accelerate or maintain a constant
speed?), for product selection (Given that I liked the trousers from shop X
last time, which new trousers shall I buy?), or in a clinical setting (Should
I administer medication A or B to the patient?). Other discrete examples yield
occurrence of events (At which day did I see you for the last time?), social
networks (Do two persons know each other?) or data compression (How many bits
do we need to store this information?). Modelling these concepts evidently
pleads for discrete models, on which we focus in this work.

The immensely grown popularity of machine learning, and in particular deep
learning, has given rise to high capacity models that have consistently been
outperforming more classical models in various domains. Whereas the first
neural networks typically comprised a stack of few fully-connected layers,
activated with non-linear functions, current deep learning models often
exhibit a (very) deep modular architecture. Specifically generative models
(e.g. the variational auto-encoder (VAE) [1] and auto-regressive pixelCNN
[2]), and Bayesian models combine such deep architectures with a form of
stochasticity. Trainable parameters of these models are generally learned in a
data-driven fashion, i.e. they are optimized over a set of training examples
by back-propagating the error for a downstream task. The probability
distributions are then treated as stochastic nodes in a computation graph,
from which a sample is drawn to compute the output, and subsequently the
error.

Sampling from resulting discrete distributions may raise challenges when the
probability mass function is unnormalized, or in case of an (exponentially)
large sampling domain. The latter is often encountered in sequence models,
where one ‘sample’ entails a full sequence, and the sampling domain grows
combinatorially for each additional sequence element. Moreover, incorporating
(discrete) stochasticity in deep learning models poses a second challenge.
Gradient computation, and therefore error backpropagation, through a
stochastic node is hampered, which is required for updating the distribution’s
parameters and all model parameters that precede this node in the network.
These two challenges, i.e. _sampling_ and _gradient estimation_ of discrete
stochasticity in deep learning models, will be the focus of this review.

Several algorithms today exist to sample from structured models with
exponentially large sampling domains (e.g. ancestral sampling), or to sample
from a set of unnormalized probabilities (e.g. Markov Chain Monte Carlo (MCMC)
methods [3], or the Gumbel-max trick [4]). The Gumbel-max trick recently found
renewed attention for use in deep learning models, thanks to the proposed
Gumbel-Softmax (GS) gradient estimator that is based on a relaxation of this
trick [5, 6]. The GS estimator (and variants thereof) have become popular
(biased) alternatives for the high-variance REINFORCE estimator [7], thanks to
their good empirical performance and straightforward implementation.

This article provides both an intuitive and mathematical understanding of the
Gumbel-max trick, reviews extensions of this trick, provides handles for
algorithm selection and corresponding design choices, and sketches a future
perspective. The content can be summarized as follows:

  * •

Background on categorical random variables, the Gumbel distribution and
inverse transform sampling is covered in Section 2.

  * •

Applications of Gumbel-based algorithms in machine learning are discussed in
Section 3.

  * •

Section 4 and 5 present technical details about Gumbel-based sampling
algorithms and gradient estimators, respectively.

  * •

Practical considerations and commonly-made design choices are presented in
Section 6.

  * •

Section 7 summarizes this review and sketches a future perspective.

##  2 Preliminaries

###  2.1 Categorical distribution

A categorical distribution is a probability distribution that assigns a
probability to $N$ distinct classes. We consider three different
parameterizations of the same distribution: _normalized probabilities_
$\bm{\pi}$, _unnormalized probabilities_ $\bm{\theta}$, or _unnormalized log-
probabilities_ (or _logits_) $\log\bm{\theta}={\bm{a}}/T$, where
$T\in\mathbb{R}_{>0}$ is a temperature parameter. We denote the
$i^{\text{th}}$ unnormalized probability - which depends on $T$ \- with
$\theta_{i;T}$. The probabilities $\pi_{i;T}$ for each class $i\in
D=\\{1,\ldots,N\\}$ are then given by:

$\pi_{i;T}=\frac{\theta_{i;T}}{\sum_{j\in
D}\theta_{j;T}}=\frac{\exp(a_{i}/T)}{\sum_{j\in
D}\exp(a_{j}/T)}=\frac{\exp(a_{i}/T)}{Z_{D}(\bm{\theta})},$ (1)

where $Z_{D}(\bm{\theta})=\sum_{j\in D}\theta_{j;T}=\sum_{j\in
D}\exp(a_{j}/T)\in\mathbb{R}_{>0}$ is the _partition function_ that normalizes
the distribution. To prevent clutter, both the parameter and subscript of $Z$,
and temperature subscript $T$ are, in the rest of this paper, omitted when
context allows it.

The temperature $T$ controls the distribution’s entropy, such that the
distribution can vary between a degenerate/deterministic ‘one-hot’
distribution for $T\rightarrow 0^{+}$ (i.e. all mass is centralized in one
class), and a Uniform distribution (all probabilities are equal) for
$T\rightarrow\infty$. In many cases, $T$ simply equals one, making ${\bm{a}}$
the unnormalized log-probabilities. Equation (1) is known as a _softmax_
function with temperature parameter $T$. Parameterized as
$\operatorname{Cat}({\bm{a}},T)$, the categorical distribution is often
referred to as the Gibbs or Boltzmann distribution.

We will use the terms _unstructured_ and _structured_ models in this work. The
former is used for what is normally referred to as a categorical distribution:
a distribution assigning a probability to a single event/class (e.g. the
probability that horse $X$ wins a horse race). The number of distribution
parameters thus equals the size of the sampling domain. A structured model, on
the other hand, places a probability on structures, i.e. different
combinations of events/classes, rather than on single events (e.g. the
probability that horses $X$, $Y$, and $Z$ cover the top-3). A sequence model
is another common example of such a structured model, in which the sampling
distribution of each element may be conditional upon the previously sampled
element.

A random variable following a categorical distribution is denoted with $I$ in
this work, and a realization (i.e. a sample from the categorical distribution)
is denoted with $I=\omega$, providing the _index_ of the sampled class. In
certain contexts, it is useful to define this sample as a one-hot embedding; a
unit vector of length $N$, with a one at index $\omega$ and zeros otherwise,
which we denote with $\mathbbm{1}_{\omega}$. Several algorithms exist to
sample from a categorical distribution. Inverse transform sampling is the most
basic approach (see Section 2.3). The Gumbel-max trick (see Section 4.1.1),
and variants thereof (see Section 4.3) are commonly-used alternatives in
machine learning applications.

###  2.2 Gumbel distribution

The Gumbel distribution [8] is an instance (type I) of the generalized extreme
value distribution111Also the Fréchet (type II) and Weibull (type III)
distribution generalize to this generalized extreme value distribution. [9],
which models optima and rare events. A Gumbel random variable - which is often
referred to in this work as ‘a Gumbel’ - is parameterized by location and
scale parameters $\mu\in\mathbb{R}$ and $\beta\in\mathbb{R}_{\geq 0}$,
respectively. The corresponding probability (PDF) and cumulative density (CDF)
functions are respectively given by

$\displaystyle f(x)$
$\displaystyle=\frac{1}{\beta}e^{-\frac{x-\mu}{\beta}}e^{-e^{-\frac{x-\mu}{\beta}}},\hskip
28.45274pt\mathrm{and}$ (2) $\displaystyle F(x)$
$\displaystyle=e^{-e^{-\frac{x-\mu}{\beta}}}.$ (3)

We denote a Gumbel distribution (as defined in Eq. 2) with
$\operatorname{Gumbel}(\mu,\beta)$ (or sometimes $\operatorname{G}(\mu,\beta)$
in short), and a random variable following this distribution with
$G_{\mu,\beta}$. To prevent clutter, the scale, or both parameters are
frequently omitted when standard settings are assumed (i.e.
$G_{\mu}:=G_{\mu,1}$ and $G:=G_{0,1}$). We often consider a set of identically
and independently distributed (i.i.d.) Gumbel variables, where we use
$G^{(i)}$ to denote an $i^{\text{th}}$ standard Gumbel. Note that, $\mu$ and
$\beta$ are _not_ the mean and variance of a Gumbel, which are given by:

$\displaystyle\mathbb{E}[\operatorname{G}_{\mu,\beta}]=\mu+\gamma\beta,\hskip
56.9055pt\mathrm{and}$ (4)
$\displaystyle\mathbb{E}[(\operatorname{G}_{\mu,\beta}-\mathbb{E}[\operatorname{G}_{\mu,\beta}])^{2}]=\frac{\pi^{2}}{6}\beta^{2},$
(5)

where $\gamma\approx 0.577$ is the Euler constant and $\pi\approx 3.14$ is the
constant pi (not to be confused with the normalized probability of the
categorical distribution), respectively.

The inverse cumulative density function (ICDF; also called quantile function)
is given by

$F^{-1}(u)=-\beta\log(-\log u)+\mu.$ (6)

From Eq. 6 it can be seen that the Gumbel distribution is closed under
_scaling and addition_ , i.e. any Gumbel variable can be generated by scaling
and shifting a standard Gumbel.

Equation (6) is used in inverse transform sampling (see Section 2.3) to
transform a sample from the Uniform distribution $\operatorname{U}(0,1)$ into
a Gumbel sample via a double (negative) logarithmic relation. Taking only one
negative logarithm of a Uniform random variable, i.e. $X=-\log U$, defines a
sample from the (standard) Exponential distribution. The Gumbel distribution
is, therefore, sometimes also referred to as the ‘Double Exponential’
distribution. Thanks to the aforementioned relations, some properties of the
Gumbel distribution are closely related to properties of the Exponential, and
Uniform/Beta distribution, on which we elaborate in Section 4.1.2.

Figure 1: Illustration of inverse transform sampling for drawing a sample from
$\operatorname{Cat}(\bm{\pi})$. A sample from the Uniform distribution is
converted to one of the categories/classes of the categorical distribution via
its inverse cumulative density function. The chance that class $i$ is being
sampled equals $\pi_{i}$; the width of the corresponding box in this
illustration.

###  2.3 Inverse transform sampling

Inverse transform sampling is the most common way to sample from a
distribution. It transforms a standard Uniform variable into another random
variable, via a transformation characterized by the random variable’s ICDF.
Figure 1 illustrates this process for sampling from a categorical
distribution: a sample form a Uniform random variable is inserted into the
ICDF of the categorical distribution in order to find the corresponding
bin/class. Inverse transform sampling is easy to implement and only requires
generation of one random variate per categorical sample. It does, however,
require the categorical distribution to be normalized. The ICDF of the Gumbel
distribution (given in Eq. 6) can equivalently be used to draw a sample from
this distribution.

##  3 Applications

In machine learning, there is significant interest in over-parameterized
modular and structured models, which often involve one or more stochastic
components. When concerned with _discrete_ stochasticity, challenges arise
regarding both _sampling_ from discrete distributions, and _gradient
estimation_ thereof.

Sampling from discrete distributions can be achieved with (among others)
inverse transform sampling, or the Gumbel-max trick [4] (see Section 4.1.1)
and extensions thereof (see Section 4.3). Gumbel-based sampling algorithms
have, for example, been used for (discrete) action selection in a multi-armed
bandit setting [10], for sampling data points in active learning [11], for
text generation in dialog systems [12], or in translation tasks [13, 14].

When the stochastic components are parameterized, and the model is trained
end-to-end with stochastic gradient-based methods, gradients need to be taken
through the stochastic sampling process and gradient estimators are therefore
called for. Finding such estimators is particularly challenging when the
stochastic components are discrete. Relaxations of the (non-differentiable)
Gumbel-max trick, notably the Gumbel-Softmax (GS) estimator and variants
therefore, have been found useful for this purpose, and will be discussed in
Section 5. In this section we review the different applications, and refer to
any variant of the GS estimator as a Gumbel-based estimator, in order to
provide a general overview without elaborating on technical details.
Applications are categorized into works that leverage discretized data, or
perform model selection from a discrete (model) space (see Fig. 2).

Figure 2: Categorization of the different applications in deep learning in
which Gumbel-based gradient estimators have been applied in order to
facilitate training via backpropagation through discrete stochastic nodes.
Visualizations in the bottom show where different models exhibit discretized
data, indicated by the dashed gray boxes.

###  3.1 Discretized data

We interpret the term _discretized data_ in a wide sense here. This category
includes models that rely on discrete representations (e.g. categorical
probability distributions in discrete latent variable models), attention
mechanisms (sub-selecting data is a discrete choice process), generative
models for inherently discrete data (e.g. GANs for text generation), and
models for data compression.

####  3.1.1 Discrete latent variable models

Latent variable models assume that ‘the world’ (represented as training data)
has originated from a set of latent (i.e. hidden) variables. While such
variables are often modelled continuously, many situations arise in which the
choice for discrete latent variables seems more appropriate (see Fig. 2a).
Such discrete latent variable models have mainly come in two flavours. On the
one hand, one can have a discrete approximation/representation of a continuous
(latent) variable (i.e. quantization, as in [15]), while on the other hand,
one could model the data with a discrete (non-degenerate) distribution in the
latent space. In the former case the straight-through estimator [16] is
typically used to facilitate gradient updates of the encoder. The latter case,
on the other hand, requires sampling from the introduced discrete distribution
to estimate the expectation in the forward pass and its gradient during
backpropagation. This approach is also possible in quantization settings, as
will be discussed in Section 3.1.4. All following works in the current sub-
section have adopted the second approach and leveraged Gumbel-based estimators
for discrete latent variable models.

VAEs with discrete priors are proposed by [17, 18, 19], and the authors of
[20] relax the discrete latents and durations of a recurrent hidden semi-
Markov Model. In some works, continuous and discrete latent variables have
been combined [21, 22, 23, 24]. The authors of [21] leverage - next to a
continuous latent variable - (only) one categorical random variable, which is
then interpreted as a clustering variable that assigns data points to discrete
classes, resulting in a combined classification- and generative model.
Similarly, [23] train a VAE with a Gaussian Mixture model as prior, where
assignment to either of the Gaussians displays a discrete process as well, and
in the same line, the authors of [25, 26] relax backpropagation of their deep
clustering algorithms. Applications of (Gumbel-based) discrete latent variable
models as described above include (among others) planning [27], syntactic
parsing [28], text modelling [29], speech modelling [30, 31], paraphrase
generation [32], recommender systems [33], drug-drug interaction modelling
[34], and event modelling [35].

####  3.1.2 Attention

Attention mechanisms are concerned with filtering incoming information,
analogously to the way we - human beings - selectively observe the world
around us. Hard attention resorts to either accepting or declining
information, while soft attention refers to placing more emphasis on certain
parts than on others. When training neural networks, only soft attention
allows gradients to flow, so hard (i.e. discrete) attention mechanisms have
been relaxed using Gumbel-based gradient estimators to enable model
optimization. While the term attention is typically used to denote selectivity
in hidden representations of deep learning models, we slightly widen the term
and present works that select part of the data (either hard or soft) anywhere
in the model (see Fig. 2b), with the main motivation to improve the model’s
performance or its interpretability. Attention has been used on latent
features to enhance interpretability of ‘black box’ neural networks [36, 37],
for hierarchical multi-scale recurrent neural networks [38], and in graph
neural networks (GNNs) [39, 40]. The authors of [40] apply attention on edges
with the specific goal of graph clustering. Applications range from (but are
not limited to) recommender systems [41], pose estimation [42], video
classification [43], event detection [44], and image synthesis [45]. The
authors of [46] apply attention on features of discrete input samples and an
attacking vocabulary, in order to acquire a scalable method for real-time
generation of adversarial examples on discrete input data. Hard attention on
(input) data points (i.e. learned subsampling) has, moreover, been leveraged
to reduce computational overhead of neural networks, e.g. on point clouds [47]
and graphs [48].

Also decision making in an environment of one or multiple agents can be
considered (hard) attention. The authors of [49] leverage this ‘attention’ to
enable end-to-end training of an agent that makes discrete decisions. In a
multi-agent environment, such decisions have concerned communication symbol
selection [50, 51] in order to learn a language to be used among agents.

####  3.1.3 Discrete GANs

Generative Adversarial Networks (GANs) [52] have been a popular class of
generative models. They contain a generator and discriminator module (see Fig.
2c). The latter’s aim is to distinguish between real and fake (generated by
the generator) samples. This setup requires differentiability of the
discriminator loss with respect to the generator’s parameters. However, in the
case of applying a GAN on discrete data, the discrete and stochastic fashion
of the generator’s output hampers direct differentiability, and has therefore
been a suitable candidate for Gumbel-based gradient estimators. Applications
of discrete GANs that have leveraged these Gumbel-based gradient estimators
include (among others) text generation [53, 54, 55, 56, 57, 58, 59], fake user
data creation for recommender systems [60], and action prediction [61].
Discrete GANs have also been combined with knowledge distillation frameworks.
The authors of [62], for example, propose to learn a policy (i.e. a trajectory
of subsequent states and actions in reinforcement learning setups) by learning
to imitate a (known) expert/teacher policy. The student model learns this
imitation by fooling a discriminator that should distinguish between both
policies. Such setups may, however, require a large number of training
iterations before the student model converges. To alleviate this, the authors
of [63] have proposed the concept of adversarial distillation in the context
of multi-label classification, in which both the teacher and student model are
jointly battling against a discriminator that should distinguish their
(discrete) label predictions from real annotations. Both models learn from
each other via distillation losses, finally resulting in a low-resource
student model that can be used during inference. Both the teacher and student
model’s output are relaxed using the GS estimator.

####  3.1.4 Data compression

Our growing demand for data has sky-rocketed data rates, hampering data
transfer, processing and storage. To alleviate these rates, data compression
comes into play. A shift is currently being made from non-data-driven
compression techniques (e.g. JPEG) to data-driven methods, where parameters of
a neural compression model are learned from a training set of data. Such
models typically encode the data in a quantized and lower-dimensional space,
after which the data are decoded again (see Fig. 2d). Quantization is a
discrete process, making it a suitable candidate to be learned using a Gumbel-
based estimator. The authors of [64, 65] propose to learn the quantization
levels, which form a quantization codebook, whereas the authors of [66] learn
to select such a codebook as a whole. Moreover, data-adaptive binarization has
been learned as well [67, 68, 69].

The amount of acquired data in the digital domain is governed by the analog-
to-digital (ADC) converter in the sensing stage. Instead of compressing in the
digital domain, the compressed sensing (CS) framework [70] introduces
compression directly at this sensing stage (see Fig. 2d). Simply reducing the
sampling frequency of an ADC introduces aliasing artifacts in case the minimal
Nyquist rate is not met. CS leverages incoherent sampling patterns (expressed
in a sensing matrix) to beat this famous Nyquist rate under certain
assumptions. The typically heuristic (and task-independent) design of this
sensing matrix, has been outperformed by a task-conditional (discrete) sensing
matrix, learned via a Gumbel-based estimator [71]. Once trained, such learned
sensing system can be implemented in hardware to reduce data rates at the
sensing stage directly. The proposed framework has been applied for medical
ultrasound imaging [72], magnetic resonance imaging [73], and MIMO antenna
arrays [74]. Moreover, it was extended to a data-conditional setup, by making
the sensing matrix conditional upon already acquired data [75]. Note that data
compression by subsampling can also be considered (hard) attention, as
discussed in Section 3.1.2. Although the focus of this data compression
section is on works that aim for data reduction rather than model enhancement.

###  3.2 Discrete model space

The number of available neural models has given rise to algorithms that not
only optimize a model on training data, but also search for a suitable model
itself. Such data-driven model selection can be leveraged with the aim to just
find better-fitting models, or to find compressed models that facilitate
implementation in dedicated hardware. We distinguish methods that learn
concepts related to model quantization, and methods that learn hyper-
parameters (often related to depth and width). The latter is often referred to
as neural architecture search (NAS) or pruning.

####  3.2.1 Model quantization

Model quantization can take place on weights, activations and/or gradients.
Joint optimization and quantization requires differentiability of the
quantization operation, making it a suitable candidate for Gumbel-based
gradient estimators. Indeed, it was shown that quantization of network weights
and activations can effectively be learned, by either learning the
quantization grid while using a pre-defined number of bits [76], or learning
the appropriate number of bits per layer [77]. The authors of [78] also learn
layer-dependent bits assignment by leveraging a bits budget to be distributed,
for quantization of weights, activations and gradients.

####  3.2.2 Neural architecture search (NAS) & pruning

Gumbel-based algorithms for discrete NAS of deep neural networks have been
leveraged for distinct goals. Most general, hyper-parameters have been learned
- rather than set heuristically - to alleviate the burden of tuning [79, 80,
81]. In case such hyper-parameters relate to the size of the model (e.g. width
or depth), NAS (often called pruning in this case) is generally used with the
aim to find small models that facilitate implementation in (dedicated)
hardware [82, 83, 84, 85, 86, 87, 88, 89]. Moreover, models have been proposed
in which the functionality is conditioned upon incoming data. This
conditioning can be applied to improve performance [90, 86, 91], achieve run-
time acceleration (i.e. reduce computational complexity during inference) [92,
93, 94, 95, 96], or enhance interpretability [90, 97, 98]. Also multi-task
models have been subject to discrete NAS, where task-dependent
gating/branching has been learned [99, 100]. Considering graphs as structured
data models, NAS has also been leveraged to find suitable graphical
representations, e.g., by optimizing the adjacency matrix [101, 98]. Moreover,
graph structures are leveraged by Tree-LSTMs [102], which were later extended
to Gumbel-Tree LSTMs, where graph edges are co-optimized instead of pre-
defined [103]. Gumbel-Tree LSTMs have notably found use in machine translation
[104, 105].

##  4 Sampling algorithms

In this section we introduce the Gumbel-max trick and its properties (Section
4.1.1), and link it to concepts known from different fields, i.e. Poisson
processes (Section 4.1.2) and Boltzmann exploration (Section 4.1.3). We then
introduce top-down sampling using the inverted Gumbel-max trick (Section 4.2),
followed by sampling algorithms that are extensions of the conventional
Gumbel-max trick (Section 4.3).

###  4.1 Gumbel-max trick

####  4.1.1 Definition and properties

The Gumbel-max trick [4] draws a sample from a categorical random variable
$I\sim\operatorname{Cat(\bm{\pi})}$. It does so, by adding i.i.d. Gumbel
(noise) samples to the unnormalized log-probabilities and selecting the index
with the maximum value, which in turn follows a Gumbel distribution. More
specifically:

$\displaystyle I$ $\displaystyle{}={}$ $\displaystyle\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i}+G^{(i)}\\}$ $\displaystyle{}\sim{}$
$\displaystyle\operatorname{Cat}(\bm{\pi})\hskip 31.2982pt\mathrm{and}$ (7)
$\displaystyle M$ $\displaystyle{}={}$ $\displaystyle\underset{i\in
D}{\operatorname{max}}\\{{\log\theta_{i}+G^{(i)}}\\}$ $\displaystyle{}\sim{}$
$\displaystyle\operatorname{Gumbel}(\log Z).$ (8)

The arguments $G_{\log\theta_{i}}:=\log\theta_{i}+G^{(i)}$ ($\forall i\in D$)
are shifted independent Gumbels, and often referred to as _perturbed logits_.
Equation (8) is known as the _max-stability_ property, which states that the
maximum is _independent_ of the argmax random variable: $M\perp\\!\\!\\!\perp
I$. The interested reader is referred to Appendix A for the proofs of Eq. 7
and (8). Section 4.1.2 provides a more intuitive understanding of the Gumbel-
max trick, by linking it to properties of well-known Poisson processes.

Scaling (or normalizing) unnormalized probabilities $\bm{\theta}$ by any
positive constant $Z$, results in a subtraction in logarithmic space:
$\log(\bm{\theta}/Z)=\log\bm{\theta}-\log Z$. Given the translation-invariance
of the argmax function, we can see that the Gumbel-max trick as defined in Eq.
7, also applies for normalized probabilities $\bm{\pi}$:

$I=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i}+G^{(i)}\\}=\underset{i\in
D}{\operatorname{argmax}}\\{\underbrace{~{}\log\pi_{i}~{}}_{\log\theta_{i}-\log
Z}+~{}G^{(i)}\\}.$

This property is a direct result of the fact that Gumbel random variables
adhere to Luce’s choice axiom [106], which states that the probability ratio
of selecting two elements from a set is independent of the other elements
present in that set. Thanks to this independence, a constant scaling of the
probabilities does thus not influence the argmax output. As a result, the
Gumbel-max trick can also be used over sub-domain $B\subseteq D$, with
unnormalized probabilities $\theta_{i}$, for $i\in B$ [107]:

$\underset{i\in
B}{\operatorname{argmax}}\\{\log\theta_{i}+G^{(i)}\\}\sim\operatorname{Cat}\Big{(}\frac{\mathbbm{1}(i\in
B)~{}\bm{\pi}}{\sum_{i\in B}\pi_{i}}\Big{)}.$

In machine learning applications, the parameters of the categorical
distribution are often learned or predicted. The Gumbel-max trick is therefore
typically used with unnormalized - rather than normalized - probabilities, as
this allows for unconstrained optimization of $\log\bm{\theta}$ (whereas
$\bm{\pi}$ is a probability vector).

The Gumbel-max trick leverages standard Gumbel noise samples. Nevertheless,
when using i.i.d. samples from $\operatorname{Gumbel}(\mu,\beta)$, with
$\mu\neq 0$ and $\beta\neq 1~{}(\geq 0)$, the distributions of the sampled
index and maximum can still be defined analytically:

$\displaystyle I^{\prime}=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i;T}+G^{(i)}_{\mu,\beta}\\}$
$\displaystyle{}\sim{}$
$\displaystyle\operatorname{Cat}\Big{(}\frac{\exp({\bm{a}}/(T\beta))}{\underset{i\in
D}{\sum}\exp(a_{i}/(T\beta))}\Big{)},$ (9) $\displaystyle
M^{\prime}=\underset{i\in
D}{\operatorname{max}}\\{{\log\theta_{i;T}+G^{(i)}_{\mu,\beta}}\\}$
$\displaystyle{}\sim{}$ $\displaystyle\operatorname{Gumbel}(\mu+\beta\log
Z^{\prime},\beta),$ (10)

with $G^{(i)}_{\mu,\beta}$ the $i^{\text{th}}$ independent Gumbel random
variable with location $\mu$, and scale $\beta$, and $Z^{\prime}=\sum_{i\in
D}\exp(a_{i}/(T\beta))$. Appendix B derives the aforementioned relations.
Interestingly, from Eq. 9 we conclude that
$I^{\prime}\sim\operatorname{Cat}({\bm{a}},T\beta)$. Changing the scale of the
independent Gumbel variates thus changes the temperature of the Boltzmann
distribution from which exact samples are being drawn (see Section 4.1.3 for a
discussion about this relation).

Figure 3: The Gumbel (double Exponential) distribution directly relates to the
Uniform/Beta and Exponential distribution via a single, respectively double,
negative logarithmic relation. The green ellipses indicate the probability
distribution of the corresponding random variable in the cell. To prevent
clutter, we here define $U:=U^{(i)}$, $X:=X^{(i)}$, and $G:=G^{(i)}$, being an
$i^{\text{th}}$ independent standard Uniform, Exponential, and Gumbel random
variable, respectively.

####  4.1.2 Link to reservoir sampling and Poisson processes

As mentioned in Section 2.2, the Uniform/Beta, Exponential and Gumbel
distribution are closely related, and can be used to express different
parameterizations of the same stochastic process. These distinct
parameterizations are typically used in different research areas: the
Uniform/Beta parameterization is notably used for weighted reservoir sampling
[108], whereas the Exponential counterpart is widely used in queueing
theory/Poisson processes [109]. The ‘double Exponential’ parameterization is
leveraged for the Gumbel-max trick [4], and its extensions. The aim of this
section is to provide a summarizing overview of the existing relations between
these distributions, and illustrate this through an intuitive example. Figure
3 relates samples and properties of the three parameterizations (one per
column), and will further be explained in this section. Some of the
connections that are presented have previously been made, e.g., both the
author(s) of [110] and [111] showed the equivalence between weighted reservoir
sampling and the Gumbel-max trick, and the author of [112] extensively
discussed the relation between Poisson processes, especially Exponential
races, and the Gumbel-max trick.

The ICDF of the Gumbel distribution - as provided in Eq. 6 \- already showed
the relation between a Gumbel and a Uniform random variable. The Gumbel
distribution also relates to the $\operatorname{Exponential}(\theta)$
distribution ($\theta\in\mathbb{R}_{\geq 0}$), which is characterized by a CDF
being $F(x)=1-e^{-\theta x}$ for $x\geq 0$ (and 0 for $x<0$), and location
$\mathbb{E}[X]=1/\theta$. Its ICDF equals
$F^{-1}(u)=\frac{-\log(1-u)}{\theta}$. Re-introducing
$X\sim\operatorname{Exponential}(1)$ as a standard Exponential random variable
with $\theta=1$, allows writing:

$\frac{-\log(1-V)}{\theta}\overset{d}{=}\frac{-\log(U)}{\theta}=\frac{X}{\theta},$
(11)

where both $V$ and $U$ are standard Uniform random variables, and
$\overset{d}{=}$ denotes equality in distribution. Taking the negative
logarithm of Eq. 11, result in:

$-\log\frac{X}{\theta}=-\log X+\log\theta=-\log(-\log U)+\log\theta,$ (12)

from which we can see that the negative logarithm transforms an Exponential
random variable into a Gumbel random variable, located at $\log\theta$ (right-
hand side of Eq. 12). Note that for $\theta=1$, this Gumbel turns into a
standard Gumbel variable with $\mu=0$. Rows I and II in Fig. 3 relate both the
standard and parameterized Uniform/Beta, Exponential and Gumbel distributions.
The green ellipses indicate the corresponding distributions.

Besides relating the distributions, we can also link some properties of the
Gumbel-max trick to properties of a Poisson process; a stochastic process
widely used for modelling the occurrence of events over time. Thanks to the
wide study of this process, many readers might have an intuitive understanding
about its properties, enabling introduction of these relations via an
intuitive example.

The Exponential distribution models the time between sequential events
occurring according to a Poisson process. The distribution’s parameter
$\theta$ represents the intensity (or rate), expressed in events per time
unit. As an example, consider a scientific poster presentation where on
average 8 students and 2 professors arrive per hour, both according to
independent Poisson processes. The Exponential distribution models the time
between subsequent arrivals of students and professors (with intensities 8 and
2, respectively). We can sample the arrival process by independently sampling
these two Exponentials. From this example we intuitively realize that
$P[I=\mathrm{student}]=\frac{\theta_{\mathrm{student}}}{\theta_{\mathrm{student}}+\theta_{\mathrm{prof.}}}$,
i.e. the probability that the first person to arrive is a student222This can
be shown by computing $P(X_{s}<X_{p})$ where $X_{s}$, $X_{p}$ are the
Exponential inter-arrival times for students/professors, respectively. equals
$\frac{8}{8+2}=0.8$. This property can be linked to the Gumbel-max trick.
Formally, given $X^{(i)}\sim\operatorname{Exponential}(1)$ i.i.d. $\forall
i\in D$, Eq. 12, and the monotonicity of the logarithmic function, the
following equalities hold:

$\displaystyle I$ $\displaystyle=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i}+G^{(i)}\\}=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i}-\log X^{(i)}\\}$
$\displaystyle=\underset{i\in
D}{\operatorname{argmin}}\\{\underbrace{~{}X^{(i)}/\theta_{i}~{}}_{\sim\operatorname{Exp.}(\theta_{i})}\\}\sim\operatorname{Cat}(\bm{\pi}).$
(13)

Thus $P[I=\omega]=\pi_{\omega}=\theta_{\omega}/\sum_{i\in D}\theta_{i}$, which
indeed coincides with our intuitively derived probability
$P[I=\mathrm{student}]$. Row III in Fig. 3 displays this relation, as well as
the relation to the Beta/Uniform distribution. The latter is known from the
field of weighted reservoir sampling [108], and can directly be deduced from
the Exponential reparameterization by taking the negative exponent of the
Exponential sample (which, in turn, switches the argmin function to argmax
thanks to $1/\exp(\cdot)$ being monotonically decreasing).

Instead of independently sampling the two Exponentials (i.e. one for students
and one for professors), we can develop an alternative but equivalent sampling
process thanks to the memoryless property of the Exponential
distribution333$P(X>t_{1}+t_{2}|X>t_{1})=P(X>t_{2})$. More specifically, we
can sample the arrivals of _people_ (students and professors) at the poster
from an Exponential with an intensity of $8+2=10$ persons per hour, and then
independently assign each arrival to either being a student or professor with
probabilities 0.8 and 0.2, respectively. The fact that this ‘merged’ Poisson
process can be sampled using an Exponential with rate $Z=\sum_{i\in
D}\theta_{i}$, relates to the max-stability property of independent Gumbel
variables. Row IV of Fig. 3 shows the relation between the three different
parameterizations. The actual value of the optimum in the Exponential
parameterization equals the shortest arrival time, i.e. the time at which the
first person arrives. Note that this time does _not_ provide any information
about the person being a student or a professor, i.e. the index and value of
the optimum are independent, analogously to the independence between $I$ and
$M$ in the Gumbel-max trick.

Figure 4: Drawing a sample from a categorical distribution can be done via
different paths. The most suitable path depends on the context in which one
would like to draw a sample. One can start from either the unnormalized log-
probabilities ${\bm{a}}$ (and a Boltzmann temperature $T$), or the normalized
probabilities $\bm{\pi}$. The green ellipses in the upper right corner of
certain nodes indicate that the node represents a random variable following
the distribution indicated in the ellipse. The different partition functions
$Z_{D}$ all normalize their corresponding node, their input is omitted for
readability reasons. The yellow square boxes are used to refer to certain
paths in the text. A Jupyter Notebook that accompanies this figure is publicly
available.555 https://github.com/iamhuijben/gumbel_softmax_sampling

####  4.1.3 Exploration vs exploitation

Sampling from a categorical distribution may in some applications adhere to an
exploration-exploitation dilemma, where exploitation of already acquired
information should be balanced with exploration of new information (e.g. when
sampling an action in reinforcement learning). Pure exploitation often results
in greedy solutions that appear sub-optimal in the long run, while
continuously exploring new directions can also cause sub-optimality and
instable or diverging behavior.

A common approach to deal with this dilemma is Boltzmann exploration [113], in
which the temperature $T$ of the Boltzmann distribution is gradually lowered
during training, therewith reducing the distribution’s entropy, and moving
from an exploration to an exploitation regime. Evidently, one can change the
Boltzmann temperature and use inverse transform sampling to draw a sample from
the categorical. On the other hand, from Eq. 9 we see that samples from a
tempered categorical distribution (i.e. $T\neq 1$) may also be drawn using the
Gumbel-max trick, either by changing the Boltzmann temperature explicitly, or
by changing the scale $\beta$ of the Gumbel distribution, from which
independent samples are being drawn. Note from Eq. 6 that such noise samples
are simply scaled (with $\beta$) samples from the standard Gumbel
distribution. Intuitively, down-scaling (i.e. setting $\beta<1$) the injected
noise in the Gumbel-max trick thus perturbs the logits in a lower extent,
therefore giving rise to a lower-entropy categorical distribution from which
we are effectively sampling. Figure 5 visually relates sampling using an
altered Boltzmann temperature (start in the bottom left and follow track
$A\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}D\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}F$) and sampling using scaled Gumbel noise (start in the bottom left
and follow track
$B\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}C\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}G$). Figure 8 in Appendix C shows results of sampling experiments
for various values of $\beta$.

This relation between the Boltzmann temperature and Gumbel noise scaling in
the Gumbel-max trick, has been used by the authors of [10], who propose
Boltzmann–Gumbel exploration (BGE), an algorithm - inspired by the Gumbel-max
trick - in which a class-dependent Gumbel noise scaling factor (i.e.
$\beta(i)$) is used to guarantee sub-linear regret in a stochastic multi-armed
bandit problem. In contrast to a class-_independent_ scaling factor, class-
_dependent_ scaling factors do _not_ admit an analytical expression for the
resultant categorical distribution. BGE was later leveraged in a recommender
system setting by the authors of [114], who more explicitly expressed the
relation between the Boltzmann temperature and Gumbel noise scaling.

A similar exploration-exploitation trade-off exists in the field of natural
text generation, where generated text should be of high quality (analogous to
exploitation), but also diverse (analogous to exploration). Next to
temperature scaling [55, 59], other algorithms have been proposed to balance
the aforementioned trade-off in this field, e.g. redistributing all mass of
the Boltzmann distribution to either the $k$ classes covering the top-$k$
highest probabilities [115], or the $p$ classes covering the top-$p$ ratio of
the total mass [116]. Note that for $\beta=0$ in Eq. 9 (i.e. no Gumbel noise
is injected), the Gumbel-max trick resorts to setting $k=1$ in the work of
[115]; i.e. all mass is redistributed to the class with the largest
probability (resulting in a zero-entropy one-hot distribution) and sampling
becomes a deterministic (argmax) function.

Figure 5: Sampling Gumbels and the corresponding argmax $I$ and maximum $M$,
can be done in a bottom-up or top-down approach. The former finds the
$\operatorname{(arg)max}$ from the perturbed logits, starting from the
Gumbels. Top-down sampling, on the other hand, starts from either $M$ and/or
$I$, and conditionally samples the corresponding perturbed logits. The maximum
and argmax are independent ($\perp\\!\\!\\!\perp$), allowing for independent
sampling of either of them in case only the other entity is known. The top-
down procedure can be applied in parallel for the entire domain $D$ (depicted)
or sequentially (not visualized).

###  4.2 Inverting the Gumbel-max trick: top-down sampling

Both the argmax and max function in Eq. 7 and Eq. 8, respectively, are
surjective/many-to-one mappings, i.e. numerous (multi)sets666A multiset or bag
is a set in which an element can occur multiple times. of perturbed logits
exist that result in the same index $I$ or maximum $M$. In certain scenarios
one can be interested in reverting this sampling process, i.e. inferring the
shifted Gumbels (i.e. perturbed logits) that may have generated a particular
sample/index or maximum. This reversion is often referred to as _top-down
sampling_ [107, 13]. To denote the standard procedure, i.e. the ‘conventional’
Gumbel-max trick in which Gumbels are sampled unconditionally, the term
_bottom-up_ sampling was introduced [107]. The notion of top-down sampling
gave rise to several applications, ranging from Gumbel-based sampling
algorithms [107, 13, 14, 117, 118], gradient estimators [119, 120, 121],
counterfactual outcome prediction in structural causal models [122], and
discrete flow models [123]. We will elaborate on the principles of top-down
sampling in this section.

We seek a set of $N$ independently perturbed logits, with a maximum value
$M=m$, located at index $I=\omega$. Given either of the two (i.e. the maximum
or index), we can - thanks to their independence - always sample the other
from $\operatorname{Cat}(\bm{\pi})$ (sampling $I$) or
$\operatorname{Gumbel}(\log Z)$ (sampling $M$). Once both the maximum and
corresponding index are known/sampled, the value of all ($N-1$) other
perturbed logits should now be restricted to be $\leq m$. Acquiring Gumbels,
given this restriction, can be achieved by drawing them from a right-truncated
Gumbel distribution [107, 124]:

$\tilde{G}_{\log\theta_{i}1,m}\sim\operatorname{TruncGumbel}(\log\theta_{i},1,m),$
(14)

with the ICDF of $\operatorname{TruncGumbel}(\mu,\beta,m)$ being:

$F^{-1}(u)=-\beta\log(\exp(\mu-\frac{m}{\beta})-\log u)+\mu.$ (15)

Analogous to the max-stability property given in Eq. 8, the maximum of a set
of independent samples from $\operatorname{TruncGumbel}(\log\theta_{i},1,m)$,
follows again a truncated Gumbel distribution:

$\underset{i\in
D}{\operatorname{max}}\\{{\tilde{G}_{\log\theta_{i},1,m}}\\}\sim\operatorname{TruncGumbel}(\log
Z,1,m).$ (16)

The top-down sampling procedure can be summarized as follows, $\forall i\in
D$:

$\displaystyle G_{\log\theta_{i},1}|\\{\omega,m\\}=\begin{cases}m&\hskip
5.69046pti=\omega\\\ -\log(\exp(-m)-\frac{\log U^{(i)}}{\theta_{i}})&\hskip
5.69046pti\neq\omega.\end{cases}$ (17)

The second case ($i\neq\omega$) expresses the ICDF of
$\operatorname{TruncGumbel}(\log\theta_{i},1,m)$. In case $m$ was sampled from
$\operatorname{Gumbel}(\log Z)$ (rather than given), $m=-\log(-\log
U^{(\omega)})+\log Z$, and $\exp(-m)$ reduces to $-\log(U^{(\omega)})/Z$.
Equivalent expressions were given by the authors of [119, App. C] (assuming
$Z=1$), and [121, Eq. (9)].

The authors of [13] observe that, when conditioning on $M=m$, we can even
directly provide an expression for all perturbed logits without explicitly
sampling $I=\omega$ first. More specifically, a set of (unconditional) Gumbels
$G_{\log\theta_{i}}\sim\operatorname{Gumbel}(\log\theta_{i})~{}\forall i\in
D$, with maximum $q=\underset{i\in
D}{\operatorname{max}}\\{G_{\log\theta_{i}}\\}$, can be transformed to
conditional Gumbels $\tilde{G}_{\log\theta_{i},1,m}$, via the respective
(inverse) CDFs. More specifically, $\forall i\in D$:

$\tilde{G}_{\log\theta_{i},1,m}:=G_{\log\theta_{i},1}|m=F_{m}^{-1}(F_{q}(G_{\log\theta_{i}})),$
(18)

where $F_{q}(\cdot)$ and $F_{m}^{-1}(\cdot)$ denote the CDF and ICDF of
truncated Gumbel distributions777The ICDF of
$\operatorname{TruncGumbel}(\mu,\beta,m)$ is provided in Eq. 15.
The CDF is expressed as
$F(x)=\exp(-\exp(-\frac{\operatorname{min}(x,m)-\mu}{\beta}))/\exp(-\exp(-\frac{m-\mu}{\beta}))$.,
both located at $\log\theta_{i}$, and truncated at $q$ and $m$, respectively.
The set $\\{\tilde{G}_{\log\theta_{i},1,m}\\}_{i\in D}$ is now ensured to have
a maximum $m$. Figure 5 visually summarizes both the bottom-up and the top-
down sampling procedure.

Rather than simultaneously sampling all conditional Gumbels for $i\neq\omega$,
they could also be acquired sequentially. This enables conditional Gumbel
sampling for applications in which one is not interested in the perturbed
logits over the entire domain $D$, because it is typically too large. The
authors of [107] note that hybrid top-down sampling formats (i.e. approaches
that are in between simultaneous and fully sequential sampling) can also be
used without loss of generality. They propose to partition the sampling space
into two (random and mutually exclusive) sub-domains, and apply the sequential
process within these domains separately. Their top-down construction algorithm
generalizes the sequential conditional sampling procedure, and is provided in
Appendix D.

###  4.3 Extended sampling algorithms

The Gumbel-max trick requires $N$ Gumbel realizations to draw one sample from
a categorical distribution of $N$ classes. This can become cumbersome when
drawing many samples, or even intractable when drawing a sample from an
exponentially large domain (i.e. large $N$). Fortunately, its theoretical
properties naturally extend to sampling algorithms with varying purposes, e.g.
sampling with or without replacement, and sampling from structured models.
This section discusses such Gumbel-based sampling algorithms, which are
summarized in Fig. 6, including conventional alternatives. The top and bottom
row display the unstructured and structured setting, as discussed in Section
4.3.1 and 4.3.2, respectively. Algorithms that require a normalized
distribution, i.e. a known partitioning function $Z$, are denoted with a
superscript $z$. Green ellipses indicate the distribution (when known), from
which algorithms return samples. Note that soft samples (first column of Fig.
6) are discussed in Section 5.

Figure 6: Overview of sampling algorithms for different scenarios: single
(discrete or soft) or multiple sample(s) from (un)structured distributions.
For all cases, both default/conventional and Gumbel-based algorithms are
given. The green ellipses indicate the distribution from which samples of the
indicated algorithms originate. In case of empty ellipses, it is unknown which
distribution the samples follow. A superscript $z$ indicates that the
algorithm requires normalized probabilities, i.e. partition function $Z$
should be known.

####  4.3.1 Unstructured distributions

Drawing a single sample from an (unstructured) categorical distribution can
among others be done using inverse transform sampling (Section 2.3), and the
Gumbel-max trick (Section 4.1.1). Figure 5 visually relates these sampling
algorithms. The $A^{*}$ sampling algorithm [107] can be interpreted as a
continuous counterpart of the Gumbel-max trick, that samples from continuous
domains (i.e. $N\rightarrow\infty$). $A^{*}$ sampling does - analogously to
the Gumbel-max trick - not require knowledge about the partition function $Z$.
It leverages top-down sampling (as explained in Section 4.2) and bounds
related to Gumbel processes [107], to diminish the number of i.i.d. Gumbels to
be drawn, while still guaranteeing exact samples from unnormalized continuous
distributions.

Sampling with replacement
Repeating a discrete sampling process, we can draw multiple (independent)
samples _with_ replacement. Depending on whether the order of samples is
considered, this gives rise to an (ordered) _sequence_ , or an (unordered)
_bag_ or _multiset_ of samples. Such sequences and bags (or certain
representations thereof) can in itself be considered samples from structured
models (see Section 4.3.2). For example, when a bag of samples is represented
as a vector of counts over the complete categorical sampling domain, this
vector has a (structured) distribution that is known as the multinomial
distribution. Formally, the counts-vector resulting from drawing $k$ samples
(with replacement) from $\operatorname{Cat}(\bm{\pi})$ is a _single_ sample
from $\operatorname{Multinomial}(\bm{\pi},k)$. Therefore, the same sampling
process may give rise to different distributions, depending on what is
considered ‘the sample’.

When repeated samples from the categorical are drawn using the Gumbel-max
trick, it may be possible to gain efficiency by reusing computations, as
proposed by the authors of [107, 125].

Sampling without replacement
When _rejecting_ (or removing) duplicates from the bag or sequence of samples
drawn with replacement, we obtain a set, respectively sequence, of _unique
samples_. This can be seen as _sampling without replacement_ , implemented
using rejection sampling, where proposals (i.e. individual samples from the
bag/sequence) get rejected if they have already been sampled. Depending on the
distribution, it may require many proposal samples to acquire a sequence or
set of $k$ unique elements (e.g. repeated sampling from a low-entropy
distribution often returns the same class). Therefore, a more common
alternative to rejection sampling is to _sequentially sample without
replacement_ by removing a sampled element from the sampling domain,
renormalize the distribution, and continue to draw the next (unique) sample
from the updated domain. The subsequent (conditional) samples are typically
sampled using standard inverse transform sampling.

If the order of the resulting unique samples is considered, this process gives
rise to a probability distribution over sequences $X$ of length $k$:

$p(X)=\prod_{i=1}^{k}\frac{p(x_{i})}{1-\sum\limits_{j<i}p(x_{j})}.$ (19)

When $k=N$, i.e. the full domain is sampled, this represents a distribution
over _permutations_ known as the Plackett-Luce model [126, 127]. When we do
not care about the order of the unique samples, we can find the probability
for a certain set of samples $\mathcal{S}$, by summing the probabilities for
all possible permutations of $\mathcal{S}$:

$p(\mathcal{S})=\sum_{X\in~{}\text{Perm}(\mathcal{S})}p(X).$ (20)

This probability distribution is a special case of the Wallenius’ noncentral
hypergeometric distribution [128, 129], and can be computed exactly in
$O(N!)$, or approximated by an integral [129, 130].

Renormalizing the distribution for sequential sampling without replacement (as
done in Eq. 19) may be difficult/intractable, especially in large domains. An
appealing alternative is to use the Gumbel-max trick on the updated domain
(being a subset of the original sampling domain), as it avoids the need for
explicit renormalization. Even more interesting, when we apply the Gumbel-max
trick repeatedly for sampling without replacement, the $N$ perturbed logits
can be _reused_ for all $k$ samples, by simply selecting the top-$k$ perturbed
logits in a single step [110, 111], see Fig. 7. This algorithm, which has
become known as Gumbel-top-$k$ sampling [14], is a strict generalization of
the Gumbel-max trick (which is the special case for $k=1$). This
parameterization of sampling without replacement allows for another derivation
of the integral to approximate Eq. 20, or to compute it exactly in $O(2^{N})$
[131].

Figure 7: Illustration of Gumbel-top-$k$ sampling. Figure adapted from [14].
Gumbel-top-$k$ perturbs all logits with i.i.d. Gumbel samples, and returns the
index of the $k$ perturbed logits with the highest values. These indices yield
$k$ independent samples without replacement from
$\operatorname{Cat}(\bm{\theta})$.

Gumbel-top-$k$ sampling is analogous to weighted reservoir sampling, as shown
by [110, 111]. The latter uses the Beta-distribution parameterization (see
Section 4.1.2), and selects the top-$k$ perturbed keys using a priority queue,
enabling sampling from streaming applications [108]. Gumbel-top-$k$ sampling
also has connections to ranking literature, as it was shown by [106] that the
Gumbel distribution is the _only_ distribution of discriminal scores in a
Thurstonian ranking model [132], inducing a ranking that satisfies Luce’s
axiom of choice [127].

Samples without replacement can be used to form statistical estimators of
functions of the underlying categorical distribution [133, 134, 135, 136,
131]. Similar ideas have been used to construct gradient estimators for the
process of sampling without replacement (see Section 5.2).

####  4.3.2 Structured distributions

The Gumbel-max trick can be seen as an instance of perturb-and-MAP [137],
which is a class of methods that transform sampling into an optimization
problem, where the sample corresponds to the optimum of a perturbed energy
function. Specifically, in the context of the Gumbel-max trick, the energies
correspond to (unnormalized) log-probabilities and the argmax operator
represents the optimization problem, i.e. ‘finding the maximum’ of the
perturbed log-probabilities/energies. Whereas finding the maximum in an
unstructured setting with a limited number of classes is hardly considered a
‘problem’, it is challenging in structured models, where there are multiple
dependent variables with a sampling domain of exponential size. Various
(approximate) sampling algorithms have been proposed, which rely on the
perturb-and-MAP principle by finding the (approximate) maximum of the
perturbed energy over all possible configurations of such structured models
[137, 138, 139, 140, 141].

Note that there is a vast amount of literature on alternatives to sample from
structured models. Standard _ancestral sampling_ can, for example, be used
when the conditional distributions for variables are normalized (i.e. the
partition function is known), whereas Markov Chain Monte Carlo (MCMC) sampling
[3] and variants thereof can be leveraged when $Z$ is unknown. In this
section, we focus on methods related to Gumbel variables, such as the perturb-
and-MAP paradigm.

The Gumbel-max trick can be seen as a _reparameterization trick_ [1, 142],
that allows to sample a categorical variable as a deterministic transformation
of independent (Gumbel) noise variables. This reparameterization property is
exploited by the authors of [118], that aim to parallelize sampling by
forecasting (reparameterized) samples from a sequential model.

Multiple samples
Whereas samples with or without replacement (from unstructured distributions)
can in itself be seen as a single sample from a structured distribution (e.g.
over sets or sequences), it is also possible to draw multiple samples (with or
without replacement) from structured distributions themselves (see bottom
right corner of Fig. 6). For example, we can draw a set (or sequence) of
(unique) sequences from a (neural) sequence model. The authors of [143, 144]
show how to efficiently compute multiple (structured) samples with replacement
from unnormalized distributions using the perturb-and-MAP paradigm. When a
structured distribution is normalized, such as in neural sequence models,
generalizations of ancestral sampling can be used to efficiently draw samples
without replacement. Examples of such algorithms are _stochastic beam search_
[13], or its generalization _ancestral Gumbel-top- $k$ sampling_ [14], which
implicitly apply Gumbel-top-$k$ sampling using top-down sampling (Section
4.2), and _UniqueRandomizer_ [117], which applies sequential sampling without
replacement using an efficient data structure to compute the necessary
renormalizations.

##  5 Gradient estimation in neural network optimization

Section 3 introduced various applications of gradient estimation of discrete
stochastic nodes in neural network optimization. Most Gumbel-based estimators
(as discussed in Section 5.2) build upon the recently introduced Gumbel-
Softmax distribution, which will be discussed in Section 5.1.

###  5.1 Gumbel-Softmax distribution

Instead of drawing hard/discrete samples from an (unstructured) categorical
distribution, one can also define soft samples, which may notably be used for
gradient estimation. To see the relation between these hard and soft samples,
we must consider the hard samples in their one-hot embedding format, i.e.
$\mathbbm{1}_{\omega}\in\\{0,1\\}^{N}$. A soft sample $S_{\lambda}$ can then
be defined as a vector of equivalent length, in which the mass is ‘spread out’
over multiple bins, rather than centered in one class. Concurrent works [5, 6]
have introduced the Gumbel-Softmax (GS; or Concrete) distribution, of which an
exact sample is a relaxation of $\mathbbm{1}_{\omega}$. We refer to [5, 6] for
(derivation of) the PDF of this distribution, here denoted with
$\operatorname{GS}(\bm{\pi},\lambda)$. More specifically, the
$i^{\mathrm{th}}$ index of soft sample $S_{\lambda}\in\\{\mathbb{R}^{N}_{\geq
0}:|S_{\lambda}|=1\\}$ is defined as:

$S_{i;\lambda}=\frac{\exp\big{(}(\log\theta_{i}+G^{(i)})/\lambda\big{)}}{\sum_{j\in
D}\exp\big{(}(\log\theta_{j}+G^{(j)})/\lambda\big{)}},$ (21)

where $\lambda$ is a temperature parameter that influences the entropy of both
the GS distribution, and the corresponding samples. Viewing these samples as
relaxations of samples from $\operatorname{Cat}(\bm{\pi})$, $\lambda$ can also
be interpreted as a parameter that defines the ‘amount of relaxation’ of soft
sample $S_{\lambda}$. In the limit of $\lambda\rightarrow 0^{+}$,
$S_{\lambda}$ converges to hard sample $\mathbbm{1}_{\omega}$, and the PDF of
the GS distribution converges to $\operatorname{Cat}(\bm{\pi})$. In fact,
drawing a GS sample can be interpreted as solving an entropy-regularized (by
$\lambda$) linear program on the probability simplex [145].

In the special case of $N=2$, i.e. when the categorical distribution converges
to the Bernoulli distribution, the softmax function in the GS distribution can
be replaced by its binary counterpart, the sigmoid function [5, 6].

Analogously to the Gumbel-max trick, $S_{\lambda}$ is invariant to scaling of
the unnormalized probabilities in $\bm{\theta}$: scaling (the entire vector
with the same constant) in probability space results in addition in log-space,
which leaves the translation-invariant softmax function in Eq. 21 unaffected.
Moreover, relaxing $\mathbbm{1}_{\omega}$, using $S_{\lambda}$, preserves
order, i.e. $S_{\lambda}$ can immediately be transformed to
$\mathbbm{1}_{\omega}$ via the argmax operation (see the upper left cell in
Fig. 6). Figure 5 visually relates hard and soft samples. Also the relation
between both temperatures $\lambda$ and $T$ (see track
$B\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}E$) is shown, which will further be discussed in Section 6.4. Figure
8 in Appendix C shows results of sampling experiments with various values of
$\lambda$. An interactive setup to repeat these experiments is provided with
this work888https://iamhuijben.github.io/gumbel_softmax_sampling.html.

Although the GS distribution has (so far) found most use as a means to relax
discrete stochastic nodes for differentiable neural network optimization, it
has recently also found use in relaxing a maximum-a-posteriori (MAP) objective
in the context of inverse problems with discrete multi-variate random
variables [146], and in solving combinatorial problems [147].

###  5.2 Gradient estimators

Most of the applications discussed in Section 3 pose modular machine learning
models that involve one or several discrete stochastic components. We
highlight the challenges posed by training such models by considering a
simplified setting with the following optimization program:

$\min_{\phi}\mathbb{E}_{X}[f(X)]\,\quad\text{ with }X\sim p_{\phi},$ (22)

where $X$ is a random variable, $p_{\phi}$ is a probability distribution
parameterized by $\phi$, and $f$ some function with range in $\mathbb{R}$. The
expectation is taken over random variable $X$ and notably depends on $\phi$.
Using gradient-based methods to optimize Eq. 22 requires an estimator999If
$p_{\theta}$ is designed to be a degenerate distribution, i.e. there is no
stochasticity (e.g. as in [15]), both the expectation and its gradient do not
need to be estimated. Therefore, we only consider cases where $p_{\theta}$ is
a non-degenerate distribution. for
$\nabla_{\phi}:=d\mathbb{E}_{X}[f(X)]/d\phi$. There are at least two groups of
estimators that could be considered, _score-function estimators_ and _pathwise
estimators_ [148].

The score-function estimator, also known as REINFORCE estimator [149, 150], is
in its simplest form given by

$\nabla_{\text{REINF}}:=f(X)\frac{\partial\log p_{\phi}(X)}{\partial\phi}.$
(23)

This estimator makes minimal assumptions on $p_{\phi}$ and $f$ and is notably
applicable when $X$ is a discrete random variable. The estimator is unbiased
[148, Section 4.3], but tends to have high variance that may render it
ineffective in practical applications. To reduce its variance, numerous
gradient estimators based on Eq. 23 have been proposed. Many of these
estimators involve control variates [151, 119, 120, 152], sometimes in
conjunction with drawing multiple samples from the discrete distribution [153,
154, 131, 155].

Pathwise estimators are also known as reparameterization gradient estimators
[156, 1]. They reparameterize the random variable $X$ using an auxiliary
random variable $\epsilon$ and a deterministic path $h(\epsilon,\phi)$, such
that $X=h(\epsilon,\phi)$ for $\epsilon\sim p(\epsilon)$. This
reparameterization can be exploited to interchange the differential and
expectation operators, giving rise to a gradient estimator of the form

$\nabla_{\text{REPARAM}}:=\frac{\partial f(X)}{\partial X}\frac{\partial
h(\epsilon,\phi)}{\partial\phi}.$ (24)

This estimator makes more assumptions on $p_{\phi}$ and $f$ [148, Section 5].
It, for example, requires $f$ to be differentiable and the reparameterization
of $X$ to be continuous. When well-defined, the pathwise estimator is unbiased
and tends to have low variance, which makes it a popular choice in practice.
Unfortunately, discrete random variables do not admit reparameterization
gradient estimators, because their reparameterizations are discontinuous (e.g.
consider the discontinuous argmax operator in the Gumbel-max trick).
Therefore, for discrete random variables, relaxed gradient estimation was
proposed [5, 6, 157]. Relaxed gradient estimators approximate a discrete
random variable with a continuous random variable that admits a
reparameterization gradient. This reparameterization gradient is then used as
a biased gradient estimator for the parameters of the discrete distribution.
Notably, the continuous random variable is only used during training to give a
gradient estimator, while at test time the model is evaluated using the
discrete random variable. For various discrete distributions, relaxed gradient
estimators have been proposed, which we now briefly survey.

####  5.2.1 Gradients for sampling unstructured distributions

When $X$ is a categorical random variable, the GS estimator [5, 6] is a
popular estimator. It approximates the categorical distribution with the GS
distribution (as presented in Section 5.1). The estimator is given by

$\nabla_{\text{GS}}:=\frac{\partial f(S_{\lambda})}{\partial
S_{\lambda}}\frac{\partial S_{\lambda}}{\partial\phi},$ (25)

where $S_{\lambda}=\operatorname{softmax}_{\lambda}\left(\phi+\bm{G}\right)$,
with $\phi=\log\bm{\theta}$ and $\bm{G}$ a vector of i.i.d. Gumbel variables
(see Eq. 21). For this estimator to be well-defined, $f$ must be defined and
differentiable on the interior of the simplex. The GS estimator reduces
variance at the expense of introducing bias: $\nabla_{\text{GS}}$ is an
unbiased estimator of
$d\mathbb{E}_{S_{\lambda}}\left[f(S_{\lambda})\right]/d\phi$, but a biased
estimator of $\nabla_{\phi}$. Empirically, the temperature parameter $\lambda$
of the GS distribution can be used to trade-off bias and variance of the
gradient estimator. Lower temperatures tend to be associated with higher
variance, but reduced bias. An important variant of the above estimator is the
straight-through GS (ST-GS) estimator [5, 121] that replaces $S_{\lambda}$
with $X$ in the first term of Eq. 25 (i.e. it uses non-relaxed computations in
the forward pass). Other work has proposed the use of alternative continuous
distributions to obtain relaxed gradient estimators for categorical
distributions [158, 157] that may lead to improvements in some applications.

####  5.2.2 Gradients for sampling structured distributions

Recent works have extended the GS estimator from categorical distributions to
structured distributions [157, and references therein]. For example, various
gradient estimators have been proposed for the case where $X$ is a subset of
fixed size [36, 111, 157]. The authors of [36] propose a relaxation that uses
repeated samples from the GS distribution to relax subset selection, and [111]
use successive applications of the softmax function. The authors of [157]
present a framework that naturally extends the Gumbel-max trick and its
relaxation to other structured variables and use it to define various relaxed
gradient estimators for subset selection. For the case where $X$ is a
permutation, the GS distribution has been generalized into the Gumbel-Sinkhorn
distribution to obtain a relaxed gradient estimator [145], and the authors of
[159] and [160] develop gradient estimators based on alternative
transformations.

##  6 Practical considerations

When using Gumbel-based (soft) sampling algorithms, several implementations
and settings might be considered. This section reviews distinct options and
commonly-made choices.

###  6.1 Parameterization of the logits

Each of the presented algorithms in Section 4 (see Fig. 6) rely on
availability of (normalized or unnormalized) probabilities of a categorical
distribution. These probabilities can either be considered trainable
parameters (i.e. they are optimized over the full training set), yielding
direct optimization, or they can be optimized indirectly, e.g. by predicting
them by a neural network. Both options will be discussed below.

####  6.1.1 Direct parameterization

When considering direct optimization of the unnormalized log-probabilities (or
logits), the number of trainable parameters typically equals the number of
classes in case of a flat/unstructured distribution. For structured
distributions, the number of trainable parameters is typically smaller than
the number of possible structures. The work of [71], e.g., distinguishes two
parameterizations for subset sampling. On the one hand $k$ samples without
replacement can be drawn from one set of logits (referred to as top$K$
sampling in their work), while on the other hand these samples could each be
drawn from their own categorical distribution (top$1$ sampling). Even though
the top$1$ parameterization allows for learning dedicated distributions for
each of the samples, this parameterization was not found to consistently
outperform the top$K$ counterpart. A similar observation was reported by [31],
where on par results were found when either sharing or not sharing learned
(via ST-GS) quantization codebooks among latent groups. Imperfect optimization
could explain these observations.

####  6.1.2 Indirect parameterization

Logits are often parameterized by a neural network when conditionality on
incoming data is required. Example applications of such data-conditionality
entail NAS [92, 93, 94], attention [39], active data acquisition [75], and
when GS is used at the output of a GAN’s generator [53]. The number of
trainable parameters that parameterize the logits now depends on the
architecture of the logit-predicting model. Despite being unrestricted in
size, such models are often designed to be simple and shallow in practice [64,
25, 39, 93, 31, 75].

Taking a different approach, the GS estimator can also be used without
explicitly predicting the logits themselves, but by inferring them from
already drawn samples from the categorical distribution [57], leveraging
conditional Gumbel sampling [107] (see Section 4.2).

####  6.1.3 Standard implementations

Independent of the logits being directly or indirectly parameterized, once the
unnormalized logits are known and one is interested in drawing one GS-sample
from a set of unnormalized logits, implementations from popular deep learning
libraries
Tensorflow101010https://www.tensorflow.org/agents/api_docs/python/tf_agents/distributions/gumbel_softmax/GumbelSoftmax
[161] and
Pytorch111111https://pytorch.org/docs/stable/generated/torch.nn.functional.gumbel_softmax.html
[162] can be used. The Pytorch implementation also facilitates the ST-GS
estimator, i.e. drawing a hard sample in the forward pass, which is not
supported by the current Tensorflow implementation. More extended Gumbel-based
sampling algorithms and corresponding gradient estimators have not yet been
implemented in standard libraries.

###  6.2 Initialization and regularization

Considerations can be made regarding the initialization and regularization of
the logits in the Gumbel-(soft)max trick. Typical initializations consider
small random values [37], e.g. drawn from a Uniform [100] or Gaussian
distribution [71], or set all logits to zero [78]. Others have based
initialization on a priori knowledge regarding the distribution [25, 72, 74],
or a desired sampling rate [92, 82, 98].

Also various regularization strategies have been proposed. In the context of
model pruning, logits often represent the (unnormalized) log-probability with
which certain elements (e.g. layers, kernels etc.) should be selected (i.e.
sampled) to be part of the model architecture. In this field, a sparsity-
promoting loss is often leveraged to penalize dense and/or large models, which
indirectly influences the logits during training [93, 88, 87, 86, 95, 98, 96].

With the goal of learning a subsampling matrix for data compression, the
authors of [71] regularize the logits more directly by promoting low-entropy
distributions. Conversely, when GS is used in discrete GANs, entropy of the
generator’s (output) distribution is often maximized rather than minimized, in
order to stabilize optimization and prevent mode collapse (which is the case
when the generator outputs non-diverse samples) [57, 43, 99]. To deal with
this diversity-quality dilemma in GANs, also distance regularization in latent
space has been leveraged [55, 61], promoting closeness of embeddings, and
therefore (indirectly) influencing the logits predicted by the generator.

Finally, regularization has also been based upon a priori knowledge, e.g.
about a prior distribution, for which the KL divergence between the learned
logits and this prior can be leveraged [50, 79], or about labels in semi-
supervised learning [30, 32]. In the latter setup, the authors of [30, 32]
penalize the categorical cross-entropy loss on the predicted probabilities in
a latent space, given some labels.

###  6.3 Setting the Gumbel-Softmax temperature $\lambda$lambda

As introduced in Section 5.1, the GS distribution introduces an additional
hyper-parameter $\lambda$. The authors of [6] show that, as long as
$\lambda\leq(N-1)^{-1}$, with $N$ being the number of classes, the GS
distribution is log-convex. In other words, it will not have a mode in the
interior of the $N$-$1$-simplex. This is a desirable property of a relaxation,
in order to not relax discreteness ‘too much’, which could result in a
considerable performance gap when subsequently evaluating the (non-relaxed)
discrete model. Despite this theoretical insight, the empirically optimal
value was, nevertheless, often found to be higher [6]. This may be a result of
the bias-variance trade-off which is present when gradient estimation is done
on GS samples. A high value of $\lambda$ induces high bias, but low variance,
while a lower value diminishes this bias, but increases the variance, and
leads to vanishing gradients for $\tau\rightarrow 0^{+}$. In line with [6],
the authors of [5] did neither find conclusive answers on how to set
$\lambda$, and used a constant value in some experiments, and an exponential
(decaying) annealing scheme in others.

Summarizing the most popular settings for $\lambda$ in practice, it can be
seen that simply fixing its value between 0.1 and 1.0 is often done both for
the GS estimator [51, 83, 87, 100] and the ST-GS estimator [56, 55, 54, 103,
97, 93, 30, 60, 86, 88, 59]. Other works have annealed the temperature with
various schemes from a value between 1.0 and 5.0 to a value in the range
0.1-0.5 [67, 83, 72, 47, 69, 31, 87, 100].

Finally, $\lambda$ could also be made trainable (or parameterized by a neural
network), providing the model the freedom to learn the optimal entropy of the
GS distribution (and the corresponding samples) [5]. The authors of [50, 38]
indeed both predict the (inverse) temperature with a (shallow) hyper-model.
Both works propose a formula that yields a temperature between 0 and 1.

###  6.4 Interaction between $\lambda$lambda and $T$T

Section 4.1.3 already explained how the Boltzmann temperature $T$ yields a
diversity-quality or exploration-exploitation trade-off, when sampling from a
Boltzmann distribution (i.e. a tempered categorical). When the GS distribution
is used to relax this categorical to enable gradient estimation, the two
different temperatures (i.e. Boltzmann temperature $T$, and GS temperature
$\lambda$) come into play. Figure 5 illustrates the relations between relaxed
and discrete samples from a categorical distribution, and shows how both
temperatures relate.

Equation (9) in Section 4.1.1 showed that sampling from a tempered categorical
can also be achieved by scaling the Gumbel noise in the Gumbel-max trick. From
track
$B\mathrel{\vbox{\hbox{\rule[0.0pt]{3.0pt}{0.0pt}}}\mkern-4.0mu\hbox{\char
41\relax}}E$ in Footnote 5 we now see that relaxing this approach with the GS
distribution, e.g. done by [163], causes an interaction effect between
$\lambda$ and $T$, yielding a GS distribution with temperature $\lambda/T$.

Instead of tuning the Boltzmann temperature to deal with the diversity-quality
dilemma in the context of discrete GANs, the authors of [58] leverage the GS
temperature $\lambda$ instead. Reasoning is as follows; a high GS temperature
$\lambda$ (low inverse temperature in their work), introduces a high bias in
the GS gradient estimator. In order to (partly) mitigate this bias, the GAN’s
generator will implicitly be pushed towards predicting a sharp distribution
(i.e. a distribution with a low implicit Boltzmann temperature $T$). Following
this reasoning, a high $\lambda$ would thus encourage low $T$, i.e.
exploitation, while a low $\lambda$ encourages exploration (large $T$), thanks
to the implicit regulation of the generator.

###  6.5 Sampled vs noise-free inference

After having trained a discrete model with the (ST-)GS estimator or variants
thereof, inference can be run in two distinct ways. Samples could either be
drawn from the trained categorical, or the class with the highest probability
can directly be selected (which is equal to the Gumbel-max trick on logits
without Gumbel perturbations). The chosen strategy has not always been
reported, but when reported mostly noise-free inference was used. Given that
sampled inference does not require adjusted code for inference, we assume that
when not reported, inference is typically done via sampling. Though, if
sampling then takes place with the value for $\lambda$ at which training
ended, this sampled inference still approaches noise-free inference in case
$\lambda$ was annealed to a low value during training [87].

In the context of NAS, the inference strategy was much more often reported
than in other applications, and it was always noise-free [82, 80, 86, 88, 96].
Given the application, this seems a logical choice, as the final goal is
typically to prune a neural network, while remaining with the best elements of
the model (indicated by the classes with the highest probabilities), rather
than to sample a model architecture. Also in the context of (hard) attention,
some authors have reported noise-free inference [37, 49]. The choice for noisy
vs noise-free inference thus typically depends on the application. When for
example the goal is to train a (discrete) generative model, it seems evident
to sample from the trained distribution, rather than to use noise-free
inference.

The authors of [86] observe boosted inference results for ensembled
predictions from sampling multiple times from the trained distribution over
channel-gates. The same authors interestingly remark that during test time,
the probabilistic gates could also be replaced by deterministic gates that are
all ’on’, in which the ST-GS estimator only serves as a (learned) regularizer
during training. This view seems analogous to the way (non-learned) dropout is
used as a regularization technique in neural network optimization.

##  7 Summary & Future perspectives

This review provides an overview of extensions and applications of the Gumbel-
max trick in machine learning. The Gumbel-max trick [4] is a method to draw an
exact sample from an unnormalized categorical distribution. It was reviewed in
Section 4.1, and depicted in Footnote 5 and 5. The Gumbel-max trick closely
links to well-known Poisson processes, and weighted reservoir sampling, as
discussed in Section 4.1.2. Figure 3 outlines the equivalences between these
three concepts.

Currently, several extensions of the Gumbel-max trick have been proposed in
the literature to create extended sampling algorithms, e.g. to sample from
structured models. These algorithms were discussed in Section 4.3, and
visually related in Fig. 6. The Gumbel-max trick gained much popularity thanks
to a proposed relaxation based on the Gumbel-Softmax (GS) distribution [5, 6]
(see Section 5.1), that enables error backpropagation for training deep neural
networks. The GS estimator and further developments in Gumbel-based gradient
estimators were discussed in Section 5.2. Applications of these estimators
were discussed in Section 3, and categorized in Fig. 2. Section 6 discussed
practical considerations and commonly-made choices regarding Gumbel-based
algorithms in a machine learning context.

Future directions can be found both from the application, as from the
theoretical perspective. At the application side, a large range of real-life
problems could benefit from highly-expressive neural networks that incorporate
a form of discrete stochasticity. Drug research, for example, has for years
been dominated by lab research, while investigations to drug-drug interactions
[34] and drug discovery [164] are currently moving forward with state-of-the-
art machine learning models. The field of logistics is another big area that
often deals with optimization over discrete spaces, making it suitable to
exploit state-of-the-art discrete stochastic models. Moreover, the every-
growing amount of acquired data poses challenges regarding their storage,
transfer and processing. Modern radar systems acquire up to several terabytes
of data per second, therewith saturating bandwidths and decelerating
processing. Compression algorithms that achieve (extremely) high compression
rates are therefore desirable. Data-driven neural compression models that
perform (discrete) data selection and/or compression (including quantization)
have already been found useful (see Section 3), but great advances in this
field are still foreseen.

From the theoretical perspective, we may expect (even more) efficient sampling
algorithms to draw one or multiple samples from (unnormalized) structured
models. We also observe a trend towards machine learning pipelines that
increasingly leverage more structured components. These include layers that
directly encode optimization problems [165, 166, 167, 168] or structured
latent variables [160, 111, 145, 157] that are tailored to a specific task at
hand. We believe that the combination of discrete solvers and differentiable
models [169, 170, 171] is promising and that the development of relaxed
(Gumbel-based) gradient estimators will continue to be a fruitful area of
research.

## Acknowledgments

I. Huijben gratefully acknowledges support from Onera Health and the project
‘OP-SLEEP’. The project ‘OP-SLEEP’ is made possible by the European Regional
Development Fund, in the context of OPZuid. W. Kool is funded by ORTEC. M.
Paulus gratefully acknowledges support from the Max Planck ETH Center for
Learning Systems. Resources used in preparing this research were provided, in
part, by the Sustainable Chemical Processes through Catalysis (Suchcat)
National Center of Competence in Research (NCCR).

## References

  * [1] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” _arXiv preprint arXiv:1312.6114_ , 2013.
  * [2] A. v. d. Oord, N. Kalchbrenner, and K. Kavukcuoglu, “Pixel recurrent neural networks,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 48, 2016, pp. 1747–1756. [Online]. Available: https://proceedings.mlr.press/v48/oord16.html
  * [3] W. K. Hastings, “Monte carlo sampling methods using markov chains and their applications,” 1970.
  * [4] E. J. Gumbel, _Statistical Theory of Extreme Values and Some Practical Applications: A series of lectures_. US Department of Commerce, 1954, vol. 33.
  * [5] E. Jang, S. Gu, and B. Poole, “Categorical reparameterization with Gumbel-softmax,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2017.
  * [6] C. J. Maddison, A. Mnih, and Y. W. Teh, “The concrete distribution: A continuous relaxation of discrete random variables,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2017.
  * [7] R. J. Williams, “Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning,” _Mach. Learn._ , vol. 8, no. 3, pp. 229–256, 1992.
  * [8] E. J. Gumbel, “Les valeurs extrêmes des distributions statistiques,” _Ann. l’institut Henri Poincaré_ , vol. 5, no. 2, pp. 115–158, 1935\.
  * [9] R. v. Mises, “La distribution de la plus grande de n valeurs,” _Rev. Math. Union Interbalcanique_ , vol. 1, pp. 141–160, 1936.
  * [10] N. Cesa-Bianchi, C. Gentile _et al._ , “Boltzmann Exploration Done Right,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2017, pp. 5094–5101.
  * [11] L. Ma, B. Ding _et al._ , “Active Learning for ML Enhanced Database Systems,” in _ACM SIGMOD Int. Conf. Manag. Data_ , vol. 17, no. 20. ACM, 2020, pp. 175–191.
  * [12] M. Firdaus, A. P. Shandeelya, and A. Ekbal, “More to diverse: Generating diversified responses in a task oriented multimodal dialog system,” _PLoS One_ , vol. 15, no. 11, 2020. [Online]. Available: https://doi.org/10.1371/journal.pone.0241271.g001
  * [13] W. Kool, H. v. Hoof, and M. Welling, “Stochastic beams and where to find them: The Gumbel-Top-k trick for sampling sequences without replacement,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_. PMLR, May 2019, pp. 3499–3508. [Online]. Available: http://proceedings.mlr.press/v97/kool19a.html
  * [14] W. Kool, H. v. Hoof, and M. Welling, “Ancestral gumbel-top-k sampling for sampling without replacement,” _Journal of Mach. Learn. Research_ , vol. 21, pp. 1–36, 2020. [Online]. Available: http://jmlr.org/papers/v21/19-985.html.
  * [15] A. Van Den Oord, O. Vinyals _et al._ , “Neural discrete representation learning,” _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , vol. 30, 2017.
  * [16] Y. Bengio, N. Léonard, and A. Courville, “Estimating or propagating gradients through stochastic neurons for conditional computation,” _arXiv preprint arXiv:1308.3432_ , 2013.
  * [17] D. Silveira, A. Carvalho _et al._ , “Topic Modeling using Variational Auto-Encoders with Gumbel-Softmax and Logistic-Normal Mixture Distributions,” in _Proceedings of the Intern. Joint Conf. on Neur. Netw. (IJCNN)_ , vol. 2018-July. IEEE, 2018, pp. 1–8.
  * [18] V. Liévin, A. Dittadi _et al._ , “Towards Hierarchical Discrete Variational Autoencoders,” in _Proceedings of the Symp. on Adv. in Approx. Bayes. Infer. (AABI)_ , 2019, pp. 1–15.
  * [19] A. Ramesh, M. Pavlov _et al._ , “Zero-shot text-to-image generation,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_. PMLR, 2021, pp. 8821–8831.
  * [20] H. Liu, L. He _et al._ , “Structured inference for recurrent hidden semi-Markov model,” in _Proceedings of the Intern. Joint Conf. on Artif. Intell. (IJCAI)_ , 2018, pp. 2447–2453.
  * [21] J. A. Figueroa and A. R. Rivera, “Is Simple Better?: Revisiting Simple Generative Models for Unsupervised Clustering,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , no. Nips, 2017, pp. 1–6.
  * [22] E. Dupont, “Learning Disentangled Joint Continuous and Discrete Representations,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2018.
  * [23] J. A. Figueroa, “Semi-supervised Learning using Deep Generative Models and Auxiliary Tasks,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , no. 1, 2019, pp. 1–9.
  * [24] C. Wang, C. Deng, and V. Ivanov, “SAG-VAE: End-to-end joint inference of data representations and feature relations,” in _Proceedings of the Intern. Joint Conf. on Neur. Netw. (IJCNN)_ , 2020, pp. 1–9.
  * [25] P. H. Chen, S. Si _et al._ , “Learning to screen for fast softmax inference on large vocabulary neural networks,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018.
  * [26] B. Gao, Y. Yang _et al._ , “Deep clustering with concrete K-means,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2020, pp. 4252–4256.
  * [27] M. Asai and A. Fukunaga, “Classical planning in deep latent space: Bridging the subsymbolic-symbolic boundary,” in _Proceedings of the Conf. on Artif. Intell. (AAAI)_ , vol. 32, no. 1, 2018.
  * [28] C. Corro and I. Titov, “Differentiable perturb-and-parse: Semi-supervised parsing with a structured variational autoencoder,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2019. [Online]. Available: http://universaldependencies.org/
  * [29] R. Baynazarov and I. Piontkovskaya, “Binary autoencoder for text modeling,” _Conf. Artif. Intell. Nat. Lang._ , no. November, pp. 139–150, 2019.
  * [30] A. Tjandra, S. Sakriani, and S. Nakamura, “End-to-end speech loss in speech chain framework via straight-through estimator,” in _Proceedings of the IEEE Intern. Conf. on Acoustics, Speech, and Signal Process. (ICASSP)_ , 2019, pp. 6281–6285.
  * [31] A. Baevski, S. Schneider, and M. Auli, “VQ-WAV2VEC: Self-supervised learning of discrete speech representations,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2020.
  * [32] Y. Fu, Y. Feng, and J. P. Cunningham, “Paraphrase Generation with Latent Bag of Words,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2019.
  * [33] P. Li and A. Tuzhilin, “Latent Multi-Criteria Ratings for Recommendations,” in _ACM Conf. Recomm. Syst._ , 2019, pp. 428–431. [Online]. Available: https://doi.org/10.1145/3298689.3347068
  * [34] Y. Dai, C. Guo _et al._ , “Drug–drug interaction prediction with Wasserstein Adversarial Autoencoder-based knowledge graph embeddings,” _Brief. Bioinform._ , pp. 1–16, 2020. [Online]. Available: https://academic.oup.com/bib/advance-article/doi/10.1093/bib/bbaa256/5943784
  * [35] M. Rezaee and F. Ferraro, “Event Representation with Sequential, Semi-Supervised Discrete Variables,” _NAACL_ , vol. June, pp. 4701–4716, 2021.
  * [36] J. Chen, L. Song _et al._ , “Learning to explain: An information-theoretic perspective on model interpretation,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 2, 2018, pp. 1386–1418.
  * [37] A. Abid, M. Fatih Balin, and J. Zou, “Concrete Autoencoders: Differentiable Feature Selection and Reconstruction,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2019.
  * [38] S. Yan, J. S. Smith _et al._ , “Hierarchical Multi-scale Attention Networks for action recognition,” _Signal Process. Image Commun._ , vol. 61, no. August 2017, pp. 73–84, 2018. [Online]. Available: https://doi.org/10.1016/j.image.2017.11.005
  * [39] Z. Kang, H. Xu _et al._ , “Learning Dynamic Graph Embedding for Traffic Flow Forecasting: A Graph Self-Attentive Method,” in _Proceedings of the IEEE Intelligent Transportation Syst. Conf. (ITSC)_ , vol. 6, no. 1. IEEE, 2019, pp. 45–52.
  * [40] D. B. Acharya and H. Zhang, “Community Detection Clustering via Gumbel Softmax,” _SN Comput. Sci._ , vol. 1, no. 5, pp. 1–11, 2020. [Online]. Available: https://doi.org/10.1007/s42979-020-00264-2
  * [41] Y. Tay, L. A. Tuan, and S. C. Hui, “Multi-pointer co-attention networks for recommendation,” _ACM SIGKDD Int. Conf. Knowl. Discov. Data Min._ , pp. 2309–2318, 2018.
  * [42] L. Zhou, Y. Chen _et al._ , “Bi-directional message passing based scanet for human pose estimation,” in _Proceedings of the IEEE Intern. Conf. on Multim. and Expo (ICME)_ , vol. 2019-July, 2019, pp. 1048–1053.
  * [43] Y. M. Huang, H. H. Tsengt, and J. T. Chien, “Stochastic fusion for multi-stream neural network in video classification,” in _Proceedings of the ASC Asia-Pacific Signal and Inf. Process. Assoc. (APSIPA)_ , no. November, 2019, pp. 69–74.
  * [44] N. T. Ngo, T. N. Nguyen, and T. H. Nguyen, “Learning to Select Important Context Words for Event Detection,” in _Proceedings of the Pacific Asia Conf. on Knowl. Disc. and Data Mining (PAKDD)_ , vol. 12085 LNAI. Springer International Publishing, 2020, pp. 756–768.
  * [45] H. Y. Tseng, H. Y. Lee _et al._ , “RetrieveGAN: Image Synthesis via Differentiable Patch Retrieval,” in _Proceedings of the Europ. Conf. on Comp. Vis. (ECCV)_ , vol. 12353, 2020, pp. 242–257.
  * [46] P. Yang, J. Chen _et al._ , “Greedy attack and gumbel attack: Generating adversarial examples for discrete data,” _J. Mach. Learn. Res._ , vol. 21, pp. 1–36, 2020.
  * [47] J. Yang, Q. Zhang _et al._ , “Modeling point clouds with self-attention and gumbel subset sampling,” in _Proceedings of the Conf. on Comp. Vis. and Pattern Recogn. (CVPR)_ , 2019, pp. 3323–3332.
  * [48] D. Bhaskar Acharya and H. Zhang, “Feature Selection and Extraction for Graph Neural Networks,” in _ACM Southeast Conf._ , 2020. [Online]. Available: https://doi.org/10.1145/3374135.3385309
  * [49] Y. Niu, H. Zhang _et al._ , “Recursive visual attention in visual dialog,” _CVPR_ , vol. June, pp. 6672–6681, 2019.
  * [50] S. Havrylov and I. Titov, “Emergence of Language with Multi-agent Games: Learning to Communicate with Sequences of Symbols,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2017.
  * [51] I. Mordatch and P. Abbeel, “Emergence of grounded compositional language in multi-agent populations,” in _Proceedings of the Conf. on Artif. Intell. (AAAI)_ , vol. 32, no. 1, 2017, pp. 1495–1502.
  * [52] I. J. Goodfellow, J. Pouget-Abadie _et al._ , “Generative adversarial nets,” in _NIPS_ , 2014.
  * [53] M. J. Kusner and J. M. Hernández-Lobato, “GANS for Sequences of Discrete Elements with the Gumbel-softmax Distribution,” _arxiv_ , 2016\. [Online]. Available: http://arxiv.org/abs/1611.04051
  * [54] B. Xu, W. Kong, and J. Chen, “Semi-supervised Image Captioning via Reconstruction,” in _Proceedings of the Intern. Conf. on Comp. Vis. (ICCV)_ , 2017, pp. 4135–4144.
  * [55] R. Shetty, M. Rohrbach _et al._ , “Speaking the Same Language: Matching Machine to Human Captions by Adversarial Training,” in _Proceedings of the Intern. Conf. on Comp. Vis. (ICCV)_ , 2017. [Online]. Available: https://goo.gl/3yRVnq
  * [56] J. Lu, A. Kannan _et al._ , “Best of Both Worlds: Transferring Knowledge from Discriminative Learning to a Generative Visual Dialog Model,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2017.
  * [57] J. Gu, D. J. Im, and V. O. K. Li, “Neural Machine Translation with Gumbel-Greedy Decoding,” _AAAI_ , vol. 32, no. 1, 2018.
  * [58] W. Nie, N. Narodytska, and A. B. Patel, “Relgan: Relational generative adversarial networks for text generation,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2019, pp. 1–20.
  * [59] T. Kojima, Y. Iwasawa, and Y. Matsuo, “Making Use of Latent Space in Language {GAN}s for Generating Diverse Text without Pre-training,” _Proc. 16th Conf. Eur. Chapter Assoc. Comput. Linguist. Student Res. Work._ , pp. 175–182, 2021. [Online]. Available: https://www.aclweb.org/anthology/2021.eacl-srw.23
  * [60] P. Chonwiharnphan, P. Thienprapasith, and E. Chuangsuwanich, “Generating Realistic Users Using Generative Adversarial Network with Recommendation-Based Embedding,” _IEEE Access_ , vol. 8, pp. 41 384–41 393, 2020.
  * [61] H. Zhao and R. P. Wildes, “On Diverse Asynchronous Activity Anticipation,” in _Proceedings of the Europ. Conf. on Comp. Vis. (ECCV)_ , vol. 12374\. Springer Science and Business Media Deutschland GmbH, Aug 2020, pp. 781–799. [Online]. Available: https://doi.org/10.1007/978-3-030-58526-6{\\\\_}46
  * [62] N. Baram, O. Ansehel _et al._ , “End-to-end differentiable adversarial imitation Learning,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 1, 2017, pp. 622–631.
  * [63] X. Wang, R. Zhang _et al._ , “Adversarial Distillation for Learning with Privileged Provisions,” _IEEE Trans. Pattern Anal. Mach. Intell._ , pp. 1–1, 2019.
  * [64] R. Shu and H. Nakayama, “Compressing word embeddings via deep compositional code learning,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018, pp. 1–13.
  * [65] W. C. Wei, C. J. Jhang _et al._ , “A Relaxed Quantization Training Method for Hardware Limitations of Resistive Random Access Memory (ReRAM)-Based Computing-in-Memory,” _IEEE J. Explor. Solid-State Comput. Devices Circuits_ , vol. 6, no. 1, pp. 45–52, 2020.
  * [66] Y. Kim, K. Kim, and S. Lee, “Adaptive Compression of Word Embeddings,” in _Proceedings of the Annual Meeting of the Assoc. for Computational Linguistics (ACL)_ , 2020, pp. 3950–3959.
  * [67] Y. H. Tsai, M. Y. Liu _et al._ , “Learning binary residual representations for domain-specific video streaming,” in _Proceedings of the Conf. on Artif. Intell. (AAAI)_ , 2018, pp. 7363–7370.
  * [68] Y. Zhang and H. Zhu, “Doc2Hash: Learning discrete latent variables for document retrieval,” in _Conf. North Am. Chapter Assoc. Comput. Linguist. Hum. Lang. Technol._ , vol. 1, 2019, pp. 2235–2240.
  * [69] Y. Yang, R. Bamler, and S. Mandt, “Improving Inference for Neural Image Compression,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2020.
  * [70] Y. C. Eldar and G. Kutyniok, _Compressed sensing: theory and applications_. Cambridge university press, 2012.
  * [71] I. A. M. Huijben, B. S. Veeling, and R. J. G. v. Sloun, “Deep probabilistic subsampling for task-adaptive compressed sensing,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2020, pp. 1–16.
  * [72] I. A. M. Huijben, B. S. Veeling _et al._ , “Learning sub-sampling and signal recovery with applications in ultrasound imaging,” _IEEE Trans. Med. Imaging_ , pp. 1–13, 2020.
  * [73] I. A. M. Huijben, B. S. Veeling, and R. J. G. v. Sloun, “Learning Sampling and Model-Based Signal Recovery for Compressed Sensing MRI,” in _Proceedings of the IEEE Intern. Conf. on Acoustics, Speech, and Signal Process. (ICASSP)_ , vol. 2020-May, Apr 2020, pp. 8906–8910.
  * [74] Y. Yang, S. Zhang _et al._ , “Deep Learning Based Antenna Selection for Channel Extrapolation in FDD Massive MIMO,” in _WCSP_ , 2020, pp. 182–187.
  * [75] H. v. Gorp, I. A. M. Huijben _et al._ , “Active deep probabilistic subsampling,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2021.
  * [76] C. Louizos, M. Reisser _et al._ , “Relaxed Quantization for Discretized Neural Networks,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , Sep 2018.
  * [77] B. Wu, Y. Wang _et al._ , “Mixed precision of convnets via differentiable architecture search,” _arXiv_ , 2018.
  * [78] G. Lacey, G. W. Taylor, and S. Areibi, “Stochastic layer-wise precision in deep neural networks,” in _Proceedings of the Conf. on Uncertainty in Artif. Intell. (UAI)_ , vol. 2, 2018, pp. 663–672.
  * [79] Y. Gal, J. Hron, and A. Kendall, “Concrete dropout,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2017.
  * [80] J. Chang, X. Zhang _et al._ , “DATA: Differentiable ArchiTecture Approximation,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , vol. 32, 2019, pp. 876–886.
  * [81] S. Xie, H. Zheng _et al._ , “SNAS: Stochastic Neural Architecture Search,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2019\.
  * [82] C. Louizos, M. Welling, and D. P. Kingma, “Learning sparse neural networks through L0 regularization,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018.
  * [83] W. Grathwohl, E. Creager _et al._ , “Gradient-based optimization of neural network architecture,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018.
  * [84] B. Wu, X. Dai _et al._ , “FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search,” in _Proceedings of the Conf. on Comp. Vis. and Pattern Recogn. (CVPR)_ , 2019.
  * [85] X. Dong and Y. Yang, “Network pruning via transformable architecture search,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2019, pp. 1–12.
  * [86] C. Herrmann, R. S. Bowen, and R. Zabih, “Channel Selection Using Gumbel Softmax,” in _Proceedings of the Europ. Conf. on Comp. Vis. (ECCV)_ , 2020, pp. 241–257.
  * [87] A. Wan, X. Dai _et al._ , “FBNetV2: Differentiable Neural Architecture Search for Spatial and Channel Dimensions,” in _Proceedings of the Conf. on Comp. Vis. and Pattern Recogn. (CVPR)_ , 2020, pp. 12 962–12 971.
  * [88] M. Kang and B. Han, “Operation-Aware Soft Channel Pruning using Differentiable Masks,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2020, pp. 5122–5131.
  * [89] L. Gonzalez-Carabarin, I. A. Huijben _et al._ , “Dynamic probabilistic pruning: A general framework for hardware-constrained pruning at different granularities,” _arXiv preprint arXiv:2105.12686_ , 2021.
  * [90] Z. Li, D. He _et al._ , “Towards binary-valued gates for robust LSTM training,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2018.
  * [91] J. Chien and C. Kuo, “Stochastic markov recurrent neural network for source separation,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2019, pp. 8072–8076.
  * [92] A. Veit and S. Belongie, “Convolutional networks with adaptive inference graphs,” in _Proceedings of the Europ. Conf. on Comp. Vis. (ECCV)_ , 2018\.
  * [93] B. E. Bejnordi, T. Blankevoort, and M. Welling, “Batch-Shaping for Learning Conditional Channel Gated Networks,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , Sep 2019.
  * [94] Y. Wang, X. Zhang _et al._ , “Dynamic Network Pruning with Interpretable Layerwise Channel Selection,” _AAAI_ , vol. 34, no. 04, pp. 6299–6306, 2020\.
  * [95] T. Verelst and T. Tuytelaars, “Dynamic convolutions: Exploiting spatial sparsity for faster inference,” in _Proceedings of the Conf. on Comp. Vis. and Pattern Recogn. (CVPR)_ , 2020, pp. 2317–2326. [Online]. Available: https://github.com/thomasverelst/dynconv.
  * [96] S. Cai, Y. Shu, and W. Wang, “Dynamic Routing Networks,” in _Proceedings of the Winter Conf. on Applications of Comp. Vis. (WACV)_ , 2021, pp. 3588–3597.
  * [97] D. Liu, H. Zhang _et al._ , “Learning to Assemble Neural Module Tree Networks for Visual Grounding,” in _Proceedings of the Intern. Conf. on Comp. Vis. (ICCV)_ , 2019.
  * [98] M. S. Schlichtkrull, N. De Cao, and I. Titov, “Interpreting graph neural networks for NLP with differentiable edge masking,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2021.
  * [99] F. J. Bragman, R. Tanno _et al._ , “Stochastic filter groups for multi-task CNNs: Learning specialist and generalist convolution kernels,” in _Proceedings of the Intern. Conf. on Comp. Vis. (ICCV)_ , 2019, pp. 1385–1394.
  * [100] P. Guo, C. Lee, and D. Ulbricht, “Learning to Branch for Multi-Task Learning,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2020.
  * [101] Z. Zhang, Y. Zhao _et al._ , “A general deep learning framework for network reconstruction and dynamics learning,” _Appl. Netw. Sci._ , vol. 4, no. 1, pp. 1–17, Dec 2019.
  * [102] K. S. Tai, R. Socher, and C. D. Manning, “Improved semantic representations from tree-structured long short-Term memory networks,” _ACL_ , vol. 1, pp. 1556–1566, 2015.
  * [103] J. Choi, K. M. Yoo, and S. G. Lee, “Learning to compose task-specific tree structures,” in _Proceedings of the Conf. on Artif. Intell. (AAAI)_ , 2018, pp. 5094–5101.
  * [104] A. Currey and K. Heafield, “Unsupervised Source Hierarchies for Low-Resource Neural Machine Translation,” in _Work. Relev. Linguist. Struct. Neural Archit. NLP_ , 2018, pp. 6–12.
  * [105] C. Su, H. Huang _et al._ , “Neural machine translation with Gumbel Tree-LSTM based encoder,” _J. Vis. Commun. Image Represent._ , vol. 71, p. 102811, Aug 2020.
  * [106] J. I. Yellott, “The relationship between Luce’s Choice Axiom, Thurstone’s Theory of Comparative Judgment, and the double exponential distribution,” _J. Math. Psychol._ , vol. 15, no. 2, pp. 109–144, 1977.
  * [107] C. J. Maddison, D. Tarlow, and T. Minka, “A* Sampling,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , vol. 27, 2014, pp. 3086–3094.
  * [108] P. S. Efraimidis and P. G. Spirakis, “Weighted random sampling with a reservoir,” _Inf. Process. Lett._ , vol. 97, no. 5, pp. 181–185, 2006.
  * [109] J. F. C. Kingman, _Poisson processes_. Clarendon Press, 1992, vol. 3.
  * [110] T. Vieira, “Gumbel-max trick and weighted reservoir sampling,” 2014, accessed on Feb. 3, 2021. [Online]. Available: http://timvieira.github.io/blog/post/2014/08/01/gumbel-max-trick-and-weighted-reservoir-sampling/
  * [111] S. M. Xie and S. Ermon, “Reparameterizable subset sampling via continuous relaxations,” in _Proceedings of the Intern. Joint Conf. on Artif. Intell. (IJCAI)_ , vol. 2019-Augus, 2019, pp. 3919–3925.
  * [112] C. J. Maddison, “A Poisson Process Model for Monte Carlo,” _Perturbations, Optim. Stat._ , pp. 1–37, 2018.
  * [113] R. S. Sutton, “Integrated architectures for learning, planning, and reacting based on approximating dynamic programming,” in _Machine learning proceedings 1990_. Elsevier, 1990, pp. 216–224.
  * [114] A. Biswas, T. T. Pham _et al._ , “Seeker: Real-Time Interactive Search,” in _Int. Conf. Knowl. Discov. Data Min._ , 2019, pp. 2867–2875. [Online]. Available: https://doi.org/10.1145/3292500.3330733
  * [115] A. Fan, M. Lewis, and Y. Dauphin, “Hierarchical neural story generation,” in _Proceedings of the Annual Meeting of the Assoc. for Computational Linguistics (ACL)_ , vol. 1, 2018, pp. 889–898.
  * [116] A. Holtzman, J. Buys _et al._ , “The curious case of neural text degeneration,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2020.
  * [117] K. Shi, D. Bieber, and C. Sutton, “Incremental sampling without replacement for sequence models,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2020, pp. 8785–8795.
  * [118] A. Wiggers and E. Hoogeboom, “Predictive Sampling with Forecasting Autoregressive Models,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2020, pp. 10 260–10 269.
  * [119] G. Tucker, A. Mnih _et al._ , “REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , vol. 30, 2017, pp. 2627–2636.
  * [120] W. Grathwohl, D. Choi _et al._ , “Backpropagation through the void: Optimizing control variates for black-box gradient estimation,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018.
  * [121] M. B. Paulus, C. J. Maddison, and A. Krause, “Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , Oct 2021.
  * [122] M. Oberst and D. Sontag, “Counterfactual off-policy evaluation with gumbel-max structural causal models,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_. PMLR, May 2019, pp. 4881–4890. [Online]. Available: http://proceedings.mlr.press/v97/oberst19a.html
  * [123] E. Hoogeboom, D. Nielsen _et al._ , “Argmax Flows and Multinomial Diffusion: Towards Non-Autoregressive Language Models,” _arxiv_ , 2021.
  * [124] C. J. Maddison and D. Tarlow, “Gumbel machinery,” 2017, accessed on Feb. 10, 2021\. [Online]. Available: https://cmaddis.github.io/gumbel-machinery
  * [125] Y. Qi, P. Wang _et al._ , “Fast Generating A Large Number of Gumbel-Max Variables,” in _Proceedings of the World Wide Web Conf. (WWW)_. ACM, 2020, pp. 796–807.
  * [126] R. Plackett, “The Analysis of Permutations,” _Appl. Stat._ , vol. 24, no. 2, pp. 193–202, 1975.
  * [127] R. D. Luce, “Individual choice behavior.” 1959.
  * [128] K. T. Wallenius, “Biased sampling; the noncentral hypergeometric probability distribution,” Stanford Univ CA Applied Mathematics and Statistics Labs, Tech. Rep., 1963.
  * [129] J. Chesson, “A non-central multivariate hypergeometric distribution arising from biased sampling with application to selective predation,” _Journal of Applied Probability_ , pp. 795–797, 1976.
  * [130] A. Fog, “Calculation methods for wallenius’ noncentral hypergeometric distribution,” _Communications in Statistics—Simulation and Computation®_ , vol. 37, no. 2, pp. 258–273, 2008.
  * [131] W. Kool, H. v. Hoof, and M. Welling, “Estimating Gradients for Discrete Random Variables by Sampling without Replacement,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , no. 2016, 2020, pp. 1–28. [Online]. Available: http://arxiv.org/abs/2002.06043
  * [132] L. L. Thurstone, “A law of comparative judgement,” _Psychological Review_ , vol. 34, pp. 278–286, 1927.
  * [133] D. Raj, “Some estimators in sampling with varying probabilities without replacement,” _Journal of the American Statistical Association_ , vol. 51, no. 274, pp. 269–284, 1956.
  * [134] M. Murthy, “Ordered and unordered estimators in sampling without replacement,” _Sankhyā: The Indian Journal of Statistics (1933-1960)_ , vol. 18, no. 3/4, pp. 379–390, 1957.
  * [135] N. Duffield, C. Lund, and M. Thorup, “Priority sampling for estimation of arbitrary subset sums,” _Journal of the ACM (JACM)_ , vol. 54, no. 6, p. 32, 2007.
  * [136] T. Vieira, “Estimating means in a finite universe,” 2017, accessed on April 8, 2021. [Online]. Available: https://timvieira.github.io/blog/post/2017/07/03/estimating-means-in-a-finite-universe/
  * [137] G. Papandreou and A. L. Yuille, “Perturb-and-MAP random fields: Using discrete optimization to learn and sample from energy models,” in _Proceedings of the Intern. Conf. on Comp. Vis. (ICCV)_ , 2011, pp. 193–200.
  * [138] D. Tarlow, R. P. Adams, and R. S. Zemel, “Randomized optimum models for structured prediction,” in _Proceedings of the Intern. Conf. on Artif. Intell. and Stats. (AISTATS)_ , vol. 22, 2012, pp. 1221–1229.
  * [139] T. Hazan, S. Maji, and T. Jaakkola, “On sampling from the Gibbs distribution with random maximum a-posteriori perturbations,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2013, pp. 1–9.
  * [140] Y. Chen and Z. Ghahramani, “Scalable discrete sampling as a multi-armed bandit problem,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 5, 2016, pp. 3691–3707.
  * [141] C. Kim, A. Sabharwal, and S. Ermon, “Exact sampling with integer linear programs and random perturbations,” in _Proceedings of the Conf. on Artif. Intell. (AAAI)_ , 2016, pp. 3248–3254.
  * [142] D. J. Rezende, S. Mohamed, and D. Wierstra, “Stochastic backpropagation and approximate inference in deep generative models,” in _”Proceedings of the Intern. Conf. on Mach. Learn. (ICML)”_. PMLR, 2014, pp. 1278–1286.
  * [143] S. Mussmann and S. Ermon, “Learning and inference via maximum inner product search,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 6, 2016, pp. 3814–3826.
  * [144] S. Mussmann, D. Levy, and S. Ermon, “Fast Amortized Inference and Learning in Log-linear Models with Randomly Perturbed Nearest Neighbor Search,” in _Conf. Uncertain. Artif. Intell._ , 2017.
  * [145] G. E. Mena, D. Belanger _et al._ , “Learning latent permutations with gumbel-sinkhorn networks,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , 2018.
  * [146] E. Beck, C. Bockelmann, and A. Dekorsy, “Concrete MAP detection: A machine learning inspired relaxation,” in _Proceedings of the IEEE Intern. ITG Workshop on Smart Antennas (WSA)_ , 2020, pp. 1–5.
  * [147] Y. Li, J. Liu _et al._ , “Gumbel-softmax-based Optimization: A simple general framework for optimization problems on graphs,” in _Int. Conf. Complex Networks Their Appl._ Springer International Publishing, 2020, pp. 879–890. [Online]. Available: http://dx.doi.org/10.1007/978-3-030-36687-2{\\\\_}73
  * [148] S. Mohamed, M. Rosca _et al._ , “Monte carlo gradient estimation in machine learning,” 2020.
  * [149] P. W. Glynn, “Likelihood ratio gradient estimation for stochastic systems,” _Communications of the ACM_ , vol. 33, no. 10, pp. 75–84, 1990.
  * [150] R. J. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning,” _Machine learning_ , vol. 8, no. 3, pp. 229–256, 1992.
  * [151] S. Tokui and I. Sato, “Evaluating the variance of likelihood-ratio gradient estimators,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , vol. 7, 2017, pp. 5244–5257.
  * [152] A. Mnih and K. Gregor, “Neural variational inference and learning in belief networks,” 2014.
  * [153] A. Mnih and D. J. Rezende, “Variational inference for monte carlo objectives,” 2016.
  * [154] W. Kool, H. v. Hoof, and M. Welling, “Buy 4 reinforce samples, get a baseline for free!” _Deep Reinf. Learn. Meets Struct. Predict. Deep. 2019 Work._ , pp. 1–14, 2019.
  * [155] M. Yin, Y. Yue, and M. Zhou, “ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_. PMLR, May 2019, pp. 7095–7104. [Online]. Available: https://github.com/ARM-gradient/ARSM.
  * [156] T. Salimans and D. A. Knowles, “Fixed-form variational posterior approximation through stochastic linear regression,” _Bayesian Analysis_ , vol. 8, no. 4, Dec 2013. [Online]. Available: http://dx.doi.org/10.1214/13-BA858
  * [157] M. B. Paulus, E. Zürich _et al._ , “Gradient Estimation with Stochastic Softmax Tricks,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2020.
  * [158] A. Potapczynski, G. Loaiza-Ganem, and J. P. Cunningham, “Invertible Gaussian reparameterization: Revisiting the Gumbel-Softmax,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2019. [Online]. Available: https://github.com/cunningham-lab/
  * [159] S. W. Linderman, G. E. Mena _et al._ , “Reparameterizing the birkhoff polytope for variational permutation inference,” _AISTATS_ , vol. 84, pp. 1618–1627, 2018.
  * [160] A. Grover, E. Wang _et al._ , “Stochastic optimization of sorting networks via continuous relaxations,” in _Proceedings of the Intern. Conf. on Learn. Repr. (ICLR)_ , Mar 2019.
  * [161] M. Abadi, A. Agarwal _et al._ , “TensorFlow: Large-scale machine learning on heterogeneous systems,” 2015, software available from tensorflow.org. [Online]. Available: https://www.tensorflow.org/
  * [162] A. Paszke, S. Gross _et al._ , “Automatic differentiation in pytorch,” 2017\.
  * [163] F. Guo, J. Boyd-Graber, and L. Findlater, “Which Evaluations Uncover Sense Representations that Actually Make Sense?” in _Lang. Resour. Eval. Conf._ , 2020, pp. 1727–1738.
  * [164] O. Wieder, S. Kohlbacher _et al._ , “A compact review of molecular property prediction with graph neural networks,” _Drug Discovery Today: Technologies_ , 2020. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S1740674920300305
  * [165] B. Amos and J. Z. Kolter, “OptNet: Differentiable optimization as a layer in neural networks,” _ICML_ , vol. 1, pp. 179–191, 2017.
  * [166] S. Gould, R. Hartley, and D. J. Campbell, “Deep Declarative Networks,” _IEEE Trans. Pattern Anal. Mach. Intell._ , vol. 8828, 2021.
  * [167] J. Djolonga and A. Krause, “Differentiable learning of submodular models,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , vol. December, no. 1, 2017, pp. 1014–1024.
  * [168] S. Tschiatschek, A. Sahin, and A. Krause, “Differentiable submodular maximization,” in _Proceedings of the Intern. Joint Conf. on Artif. Intell. (IJCAI)_ , vol. July, 2018, pp. 2731–2738.
  * [169] V. Niculae, A. F. T. Martins _et al._ , “SparseMAP: Differentiable Sparse Structured Inference,” in _Proceedings of the Intern. Conf. on Mach. Learn. (ICML)_ , 2018, pp. 3799—-3808.
  * [170] G. Lorberbom, A. Gane _et al._ , “Direct optimization through arg max for discrete variational auto-encoder,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2019, pp. 6203–6214.
  * [171] Q. Berthet, M. Blondel _et al._ , “Learning with differentiable perturbed optimizers,” in _Proceedings of the Conf. on Neur. Inf. Process. Syst. (NIPS)_ , 2020, pp. 9508—-9519.

|  Iris A. M. Huijben (Student Member, IEEE) received her B.Sc. degree in
Psychology & Technology, and her M.Sc. degree in Electrical Engineering (both
_cum laude_), in which she enrolled as a PhD student afterwards. Both degrees
were acquired from Eindhoven University of Technology, the Netherlands, where
she is also pursuing her PhD currently. Her research interests include
representation learning, compressed sensing and (biophysiological) signal
processing.
---|---
|  Wouter Kool received his M.Sc. degrees in Business Analytics and Operations
Research (both _cum laude_) from VU University in Amsterdam, The Netherlands.
He is currently an Operations Research Engineer at ORTEC and a PhD student at
the University of Amsterdam. His research focuses on theoretical and practical
aspects of machine learning applied to combinatorial optimization.
---|---
|  Max B. Paulus is a PhD student at ETH Zurich, a doctoral fellow of the Max
Planck ETH Center for Learning Systems and an ELLIS PhD Student. He holds a
M.Sc. degree in statistics from ETH Zurich and a B.Sc. degree (Hons) in
economics from the University of Cambridge. His research interests lie at the
intersection of machine learning and optimization and include latent variable
models, interpretability in machine learning, as well as machine learning for
integer programming.
---|---
|  Ruud J. G. van Sloun (Member, IEEE) received the B.Sc. and M.Sc. degrees
(_cum laude_) in electrical engineering and the Ph.D. degree (_cum laude_)
from the Eindhoven University of Technology, Eindhoven, The Netherlands, in
2012, 2014, and 2018, respectively. Since then, he has been an Assistant
Professor with the Department of Electrical Engineering at the Eindhoven
University of Technology and since January 2020 a Kickstart-AI fellow at
Philips Research, Eindhoven. From 2019-2020 he was also a Visiting Professor
with the Department of Mathematics and Computer Science at the Weizmann
Institute of Science, Rehovot, Israel. He is an NWO Rubicon laureate and
received a Google Faculty Research Award in 2020. His current research
interests include artificial intelligence and deep learning for front-end
(ultrasound) signal processing, model-based deep learning, compressed sensing,
ultrasound imaging, and probabilistic signal and image analysis.
---|---

##  Appendix A Proof Gumbel-max trick and max-stability

The Gumbel-max trick is used to draw an exact sample $\omega$ from a
categorical distribution $\operatorname{Cat}(\bm{\theta})$, defined by its
unnormalized probabilities $\bm{\theta}\in\mathbb{R}_{\geq 0}^{N}$. The
Gumbel-max trick [4] takes the following form:

$I=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i}+G^{(i)}\\}\sim{}\operatorname{Cat}(\bm{\pi}),$
(26)

with $G^{(i)}\sim\operatorname{G}(0,1)$ i.i.d. $\forall i\in D$. We will now
proof that $P[I=\omega]=\pi_{\omega}$, i.e. sampling using the Gumbel-max
trick results in exact samples from $\operatorname{Cat}(\bm{\pi})$, with
$\bm{\pi}=\bm{\theta}/\sum_{i\in D}\theta_{i}.$ We use $f_{\omega}(\cdot)$ to
denote the PDF of $\operatorname{Gumbel}(\log\theta_{\omega})$, and introduce
sub-domain $D^{\prime}=D\backslash\\{\omega\\}$.

We start the proof with intuitively understanding that $I=\omega$ can only
hold iff $G_{\log\theta_{\omega}}>G_{\log\theta_{i}}~{}\forall i\in
D^{\prime}$. By independence of the (shifted) Gumbel variables, we can
factorize the probability that all $G_{\log\theta_{i}}~{}(\text{with}~{}i\in
D^{\prime})$, are smaller than $M:=G_{\log\theta_{\omega}}$:

$\displaystyle P[I=\omega]=\mathbb{E}_{M}\bigg{[}\prod_{i\in
D^{\prime}}p(G_{\log\theta_{i}}<M)\bigg{]}$
$\displaystyle=\int_{-\infty}^{\infty}f_{\omega}(m)\prod_{i\in
D^{\prime}}p(G_{\log\theta_{i}}<m)~{}dm$
$\displaystyle=\int_{-\infty}^{\infty}f_{\omega}(m)\prod_{i\in
D^{\prime}}e^{-e^{\log\theta_{i}-m}}~{}dm$
$\displaystyle=\int_{-\infty}^{\infty}f_{\omega}(m)~{}e^{-\sum_{i\in
D^{\prime}}e^{\log\theta_{i}-m}}~{}dm$
$\displaystyle=\int_{-\infty}^{\infty}e^{\log\theta_{\omega}-m-e^{\log\theta_{\omega}-m}}e^{-\sum_{i\in
D^{\prime}}e^{\log\theta_{i}-m}}~{}dm$
$\displaystyle=\int_{-\infty}^{\infty}e^{\log\theta_{\omega}-m}e^{-\sum_{i\in
D}e^{\log\theta_{i}-m}}~{}dm$
$\displaystyle=\int_{-\infty}^{\infty}\theta_{\omega}e^{-m}e^{-e^{-m}\sum_{i\in
D}\theta_{i}}~{}dm.$
$\displaystyle\mathrm{From~{}\lx@cref{creftype~refnum}{eq:normprobs_vs_unnorm_logits}:}~{}Z=\sum_{i\in
D}\theta_{i}\mathrm{~{}and~{}}\theta_{\omega}=\pi_{\omega}Z:$ $\displaystyle
P[I=\omega]=\pi_{\omega}Z\int_{-\infty}^{\infty}e^{-m}e^{-Ze^{-m}}~{}dm.$
$\displaystyle\mathrm{Using}~{}\int_{-\infty}^{\infty}e^{-m}e^{-Ze^{-m}}~{}dm=\frac{1}{Z}:$
$\displaystyle P[I=\omega]=\pi_{\omega}.$ (27)

The max-stability property yields:

$M=\underset{i\in
D}{\operatorname{max}}\\{{\log\theta_{i}+G^{(i)}}\\}\sim{}\operatorname{Gumbel}(\log
Z).$ (28)

To proof that the maximum of a set of shifted Gumbels follows a Gumbel
distribution with location $\log Z$, we formulate the joint distribution of
the (independent) argmax and maximum. Since the (mathematical) steps are very
similar to the proof from Appendix A, we deliberately omit some of them to
prevent unnecessary clutter. Formally:

$\displaystyle p(M,I)=f_{\omega}(m)\prod_{i\in
D^{\prime}}p(G_{\log\theta_{i}}<m)$
$\displaystyle=\theta_{\omega}~{}e^{-m}e^{-e^{-m}\sum_{i\in D}\theta_{i}}$
$\displaystyle=\pi_{\omega}Z~{}e^{-m}e^{-Ze^{-m}}$
$\displaystyle=\pi_{\omega}~{}f_{\log Z}(m)$
$\displaystyle=P[I=\omega]~{}p(M),$ (29)

where $p(M)=f_{\log Z}(m)$ denotes the PDF of a Gumbel variable located at
$\log Z$, which thus implies that $M\sim\operatorname{Gumbel}(\log Z)$.

##  Appendix B Gumbel-max trick with non-standard Gumbel samples

Given a set of i.i.d. standard Gumbel random variables $G^{(i)}$, their scaled
and shifted counterparts $\beta G^{(i)}+\mu\sim\operatorname{G}(\mu,\beta)$
i.i.d. $\forall i\in D$, and the translation and scale _invariance_ of the
argmax function, we can write:

$\displaystyle I^{\prime}$ $\displaystyle=\underset{i\in
D}{\operatorname{argmax}}\\{\log\theta_{i;T}+\beta G^{(i)}+\mu\\}$
$\displaystyle=\underset{i\in D}{\operatorname{argmax}}\\{a_{i}/T+\beta
G^{(i)}+\mu\\}$ $\displaystyle=\underset{i\in
D}{\operatorname{argmax}}\\{a_{i}/T+\beta G^{(i)}\\}$
$\displaystyle{}\backslash\backslash{}$ translation invariance
$\displaystyle=\underset{i\in
D}{\operatorname{argmax}}\\{a_{i}/(T\beta)+G^{(i)}\\}.$
$\displaystyle{}\backslash\backslash{}$ scale invariance

Given the Gumbel-max trick in Eq. 26, we conclude that
$I^{\prime}\sim\operatorname{Cat}(\frac{\exp({\bm{a}}/(T\beta))}{\sum_{i\in
D}\exp(a_{i}/(T\beta))})=\operatorname{Cat}({\bm{a}},T\beta)$.

Analogously, using the translation and scale _equivariance_ of the max
function, we can write:

$\displaystyle M^{\prime}$ $\displaystyle=\underset{i\in
D}{\operatorname{max}}\\{\log\theta_{i}+\beta G^{(i)}+\mu\\}$
$\displaystyle=\underset{i\in D}{\operatorname{max}}\\{a_{i}/T+\beta
G^{(i)}+\mu\\}$ $\displaystyle=\underset{i\in
D}{\operatorname{max}}\\{a_{i}/T+\beta G^{(i)}\\}+\mu$
$\displaystyle{}\backslash\backslash{}$ translation equiv.
$\displaystyle=\beta\cdot\underset{i\in
D}{\operatorname{max}}\\{a_{i}/(T\beta)+G^{(i)}\\}+\mu.$ $\displaystyle\hskip
8.5359pt{}\backslash\backslash{}$ scale equiv.

Given the max-stability, scaling and shifting properties of Gumbels, we
conclude that $M^{\prime}\sim\operatorname{Gumbel}(\mu+\beta\log
Z^{\prime},\beta)$, with $Z^{\prime}=\sum_{i\in D}\exp(a_{i}/(T\beta))$.

##  Appendix C Sampling experiments

Figure 8 shows the results of sampling experiments with varying values for
Gumbel noise scale $\beta$ and GS temperature $\lambda$. Experiments were run
with the demonstration setup that is provided with this paper, and can be
found at https://iamhuijben.github.io/gumbel_softmax_sampling.html.

Figure 8: A categorical distribution with five classes was randomly created
(2nd row, left), from which 1 Gumbel-max sample results in a one-hot vector
(2nd row, right). When multiple (hard) samples are drawn, the (Monte-Carlo)
histogram of the one-hot encoded Gumbel-max samples takes the form of the
categorical distribution (1st row, middle). On the other hand, if the Gumbel
noise’s scale $\beta$ is altered from its default setting of 1, the
distribution from which we effectively sample has lower entropy for $\beta<1$
(1st row, left), or higher entropy (i.e. more uniform) for $\beta>1$ (1st row,
right). One GS sample ‘spreads’ the mass over the different classes (as
compared to the Gumbel-max sample), and varies for GS temperature $\lambda$
(3rd row). The expectation of multiple GS samples converges to the categorical
distribution for $\lambda\rightarrow 0^{+}$, while it becomes more uniform for
higher values of $\lambda$ (4rd row).

##  Appendix D Top-down sampling

Algorithm 1 provides the top-down construction algorithm from [107], rewritten
in discrete sampling space and with notations as used in the rest of this
paper.

Algorithm 1 Top-Down construction

Sampling space $D\in\\{1,\ldots,N\\}$, unnormalized probabilities
$\bm{\theta}\in\mathbb{R}_{\geq 0}^{N}$, partition function
$Z_{D}(\bm{\theta})=\sum_{i\in D}\theta_{i}$, queue $Q$

$D_{1}\leftarrow D$

$m_{1}\sim\operatorname{Gumbel}(\log Z_{D_{1}}(\bm{\theta}))$

$\omega_{1}\sim\operatorname{Cat}(\bm{\theta}/Z_{D_{1}}(\bm{\theta}))$

$Q.push(1)$

$k\leftarrow 1$

while $!Q.empty()$ do

$j\leftarrow Q.pop()$

$L,R\leftarrow partition(D_{j}-\\{\omega_{j}$})

for $C\in\\{L,R\\}$ do

if $C\neq\emptyset$ then

$k\leftarrow k+1$

$D_{k}\leftarrow C$

$m_{k}\sim\operatorname{TruncGumbel}(\log Z_{D_{k}}(\bm{\theta}),1,m_{j})$

$\omega_{k}\sim\operatorname{Cat}(\frac{\mathbbm{1}(i\in
D_{k})\bm{\theta}}{Z_{D_{k}}(\bm{\theta})})$

$Q.push(k)$

yield ($m_{k},\omega_{k}$)

end if

end for

end while