変分オートエンコーダ(VAE)

1. 概要

変分オートエンコーダ(VAE: Variational AutoEncoder)[Kingma and Welling, 2014], [Kingma and Welling, 2019]は,潜在変数付き確率グラフィカルモデル形式の,深層生成モデルである.

オートエンコーダのボトルネック部に,低次元ベクトルの潜在変数$\bm{z}$が正規分布として学習できる構造の,(連続)潜在変数付きグラフィカルモデルとして,変分オートエンコーダは設計されている.変分ベイズ(Variational Bayes)の目的関数を応用して学習を行うことから,『変分』オートエンコーダと名付けられた.

「グラフィカル生成モデル構造を持っていながら,ニューラルネットなのでSGDで学習できる」というのが,VAEの1番の特徴である.よって,他のニューラルネットワーク構造とも容易に組み合わせられる上に,(それまで深層学習には無かった)生成モデルであるのにディープモデルなので,スケール化させることも容易である.また,従来のグラフィカル生成モデルの発想を持って来れるので,(古くからの生成モデルの研究者も含めて)VAEを元にした新しい生成モデルを設計しやすい利点もある.

VAEでは reparametrization trickという学習テクニックを導入したおかげで,中間に正規分布を挟んでいるにも関わらずEnd-to-End学習が可能となっている.従来のベイズグラフィカルモデルで,大量のデータから潜在変数も学習する際は,MCMC法など計算コストの高い,サンプリング方式のベイズ推論手法が必須であった.それが変分オートエンコーダでは,ただバックプロパゲーションを実施するだけで(グラフ設計の苦労がいらずに)学習できるようになったのが,潜在変数付きグラフィカルモデルの視点でパラダイムシフトであった.

MNISTデータセットのような潜在的な分布$p(\bm{z})$が存在するデータセットからVAEを学習すると,元の高次元データの特徴ベクトルが近いサンプル同士が近くに配置された,ガウス分布の潜在変数ベクトル空間$\bm{z}$を教師なし学習で獲得できる.テスト時は,ユーザーが指定した$\bm{z}$の値を入力としてデコーダを順伝搬すると,$\bm{z}$の値に対応した生成サンプル$\bm{x}^{\prime}$を生成することができる.

1.1 記事の構成

2節以降は,オリジナルのVAEの提案 [Kingma and Welling, 2014][Kingma and Welling, 2019] を中心として,以下の構成によりVAEの基本的な仕組みと特徴を説明していく:

  • 2 VAEの構成
  • 3 VAEのグラフィカル構造
  • 4 VAE の学習
  • 5 VAEの特徴と発展・応用

VAEの基本モデルについての記事なので,VAEのその後の応用や発展については,3.1節や5.2節などで少し述べている程度であることに注意されたい.

2 VAEの構成

2.1 問題設定とモデルの概要

VAEの学習の目的は,大規模データセット$\mathcal{D} = \{\bm{x}^1,\bm{x}^2,\ldots,\bm{x}^N\}$の分布を模した生成モデルを学習することである.潜在変数グラフィカルモデルや非線形次元削減においてよく用いられてきた「高次元データ群の分布は,低次元多様体$\bm{z}$上に埋め込み可能である」という仮定を,VAEでも採用する.よって,データ群$\mathcal{D}$のなかに潜在的に存在する低次元潜在変数$\bm{z}$を推論(Inference)し,その潜在変数からデータセットの分布に近いデータを生成することを,VAEの目的とする.つまり,データ$\bm{x}$の背後に潜在変数$\bm{z}$とその分布$p_{\theta}(\bm{z})$が存在すると仮定する.さらに,潜在変数からデータの生成を行う生成用分布$p_{\theta}(\bm{x}|\bm{z})$も存在すると仮定する.

VAEでは,この生成グラフィカルモデルとして,変数間の依存関係をニューラルネットで学習する深層潜在変数モデル(Deep Latent-Variable Model)を採用している.このモデル下で,$p_{\theta}(\bm{z})$および$p_{\theta}(\bm{x}|\bm{z})$を,教師なし学習する問題設定である.

