Categorical Reparameterization with Gumbel-Softmax

https://arxiv.org/abs/1611.01144
2016

カテゴリー変数は、世界の離散的な構造を表現するための自然な選択です。しかし、確率的ニューラルネットワークでは、サンプルをバックプロパゲートすることができないため、カテゴリー型の潜在変数を使用することはほとんどない。本研究では、カテゴリー分布からの非微分サンプルを、新しいGumbel-Softmax分布からの微分サンプルに置き換える効率的な勾配推定法を提案する。この分布は、カテゴリ分布にスムーズにアニーリングできるという本質的な特性を持っている。我々のGumbel-Softmax推定量は、カテゴリー潜在変数を用いた構造化出力予測や教師なし生成モデリングのタスクにおいて、最先端の勾配推定量を上回り、半教師付き分類では大幅なスピードアップが可能であることを示している。

1 はじめに

離散的な確率変数を用いた確率的ニューラルネットワークは，教師なし学習，言語モデリング，注意メカニズム，強化学習などの領域で遭遇する分布を表現するための強力な手法である。例えば、離散変数は、明確な意味クラス（Kingmaら、2014年）、画像領域（Xuら、2015年）、および記憶位置（Gravesら、2014年、Gravesら、2016年）に対応する確率的な潜在的表現を学習するために使用されている。離散的な表現は、多くの場合、連続的な類似物よりも解釈しやすく（Chen et al. しかし、離散変数を持つ確率的ネットワークは、バックプロパゲーションアルゴリズム（パラメータ勾配の効率的な計算を可能にするが、非分化層には適用できない）のため、訓練が困難である。確率的勾配推定に関する先行研究では、従来、モンテカルロ分散削減技術で補強されたスコア関数推定量（Paisley et al.2012 Mnih & Gregor, 2014 Gu et al.2016 Gregor et al.2013）、またはベルヌーイ変数の偏った経路微分推定量（Bengio et al.2013）のいずれかに焦点が当てられていました。しかし、カテゴリー変数に特化して定式化された既存の勾配推定量はありません。

本研究の貢献は以下の3点である。

1.カテゴリー変数のサンプルを近似できるシンプレックス上の連続分布であるGumbel-Softmaxを導入し、そのパラメータ勾配を再パラメータ化トリックによって簡単に計算することができる。

Gumbel-Softmaxは、ベルヌーイ変数とカテゴリー変数の両方において、すべての単一サンプル勾配推定量よりも優れていることを実験的に示します。
この推定量を用いて、観測されていないカテゴリー潜在変数に対するコストのかかる周辺化を行わずに、半教師付きモデル（例えばKingma et al.2014）を効率的に訓練できることを示す。

本論文の実用的な成果は、ニューラルネットワークに統合し、標準的なバックプロパゲーションを用いて学習することができる、カテゴリー変数のためのシンプルで微分可能な近似サンプリングメカニズムである。

2 ガンベル・ソフトマックス分布

まず、ガンベル・ソフトマックス分布の定義から始めます。これは、カテゴリ分布からのサンプルを近似できるシンプレックス上の連続分布です。 zをクラス確率π1, π2, ...πkを持つカテゴリ変数とします。本論文では、カテゴリ分布のサンプルは、(k - 1)次元のsimplex、∆^{k-1}のコーナーに位置するk次元のワンショット・ベクトルとしてエンコードされていると仮定する。これにより、これらのベクトルの要素ごとの平均E_p[z] = [π1, ..., πk]などの量を定義することができます。 Gumbel-Maxトリック(Gumbel, 1954; Maddison et al., 2014)は，クラス確率πのカテゴリカル分布からサンプルzを引き出す効率的な方法を提供する

ここでg1...gkはGumbel(0, 1)^1から抽出したi.i.d.サンプルである。 arg maxの連続的で微分可能な近似値としてsoftmax関数を使用し，k次元のサンプルベクトルy∈Δ^{k-1}を生成します．

ガンベル・ソフトマックス分布の密度（付録Bで導出）は

