Gumbel-Softmax: Complete Guide — Differentiable Discrete Sampling

Summary — Gumbel-Softmax is a trick that allows training machine learning models involving discrete sampling through a differentiable mechanism. Proposed by Jang, Gu, and Poole in 2016 and Maddison et al. the same year, it solves a fundamental problem: you cannot backpropagate through a categorical sampling operation. Gumbel-Softmax replaces this operation with a continuous, differentiable approximation, enabling end-to-end training of models with discrete latent variables.

Mathematical Principle

1. The problem: non-differentiable sampling

To sample from a categorical distribution with probabilities π = [π₁, π₂, …, π_K], we traditionally use:

k = argmax_i (log(π_i) + g_i)  where g_i ~ Gumbel(0, 1)

Where g_i = -log(-log(u_i)), u_i ~ Uniform(0, 1). This is the Gumbel-Max trick, which gives exact samples. The problem: argmax has no gradient.

2. Softmax relaxation

Instead of argmax, we use softmax with a temperature parameter τ (tau):

y_i = exp((log(π_i) + g_i) / τ) / sum_j exp((log(π_j) + g_j) / τ)

• When τ → 0: y approaches one-hot (like argmax)
• When τ → ∞: y approaches the uniform distribution
• For intermediate τ: y is a smooth approximation of one-hot

And most importantly: softmax is differentiable. The gradients of L with respect to log(π) flow through the softmax:

dL/d(log(π_i)) = sum_j dL/dy_j · dy_j/d(log(π_i))

3. Hard Gumbel-Softmax

To get one-hot samples while passing gradients in forward mode, we use the straight-through estimator:

# Forward: one-hot
y_hard = one_hot(argmax(y))
# Backward: gradients from y (soft)
y = y_hard - y.detach() + y

This gives discrete outputs in forward mode (useful for text generation) while having correct gradients in backward mode.

4. Temperature annealing

In practice, we start with a high temperature (τ = 1.0, smooth distribution) and gradually reduce it during training:

τ(t) = max(τ_min, τ_0 · exp(-rate · t))

This allows initial exploration (soft softmax) and then convergence toward discrete decisions (one-hot).

Intuition

Imagine you have to choose between several paths in a forest, but you want to be able to go back and explore other options.

Gumbel-Softmax is like a smart GPS that picks a path but always leaves a small door open to go back. The temperature τ controls this trade-off:

• High temperature (τ = 5.0): the GPS is indecisive, it takes all paths at once with similar weights. This is maximum exploration.
• Low temperature (τ = 0.1): the GPS is firmly committed to a single path. This is exploitation.

By starting soft and gradually hardening, the explorer (the model) first explores all options then specializes in the best ones.

Python Implementation

1. Gumbel-Softmax from scratch

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

def sample_gumbel(shape, eps=1e-20):
    """Sample from Gumbel(0, 1)."""
    U = torch.rand(shape)
    return -torch.log(-torch.log(U + eps) + eps)

def gumbel_softmax(logits, temperature=1.0, hard=False, dim=-1):
    """Differentiable Gumbel-Softmax sampling."""
    gumbel_noise = sample_gumbel(logits.shape)
    y = logits + gumbel_noise
    y_soft = F.softmax(y / temperature, dim=dim)

    if hard:
        index = y_soft.argmax(dim=dim, keepdim=True)
        y_hard = F.one_hot(index, num_classes=y_soft.shape[dim]).float().squeeze(dim)
        y = y_hard - y_soft.detach() + y_soft
        return y
    return y_soft

# Demo: differentiable categorical sampling
logits = torch.tensor([[2.0, 1.0, 0.1],
                       [0.5, 3.0, 1.0]], requires_grad=True)

for tau in [5.0, 1.0, 0.5, 0.1]:
    sample = gumbel_softmax(logits, temperature=tau, hard=False)
    print(f'tau={tau}: {sample[0].detach().numpy().round(3)}')
    loss = sample.sum()
    loss.backward()
    print(f'  Gradient on logits[0]: {logits.grad[0].numpy().round(3)}')
    logits.grad.zero_()

2. VAE with discrete latent variables via Gumbel-Softmax

class DiscreteVAE(nn.Module):
    """VAE with categorical latent space via Gumbel-Softmax."""
    def __init__(self, input_dim=784, latent_categories=32, latent_dim=10):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256), nn.ReLU(),
            nn.Linear(256, latent_categories * latent_dim)
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_categories * latent_dim, 256), nn.ReLU(),
            nn.Linear(256, input_dim), nn.Sigmoid()
        )
        self.latent_categories = latent_categories
        self.latent_dim = latent_dim

    def encode(self, x):
        h = self.encoder(x)
        h = h.view(-1, self.latent_dim, self.latent_categories)
        return F.log_softmax(h, dim=2)

    def sample_latent(self, log_probs, temperature, hard=False):
        samples = []
        for i in range(self.latent_dim):
            log_p = log_probs[:, i, :]
            s = gumbel_softmax(log_p, temperature=temperature, hard=hard, dim=1)
            samples.append(s)
        return torch.stack(samples, dim=1)

    def forward(self, x, temperature=1.0, hard=True):
        log_probs = self.encode(x)
        z = self.sample_latent(log_probs, temperature, hard=hard)
        x_recon = self.decode(z.view(z.size(0), -1))
        return x_recon, log_probs