ここで,$p_{\theta}(\bm{z}|\bm{x})$は,ベイズ法則(Bayes’s rule)を用いると,以下のように分解できる:

\[p_{\theta}(\bm{z}|\bm{x}) = p_{\theta}(\bm{x}|\bm{z}) p_{\theta}(\bm{z}) \tag{2.1}\]

この右辺の$p_{\theta}(x)$を以下のように周辺化して計算すると,左辺の値を求めることができる:

\begin{equation}p(\bm{x}) = \int p_{\theta}(\bm{x}|\bm{z}) p_{\theta}(\bm{z}) \tag{2.2}\end{equation}

ここで,事後分布$p_{\theta}(\bm{z}|\bm{x})$は毎回の学習ループで推論を行う必要があるのだが,この際に正確な(exact)な$p_{\theta}(\bm{z}|\bm{x})$を用いようとすると,データ数$N$が大きいほど大量のサンプリングが必要で $\bm{z}$の候補全てについて周辺化するにも計算量が多くなり,扱いづらい(intractable)モデルとなる.

そこでVAEでのエンコーダによる推論は,元のExactな分布を用いる代わりに,近似分布 $q_{\phi}(\bm{z}|\bm{x})$を推定する変分推論で効率的に学習を行う,Auto-Encoding Variational Bayes(AEVB) [Kingma and Welling, 2014]という学習アルゴリズム採用している(4節).

2.2 VAE のネットワーク構造

変分オートエンコーダ(VAE: Variational AutoEncoder) の構造
図1 変分オートエンコーダ(VAE: Variational AutoEncoder) の構造.

VAEは,低次元ベクトルの潜在変数$\bm{z}$(連続変数から構成)をボトルネックとして中間に挟む形で,以下の2つのネットワークによりオートエンコーダ構造を成している(図1):

  • エンコーダ $q_{\phi}(\bm{z}|\bm{x})$:
    • 潜在変数の推論ネットワーク.データ$\bm{x}$から潜在変数$\bm{z}$の確率分布を推論する.
    • 具体的には,$\bm{z}$をサンプリングするための,平均ベクトル$\bm{\mu}$と標準偏差ベクトル$\bm{sigma}$を結合したベクトルを推定する.
    • 学習の際にオートエンコーダを構成する補助役である.データの生成を行うテスト時には使用しないが,実データから潜在変数$\bm{z}$の空間に射影することも可能であるので,テスト時にも便利に使用できる.
  • デコーダ $p_{\theta}(\bm{x}|\bm{z})$:
    • 生成モデルのネットワーク.
    • 潜在変数$\bm{z}$を入力として,データ$\bm{x}^{\prime}$を生成(再構成)する.

$\bm{z}$の事前分布には0~1の値を取る正規分布を仮定したい.その正規分布からサンプリングされた値$\bm{z}$をデコーダに入力すると,学習に用いたデータセットの分布に基づいた,$\bm{z}$に対応するランダム変数$bm$を生成する.また,潜在変数$z$を推論するエンコーダ$q_{\phi}(\bm{z}|\bm{x})$(条件付き近似分布)も,ニューラルネットワーク関数化されている(その学習原理については3.2節).

中間部の低次元潜在表現$\bm{z}$を正規分布からサンプリングする部分は,本来は微分不可能であるので,2個のネットワークは同時に学習できない.そこを,VAEでは Reparametrization Trick により,グラフィカルモデル構造を保ちながらも2つのネットワークをEnd-to-Endで学習することができる (※ VAEのキーアイデアである).

GANも提案されて以降数年は,これら2つの深層生成モデルが中心選手として,双方とも発展していた.両者ともに「大量の画像データを用いて画像生成モデルを学習できる」という点が,それまでに存在しなかった利点である.

5.2節や,親記事(深層生成モデル (Deep Generative Model))でも述べるように,GANとVAEは相互補完的な関係も持っており,片方の弱点をもう片方だけが実現できる長所がある.よって,同時に両者を使用するモデルも多く提案されている.

