Leitura de longo prazo · 《Categorical Reparameterization with Gumbel-Softmax》(2016)
In stochastic neural networks, categorical variables are the natural choice for representing discrete structures. However, because backpropagation cannot be performed through samples, such networks rarely use categorical latent variables. The Gumbel-Softmax estimator proposed by Jang et al. aims to address this obstacle in gradient computation.
The theoretical foundation of this method stems from the Gumbel-Max trick, which provides a simple and efficient way to draw samples from a categorical distribution with class probabilities. The authors introduce a new continuous distribution, the Gumbel-Softmax distribution, which is a continuous distribution on the simplex that can approximate categorical samples, and whose parameter gradients can be easily computed via the reparameterization trick.
A key characteristic of the Gumbel-Softmax distribution is its ability to smoothly anneal into a categorical distribution. As the softmax temperature τ approaches 0, samples from the Gumbel-Softmax distribution become one-hot encoded, and the distribution becomes identical to the categorical distribution. This interpolation capability allows models to flexibly adjust the degree of discreteness during training.
In practical applications, there is a trade-off: at low temperatures, samples are close to one-hot but have high gradient variance; at high temperatures, samples are smooth but have low gradient variance. Therefore, in practice, it is common to start at a high temperature and anneal down to a small non-zero temperature. If the temperature is treated as a learnable parameter rather than a fixed schedule, this can be interpreted as entropy regularization, allowing the distribution to adaptively adjust its "confidence".
For scenarios where discrete values must be sampled (such as action spaces in reinforcement learning), the paper proposes the Straight-Through (ST) Gumbel Estimator. It uses arg max to discretize y during the forward pass, but uses the continuous approximation to estimate gradients during the backward pass. This allows samples to maintain sparsity even at high temperatures.
Experiments show that the Gumbel-Softmax estimator outperforms existing single-sample gradient estimators in structured output prediction and unsupervised generative modeling tasks. Additionally, it can be used to efficiently train semi-supervised models, avoiding expensive marginalization operations over unobserved categorical latent variables.
"Categorical Reparameterization with Gumbel-Softmax" was written by Eric Jang, Shixiang Gu, and Ben Poole, published in 2016 (ICLR 2017). It is suitable reading for machine learning practitioners researching variational autoencoders, reinforcement learning, and any field involving optimization with discrete latent variables. It is recommended to start with Section 2 regarding the definition of the distribution and the analysis of temperature effects in Figure 1.