# Training with annealing
model = DiscreteVAE()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(50):
    temperature = max(0.1, 1.0 * np.exp(-0.05 * epoch))
    model.train()
    for batch in dataloader:
        x_recon, log_probs = model(batch, temperature=temperature, hard=True)
        recon_loss = F.binary_cross_entropy(x_recon, batch.view(batch.size(0), -1), reduction='sum')
        q = F.softmax(log_probs, dim=2)
        kl = (q * (log_probs - np.log(1.0 / model.latent_categories))).sum()
        loss = recon_loss + kl
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch} | τ: {temperature:.3f}')

Hyperparameters

Advantages

1. Differentiable: Enables gradient descent training of models with discrete latent variables, previously impossible without REINFORCE.
2. More stable than REINFORCE: Much lower variance than score function estimators (REINFORCE/RELAX).
3. Exploration/exploitation control: Temperature provides a natural mechanism to transition from exploration to exploitation.
4. Easy to implement: PyTorch provides F.gumbel_softmax() natively. Just a few lines needed.

Limitations

1. Gradient bias: The gradient is biased (especially with hard=True). At low temperature, the bias is significant.
2. Temperature sensitivity: Annealing that is too fast can trap the model in suboptimal solutions.
3. Quality vs REINFORCE: In some cases, REINFORCE with baseline can give better final results despite higher variance.

4 concrete use cases

1. VAE with discrete latent space

Unlike classical VAEs with continuous latents (Gaussians), a VAE with discrete latents learns a structured categorical latent space. Each latent dimension can represent a discrete attribute like “the digit is handwritten or typed” on MNIST.

2. Text generation with differentiable word selection

In text generation models, Gumbel-Softmax enables selecting words in a differentiable way during training. This allows direct optimization of non-differentiable metrics like BLEU or ROUGE.

3. Neural network with selective attention

Instead of soft attention over all tokens, Gumbel-Softmax can be used to differentiably select a subset of tokens. This creates sparse models that only read the relevant parts of the text, saving computation.

4. Neuron quantization

In model compression, Gumbel-Softmax can be used to learn optimal weight quantization: instead of forcing weights to discrete values, the quantization centroids are learned in a differentiable way.

Conclusion

Gumbel-Softmax is a fundamental contribution to deep learning: it opened the door to training models with discrete decisions, a domain previously dominated by high-variance methods like REINFORCE.

Although gradient bias remains a challenge, progressive temperature annealing and improved variants have made this technique reliable enough for real-world applications in text generation, quantization, and selective attention.

See also

• Calculer la Longueur de l’Union de Segments en Python: Guide Complet et Code Optimisé
• Découvrez les Secrets du Pandigital Prime avec Python : Guide Complet et Astuces de Programmation
• #135 VAE — The classical VAE with continuous latents, precursor to the discrete VAE with Gumbel-Softmax.
• #073 MLP — The basic architecture used in encoders/decoders.
• #098 Transformer — The architecture where Gumbel-Softmax can be used for selective attention.

Comparison with alternatives

Gumbel-Softmax vs REINFORCE

Gumbel-Softmax vs Relaxed Bernoulli

The binary version of Gumbel-Softmax (two categories) is called Relaxed Bernoulli:
y = sigmoid((log(π/(1-π)) + g₁ - g₂) / τ)
This is mathematically equivalent to Gumbel-Softmax with K=2, but more efficient since only one Gumbel sample is needed instead of two.

Extensions and variants

1. Gumbel-Softmax with corrected bias

To reduce gradient bias, Gumbel-Softmax can be combined with a correction term:
∇L_corrected = ∇L_gs + β · (∇L_reinforce - ∇L_gs_baseline)
Where β increases progressively to compensate for the remaining bias.

2. Categorical Reparameterization with Gompertz

An alternative to Gumbel is to use a relaxation based on the Gompertz distribution, which offers more stable gradients for certain architectures.

3. Improved Straight-Through

The standard straight-through estimator (y_hard – y_soft.detach() + y_soft) can be improved with a scale-aware version:
y = scale_factor · (y_hard - y_soft).detach() + y_soft
Where scale_factor is dynamically adapted to compensate for the magnitude difference between soft and hard gradients.

—