While studying the MADDPG paper, I encountered the use of the Gumbel-Softmax estimator in discrete communication environments. Further exploration revealed that this technique has broad applications, including various GANs in deep learning and algorithms like A2C and MADDPG in reinforcement learning. The Gumbel-Softmax Trick can be applied whenever reparameterization is needed on discrete distributions.
This article reflects my understanding based on several resources listed below. For deeper insights, readers are encouraged to refer to the original materials.
- The Humble Gumbel Distribution
- The Gumbel-Max Trick for Discrete Distributions
- The Gumbel-Softmax Trick for Inference of Discrete Variables
- How to Understand the Gumbel-Max Trick?
The article begins with an intuitive introduction, explaining where and how to apply the Gumbel-Softmax Trick, followed by practical coding examples to observe Gumbel distribution behavior, and concludes with mathematical proof.
Table of Contents
- I. Application Areas of the Gumbel-Softmax Trick
- Problem Statement
- Reparameterization Trick
- Gumbel-Softmax Trick
- II. Sampling Behavior of the Gumbel Distribution
- III. Mathematical Proof
I. Application Areas of the Gumbel-Softmax Trick
Problem Statement
In reinforcement learning, discrete action spaces often result in network outputs of one-hot vectors representing actions. For instance, four actions—up, down, left, right—are represented as [1,0,0,0], [0,1,0,0], etc., respectively. The selected action corresponds to the index with the highest value, i.e., \(\arg\max(v)\).
For example, given output vector \(v = [-20, 10, 9.6, 6.2]\), the second dimension has the maximum value (10), so the action is [0,1,0,0]. However, this approach does not allow gradient computation, which prevents network updates.
Typically, applying softmax normalizes the vector, enabling gradient calculation while preserving probabilistic interpretations. The softmax function is defined as:
[\sigma(z_i)=\frac{e^{z_i}}{\sum\limits_{j=1}^Ke^{z_j}} ]
Applying softmax to \(v = [-20, 10, 9.6, 6.2]\) gives \(\sigma(v) = [0, 0.591, 0.396, 0.013]\). While this maintains the same selection outcome, it also amplifies the differences in probabilities, aiding in training towards one-hot outputs—a common practice in classification tasks.
Yet, there's another issue: the resulting probability vector doesn't truly reflect probability in a meaningful way. For instance, both \([0, 0.591, 0.396, 0.013]\) and \([0, 0.9, 0.1, 0]\) yield the same choice, yet their probabilistic interpretations differ significantly. Thus, a method is needed that both selects actions and respects probability semantics.
A straightforward solution is sampling according to the probability distribution, such as using np.random.choice. However, this raises the question: how do we compute gradients? Without gradients, backpropagation for network updates becomes impossible.
Here, the reparameterization trick comes into play. It moves the sampling step out of the computational graph, allowing full gradient computation. This concept is elaborated upon in detail elsewhere. In short, reparameterization enables sampling in scenarios where gradients are required, such as in VAEs and GANs.
Reparameterization Trick
Standard autoencoders typically consist of input/output networks with an intermediate feature vector. VAEs, however, work differently by extracting distributional parameters instead of direct features. During decoding, samples are drawn from these distributions.
Reparameterization addresses the gradient issue by transforming random variables into deterministic nodes. For instance, if \(z \sim N(x, \phi^2)\), then \(z = x + \epsilon \cdot \phi\) where \(\epsilon \sim N(0,1)\). This reformulation removes randomness from the computation graph, enabling backpropagation while keeping \(\epsilon\) fixed during training.
Gumbel-Softmax Trick
Similar principles apply to discrete distributions. First, sampling must produce meaningful discrete outcomes rather than just selecting the highest probability. Second, gradients must be computable. The Gumbel-Softmax Trick achieves this through:
For an \(n\)-dimensional probability vector \(\pi\), add independent standard Gumbel noise to each element before taking the argmax:
[x_{\pi}=\arg\max(\log(\pi_i)+G_i) ]
Where \(G_i\) follows the standard Gumbel distribution with CDF \(F(x)=e^{-e^{-x}}\). Since argmax is non-differentiable, the Softmax function replaces it, forming the Gumbel-Softmax Trick. Gumbel samples can be generated from uniform distribution via \(G_i = -\log(-\log(U_i)), U_i \sim U(0,1)\).
Implementation steps:
- Generate \(n\) independent uniform samples \(\epsilon_1, ..., \epsilon_n\).
- Compute \(G_i = -\log(-\log(\epsilon_i))\).
- Add to logits: \(v' = [v_1 + G_1, v_2 + G_2, ..., v_n + G_n]\).
- Apply softmax with temperature parameter \(\tau\):
[\sigma_{\tau}(v'_i)=\frac{e^{v'i/\tau}}{\sum\limits{j=1}^ne^{v'_j/\tau}} ]
This introduces stochasticity, enhancing exploration in RL and simulating distribution samples in DL tasks.
II. Sampling Behavior of the Gumbel Distribution
Why does using Gumbel distribution generate samples that accurately mimic discrete distributions? The following code demonstrates this phenomenon.
The probability density function of the Gumbel distribution is:
[p(x)=\frac{1}{\beta}e^{-z-e^{-z}} ]
Where \(z = \frac{x-\mu}{\beta}\).
Gumbel distributions model extreme values. Consider drinking water: if you drink multiple times daily, each volume follows a normal distribution. The maximum volume per day follows a Gumbel distribution. This applies generallly to exponential family distributions.
Here’s Python code for visualizing this:
from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as plt
mean_hunger = 5
samples_per_day = 100
n_days = 10000
samples = np.random.normal(loc=mean_hunger, size=(n_days, samples_per_day))
daily_maxes = np.max(samples, axis=1)
def gumbel_pdf(prob, loc, scale):
z = (prob - loc) / scale
return np.exp(-z - np.exp(-z)) / scale
def plot_maxes(daily_maxes):
probs, hungrers, _ = plt.hist(daily_maxes, density=True, bins=100)
plt.xlabel('Volume')
plt.ylabel('Probability of Volume being daily maximum')
(loc, scale), _ = curve_fit(gumbel_pdf, hungrers[:-1], probs)
plt.plot(hungrers, gumbel_pdf(hungrers, loc, scale))
plt.figure()
plot_maxes(daily_maxes)
To demonstrate Gumbel distribution superiority, consider a multinomial distribution with 7 categories. Compare sampling methods:
n_cats = 7
cats = np.arange(n_cats)
probs = np.random.randint(low=1, high=20, size=n_cats)
probs = probs / sum(probs)
logits = np.log(probs)
def plot_probs():
plt.bar(cats, probs)
plt.xlabel("Category")
plt.ylabel("Probability")
plt.figure()
plot_probs()
Direct sampling using np.random.choice:
n_samples = 1000
def plot_estimated_probs(samples, ylabel=''):
n_cats = np.max(samples) + 1
estd_probs, _, _ = plt.hist(samples, bins=np.arange(n_cats+1), align='left', edgecolor='white', density=True)
plt.xlabel('Category')
plt.ylabel(ylabel + 'Estimated probability')
return estd_probs
def print_probs(probs):
print(' '.join(['{:.2f}'] * len(probs)).format(*probs))
samples = np.random.choice(cats, p=probs, size=n_samples)
plt.figure()
plt.subplot(1,2,1)
plot_probs()
plt.subplot(1,2,2)
estd_probs = plot_estimated_probs(samples)
plt.tight_layout()
print('Original probabilities:\t\t', end='')
print_probs(probs)
print('Estimated probabilities:\t', end='')
print_probs(estd_probs)
Output:
Original probabilities: 0.11 0.05 0.12 0.21 0.12 0.26 0.14
Estimated probabilities: 0.12 0.04 0.12 0.23 0.10 0.26 0.13
Next, compare with Gumbel, Normal, and Uniform noise:
def sample_gumbel(logits):
noise = np.random.gumbel(size=len(logits))
sample = np.argmax(logits + noise)
return sample
def sample_uniform(logits):
noise = np.random.uniform(size=len(logits))
sample = np.argmax(logits + noise)
return sample
def sample_normal(logits):
noise = np.random.normal(size=len(logits))
sample = np.argmax(logits + noise)
return sample
gumbel_samples = [sample_gumbel(logits) for _ in range(n_samples)]
uniform_samples = [sample_uniform(logits) for _ in range(n_samples)]
normal_samples = [sample_normal(logits) for _ in range(n_samples)]
plt.figure(figsize=(10, 4))
plt.subplot(1,4,1)
plot_probs()
plt.subplot(1,4,2)
gumbel_estd_probs = plot_estimated_probs(gumbel_samples, 'Gumbel ')
plt.subplot(1,4,3)
normal_estd_probs = plot_estimated_probs(normal_samples, 'Normal ')
plt.subplot(1,4,4)
uniform_estd_probs = plot_estimated_probs(uniform_samples, 'Uniform ')
plt.tight_layout()
print('Original probabilities:\t\t', end='')
print_probs(probs)
print('Gumbel Estimated probabilities:\t', end='')
print_probs(gumbel_estd_probs)
print('Normal Estimated probabilities:\t', end='')
print_probs(normal_estd_probs)
print('Uniform Estimated probabilities:', end='')
print_probs(uniform_estd_probs)
Results show Gumbel noise performs best, followed by Normal, and worst for Uniform.
III. Mathematical Proof
Why does adding Gumbel noise yield accurate results? Here’s the formal derivation.
Assume a \(K\)-dimensional output vector \(x_k\) with softmax probabilities:
[\pi_k=\frac{e^{x_k}}{\sum^K_{k'=1}e^{x_{k'}}} ]
Alternatively, add independent standard Gumbel noise \(G_k\) to each \(x_k\), then select the maximum:
[z_k = x_k + G_k \Rightarrow \text{select } \arg\max(z_k) ]
This produces the same probabilities \(\pi_k\).
Gumbel PDF and CDF: [f(z;\mu)=e^{-(z-\mu)-e^{-(z-\mu)}} ] [F(z;\mu)=e^{-e^{-(z-\mu)}} ]
Let \(z_k = x_k + G_k\), where \(G_k \sim Gumbel(0,1)\). The goal is to prove:
[P(z_k \ge z_{k'} \text{ for all } k' \ne k | {x_{k'}}) = \pi_k ]
The cumulative distribution function is:
[P(z_k \ge z_{k'} \text{ for all } k' \ne k) = \prod_{k' \ne k} e^{-e^{-(z_k - x_{k'})}} ]
Integrating over \(z_k\):
[P(z_k \ge z_{k'} \text{ for all } k' \ne k) = \int \prod_{k' \ne k} e^{-e^{-(z_k - x_{k'})}} \cdot e^{-(z_k - x_k) - e^{-(z_k - x_k)}} dz_k ]
After simplification, the integral evaluates to:
[\frac{e^{x_k}}{\sum_{k'=1}^K e^{x_{k'}}} ]
This matches the softmax expression, proving equivalence.