Gumbel-Softmax：用连续分布近似离散采样以解决反向传播难题

在随机神经网络中，类别变量是表示离散结构的自然选择。然而，由于无法通过样本进行反向传播，这类网络很少使用类别潜在变量。Jang等人提出的Gumbel-Softmax估计器，旨在解决这一梯度计算障碍。

该方法的理论基础源于Gumbel-Max技巧，它提供了一种从具有类概率的类别分布中抽取样本的简单高效方式。作者引入了一种新的连续分布，即Gumbel-Softmax分布，它是单纯形上的连续分布，能够近似类别样本，且其参数梯度可通过重参数化技巧轻松计算。

Gumbel-Softmax分布的关键特性在于它可以平滑地退火为类别分布。随着softmax温度τ趋近于0，来自Gumbel-Softmax分布的样本变为one-hot编码，该分布变得与类别分布相同。这种插值能力使得模型可以在训练过程中灵活调整离散程度。

在实际应用中，存在一个权衡：低温下样本接近one-hot但梯度方差大，高温下样本平滑但梯度方差小。因此，实践中通常从高温开始并退火至非零的小温度。若将温度作为学习参数而非固定调度，这可以解释为熵正则化，使分布自适应调整“置信度”。

对于必须采样离散值的场景（如强化学习中的动作空间），论文提出了Straight-Through (ST) Gumbel Estimator。它在前向传播中使用arg max离散化y，但在后向传播中使用连续近似来估算梯度。这使得即使在高温度下，样本也能保持稀疏性。

实验表明，Gumbel-Softmax估计器在结构化输出预测和无监督生成建模任务中优于现有的单样本梯度估计器。此外，它还能用于高效训练半监督模型，避免了对未观测类别潜在变量进行昂贵的边缘化操作。

《Categorical Reparameterization with Gumbel-Softmax》由Eric Jang、Shixiang Gu和Ben Poole撰写，发表于2016年（ICLR 2017）。适合研究变分自编码器、强化学习及任何涉及离散潜在变量优化的机器学习从业者阅读。建议从第2节关于分布定义及图1的温度效应分析入手。