Gumbel–Softmax Relaxation

Worked Example

Consider a categorical distribution with 3 classes, parameterized by unnormalized log-probabilities (logits). The standard Gumbel-Max trick states that a discrete sample can be drawn by adding Gumbel noise to each logit and taking the argmax.

Let the logits be L = [2, 0, -2] (favoring class A), and we draw standard uniform variables U = [exp(-1), exp(-2), exp(-0.5)].

The Gumbel noise for each class is g = -ln(-ln U). So: g_A = -ln(-ln(exp(-1))) = -ln(1) = 0 g_B = -ln(-ln(exp(-2))) = -ln(2) = -0.6931471805599453 g_C = -ln(-ln(exp(-0.5))) = -ln(0.5) = 0.6931471805599453 We add the noise to the logits: L_A + g_A = 2 + 0 = 2 L_B + g_B = 0 - 0.6931471805599453 = -0.6931471805599453 L_C + g_C = -2 + 0.6931471805599453 = -1.3068528194400547 The true discrete sample is argmax([2, -0.6931471805599453, -1.3068528194400547]) = A. With Gumbel-Softmax relaxation at a temperature τ = 1, we apply softmax: exp(2/1) = 7.38905609893065 exp(-0.6931471805599453/1) = 0.5 exp(-1.3068528194400547/1) = 0.2706705664732254 sum = 7.38905609893065 + 0.5 + 0.2706705664732254 = 8.159726665403876 Relaxed A = 7.38905609893065 / 8.159726665403876 = 0.9055519139212392 Relaxed B = 0.5 / 8.159726665403876 = 0.0612765623779938 Relaxed C = 0.2706705664732254 / 8.159726665403876 = 0.033171523700767036

As the temperature τ → 0, this relaxed vector approaches the one-hot encoding of the argmax result [1, 0, 0]. The expected result is Relaxed A = 0.9055519139212392.

Experiment

Observe the effect of the temperature parameter in the interactive Gumbel-Softmax module.

1. Start with Logits set to Favor A, Noise set to Tilted draw, and Temperature set to 1. The relaxed probabilities spread across the classes, but strongly favor A. The expected numerical result Relaxed A is 0.9055519139212392.
2. Change the temperature to 0.25. Notice how the point on the simplex sharply converges towards the corner representing class A. The entropy of the relaxed probabilities drops significantly.
3. If you set logits to "Equal" and noise to "Equal draw", all classes tie. As temperature decreases, the relaxed output stays uniformly distributed because of the exact tie, ensuring we do not falsely promise a one-hot limit when an argmax is ambiguous.

Limitations

While the Gumbel-Softmax estimator enables reparameterized gradients, it is biased. The expectation over the relaxed Gumbel-Softmax distribution is only equal to the categorical distribution expectation in the limit as temperature approaches zero.

Additionally, finding a good temperature schedule can be difficult. A high temperature leads to low variance but high bias, while a very low temperature provides low bias but high gradient variance, and can cause numerical instability.

References

• Jang, E., Gu, S., & Poole, B. (2016). Categorical Reparameterization with Gumbel-Softmax.
  https://arxiv.org/abs/1611.01144

This paper introduces the Gumbel-Softmax distribution, a continuous distribution over the simplex that can approximate categorical samples. It establishes the foundations for differentiable sampling of discrete variables, overcoming a major challenge in training models with stochastic discrete nodes.