この分布はMaddisonら（2016）によって独自に発見されたもので、コンクリート分布と呼ばれています。ソフトマックス温度τが0に近づくと、ガンベルソフトマックス分布からのサンプルはワンショットになり、ガンベルソフトマックス分布はカテゴリー分布p(z)と同一になります。

2.1 GUMBEL-SOFTMAX ESTIMATOR

Gumbel-Softmax分布はτ>0に対して滑らかであるため、パラメータπに対する勾配∂y/∂πが明確に定義されています。したがって、カテゴリーサンプルをGumbelSoftmaxサンプルで置き換えることにより、バックプロパゲーションを用いて勾配を計算することができます（セクション3.1参照）。ここでは、学習中に非微分のカテゴリー・サンプルを微分可能な近似値に置き換えるこの手順を、Gumbel-Softmax estimatorと呼びます。 Gumbel-Softmax推定値は微分可能ですが、温度が0でない場合の対応するカテゴリ分布からのサンプルとは同一ではありません。学習のためには、サンプルがワンhotに近いがグラデーションの分散が大きい小さな温度と、サンプルが滑らかだがグラデーションの分散が小さい大きな温度の間でトレードオフが生じます（図1）。実際には，高い温度から始めて，小さくてもゼロではない温度までアニールします．実験の結果、ソフトマックス温度τは、様々なスケジュールに従ってアニールしても、十分な性能を発揮することが分かりました。 τが（固定スケジュールでアニールされるのではなく）学習パラメータである場合、このスキームはエントロピー正則化（Szegedy et al., 2015; Pereyra et al., 2016）と解釈することができ、ガンベル-ソフトマックス分布は学習プロセス中に提案されたサンプルの「信頼性」を適応的に調整することができます。

2.2 STRAIGHT-THROUGH GUMBEL-SOFTMAX ESTIMATOR

ワンショットベクトルの連続的な緩和は、隠れた表現の学習やシーケンスモデリングなどの問題に適しています。離散的な値をサンプリングするように制約されているシナリオ（例えば、強化学習のための離散的なアクション空間や量子化された圧縮など）では、arg maxを使ってyを離散化しますが、バックワードパスでは∇θz≒∇θyを近似することで連続的な近似を使用します。これは、Bengio et al. (2013)で説明されている偏った経路微分推定法を彷彿とさせるため、Straight-Through (ST) Gumbel Estimatorと呼びます。 ST Gumbel-Softmaxは、温度τが高い場合でもサンプルを疎にすることができます。

3 RELATED WORK

本節では、離散変数に対する既存の確率的勾配推定技術をレビューする（図2に図示）。分布がパラメータθに依存する離散的なランダム変数zと、コスト関数f(z)を持つ確率的計算グラフ(Schulman et al., 2015)を考える。目的は、勾配降下によって期待コストL(θ) = E_{z∼p_θ(z)} [f(z)]を最小化することであり、そのためには∇θ E{z∼p_θ(z)} [f(z)]を推定する必要があります。

3.1 経路的勾配推定法

再パラメータ化可能な分布では、パラメータθと独立確率変数 eps, の決定論的関数 g として標本 z を計算することができ、 z = g(θ, eps) となります。 fからθへの経路勾配は，確率的なノードに遭遇することなく計算することができます．

