This document proposes the Gumbel-Softmax distribution as a way to differentially sample from a categorical distribution. It describes how categorical variables are non-differentiable to train with backpropagation. The REINFORCE algorithm uses the likelihood ratio to estimate gradients but has high variance. Gumbel-Softmax approximates the categorical with a continuous relaxation using the Gumbel-Max trick, allowing gradients to flow. It shows this continuous relaxation performs similarly to the discrete categorical distribution but is differentiable, enabling lower-variance training of models with categorical latent variables.

1. Categorical Reparameterization
with Gumbel-Softmax
Eric Jang et al., ICLR 2017
A summary by
Jin-Hwa Kim
BI Lab. SNU

2. Table of Contents
• Problem definition
• Score function estimator (a policy gradient method, REINFORCE)
• Gumbel-Softmax
• Implementations

3. Problem Definition
Feed-forward Sampling
Linear
Softmax
Cross Entropy Loss
Input x
Linear
Softmax
Input x
label
Sampler
Environment
reward𝜕ℒ/ 𝜕𝑝
𝜕𝑝/ 𝜕𝑜
𝜕𝑝/ 𝜕𝑜
𝜕𝑜/ 𝜕𝑥
𝜕𝑜/ 𝜕𝑥
Sampling is a selection following
a given probability distribution!
Non-differentiable

4. Score Function Estimator
• A policy gradient method
• REINFORCE (Williams, 1992)

5. Score Function
• Assume that we know the gradient
• Likelihood ratio estimator (Glynn, 1990) uses the following identity
score function
Use MC Sampling
Gradient of
policy networks
0.5
0.2
0.3
a b c

6. Monte-Carlo Sampling
• Update parameters by stochastic gradient descent
i-th action
For the other actions,
gradients are zero.
Since 𝜋 𝜃
(𝑖)
is sampled following a probabilistic distribution 𝜋 𝜃,
we can use 𝛻𝜃log𝜋 𝜃
(𝑖)
in the batch learning.

7. Monte-Carlo Policy Gradient (REINFORCE)
• REINFORCE only require that 𝜋 𝜃 is continuous in 𝜃
• The variance of SF is linearly proportional to sampling space (Rezende et
al., 2014a)
score function
reward baseline
MSE to reduce variance
batch learning using MC sampling

8. Reducing Variance using a Baseline
• Reduce variance without changing expectation

10. Cross Entropy vs. REINFORCE
Cross Entropy REINFORCE
pq r
𝜋
similar to cross entropy

11. Cross Entropy vs. REINFORCE
Cross Entropy REINFORCE
pq r
𝜋
Likewise to cross entropy…
Input of
SoftMax
learning rate
advantage

12. Derivative w.r.t. Input of SoftMax

13. Comparison
Feed-forward Sampling
Linear
Softmax
Loss
Input x
Linear
Softmax
Input x
label
Sampler
Environment
reward
𝜕ℒ/ 𝜕𝑝
𝜕𝑝/ 𝜕𝑜
𝜕𝑝/ 𝜕𝑜
𝜕𝑜/ 𝜕𝑥
𝜕𝑜/ 𝜕𝑥
REINFORCE
Linear
Softmax
Input x
Sampler
Environment
reward
𝜕𝑝/ 𝜕𝑜
𝜕𝑜/ 𝜕𝑥
𝜕log𝑝/ 𝜕𝑝

14. Gist of Gumbel-Softmax
• A efficient gradient estimator that
• replace the non-differentiable sample from a categorical distribution
• with a differentiable sample from the Gumbel-Softmax distribution

15. Distribution vs. Sampling
A probability distribution Sampling
0.5
0.2
0.3
a b c
a
b
c
Sampling
discrete values

16. Distribution vs. Sampling
A probability distribution Sampling
0.5
0.2
0.3
a b c
a
b
c
a b c
a b c
a b c
Continuous relaxations
of one-hot vectors
Sampling
discrete values
Gumbel-Max Trick

17. Gumbel-Max Trick
• To draw samples z from a categorical distribution with class probabilities 𝜋:
• 𝑔𝑖 is a i.i.d samples drawn from Gumbel(0,1)
• Gumbel(0,1) = − log − log 𝑢 , 𝑢 ~ Uniform(0,1)
• Gumbel, 1954 & Maddison et al., NIPS 2014
• Using Multinoulli distribution:
Inverse transform sampling

18. max(p)

19. • To draw samples z from a categorical distribution with class probabilities 𝜋:
• To generate k-dimensional sample vectors y
• 𝜏 → 0: one-hot & categorical distribution
Softmax Replaces Argmax

20. Straight-Through Gumbel-Softmax Estimator
• Use argmax in forward pass
• But use the continuous approximation in backward pass

22. Categorical VAE Result
Test loss (negative variational lower bound) on binarized MNIST VAE with categorical latent variables (784 – (20x10) – 200).

23. REINFORCE
Torch7 PyTorch
https://github.com/pytorch/examples/blob/master/reinforcement_learning/reinforce.py
Just for a reference, not runnable code

24. Gumbel-Softmax
Torch7 PyTorch
https://github.com/pytorch/examples/blob/master/reinforcement_learning/reinforce.py
Just for a reference, not runnable code
https://github.com/Kaixhin/Autoencoders/blob/master/models/CatVAE.lua

25. Any question?