Gumbel-softmax
In plain wordsA mathematical trick that lets a neural network make a choice among separate options, such as picking one word or one category, while still learning through small, smooth adjustments.
A continuous relaxation of sampling from a categorical distributioncategorical distribution, n.: The probabilities for a single pick among a fixed set of options, like a weighted die. A language model's guess about the next word is one of these, spread over its whole vocabulary. that makes discrete choices differentiable, so that networks with discrete latent variables can be trained by backpropagation through the reparameterization trickreparameterization trick, n.: A trick that rewrites a random draw as fixed numbers plus separate, outside randomness, so that a neural network can still learn through a step that involves chance.. It builds on the Gumbel-max trick: adding independent Gumbel(0, 1) noise to the log-probabilities and taking the arg max yields an exact categorical sample. Gumbel-softmax replaces the arg max with a softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score. at temperaturetemperature, n.: A setting that controls how adventurous an AI model's word choices are. Low temperature makes it pick the safest, most likely words; high temperature makes it more varied and surprising, but more error-prone. τ, giving a point on the probability simplex that approaches a one-hot vector as τ falls toward zero; low temperatures give nearly discrete samples but high-variance gradients, high ones smooth samples but biased gradients, so τ is often annealed. In the straight-through variant, the forward pass uses the hard one-hot sample and the backward pass the relaxed gradient. Uses include discrete variational autoencodersvariational autoencoder, n.: A neural network that learns to squeeze data into a compact, slightly fuzzy code and rebuild it, arranged so that decoding random codes produces new, plausible examples. (such as the image tokenizer of the original DALL·EDALL·E, prop. n.: OpenAI's pioneering series of image generators, named after the painter Salvador Dalí and the robot WALL·E, which showed in 2021–2023 that AI could draw almost anything described in words.), neural architecture searchneural architecture search, n.: Using an automated search, rather than human trial and error, to design the structure of a neural network: which layers to use, how many, and how to connect them., hard attention and learned routing.
We pick the codebook index with straight-through Gumbel-softmax and anneal the temperature from 1.0 to 0.1.
Origin
Introduced simultaneously in November 2016 by Eric Jang, Shixiang Gu and Ben Poole, who named it Gumbel-softmax, and by Chris Maddison, Andriy Mnih and Yee Whye Teh, who called it the Concrete (continuous relaxation of discrete) distribution; both papers appeared at ICLR 2017. The underlying Gumbel-max trick is much older: the machine-learning literature credits it to the statistician Emil Julius Gumbel's 1954 lectures on extreme values, and the same fact underlies the logit models of discrete choice in 1970s econometrics.
Usage note
The gradient is biased with respect to the true discrete objective; unbiased alternatives such as REINFORCEREINFORCE, n.: The simplest way to improve an AI's strategy by trial and error: after each attempt, make every choice in it more likely if the attempt went well and less likely if it went badly.-style score-function estimators have higher variance. Compare the plain straight-through estimatorstraight-through estimator, n.: A shortcut that lets networks learn even when they contain steps with no useful slope, such as rounding: during learning, the computer simply pretends the step passed its input straight through., which passes gradients through a hard decision as if it were the identity.
Related
- reparameterization trickreparameterization trick, n.: A trick that rewrites a random draw as fixed numbers plus separate, outside randomness, so that a neural network can still learn through a step that involves chance.
- categorical distributioncategorical distribution, n.: The probabilities for a single pick among a fixed set of options, like a weighted die. A language model's guess about the next word is one of these, spread over its whole vocabulary.
- straight-through estimatorstraight-through estimator, n.: A shortcut that lets networks learn even when they contain steps with no useful slope, such as rounding: during learning, the computer simply pretends the step passed its input straight through.
- softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score.
- temperaturetemperature, n.: A setting that controls how adventurous an AI model's word choices are. Low temperature makes it pick the safest, most likely words; high temperature makes it more varied and surprising, but more error-prone.
- variational autoencodervariational autoencoder, n.: A neural network that learns to squeeze data into a compact, slightly fuzzy code and rebuild it, arranged so that decoding random codes produces new, plausible examples.
- VQ-VAEVQ-VAE, n.: A network that compresses an image or sound into a short list of entries from a learned 'codebook' — like describing a picture with a fixed set of tiles — which other models can then generate.
- REINFORCEREINFORCE, n.: The simplest way to improve an AI's strategy by trial and error: after each attempt, make every choice in it more likely if the attempt went well and less likely if it went badly.
Sources
- Jang, E., Gu, S. & Poole, B. (2017). Categorical Reparameterization with Gumbel-Softmax. ICLR 2017.
- Maddison, C. J., Mnih, A. & Teh, Y. W. (2017). The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. ICLR 2017.