Elamur Research
☰
home / papers / 2016 / Categorical Reparameterization with Gumbel-Softmax

Categorical Reparameterization with Gumbel-Softmax

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive licence to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/1611.01144

Summary

Neural networks with discrete latent variables are hard to train because sampling from a categorical distribution blocks backpropagation. Earlier workarounds used either high-variance score-function (REINFORCE-style) estimators or biased straight-through tricks built for binary units. The authors introduce the Gumbel-Softmax distribution. It adds Gumbel noise to the log-probabilities and replaces the argmax of the Gumbel-Max trick with a softmax controlled by a temperature. The result is a continuous, reparameterizable relaxation on the simplex that becomes a true categorical distribution as the temperature approaches zero, so ordinary backprop can compute gradients. During training, the temperature is usually annealed from high to low. A straight-through variant uses hard one-hot samples in the forward pass and relaxed gradients in the backward pass. On MNIST benchmarks, the estimator beat other single-sample estimators such as SF, DARN, MuProp, NVIL and ST. The benchmarks covered structured output prediction, VAEs with Bernoulli or categorical latents, and semi-supervised classification. The method gives a simple, general way to use discrete choices in VAEs, attention and RL policies.

Key points

  • Temperature trade-off: a low τ gives nearly one-hot samples but noisier gradients, while a high τ gives smooth samples with lower gradient variance. If τ is learned instead of scheduled, it acts like entropy regularization.
  • Straight-Through (ST) Gumbel-Softmax: discretizes samples with argmax but backpropagates through the relaxed sample. It is meant for settings that need truly discrete values, such as discrete action spaces or quantized compression.
  • Semi-supervised speedup: in the semi-supervised model of Kingma et al. (2014), it removes the costly sum over all values of the unobserved class label. The paper reports about a 2× speedup for 10 classes and close to 10× for 100 classes, with comparable accuracy. These figures come from my knowledge of the paper's results section, which was missing from the text provided.
  • Same as Concrete: Maddison et al. (2016) found the same distribution independently and called it the "Concrete" distribution. The paper also derives its exact density in closed form (Appendix B).
  • Baseline scope: VIMCO was left out of the comparisons because it targets multi-sample objectives. The authors note that Gumbel-Softmax extends easily to that setting.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Other

Cited by

508 papers of our library cite it (of the 140673 whose bibliographies we read); 12 of them have a page here, 496 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.1611.01144 · arXiv 1611.01144

@misc{jang2016categorical,
  title = {{Categorical Reparameterization with Gumbel-Softmax}},
  author = {Jang, Eric and Gu, Shixiang and Poole, Ben},
  year = {2016},
  doi = {10.48550/arxiv.1611.01144},
  eprint = {1611.01144},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.1611.01144}
}

References

28 references, 6 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013. find
  2. Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. CoRR, abs/1606.03657, 2016. — [arXiv:1606.03657](https://arxiv.org/abs/1606.03657) find
  3. J. Chung, S. Ahn, and Y. Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint arXiv:1609.01704, 2016. find
  4. P. W Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75–84, 1990. find
  5. A. Graves, G. Wayne, M. Reynolds, T. Harley, I. Danihelka, A. Grabska-Barwińska, S. G. Colmenarejo, E. Grefenstette, T. Ramalho, J. Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626):471–476, 2016. find
  6. Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, abs/1410.5401, 2014. find
  7. K. Gregor, I. Danihelka, A. Mnih, C. Blundell, and D. Wierstra. Deep autoregressive networks. arXiv preprint arXiv:1310.8499, 2013. — [arXiv:1310.8499](https://arxiv.org/abs/1310.8499) find
  8. S. Gu, S. Levine, I. Sutskever, and A Mnih. MuProp: Unbiased Backpropagation for Stochastic Neural Networks. ICLR, 2016. find
  9. E. J. Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures. Number 33. US Govt. Print. Office, 1954. find
  10. D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013. find our card
  11. D. P. Kingma, S. Mohamed, D. J. Rezende, and M. Welling. Semi-supervised learning with deep generative models. In Advances in Neural Information Processing Systems, pp. 3581–3589, 2014. find
  12. H. Larochelle and I. Murray. The neural autoregressive distribution estimator. In AISTATS, volume 1, pp. 2, 2011. find
  13. C. J. Maddison, D. Tarlow, and T. Minka. A* sampling. In Advances in Neural Information Processing Systems, pp. 3086–3094, 2014. find
  14. C. J. Maddison, A. Mnih, and Y. Whye Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. ArXiv e-prints, November 2016. find our card
  15. A. Mnih and K. Gregor. Neural variational inference and learning in belief networks. ICML, 31, 2014. find
  16. A. Mnih and D. J. Rezende. Variational inference for monte carlo objectives. arXiv preprint arXiv:1602.06725, 2016. — [arXiv:1602.06725](https://arxiv.org/abs/1602.06725) find
  17. J. Paisley, D. Blei, and M. Jordan. Variational Bayesian Inference with Stochastic Search. ArXiv e-prints, June 2012. find
  18. Gabriel Pereyra, Geoffrey Hinton, George Tucker, and Lukasz Kaiser. Regularizing neural networks by penalizing confident output distributions. 2016. find our card
  19. J. W Rae, J. J Hunt, T. Harley, I. Danihelka, A. Senior, G. Wayne, A. Graves, and T. P Lillicrap. Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes. ArXiv e-prints, October 2016. find
  20. T. Raiko, M. Berglund, G. Alain, and L. Dinh. Techniques for learning binary stochastic feedforward neural networks. arXiv preprint arXiv:1406.2989, 2014. — [arXiv:1406.2989](https://arxiv.org/abs/1406.2989) find
  21. D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint arXiv:1401.4082, 2014a. find our card
  22. D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of The 31st International Conference on Machine Learning, pp. 1278–1286, 2014b. find our card
  23. J. T. Rolfe. Discrete Variational Autoencoders. ArXiv e-prints, September 2016. find
  24. R. Salakhutdinov and I. Murray. On the quantitative analysis of deep belief networks. In Proceedings of the 25th international conference on Machine learning, pp. 872–879. ACM, 2008. find
  25. J. Schulman, N. Heess, T. Weber, and P. Abbeel. Gradient estimation using stochastic computation graphs. In Advances in Neural Information Processing Systems, pp. 3528–3536, 2015. find
  26. C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna. Rethinking the inception architecture for computer vision. arXiv preprint arXiv:1512.00567, 2015. find
  27. R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229–256, 1992. find
  28. K. Xu, J. Ba, R. Kiros, K. Cho, A. C. Courville, R. Salakhutdinov, R. S. Zemel, and Y. Bengio. Show, attend and tell: Neural image caption generation with visual attention. CoRR, abs/1502.03044, 2015. find our card