[论文解读] Generalized Gumbel-Softmax Gradient Estimator for Generic Discrete Random Variables

Weonyoung Joo, Dongjun Kim+2 more|arXiv (Cornell University)|Mar 4, 2020

Advanced Graph Neural Networks被引用 8

一句话总结

本文提出广义Gumbel-Softmax（GenGS），一种新颖的重参数化梯度估计器，通过结合截断、线性变换和连续松弛，将Gumbel-Softmax技巧扩展至伯努利和多项分布之外的通用离散分布。在变分自编码器（VAEs）和主题模型上的实验表明，与现有估计器相比，GenGS在更低困惑度和更稳定的训练中实现了最先进性能。

ABSTRACT

Estimating the gradients of stochastic nodes in stochastic computational graphs is one of the crucial research questions in the deep generative modeling community, which enables the gradient descent optimization on neural network parameters. Stochastic gradient estimators of discrete random variables are widely explored, for example, Gumbel-Softmax reparameterization trick for Bernoulli and categorical distributions. Meanwhile, other discrete distribution cases such as the Poisson, geometric, binomial, multinomial, negative binomial, etc. have not been explored. This paper proposes a generalized version of the Gumbel-Softmax estimator, which is able to reparameterize generic discrete distributions, not restricted to the Bernoulli and the categorical. The proposed estimator utilizes the truncation of discrete random variables, the Gumbel-Softmax trick, and a special form of linear transformation. Our experiments consist of (1) synthetic examples and applications on VAE, which show the efficacy of our methods; and (2) topic models, which demonstrate the value of the proposed estimation in practice.

研究动机与目标

• 解决深度生成模型中针对泊松、几何、二项和负二项等通用离散分布缺乏有效梯度估计器的问题。
• 将重参数化技巧从伯努利和多项分布扩展至更广泛的离散分布族。
• 实现通过包含非多项式离散潜变量的随机计算图进行稳定且高效的反向传播。
• 提供一种具有概率基础且可实现的方法，支持显式与隐式后验推断。

提出的方法

• 提出广义Gumbel-Softmax（GenGS）估计器，通过截断无限支撑，对任意离散分布进行重参数化。
• 应用线性变换将截断后的离散分布映射为与Gumbel-Softmax松弛兼容的形式。
• 利用Gumbel-Softmax技巧创建离散采样过程的连续松弛，从而实现梯度反向传播。
• 在有限层级进行截断，使支撑计算上可行，同时保持分布保真度。
• 通过变换后分布的结构，支持显式与隐式后验推断。
• 利用温度退火和多样本估计提升深度模型中的训练稳定性。

实验结果

研究问题

• RQ1 Gumbel-Softmax技巧能否被推广至伯努利和多项分布之外的离散分布，如泊松和几何分布？
• RQ2 如何有效重参数化具有无限支撑的离散分布，以实现深度生成模型中的可微优化？
• RQ3 所提出的GenGS估计器在训练稳定性和模型似然性方面是否优于现有梯度估计器？
• RQ4 GenGS能否在复杂模型（如堆叠变分自编码器和主题模型）中有效应用？

主要发现

• 在多层堆叠主题建模中，GenGS在20Newsgroups数据集（562 ± 11.5）和RCV1数据集（576 ± 18.8）上均达到最低困惑度，优于所有基线方法。
• 在1-Stacked 20Newsgroups设置中，GenGS的困惑度为759 ± 13.1，显著低于次优方法StoRB∗的919 ± 16.2。
• 在多层堆叠的RCV1案例中，GenGS是唯一实现稳定收敛的估计器，而其他方法如RF*、NVIL和MuProp均未能达到最优性能。
• GenGS中采用的温度退火与多采样策略显著提升了深度层次模型中的训练稳定性和似然优化。
• GenGS在合成VAE实验与真实世界主题建模应用中均表现出一致的优越性。