[论文解读] Straight-Through Estimator as Projected Wasserstein Gradient Flow

Pengyu Cheng, Chang Liu+4 more|arXiv (Cornell University)|Oct 5, 2019

Advanced Neuroimaging Techniques and Applications参考文献 25被引用 6

一句话总结

本文通过将直通（Straight-Through, ST）估计器解释为在离散分布上的投影Wasserstein梯度流（pWGF），为其建立了理论基础。该文提出了一种基于新型pWGF的优化方案，在泊松等无限支撑分布上实现了更优的收敛性，相较于现有方法在伯努利与泊松隐变量模型上表现更优，且与ST方法性能相当。

ABSTRACT

The Straight-Through (ST) estimator is a widely used technique for back-propagating gradients through discrete random variables. However, this effective method lacks theoretical justification. In this paper, we show that ST can be interpreted as the simulation of the projected Wasserstein gradient flow (pWGF). Based on this understanding, a theoretical foundation is established to justify the convergence properties of ST. Further, another pWGF estimator variant is proposed, which exhibits superior performance on distributions with infinite support,e.g., Poisson distributions. Empirically, we show that ST and our proposed estimator, while applied to different types of discrete structures (including both Bernoulli and Poisson latent variables), exhibit comparable or even better performances relative to other state-of-the-art methods. Our results uncover the origin of the widespread adoption of the ST estimator and represent a helpful step towards exploring alternative gradient estimators for discrete variables.

研究动机与目标

• 为在离散变分推断中广泛使用但理解不足的直通（ST）梯度估计器提供理论依据。
• 在最小化离散分布上期望代价的背景下，形式化ST与投影Wasserstein梯度流（pWGF）之间的联系。
• 基于pWGF开发一种新型梯度估计器，实现对泊松等具有无限支撑分布的更快收敛。
• 在伯努利与泊松隐变量模型上对所提方法进行实证验证，与最先进基线方法进行性能比较。

提出的方法

• 该方法将ST解释为在2-Wasserstein空间中模拟投影Wasserstein梯度流（pWGF），其中梯度在松弛的连续空间中计算，随后投影回离散分布族。
• 在每次迭代中，算法首先通过Wasserstein梯度流更新样本：$\tilde{\bm{z}}_n = \bm{z}_n + \varepsilon \nabla f(\bm{z}_n)$，生成临时分布$\tilde{\mu}_k$。
• 随后通过最小化2-Wasserstein距离$W^2(\mu, \tilde{\mu}_k)$，将分布$\tilde{\mu}_k$投影回离散族$\mathcal{M}$。
• 在分布可分解的假设下，Wasserstein距离最小化可简化为在每个维度上最小化边缘距离，从而实现高效计算。
• 对于伯努利分布，通过基于累积分布函数的阈值规则，解析推导出最优投影。
• 通过将投影步骤适配于处理无界支撑，将该方法扩展至泊松等无限支撑分布，从而实现更快收敛。

实验结果

研究问题

• RQ1 直通估计器能否通过连续优化框架获得理论解释？
• RQ2 投影Wasserstein梯度流（pWGF）是否可作为ST估计器的有效且可推广的解释？
• RQ3 基于pWGF的估计器是否在离散隐变量上优于现有梯度估计器，尤其是在无限支撑分布上？
• RQ4 所提pWGF估计器在伯努利与泊松隐变量模型上的性能，与ST及其他最先进方法相比如何？

主要发现

• 直通估计器被理论证明为投影Wasserstein梯度流（pWGF）的模拟，为其经验成功提供了严谨基础。
• 在伯努利隐变量模型上，所提pWGF估计器性能与ST相当或更优，在使用非线性架构的MNIST数据集上测试ELBO达到104.6。
• 在泊松隐变量模型上，所提pWGF变体相较于ST及其他基线方法表现出更快收敛，尤其得益于其针对无限支撑分布的设计。
• 在MNIST数据集上，pWGF方法在所有测试架构中均优于Gumbel-Softmax、RELAX与REBAR，线性架构下测试ELBO达129.2。
• 理论分析证明，对于可分解分布，两个分布之间的2-Wasserstein距离等于各维度上平方边缘距离之和。
• 对于伯努利分布，pWGF下的最优投影被证明等价于基于累积分布函数的阈值规则，其中$W^2(\text{Bern}(p), \text{Bern}(q)) = |p - q|$。