3. VAEのグラフィカル構造と潜在変数

3.1 確率的グラフィカルモデルとしての構造

変分オートエンコーダ(VAE) のグラフィカルモデル構造
図2:変分オートエンコーダ(VAE)の確率的グラフィカルモデル構造

VAEは図2で示されるグラフィカルモデルに沿って,正規分布からサンプリングした$\bm{z}$を入力として受け付けるデコーダネットワークを用いて$\bm{x}$を生成するモデルとなっている.

ベイズ的には以下の構成である:

  • 事前分布:$p_{\theta}(\bm{z})$
  • 尤度:$p_{\theta}(\bm{x}|\bm{z}) $
  • 事後分布:$p_{\theta}(\bm{z}|\bm{x}) \approx q_{\phi}(\bm{z}|\bm{x})$$

VAEではデコーダ$p_{\theta}(\bm{x}|\bm{z})$をニューラルネットワークで実装しており,$\theta$はそのエンコーダのパラメータである.また,学習時にベイズ推定を行うために使用する($p(\bm{z}|\bm{x})$を近似した分布$q_{\phi}(\bm{z}|\bm{x})$もニューラルネットワークで実装しており,$\theta$はこのデコーダの(変分)パラメータである.

VAEは,図1のように,確率分布中間変数$\bm{z}$をボトルネックとするAutoencoders構造である.2.2節で述べたように,$\bm{z}$は事前分布が$N(0,\bm{I})$ (原点中心の等方性正規分布)であり,学習を通じてそのががウシアン事前分布上の点$\bm{z}$に,各学習データが埋め込まれていく.

3.1 もつれほどき(disentanglement)

VAEの学習を通じて得られる,潜在変数間のもつれほどき(disentanglement)について,簡単にだけ触れておきたい(原理や最近の動向についてまでは,この記事では述べない).

元論文での,顔画像から学習したVAEの,潜在変数2変数を変化させて生成した顔画像群の可視化図 ( [Kingma and Welling, 2019] Fig.4 (a) Learned Frey Face manifold )のように,VAEの$\bm{z}$の各変数$z_1, z_2,\ldots, z_D$は,お互いがもつれほどきされた(disentangled),独立性の高い表現になりやすい.したがって,各変数の値を調節して生成することで,それらの独立した潜在成分の混合具合をコントロールした画像生成が行いやすい点がVAEのメリットの1つであると言える.

例えば,古典的な線形の潜在変数モデルのPCAを同じデータセットから学習させた場合,潜在変数には正規直行基底を仮定しており,PCAで学習できる正規直行基底レベルまでしか,各変数間のもつれほどきは起こらない.しかし,VAEだと,非線型かつ深い(階層的な)潜在表現を,教師なし学習を通して獲得できる.

β-VAEという手法の登場以降,潜在変数のもつれほどき具合を,学習パラメータ$\beta$により明示的にコントロールを行うようにもなってきている.また関連する話として,2つのEncoderを用意して明示的に2成分にもつれほどきを行う研究が,画像生成や画像変換界隈で盛んであるが,この記事では述べない(例「Pose v.s Appearance」のもつれほどき,「顔のIdentityに関わる成分 v.s. 顔のIdentityは除去した成分」 のもつれほどきなど).

4 VAE の学習: Auto-Encoding Variational Bayes (AEVB)

VAEの学習アルゴリズムはAuto-Encoding Variational Bayes(AEVB) [Kingma and Welling, 2014] と呼ばれる.VAEの学習では,変分ベイズ推定(Variational Bayesian Inference)による潜在変数の推定も同時に行われる.そのベイズ推定の過程がニューラルネットワーク化されたSGD学習になっているので,旧来の潜在変数グラフィカルモデルよりもスケールしやすい.

AEVBによる学習の目的関数は,周辺尤度の対数 $\log p_{\theta}(\bm)$ の変分下界(variational lower bound)に相当するEvidence Lower BOund (ELBO)である.ELBOの最大化を用いると,近似解である周辺尤度の最大化を通して,VAEをSGDで学習できる.

4.1 損失関数:ELBOの構成

VAEでは,周辺尤度$p_{\theta}(\bm{x})$最大化による学習を通じて,デコーダのパラメータ$\theta$,およびエンコーダの変分パラメータ$\phi$を同時に学習していく.

以下の式が,目的関数である周辺尤度の対数を,2項に分解したものである(※この式の変形方法の詳細は省略.各文献を参考のこと):

\[
\log p_{\theta}(\bm{x}) = \underbrace{ \mathcal{L}_{\theta,\phi}(\bm{x})}_{ELBO} +\mathcal{D}_{KL}( q_{\phi}(\bm{z}|\bm{x}) || p_{\theta}(\bm{z}|\bm{x}) ) \tag{4.1}
\]

ここで右辺第2項は,真の分布$p_{\theta}(\bm{z}|\bm{x})$と近似分布$q_{\phi}(\bm{z}|\bm{x})$の間のKLダイバージェンスであり,以下のように非負の値を取り,値が0に近くなるほど近似分布は真の分布に近づく:

\begin{equation}
\mathcal{D}_{KL}(q_{\phi}(\bm{z}|\bm{x}) || p_{\theta}(\bm{z}|\bm{x})) \geq 0 \tag{4.2}
\end{equation}

また,右辺第1項はEvidence Lower Bound (ELBO)と呼ばれる:

\begin{equation}
\mathcal{L}(\bm{x},\theta,\phi)=\mathbb{E}_{q_{\phi}(\bm{z}|\bm{x})} [\log p_{\theta}(\bm{x},\bm{z})-q_{\phi}(\bm{z},\bm{x})] \tag{4.3}
\end{equation}

よって式(4.2)中に式(4.1)を適用することで,以下のようにELBOは対数尤度の下界に相当しているとわかる:

\begin{equation}
\mathcal{L}_{\theta,\phi}(\bm{x}) = \log p_{\theta}(\bm{x}) – \mathcal{D}_{KL}(q_{\phi}(\bm{z}|\bm{x})||p_{\theta}(\bm{z}|\bm{x})) \leq 0 \tag{4.4}
\end{equation}

すわなち,VAEの目的であるELBOの最大化は,周辺尤度を最大化することに相当していることが導出できた.

4.2 再パラメータ化トリック(Reparametrization Trick)

各近似事後分布の$q_{\phi}(\bm{z}|\bm{x})$を用いて$\bm{z}$の正規分布のパラメータ(平均と標準偏差)を推定した際に,そのパラメータから$\bm{z}$をサンプリングする必要がある.しかし,サンプリング処理は微分不可能であるので,このままだとVAE全体にELBOからの損失をバックプロパゲーションすることはできない.

そこで,VAEの論文では再パラメータ化トリック(Reparametrization Trick)が提案された.中間層$\bm{z}$の世紀分布のサンプリングを,外の別の確率変数$\epsilon$側へ逃した以下の表現に再パラメータ化することで,$\bm{z}$を決定的変数に変更した:

\begin{equation}\bm{\epsilon} \sim \mathcal{N}(0,\bm{I}) \tag{4.5}\end{equation}

\begin{equation}\bm{z} = \bm{\mu} + \bm{\sigma} \odot \epsilon \tag(4.6)\end{equation}

この再パラメータ化(4.5),(4.6)により,VAEはエンコーダとデコーダの間が$\bm{z}$で繋げられるようになり,VAE全体が逆伝播による一括して学習することが可能となった.確率的ノイズ変数$\epsilon$は,順伝搬時のみ別途サンプリングし,サンプリングした$\epsilon \sim p(\epsilon)$が$\bm{z}$で確率変数的な値を表現する役割を担当できる

こうしてReparametrization Trickの導入により,グラフィカルモデルとしてのVAEの挙動が$\bm{z}$が確率変数である場合と等価なまま,end-to-end化を達成でき,SGDだけで学習できるようになった

5. VAE の特徴と発展・応用

2~4節までは,オリジナルのVAEの仕組みを紹介してきた.最後にこの5節でVAEの特徴を整理し(5.1節),その発展・応用にもさわりだけ少し触れておきたい(5.2節).

5.1 VAE の特徴

ここまでの話をまとめつつ,VAEの特徴を長所・短所に分けて,以下にまとめる:

長所

  • 旧来の潜在変数付きグラフィカル生成モデルと異なり,SGDを用いたバッチ学習が可能なので,大量のデータから学習できスケールしやすい.
  • 生成モデルとしてシンプルなグラフィカルモデルを使用しており取り扱いやすい.
  • $q(\bm{z}|\bm{x})$の推論が可能なので,色々なタスクにおいてこの表現学習の仕組みを応用しやすい.
短所
  • 伝統的な尤度最大化でありその下界を用いること自体は悪くはない.ただ,自己回帰モデルやGANなど他の深層生成モデルと比べると,下界なので甘い評価関数による最適化である.
  • 分布の近似および変数のサンプリングを行うので,シャープなサンプルを生成できる(決定的モデルの)GANと比べると,生成サンプルはボケがあって質的に劣る(※ しかし最近はVQ-VAEなど,綺麗な画像が生成できるVAEも登場している)

5.2 VAE の発展・応用

(教師なしで)潜在空間付き生成モデルを学習できるというVAEの性質上, シーケンス形式のデータからそのデータ群の性質をとらえたい対象として,近年VAEは音響データへの応用が多く見うけられる.(自己回帰モデルにおいても,Wavenetの登場以降は『音響データ向けの応用』が盛んである)

KingmaとWellingによる解説 [Kingma and Welling, 2019] でも触れられているように,VAEとGANは相互補完的な関係にある.従ってVAE-GAN [Larsen et al., 2016]のようにVAEのDecoderとGANのGeneratorを同一のものとみなし2者を組み合わせると,お互いの短所を補った良いモデルを作りやすい.

また,他によく研究されている応用モデルとして,クラスラベルで条件付けしたVAEである Conditional VAE (CVAE) [Sohn et al., 2015] が挙げられる.CVAEでは,クラスラベル付きの他クラスデータセットを学習することで,クラスラベルに条件付けされたEncoderとGeneratorを学習するVAEである.多クラスラベルデータセットからクラスごとの分布の違いを把握できている単一の生成モデル学習できるので,多クラス識別問題向けの画像データセットからもVAEを教師あり学習できるようになった.

5. まとめ

この記事では,深層生成モデルの代表的モデルである変分オートエンコーダ(VAE)について,オリジナルのKingmaとWellingの提案についてまとめた.

変分オートエンコーダは,ELBOの最大化により,変分ベイズ学習を行う確率グラフィカルモデル型の潜在変数付き深層生成モデルである.中間層$\bm{z}$が多次元正規分布をサンプリングするので,そのままだとEnd-to-Endに逆伝播できない.そこで,再パラメータ化トリックを導入して,平均と標準偏差はサンプリングしないで済むようにした.こうしてVAEはEnd-to-End学習が可能となったことで学習しやすく,かつ潜在分布がExplicitな正規分布として学習できる,使い良い深層生成モデルとなった.

最後の,5節ではVAEの特徴をまとめ,VAEの発展と応用も少し述べた.

References

  • [Kingma and Welling, 2014] Auto-encoding variational bayes. In ICLR, 2014
  • [Kingma and Welling, 2019] D. P. Kingma and M. Welling. (2019). An introduction to variational autoencoders. Foundations and Trends® in Machine Learning, 12(4), 307-392.
  • [Larsen et al., 2016] A. B. L. Larsen, S. K. Sønderby, H. Larochelle, and O. Winther. Autoencoding beyond pixels using a learned similarity metric. In ICML, 2016.
  • [Sohn et al., 2015] K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep conditional generative models. In, NIPS 2015.

参照外部リンク

関連記事