例えば，正規分布z ∼ N (μ, σ)は，μ + σ * N (0, 1)と書き換えることができ，∂z/∂μと∂z/∂σを簡単に計算することができます。この再パラメータ化のトリックは、バックプロパゲーションを用いた連続潜在変数を持つ変分オートエンコーダのトレーニングによく適用されています(Kingma & Welling, 2013; Rezende et al., 2014b)。図2に示すように、我々はGumbel-Softmax推定量の構築において、このようなトリックを利用する。偏った経路微分推定量は、zがリパラメータ可能でない場合でも利用することができます。一般的には、∇_θ z≒∇_θ m(θ)と近似することができ、ここでmは確率的サンプルの微分可能な代理変数である。平均パラメータθを持つベルヌーイ変数の場合，Straight-Through (ST) estimator (Bengio et al., 2013) は，m = µ_θ(z)と近似し，∇_θ m = 1を意味します。 k = 2 (Bernoulli)の場合、ST Gumbel-SoftmaxはChungら(2016)が提案したslope-annealed Straight-Through estimatorに似ていますが、勾配を決定するためにハードシグモイドの代わりにソフトマックスを使用します。 Rolfe (2016) は、各バイナリ潜在変数が連続混合モデルをパラメータ化する代替アプローチを検討している。再パラメータ化勾配は、連続変数をバックプロパゲーションし、バイナリ変数をマージンアウトすることで得られる。 ST推定量の制限として、サンプルに依存しない平均に関してバックプロパゲーションを行うと、フォワードパスとバックワードパスの間に不一致が生じ、分散が高くなる可能性がある。

Gumbel-Softmaxは、各サンプルyが対応する離散サンプルzの微分可能なプロキシであるため、この問題を回避することができます。

3.2 SCORE FUNCTION-BASED GRADIENT ESTIMATORS SCORE FUNCTION ESTIMATORS

score function estimator （SF、REINFORCE（Williams, 1992）やLikelihood Ratio estimator（Glynn, 1990）とも呼ばれる）は、恒等式∇_θ p_θ(z) = p_θ(z) ∇_θ log p_θ(z)を用いて、以下の不偏推定量を導き出します。

SFは、pθ(z)がθで連続であることだけを必要とし、fやサンプルzをバックプロパゲーションする必要はありません。しかし、SFは分散が大きく、結果的に収束が遅いという問題があります。特に、SFの分散はサンプルベクトルの次元数に比例しており(Rezende et al., 2014a)、カテゴリー分布に使用することは特に困難です。スコア関数推定量の分散は、学習信号fから制御変量b(z)を差し引き、その解析的期待値μ_b = E_z [b(z)∇_θ log p_θ(z)]を戻し加算して推定量を不偏にすることで低減できる

制御変量を利用した最近の確率的勾配推定法を簡単にまとめた。これらの手法の詳細については、Gu et al (2016)を読者に案内する。

