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

ひとこと要約

本稿では、組み合わせ的離散分布の構造的緩和を統一的に扱うための確率的ソフトマックストリック（SSTs）を導入する。これは、Gumbel-Softmaxトリックを一般化したものであり、摂動モデルにおいて強い凸正則化子と確率的線形目的関数を組み合わせることで、サブセット選択、スパニングツリー、アバレスセンスなどの複雑な組み合わせ的空間における低分散で微分可能な勾配推定を可能にする。これにより、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トリックをワンホットベクトルを超えて、より豊かな組み合わせ的集合へと拡張する一般化可能で微分可能なフレームワークを提供すること。

提案手法

• 本手法は、確率的線形目的関数と強い凸正則化子を組み合わせることで、Gumbel-Softmax の一般化として確率的ソフトマックストリック（SSTs）を導入する。
• SSTs は、有限集合 𝒳 上の離散分布が、𝒳 上で最適化される確率的効用 U によって誘導されることを前提とする摂動モデルフレームワーク内に構築される。
• 緩和は、𝒳 の凸包内での離散的 argmax 解を連続的に近似する凸最適化問題を解くことで達成される。
• フレームワークは、スパニングツリーのエッジインジケータなどの構造的表現をサポートし、完全なワンホット符号化よりもスケーラビリティを向上させる。
• 連続的緩和を介してバックプロパゲーションによる再パrameter化勾配を可能にし、組み合わせ的潜在変数を含むモデルのエンドツーエンドトレーニングを可能にする。
• 正則化子がGumbel分布の負のログットである場合、本手法はGumbel-Softmax を特別なケースとして包含する。

実験結果

リサーチクエスチョン

• RQ1 スパニングツリー やサブセット選択のような構造的組み合わせ的分布へ、Gumbel-Softmaxトリックを一般化することは可能か？
• RQ2 多様な組み合わせ的空間にわたる緩和勾配推定のための統一的フレームワークをどのように設計できるか？
• RQ3 構造的緩和を用いることで、非構造的ベースラインと比較して、変分モデルにおける潜在構造の発見性が向上するか？
• RQ4 SSTs は、ニューラル関係推論（NRI） や L2X といったモデルにおける性能と解釈可能性を向上させることができるか？
• RQ5 異なる正則化子と構造的表現が、勾配の分散とモデル収束に与える影響は何か？

主な発見

• SSTs により、サブセット選択、スパニングツリー、アバレスセンスなどの組み合わせ的分布における低分散で微分可能な勾配推定が可能になった。
• 本フレームワークは Gumbel-Softmax を一般化しており、従来、構造的組み合わせ的空間に利用可能ではなかった新しい緩和を含む。
• NRI および L2X フレームワークにおいて、SSTs を用いたモデルは、非構造的ベースラインと比較して、より正確で解釈可能な潜在構造を発見した。
• SSTs を用いてトレーニングされたモデルは、特に真の潜在構造がモデルのインダクティブバイアスと一致する場合、下流タスクで優れた性能を達成した。
• エッジインジケータなどの構造的表現の使用は、完全な組み合わせ的列挙と比較して、スケーラビリティと効率を顕著に向上させた。
• ランダムサーチによるハイパーパramータチューニングにより、最適な学習率、重み減衰、温度パラメータが特定され、最良のモデルはバリデーション性能に基づいて選別された。