Gumbel-softmax distribution
The gumbel-softmax distribution is a continuous probability distribution on the interior of a probability simplex. It provides a differentiable relaxation of a categorical distribution by combining the Gumbel-max trick with a temperature-scaled softmax function. The same family is also known as the concrete distribution, although the two names arose from presentations emphasizing different mathematical aspects of the construction.
A gumbel-softmax random vector has nonnegative coordinates that sum to one. At finite positive temperature, its realizations are continuous rather than one-hot categorical vectors. As the temperature approaches zero, the distribution concentrates on the simplex vertices and converges to the corresponding categorical distribution. This limiting relationship permits discrete random choices to be represented by differentiable computations during the optimization of probabilistic and machine-learning models.
Historical development
The modern formulation emerged in 2016 from independent work on stochastic computation graphs containing discrete variables. Eric Jang, Shixiang Gu, and Ben Poole introduced the term “gumbel-softmax” in their treatment of differentiable categorical sampling. Their formulation centered on the reparameterized random vector and its use in gradient-based estimation.
A parallel density-oriented development by Chris J. Maddison, Andriy Mnih, Yee Whye Teh, and You Watanabe named the same distribution the “concrete distribution.” This treatment derived its density on the simplex and examined its relationship to discrete random variables. The names therefore refer to the same basic family, although “gumbel-softmax” commonly denotes the sampling construction and “concrete” commonly denotes the associated continuous distribution.
Both formulations belong to the broader development of reparameterization methods for stochastic gradient estimation. Earlier reparameterized estimators had been especially direct for continuous location-scale families, whereas categorical variables lack an ordinary pathwise derivative. The gumbel-softmax construction replaces the discontinuous categorical selection with a temperature-controlled continuous map.
Construction
Let a categorical distribution have (K) outcomes with positive unnormalized parameters
[ \alpha_1,\ldots,\alpha_K, ]
so that the probability of outcome (i) is
[ \pi_i=\frac{\alpha_i}{\sum_{j=1}^{K}\alpha_j}. ]
Independent standard Gumbel distribution variables can be generated from independent uniform variables (U_i\sim\operatorname{Uniform}(0,1)) through
[ G_i=-\log\left(-\log U_i\right). ]
The Gumbel-max construction defines the discrete index
[ I=\operatorname*{arg,max}_{1\leq i\leq K} \left(\log\alpha_i+G_i\right). ]
The resulting index satisfies
[ \Pr(I=i)=\pi_i. ]
Consequently, the one-hot vector associated with (I) is an exact sample from the specified categorical distribution. The map from the perturbed logits to the maximizing index is discontinuous, however, and therefore does not provide an ordinary pathwise derivative with respect to the parameters.
For a temperature parameter (\tau>0), the gumbel-softmax relaxation replaces the maximum operation by a softmax transformation:
[ Y_i= \frac{ \exp\left((\log\alpha_i+G_i)/\tau\right) }{ \sum_{j=1}^{K} \exp\left((\log\alpha_j+G_j)/\tau\right) }. ]
The random vector
[ Y=(Y_1,\ldots,Y_K) ]
lies in the interior of the simplex
[ \Delta^{K-1}
\left{ y\in\mathbb{R}^{K}: y_i>0,\ \sum_{i=1}^{K}y_i=1 \right}. ]
For fixed Gumbel perturbations without a tie, decreasing the temperature makes the largest perturbed logit dominate the normalized exponential expression. The vector then converges to the one-hot encoding of the Gumbel-max index:
[ \lim_{\tau\to 0^+}Y_i
\begin{cases} 1, & i=I,\ 0, & i\neq I. \end{cases} ]
Ties occur with probability zero because the Gumbel variables have continuous distributions. At the opposite limit, the influence of all finite perturbed logits vanishes relative to the temperature, giving
[ \lim_{\tau\to\infty}Y_i=\frac{1}{K}. ]
Thus, temperature changes the geometric concentration of the distribution while leaving the zero-temperature categorical probabilities determined by the relative values of the (\alpha_i).
Probability density
For (y) in the interior of the ((K-1))-dimensional simplex, the concrete density with parameters (\alpha_1,\ldots,\alpha_K) and temperature (\tau) is
[ p_{\alpha,\tau}(y)
\Gamma(K)\tau^{K-1} \left( \sum_{i=1}^{K}\alpha_i y_i^{-\tau} \right)^{-K} \prod_{i=1}^{K} \alpha_i y_i^{-\tau-1}, ]
where (\Gamma) is the gamma function. The density is defined relative to the natural ((K-1))-dimensional measure on the simplex rather than (K)-dimensional Lebesgue measure, because the coordinate-sum constraint confines the distribution to a lower-dimensional affine subspace.
Multiplying every (\alpha_i) by the same positive constant leaves the density unchanged. The parameters are therefore identifiable only through their ratios, in agreement with the corresponding categorical probabilities. They are often represented by unconstrained logits (\ell_i), with
[ \alpha_i=\exp(\ell_i). ]
Adding the same constant to every logit also leaves the distribution unchanged.
The concrete distribution is not a Dirichlet distribution. Both families are supported on a simplex, but their constructions and boundary behavior differ. A Dirichlet random vector is commonly obtained by normalizing independent gamma variables, whereas a concrete random vector is obtained by applying a softmax transformation to independently Gumbel-perturbed logits.
Differentiation and gradient estimation
For fixed Gumbel noise, the relaxed sample is differentiable with respect to the logits. Writing (Y_i) as a function of (\ell_j=\log\alpha_j), its derivative is
[ \frac{\partial Y_i}{\partial \ell_j}
\frac{1}{\tau} Y_i\left(\delta_{ij}-Y_j\right), ]
where (\delta_{ij}) is the Kronecker delta. This derivative enables the use of automatic differentiation through a sampled random vector.
If an objective (f) is evaluated on the relaxed vector, the expectation
[ L_\tau(\alpha)
\mathbb{E}_{G}\left[f(Y(\alpha,G,\tau))\right] ]
has a pathwise gradient under the usual differentiability and integrability conditions:
[ \nabla_\alpha L_\tau
\mathbb{E}_{G} \left[ \nabla_Y f(Y) \frac{\partial Y}{\partial\alpha} \right]. ]
A Monte Carlo estimate of this expression is unbiased for the gradient of the relaxed objective (L_\tau). It is not generally unbiased for the gradient of the original discrete objective, because evaluating (f) on a simplex-interior vector differs from evaluating it on a categorical one-hot vector.
Temperature mediates the discrepancy between the relaxed and discrete objectives. Low temperatures produce samples close to simplex vertices, but the resulting derivatives can become highly concentrated around transitions between competing categories. Higher temperatures produce smoother samples, while also increasing the difference between those samples and categorical outcomes. The relation between temperature and estimator variance depends on the objective, parameterization, and surrounding computational graph rather than following a universal monotonic rule.
Straight-through form
The straight-through gumbel-softmax estimator uses a discrete vector in the forward computation and the continuous relaxation in the backward computation. Given a relaxed sample (Y), the forward value is
[ Z=\operatorname{one_hot} \left( \operatorname*{arg,max}_{i}Y_i \right). ]
During differentiation, the derivative assigned to (Z) is the derivative of (Y). This convention combines an exactly one-hot forward representation with a nonzero surrogate derivative.
The resulting quantity is generally a biased estimator of the gradient of the discrete expected objective. It also need not equal the exact gradient of a scalar objective whose forward value is the displayed discrete computation. It is therefore classified as a surrogate-gradient estimator rather than an exact application of the pathwise derivative rule.
Role in latent-variable models
In a latent-variable model, a categorical latent variable can determine which component, state, or symbolic alternative contributes to an observation. Direct sampling creates a discontinuity between the latent distribution’s parameters and the selected outcome. Replacing the categorical sample with a gumbel-softmax vector changes the latent space from a finite set of vertices to the simplex interior.
Within a variational autoencoder, this replacement permits the parameters of an approximate categorical posterior to participate in a reparameterized reconstruction term. The resulting model optimizes a relaxed objective whenever downstream computations consume the continuous vector. If the forward computation instead consumes the straight-through one-hot form, optimization uses the corresponding surrogate gradient.
The relaxation does not remove the probabilistic distinction between discrete and continuous latent variables. At every positive temperature, the sample has a continuous density and assigns probability zero to each simplex vertex. Exact categorical behavior occurs in the Gumbel-max sample or in the zero-temperature limit, not in the finite-temperature distribution itself.
Terminology
“Gumbel-softmax distribution” emphasizes the transformation used to generate the random vector. “Concrete distribution” emphasizes a continuous random variable whose temperature-dependent realizations approach discrete outcomes. In mathematical treatments, both terms ordinarily denote the same distributional family.
The phrase “gumbel-softmax estimator” has a broader meaning than “gumbel-softmax distribution.” It may denote a pathwise gradient estimator for a relaxed objective, or it may denote the straight-through surrogate that uses different forward and backward representations. Distinguishing the distribution from the estimator identifies whether a statement concerns probability law, limiting behavior, or gradient computation.