[论文解读] Gradient Estimation with Stochastic Softmax Tricks

Max B. Paulus, Dami Choi+3 more|arXiv (Cornell University)|Jun 15, 2020

Generative Adversarial Networks and Image Synthesis参考文献 77被引用 24

一句话总结

本文提出了随机Softmax技巧（SSTs），这是一种用于组合离散分布结构化松弛的统一框架，推广了Gumbel-Softmax技巧。通过在扰动模型中结合强凸正则化项与随机线性目标，SSTs能够在复杂组合空间（如子集选择、生成树和有向生成树）上实现低方差、可微分的梯度估计，从而在NRI和L2X等模型中提升了性能并更好地发现了潜在结构。

ABSTRACT

The Gumbel-Max trick is the basis of many relaxed gradient estimators. These estimators are easy to implement and low variance, but the goal of scaling them comprehensively to large combinatorial distributions is still outstanding. Working within the perturbation model framework, we introduce stochastic softmax tricks, which generalize the Gumbel-Softmax trick to combinatorial spaces. Our framework is a unified perspective on existing relaxed estimators for perturbation models, and it contains many novel relaxations. We design structured relaxations for subset selection, spanning trees, arborescences, and others. When compared to less structured baselines, we find that stochastic softmax tricks can be used to train latent variable models that perform better and discover more latent structure.

研究动机与目标

• 解决在深度学习中对大规模组合状态空间上的离散分布实现可扩展、低方差梯度估计的挑战。
• 将现有的松弛梯度估计器（如Gumbel-Softmax）统一在一个适用于结构化组合分布的单一框架中。
• 通过扰动模型设计针对组合对象（如子集、生成树和有向生成树）的新颖结构化松弛方法。
• 通过支持对底层组合结构的更好发现，提升潜在变量模型的性能与可解释性。
• 提供一个通用的、可微分的框架，将Gumbel-Max技巧从one-hot向量扩展到更丰富的组合集合。

提出的方法

• 该方法通过将随机线性目标与强凸正则化项结合，将SSTs作为Gumbel-Softmax的推广，生成连续且可微分的解。
• SSTs基于扰动模型框架构建，其中有限集合𝒳上的离散分布由在𝒳上优化随机效用U而诱导得出。
• 通过求解一个凸优化问题，连续逼近𝒳凸包内的离散argmax解，实现松弛。
• 该框架支持结构化表示（例如生成树的二值边指示变量），而非完整的one-hot编码，从而提升可扩展性。
• 该方法通过反向传播实现重参数化梯度，支持具有组合潜在变量的模型端到端训练。
• 当正则化项为Gumbel分布的负对数几率时，该方法将Gumbel-Softmax作为特例推广。

实验结果

研究问题

• RQ1 我们能否将Gumbel-Softmax技巧推广到如生成树和子集选择等结构化组合分布？
• RQ2 如何设计一个统一框架，用于在多样化组合空间上实现松弛梯度估计？
• RQ3 与非结构化基线相比，使用结构化松弛是否能带来更好的潜在结构发现？
• RQ4 SSTs能否在Neural Relational Inference（NRI）和L2X等模型中提升性能与可解释性？
• RQ5 不同正则化项与结构化表示对梯度方差与模型收敛性有何影响？

主要发现

• SSTs实现了对子集选择、生成树和有向生成树等组合分布的低方差、可微分梯度估计。
• 该框架推广了Gumbel-Softmax，并包含了此前在结构化组合空间中未实现的新颖松弛方法。
• 在NRI和L2X框架中，使用SSTs的模型比非结构化基线发现了更准确且更具可解释性的潜在结构。
• 采用SSTs训练的模型在下游任务中取得了更强的性能，尤其当真实潜在结构与模型归纳偏差匹配时。
• 使用结构化表示（如边指示变量而非one-hot向量）显著提升了可扩展性与效率，优于完整的组合枚举。
• 通过随机搜索进行超参数调优，确定了最优学习率、权重衰减和温度值，最佳模型基于验证性能选取。