[论文解读] Categorical Reparameterization with Gumbel-Softmax

Eric Jang, Shixiang Gu+1 more|arXiv (Cornell University)|Nov 3, 2016

Neural Networks and Applications被引用 3,222

一句话总结

引入 Gumbel-Softmax，一种用于分类变量的可微连续放松，可以退火为真正的独热样本，支持通过离散选择进行反向传播并加速半监督学习。

ABSTRACT

Categorical variables are a natural choice for representing discrete structure in the world. However, stochastic neural networks rarely use categorical latent variables due to the inability to backpropagate through samples. In this work, we present an efficient gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample from a novel Gumbel-Softmax distribution. This distribution has the essential property that it can be smoothly annealed into a categorical distribution. We show that our Gumbel-Softmax estimator outperforms state-of-the-art gradient estimators on structured output prediction and unsupervised generative modeling tasks with categorical latent variables, and enables large speedups on semi-supervised classification.

研究动机与目标

• 在神经网络中动机化使用离散潜变量以实现可解释且高效的表示。
• 提供一个可微的梯度估计器用于分类样本以实现反向传播。
• 引入 Gumbel-Softmax 作为可迭代到分类分布的连续放松。
• 在结构化预测和变分建模任务上证明优于先前估计器的性能提高。
• 展示 Gumbel-Softmax 在具有大量类别的半监督分类中实现更快的训练速度。

提出的方法

• 在 simplex 上定义 Gumbel-Softmax 分布，作为对分类样本的可微放松。
• 使用 Gumbel-Max 技巧再加上带温度 tau 的 softmax 从 logits 和 Gumbel 噪声中获得 y_i（y_i 的方程）。
• 推导 Gumbel-Softmax（concrete）分布的密度 p_pi,tau(y)。
• 提出 Gumbel-Softmax 估计器，允许通过重参数化（路径导数）对 y 进行反向传播。
• 引入一个 Straight-Through 变体，在前向传播中对 y 做离散化，但对梯度计算时使用连续的 y。
• 解释 tau 的退火调度及其作为熵正则化的解释。

实验结果

研究问题

• RQ1 与现有估计器相比，Gumbel-Softmax 是否能为分类变量提供低方差的可微分梯度？
• RQ2 Gumbel-Softmax 估计器在结构化输出预测和变分自编码器任务中的表现相对于现有梯度估计器如何？
• RQ3 这种方法是否能够在大规模类别数量下实现更快、可扩展的半监督学习？
• RQ4 对 softmax 温度的退火对训练稳定性和性能有何影响？

主要发现

• Gumbel-Softmax 在伯努利和分类潜变量上均优于单样本梯度估计器。
• ST Gumbel-Softmax 往往达到或超过其他方法的性能，Gumbel-Softmax 通常在各任务中表现最佳。
• 在具有大量类别的半监督分类中，与边际化推断相比，该方法显著加速了训练。
• 温度 tau 的退火提供了偏差（趋向离散性）与梯度方差之间的可控权衡，实用的调度方法效果良好。
• 在基于 MNIST 的任务（SBN 和 VAE）上，对数似然指标（nats）有所提升，在半监督设置中更快收敛且不牺牲准确性。
• 该方法实现了对离散潜变量的反向传播，使具有大规模分类空间的模型的可扩展学习成为可能。