[論文レビュー] 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

ひとこと要約

本稿は、離散分布上の投影 Wasserstein 勾配流（pWGF）として解釈することで、Straight-Through（ST）推定器の理論的基盤を確立する。本稿では、ポアソンのような無限台付き分布に対して収束性を向上させる、新しい 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.

研究の動機と目的

• 離散的変分推論において広く使われているが、その理解が不十分な Straight-Through（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$ が生成される。
• その後、分布 $\tilde{\mu}_k$ は、2-Wasserstein 距離 $W^2(\mu, \tilde{\mu}_k)$ を最小化することで、離散族 $\mathcal{M}$ に射影される。
• 因子化された分布を仮定すると、Wasserstein 距離の最小化は各次元における周辺距離の最小化に帰着され、計算が効率的に行える。
• ベルヌーイ分布の場合、累積分布関数に基づくしきい値ルールを用いて、最適な射影を解析的に導出できる。
• 無限台付き分布（例：ポアソン）に対しては、射影ステップを無限大の台を持つ場合に適応させることで、収束性が向上する。

実験結果

リサーチクエスチョン

• RQ1 Straight-Through 推定器は、連続的最適化フレームワークを用いて理論的に正当化可能か？
• RQ2 投影 Wasserstein 勾配流（pWGF）は、ST 推定器の解釈として妥当かつ一般化可能か？
• RQ3 pWGF に基づく推定器は、特に無限台付き分布において、既存の勾配推定器を上回る性能を示すか？
• RQ4 提案手法の pWGF 推定器は、ベルヌーイおよびポアソンの潜在変数モデルにおいて、ST や他の最先端手法と比較してどのように性能を発揮するか？

主な発見

• Straight-Through 推定器は、理論的に投影 Wasserstein 勾配流（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|$ が成り立つ。