Neural Networks gone wild! They can sample from discrete distributions now!

Abstraction: Gumbel-softmax trick enables gradient flow through discrete sampling

Key points:

• The reparameterization trick allows gradients to propagate through continuous distributions (e.g., Gaussian) but breaks down for discrete distributions because argmax is non-differentiable
• The Gumbel-max trick samples from a discrete distribution by adding i.i.d. Gumbel noise to logits and taking argmax
• Gumbel-softmax replaces argmax with softmax, making the operation differentiable; a temperature parameter τ controls approximation sharpness (τ→0 approaches one-hot, τ→∞ approaches uniform)
• High temperature gives low-variance gradients but poor approximation; practical advice is to anneal temperature from large to small during training
• Demonstrated with a toy GAN that learns a 5-class discrete distribution using TensorFlow's RelaxedOneHotCategorical; the generator's learned probabilities converge to the true distribution within ~200 epochs

Connections: Generative Adversarial Networks · Reparameterization Trick · Discrete Distributions · Gumbel Softmax

Source: https://anotherdatum.com/gumbel-gan.html