NVIL (Mnih & Gregor, 2014) は2つのベースラインを使用する。(1）学習信号をセンタリングするためのfの移動平均¯fと、（2）f - ¯f（センタリングされた学習信号自体の制御変量）にフィットした1層のニューラルネットワークによって計算された入力依存のベースライン。最後に、分散正規化は、学習信号をmax(1, σf )で分割します。σ 2 fは、Var[f]の移動平均です。
DARN (Gregor et al., 2013) は b = f(¯z) + f'(¯z)(z - z¯) を使用します。ベースラインは f(¯z) から f(z) の一次テイラー近似に対応します。 z はベルヌーイ変数に対して 1/2 となるように選択され、推定式の補正項 µb を無視しているため、非二次 f に対して推定値が偏ります。
MuProp (Gu et al., 2016) もベースラインを一次のテイラー展開としてモデル化しています：b = f(¯z) + f'(¯z)(z - z¯) and µb = f'(¯z)∇θ Ez [z]. 離散的なサンプリングによるバックプロパゲーションを克服するために、ベースラインの計算と関連する勾配の導出には、f(z)の代わりに平均場近似f{MF}(μθ(z))が使用されています。
VIMCO（Mnih & Rezende, 2016）は、マルチサンプル目的の勾配推定器であり、他のサンプルの平均b = 1/m sum_{j !=i} f(z_j) を使用して、各サンプルzi ∈ z{1:m}のベースラインを構築する。Gumbel-Softmaxはマルチサンプル目的に簡単に拡張できるが、我々はシングルサンプル目的の推定量を比較しているため、VIMCOを実験から除外している。

3.3 半教師付きジェネレーティブ・モデル

半教師付き学習は，ラベル付きデータ(x, y) ∼ D_L とラベルのないデータ x ∼ D_U の両方から学習する問題を考える．半教師付き分類のために，Kingmaら（2014）は，潜在状態がガウス型の「スタイル」変数zとカテゴリー型の「セマンティッククラス」変数yの共同分布である変分自動符号化器（VAE）を提案している（図6，付録）。 VAEの目的は、生成モデルの下で観測値の対数尤度に対する変分下界を最大化することで、識別ネットワークq_φ(y|x)、推論ネットワークq_φ(z|x, y)、生成ネットワークp_θ(x|y, z)をエンド・ツー・エンドで学習することです。ラベル付きデータの場合、クラスyは観測されているので、推論はz ∼ q(z|x, y)に対してのみ行われます。ラベル付きデータの変分下界は次のように与えられる。

ラベルのないデータでは、カテゴリ分布が再パラメータ化できないため、困難が生じる。 Kingmaら(2014)は、すべてのクラスでyを周辺化することで、ラベルのないデータでは、各yに対するq_φ(z|x, y)を推論するようにしています。ラベルのないデータに対する下界は

完全な最大化の目的は

ここでαは生成目的と識別目的の間のスカラートレードオフである。このアプローチの限界は、クラス数が多いモデルでは、k個のクラス値すべてに対する周辺化が法外に高価になることです。 D, I, Gがそれぞれq_φ(y|x)、q_φ(z|x, y)、p_θ(x|y, z)からのサンプリングの計算コストであるとすると、教師なし目的の学習には、前進／後退の各ステップにO(D + k(I + G))が必要となる。一方、Gumbel-Softmaxでは、y ∼ q_φ(y|x)をバックプロパゲートして1サンプルの勾配推定を行うことができ、トレーニングステップごとにO(D +I +G)のコストを実現しています。学習速度の実験比較を図5に示す。

4 実験結果

最初の実験では、Gumbel-SoftmaxとST Gumbel-Softmaxを他の確率的勾配推定量と比較します。 Score-Function (SF)、DARN、MuProp、Straight-Through (ST)、そしてSlope-Annealed STです。各推定器は2つのタスクで評価されます。 (1)構造化出力予測， (2)生成モデルの変分学習．

学習と評価には，確率的勾配推定量を評価する際に一般的な手法である固定二値化を行ったMNISTデータセットを用いた(Salakhutdinov & Murray, 2008; Larochelle & Murray, 2011)．学習率は{3e-5, 1e-5, 3e-4, 1e-4, 3e-3, 1e-3}から選択し，MNISTの検証セットを用いて各推定量に最適な学習率を選択し，テストセットでの性能を報告した． Gumbel-Softmax分布から抽出されたサンプルは，学習時には連続的であるが，評価時には1hotのベクトルに離散化される．また，SF，DARN，MuPropで競争力のある性能を得るためには，分散の正規化が必要であることがわかりました．二値（ベルヌーイ）ニューラルネットワークにはシグモイド活性化関数を，カテゴリ変数にはソフトマックス活性化関数を使用しました．モデルの学習には，運動量0.9の確率的勾配降下法を用いた．

4.1 STOCHASTIC BINARY NETWORKSによる構造化出力予測

構造化出力予測の目的は、画像の上半分(14×28)が与えられたときに、28×28のMNISTの数字の下半分を予測することである。これは、ストキャスティック・バイナリ・ネットワーク（SBN）を学習するための一般的なベンチマークです（Raiko et al., 2014; Gu et al., 2016; Mnih & Rezende, 2016）。この条件付き生成モデルの最小化目的は、尤度目的の重要度サンプリング推定値である

E_{h∼pθ}(hi| x{upper}) 1/m sum^m_{i=1} log pθ(x{lower}|h_i)

ここでm = 1は学習に、m = 1000は評価に使用されます。我々は，200ユニットずつの2つの隠れ層を持つSBNを学習した．これは，200個のベルヌーイ変数（392-200-200-392と表記），または2値化された活性化を持つ20個のカテゴリー変数（各クラスは10個）（392(20×10)-(20×10)-392と表記）に対応する．図3に示すように、STのGumbel-Softmaxは、ベルヌーイ変数では他の推定量と同等であり、カテゴリー変数では優れています。一方、Gumbel-Softmaxは、Bernoulli変数とCategorical変数の両方で他の推定量よりも優れています。このタスクでは、ソフトマックスの温度をアニーリングする必要がないことがわかり、固定のτ=1を使用しました。

4.2 変分オートエンコーダによる生成モデル化

変分オートエンコーダ(Kingma & Welling, 2013)を学習するが，ここでの目的は，2値のMNIST画像の生成モデルを学習することである．今回の実験では，潜在変数を200個のベルヌーイ変数または20個のカテゴリー変数（20×10）を持つ単一の隠れ層としてモデル化した．学習目的には、Gumbel-Softmax事前分布ではなく、学習されたカテゴリー事前分布を用いている。そのため、学習時の最小化目的は、サンプルが離散的でない場合、もはや変分束ではありません。実際には、この目的を温度アニーリングと組み合わせて最適化しても、検証セットやテストセットで実際の変分限界を最小化できることがわかりました。構造化出力予測タスクと同様に、評価にはm = 1000のマルチサンプル境界を使用します。温度は、グローバルトレーニングステップtのスケジュールτ = max(0.5, exp(-rt))を用いてアニーリングされ、τはNステップごとに更新されます。 N∈{500, 1000}とr∈{1e-5, 1e-4}は，検証セットで最も性能の高い推定量を選択し，テスト性能を報告するためのハイパーパラメータである。図4に示すように、カテゴリー変数では、ST Gumbel-Softmaxが他の推定量を上回り、ベルヌーイ変数とカテゴリー変数の両方では、Gumbel-Softmaxが他の推定量を大幅に上回っています。

4.3 GENERATIVE SEMI-SUPERVISED CLASSIFICATION

我々は、Gumbel-Softmax推定量をバイナリMNISTデータセットの半教師付き分類に適用する。オリジナルの周辺化ベースの推論アプローチ(Kingma et al., 2014)と、Gumbel-SoftmaxおよびST Gumbel-Softmaxを用いた単一サンプルの推論を比較する。 100個のラベル付き例（10個のクラスのそれぞれに均等に分布）と50,000個のラベルなし例からなるデータセットで学習を行い、ミニバッチごとにラベルなし例を動的に2値化した。識別モデルqφ(y|x)と推論モデルqφ(z|x, y)は，それぞれReLU活性化関数を持つ3層の畳み込みニューラルネットワークとして実装されている．生成モデルpθ(x|y, z)は、ReLU活性化関数を持つ4層の畳み込みトランスポーズ・ネットワークです。実験の詳細は付録Aに記載されています。 α = {0.1, 0.2, 0.3, 0.8, 1.0}の複数の値に対して推定器を学習・評価し，テストセットのラベルなし分類結果のうち，各推定器で最も優れたものを選択し，表2に報告した． 2000ステップごとに更新されるτ = max(0.5, exp(-3e-5 - t))のアニーリング・スケジュールを用いた。 Kingmaら(2014)では、潜在的な状態に関する推論は、yを周辺化し、qφ(z|x, y)からのサンプリングに再パラメータ化のトリックを使用することで行われる。しかし、このアプローチは、クラスの数に対して線形にスケールする計算コストを持っています。 Gumbel-Softmaxは、結合qφ(y, z|x)からの単一サンプルを直接バックプロパゲートすることができ、生成性能や分類性能を損なうことなく、学習の大幅な高速化を実現しています。(表2、図5)

5 DISCUSSION

本研究の主な貢献は、再パラメータ化可能なGumbel-Softmax分布であり、その対応する推定量はカテゴリ分布の低分散経路微分勾配を提供する。本研究では、Gumbel-SoftmaxとStraight-Through Gumbel-Softmaxが、構造化された出力予測と変分オートエンコーダーのタスクに有効であり、ベルヌーイ潜在変数とカテゴリー潜在変数の両方に対する既存の確率的勾配推定量を上回ることを示す。最後に、Gumbel-Softmaxは、離散潜在変数に対する推論を劇的に高速化します。

e4exp / paper_manager_abstract