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

1. VAE(オートエンコーダ)とは [概要]

変分オートエンコーダ(VAE, Variational AutoEncoder),潜在変数付き確率グラフィカルモデルとして構成されている,オートエンコーダの発展型の,高次元データ向け深層生成モデルである [Kingma and Welling, 2014].中間のボトルネック層を潜在変数化した「潜在変数付きEncoder-Decoder」を学習できるゆえ,ベイズ的な解釈性の高い潜在変数も同時に学習できる利点がある生成 or 変換モデルである.

この記事では,変分エンコーダの基礎と学習方法,およびその応用について述べる.

VAEは,再パラメータ化トリック(4.3節)というテクニックの導入により,SGDのみで潜在変数付きEncoder-Decoderを実現している点がVAEの一番の工夫点である.従って,潜在変数も推論できるのに,学習時の計算効率性もよく,しかも他のニューラルネットと同様に同じ仕組みの「SGD」のみで,VAEも学習できるメリットがある.

画像や音の変換モデルとしてEncoder-Decoderネットワークを学習する場合に,VAEとして設計・学習しすると,中間のボトルネック層部分を,潜在確率変数(正規分布)として,可視化および,低次元埋込み空間化でき,デバッグ性・解釈性がよくなる.

この記事の内容が盛りだくさんであるり,2節以降の詳細に入る前に,読者のVAEへの俯瞰度を,先に高めておきたい.そこで,1節のこれ以降では,もうすこしだけ概要をブレイクダウンした「少しくわしめの概要(1.1節)」と,「MNISTでの例(1.2節)」を先に提示しておきたい.

1.1 少し詳しめのVAEの概要

1.1.1 潜在変数の推論付きの自己符号化と,その利点

VAEでは,画像などの高次元データ群から,データの再構成損失関数に沿った「正規分布の潜在変数空間$\bm{z}$」を推論するエンコーダ$q(\bm{z}|\bm{x})$と,その潜在変数から綺麗なサンプルを再構成するデコーダ$p(\bm{x}|\bm{z})$を,一体化したオートエンコーダモデルとしてend-to-end学習する.

学習後のテスト時は,ユーザーが指定した$\bm{z}$の値を入力としてデコーダを順伝搬すると,$\bm{z}$の値に対応した生成サンプル$\bm{x}^{\prime}$を生成することができる.

この構成により,生成モデルであるデコーダ$p(\bm{x}|\bm{z})$を学習するのと同時に,データの生成もとである$\bm{z}$の低次元空間上に,各サンプルをデータの距離をもととした位置への埋め込みであるエンコーダ$q(\bm{z}|\bm{x})$も学習でき,デバッグ性の高い生成モデルを学習できる.

また,VAEで学習された$\bm{z}$を構成する各次元の潜在変数$\bm{z}_i$は,互いにもつれほどき(disentanglement)されやすく,個々の変数$\bm{z}$が,生成サンプルの「変化軸」として個別に独立して機能しやすくなりやすい特性がある.

VAEは,オートエンコーのボトルネック部である$\bm{z}$を正規分布としたときに,潜在変数$\bm{z}$と観測データ$\bm{x}$の間の写像関数をニューラルネットとして学習可能な「(連続)複数潜在変数付きのグラフィカルモデル」である.変分ベイズ(Variational Bayes)の目的関数を応用して学習を行うことから,変分オートエンコーダと名付けられた.

潜在変数付きグラフィカルモデル構造を持っていながら,ニューラルネットで構成してSGDで学習できる」点が,VAEの1番の特徴・長所である.これにより,他のニューラルネットワーク構造とも組み合わせやすい.また,深層生成モデルなので,過去の複数潜在変数付きグラフィカル生成モデルよりも,スケール化させることが容易である.加えて,従来のグラフィカル生成モデルの発想を持って来やすい.その意味で,VAEを元にした,新たな生成モデルを設計しやすい利点もある.

VAEは,Auto-Encoding Variational Bayes(AEVB) というアルゴリズムによる学習を行う(4節).VAEでは潜在(正規分布)変数も同時に学習(推論)するので,AEVBも「潜在変数あり学習」の1種類である.AEVBは,潜在変数の最適化と,非潜在変数の学習を繰り返すゆえ,伝統的なEMアルゴリズムに似た学習アルゴリズムでVAEも学習するのだと「一般化」して同一視・比較ができていると,VAEについて深く理解しやすい(=自分でVAEを応用できるになる).

1.1.2 確率モデルであるが,SDGで済む利点

オートエンコーダ,VAEは,オートエンコーダの中間部に,確率的な正規分布の潜在変数を含んでおり,SGDで一気に学習することはそのままではできない.

VAE以前の,かつての潜在ランダム変数付き確率的グラフィカルモデルの場合は(例:LDA(潜在ディリクレ配分)など),潜在変数の事後分布候補を,大量に(ランダム)サンプリングする必要があった.MCMC法などサンプリング計算コストの高いベイズ推論手法が必須であった.

しかしVAEでは,再パラメータ化トリック(reparametrization trick) (4.2節)を導入したことで,end-to-end学習が可能となっているのが一番のポイントである.これにより,VAEでは,バックプロパゲーション(=SGD)だけで,変分ベイズ推論学習が可能とになった.こうして「潜在変数付き」深層生成モデルを学習するハードルがVAEで一気に下がり,パラダイムシフトがおこった.ニューラルネットにVAEも利用することで,「もつれほどきもされた低次元な事前分布への,複雑な埋込みEncoder」の学習が容易に可能にになった,とも言い換えることができる.

ちなみに同時期に出てきたGAN も生成モデルや変換モデルでよく使われていく.GANとは異なり,VAEは潜在変数を予測するEncoderも同時に学習できるのが最も異なる点である.以下の親記事「深層生成モデル」などにも,書いたように,GANとVAEは,相互補完的な側面がある対照的技術であるので,両者の融合(良いとこ取り)がその後盛んになっていく(※ この記事では対象外).

1.1.3 またあまり使われなくなってきた潜在分布推論

[2022 5月追加] 複雑な確率分布を推論しようとすると,どうしても変分ベイズやMCMCを使わざるを得ず, よって,VAEにより,こんなに手軽にデータや事前確率の(ベイズ)分布を学習できるようになるとは,ディープラーニング流行後の,予想外の展開の1つであったとも言える .

古典的な画像特徴の分布化は,線形のPCA(顔認識での固有顔)や部分空間法などがかつて主流であった.また,ビジョン幾何での行列音階,音源分離での行列分解などでも,ICAやNMFなどの線形モデリングや行列分解が主流であった.よって,それを複雑な分布にでき,なおかつ手軽に分布を得られるVAEは,よく応用されるようになったわけである.

ただし,その後,特にTransformerが人気になって以降は,NLPやコンピュータビジョンの中心をなす基盤モデル周辺だと,またも「確率分布推論」の利用頻度は減り,もう忘れ去られがちである.生成モデルにおいても,VQ-VAEの登場で,潜在変数もトークン化してしまったことも大きいかもしれない.

あと,2020頃以降,Google T5のように,複数タスクをまぜこぜにした言語変換モデルが,画像認識やV&L事前学習でも登場し,席巻しはじめている.そうしたT5系の「マルチタスクまぜこせ変換モデル」には,潜在確率分布を配置できなくなってる点も,潜在分布が使われなくなっている理由であろう.複数タスクまぜこぜで学習する時点で,(Explicitな)潜在分布はもう配置しづらいのである.

1.2 VAEの動機と生成画像例

VAEが学習した2次元多様体からの生成画像
図1VAEが生成した生成画像群の可視化

例えば,MNIST データセットのように,潜在的な分布$p(\bm{z})$が存在するデータセットからVAEを学習すると,元の高次元データの特徴ベクトルが近いサンプル同士が近くに配置された,中央0のガウス分布による低次元多様体$\bm{z} $を,教師なし学習で獲得できる.MNISTの場合,\bm{z}は10次元ベクトルである(図1-a).また,顔画像データセットのように,データ内に向きや表情などの変化を含むデータセットからVAEを学習すると,VAEの構造と学習方法から出やすいもつれほどき(disentanglement)の効果により,独立した変化が各潜在変数に(教師なしでも)表出しやすい性質を持つ(図1-b).

Kingma氏によるデモサイトでは,VAEの画像生成の潜在変数を調整しての画像生成デモを,ブラウザ上で試すことができる.このサイトで実物を触ってもらうとわかるように,VAE開発の動機は「10次元程度の潜在変数を操るだけで,それの潜在変数の成分が混ざった綺麗な画像を生成できる深層生成モデルを開発する」ことにあった.

1.3 記事の構成

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)}\}$の分布を模した,生成モデルを学習することである.

2.1.1 潜在変数モデルの仮定

従来の潜在変数付きグラフィカルモデルや非線形次元削減でよく用いられてきた「高次元データ群の分布は,低次元多様体$\bm{z}$上に埋め込み可能である」という仮定をVAEでも採用する.

元の高次元データセット$\mathcal{D}$の中に潜在的に存在する低次元隠れ変数$\bm{z}$を推論(Inference)し,逆にその潜在変数からデータセットの分布に近いデータ$\bm{x}$を生成することを,VAEの目的とする.

すなわち,データ$\bm{x}$の背後に,各データサンプルに対応する低次元のシンプルな分布$p_{\theta}(\bm{z})$が存在すると仮定する.さらに,潜在変数からデータの生成を行う生成用分布$p_{\theta}(\bm{x} \mid \bm{z})$も存在するとしてグラフィカルモデル構造を設計している(3節).一方,元のデータ$\bm{x}$には高次元データ(画像など)を用いるので,その経験的分布$q_{\mathcal{D}}(\bm{x})$は複雑になりがちである.

2.1.2 VAE の学習の概要

VAEは,生成グラフィカルモデルの変数間の依存関係をニューラルネットで学習する深層潜在変数モデル(Deep Latent-Variable Models, DLVM)を採用している.DLVMである$p_{\theta}(\bm{x}, \bm{z})$および推論モデル$q_{\phi}(\bm{z} \mid \bm{x})$を,同時に学習するモデルである.この時,AEVBアルゴリズム(4節)を用いると,SGDを用いて効率的に学習できる.

この節では確率生成モデルとしての学習で構成される確率分布間の関係の大枠を先に述べる.従って,グラフィカルモデルとしての,詳細なサンプルの生成過程は3節にて述べる.

VAE のDeep Latent variable modelsとしての構造

VAEは$\bm{z}$から$ \bm{x}$が生成するDLVM構造である.

VAEのDLVMにおける$\bm{x}$と$\bm{z}$の関係は,ベイズの法則における以下の3つの$\bm{z}$の分布から定義できる:

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

DLVMでは,生成モデル全体を表す結合分布$p_{\theta}(\bm{x} , \bm{z})$を,次のように因数分解するグラフィカルモデル構造のベイズ生成モデルとなっている(詳細は3節):

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

データ点を生成するには,まず事前分布から$\bm{z}$をサンプリングし,その後尤度分布から$\bm{x}$をサンプルする.構造的には非常にシンプルなモデルである.

対数尤度目的関数の周辺化

生成モデルでは以下のように,データ分布の対数尤度を用いて分布のパラメータを推定することが標準的である:

\[
\theta^{\ast} = \arg \max_{\theta} \sum_{i}^{N} \log p_{\theta}(\bm{x}^{(i)}) \tag{2.2}
\]

ここで,式(2.2)中のエビデンス項$p_{\bm{\theta}}(\bm{x}^{(i)})$を,$\bm{z}$について周辺化することで周辺尤度として以下のように計算することができる:

\begin{align}
p_{\bm{\theta}}(\bm{x}^{(i)})
&= \int p_{\theta}(\bm{x} , \bm{z}) d \bm{z}\\
&= \int p_{\bm{\theta}}(\bm{x}^{(i)} \mid \bm{z}) p_{\bm{\theta}}(\bm{z}) d \bm{z} \tag{2.3}
\end{align}

また潜在変数$\bm{z}$も推論したいが,ここではベイズの法則を用いて事後分布を因数分解したものを用いることができる:

\begin{equation}
p_{\theta}(\bm{z} \mid \bm{x})
= \frac{p_{\theta}(\bm{x} \mid \bm{z})p_{\theta}(\bm{z})}{p_{\theta}(\bm{x})}
\end{equation}

オートエンコーダとしての学習

さて,式(2.2)からVAEの学習を行いたいが,そうすると式(2.3)を毎回の学習ループで推論を行う必要が出てくる.この際に,$p_{\theta}(\bm{x}^{(i)} \mid \bm{z})$をサンプリング法(MCMCなど)を用いて正確(exact)なベイズ推論をしようとすると,データ数$N$が大きいほど大量のサンプリングが必要であり,$\bm{z}$の候補全てについて周辺化しようにも計算量が多く取り扱いきれない(intractable)モデルとなる.

そこでVAEでは,エンコーダによる推論には,元のExactな事後分布を用いる代わりに,近似分布 $q_{\bm{\phi}}(\bm{z} \mid \bm{x})$ を利用した変分ベイズ推論を採用する:

\[q_{\bm{\phi}}(\bm{z} \mid \bm{x})\approx p_{\bm{\theta}}(\bm{z} \mid \bm{x})\tag{2.4}\]

そして,この$p_{\bm{\theta}}(\bm{x}^{(i)} \mid \bm{z})$の事前分布を,$K$次元の多次元正規分布としてモデル化するのが,変分オートエンコーダーの基本的なアイデアである.

\begin{align}
\bm{z} &\sim p(\bm{z})= \mathcal{N}(0,\bm{I})\\
\bm{x} \mid \bm{z} &\sim p_{\bm{\theta}}(\bm{z} \mid \bm{x}) \approx q_{\bm{\phi}}(\bm{z} \mid \bm{x})\tag{2.5}
\end{align}

そして,近似分布の式(2.4)を,エンコーダとして効率的に変分ベイズ推論で計算し,なおかつ生成モデルも学習を行うアルゴリズムとして,Auto-Encoding Variational Bayes(AEVB)アルゴリズムを,VAEでは使用する(4節).

2.2 VAE のネットワーク構造

変分オートエンコーダ(VAE: Variational AutoEncoder) の構造
図2 変分オートエンコーダ(VAE) のネットワーク構造.

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

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

$\bm{z}$の事前分布には0~1の値を取る正規分布を仮定したい.その正規分布からサンプリングされた値$\bm{z}$をデコーダに入力すると,学習に用いたデータセットの分布に基づいた,$\bm{z}$に対応するランダム変数$bm$を生成する.

また,潜在変数$z$を推論するエンコーダ$q_{\phi}(\bm{z} \mid \bm{x})$(条件付き近似分布)も,ニューラルネットワーク関数化されている(その学習原理については3.2節).

中間部の潜在表現$\bm{z}$を正規分布からサンプリングする部分は,本来は微分不可能であり,2個のネットワークは同時に学習できない.しかし,VAEでは 再パラメータ化トリック(4.2節)により,ボトルネック部がランダム正規分布のように振る舞いながらも$\bm{z}$を決定論的変数に変更する再パラメータ化を行うことで,2つのネットワークをEnd-to-Endで一括に学習することができる (※ VAEを実現したキーアイデア).

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

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

3. VAE のグラフィカルモデル構造

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

VAEのグラフィカルモデル構造
図3:VAEのグラフィカルモデル構造

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

VAEではデコーダ$p_{\theta}(\bm{x} \mid \bm{z})$をニューラルネットワークで実装しており,この$p_{\theta}$の確率分布により,データ$\bm{x}$と潜在変数ベクトル$\bm{z}$の間を写像する.

テスト時には,学習済みのパラメータ $\theta^{\ast}$ が得られており,サンプル$\bm{x}$を生成するためには以下の手順をもちいる:

  • 事前分布$p_{\bm{\theta}^{\ast}}(\bm{z})$から$\bm{z}$をサンプリング.
  • 条件付き分布 $p_{\bm{\theta}}(\bm{x} \mid \bm{z})$に相当するデコーダネットワークから$\bm{x}$を生成する.

また,変分ベイズ推定で用いる推論用の近似分布$q_{\bm{\phi}}(\bm{z} \mid \bm{x})$もエンコーダとして実装する.$\phi$はこの近似エンコーダの(変分)パラメータである.

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

3.2 もつれほどき(disentanglement)

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

元論文での,Fray Faceデータセットで学習した多様体の生成画像可視化図のように(図1-b),VAEの$\bm{z}$の各変数$z_1, z_2,\ldots, z_D$は,お互いがもつれほどきされた(disentangled)独立性の高い表現になりやすい.図1-bの,縦方向$z_1$は「顔の表情変化」におおよそ対応しており,横方向$z_2$は「顔の左右向き(姿勢)変化」におおよおそ対応している.したがって,各変数の値を調節して生成することで,「独立した潜在成分の混合具合をコントロールした画像生成が行いやすい」点がVAEの利点であると言える.

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

β-VAEという手法の登場以降,潜在変数のもつれほどき具合を,損失関数間の重み付けパラメータ$\beta$により明示的にコントロールを行うようにもなってきている.また関連する話として,2つのEncoderを用意して,もつれほどきされた成分同士に明示的に個別の潜在変数+エンコーダへ分離する研究が,画像生成や画像変換界隈で盛んである(例「Pose v.s Appearance」のもつれほどき,「顔のIdentityに関わる成分 v.s. 顔のIdentityは除去した成分」 のもつれほどきなど).

4. VAE の学習: 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, 証拠下界)である (4.1節).ELBOの最大化を用いると,近似解である周辺尤度の最大化を通して,VAEの近似エンコーダとデコーダのパラメータを同時に学習できる(4.2節).しかし,真ん中に正規分布である$p(\bm{z})$が挟まっているままだと,エンコーダからデコーダまでEnd-to-EndにSGDで逆伝播することができない.そこで,AEVBでは再パラメータ化トリックを用いて,$\bm{z}$を決定論変数に変えることにより,End-to-End学習を可能としている(4.3節)

4.1 対数尤度から導かれる損失関数の構成

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

以下の式が,VAEの目的関数である「周辺尤度の対数 $\log p_{\bm{\theta}}(\bm{x})$」を,2項に分解したものである(※この式の変形方法の詳細は省略.参考文献等を参考のこと):

\[
\log p_{\bm{\theta}}(\bm{x})
= \underbrace{\mathcal{L}_{\bm{\theta},\bm{\phi}}(\bm{x})}_{\text{ELBO}}
+\underbrace{\mathcal{D}_{KL}( q_{\bm{\phi}}(\bm{z} \mid \bm{x}) || p_{\bm{\theta}}(\bm{z} \mid \bm{x}))}_{\text{KLダイバージェンス}} \tag{4.1}
\]

VAEは,この対数尤度の下界(ELBO)を代理目的関数として使用し,SGDで近似学習する.2項について,意味をそれぞれ見ていきたい.

4.1.1 KLダイバージェンス項

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

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

4.1.2 ELBO項

式(4.1)の右辺第1項は,変分下界であるEvidence Lower Bound (ELBO)である:

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

KLダイバージェンスの非負性から,式(4.3)中に,式(4.1)を適用することで,以下のようにELBOは対数尤度の下界に相当しているとわかる:

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

ここで,KLダイバージェンス$\mathcal{D}_{KL}(q_{\bm{\phi}}(\bm{z} \mid \bm{x}) || p_{\bm{\theta}}(\bm{z} \mid \bm{x}))$は,以下の2つの距離計量を同時に定めていることがわかる

  1. 真の事後分布からの,近似事後分布へのKLダイバージェンス.
  2. ELBOと対数尤度のギャップ(下界のtightness).KLダイバージェンスに関して近似分布が真の分布を表しているほど,このギャップは小さい値となる.

こうして,VAEの目的関数は「ELBOの最大化」であり,それは「周辺尤度の下界の最大化」に相当していることが導出できた.(最大化を行うので,SGDではなくStochastic Gradient Ascentを行う点に注意).

よって,ELBOを用いてSGDによりVAEを最適化する方法を次に述べていく.

4.2 SGVD: SGDを通したELBO最適化

4.1節冒頭でも述べたように,式(4.3)のELBO $\mathcal{L}(\bm{x},\bm{\theta},\bm{\phi})$(の最大化)の利点は,SGD(stochastic gradient descent)により,パラメータ$\bm{\theta}$,およびエンコーダの変分パラメータ$\bm{\phi}$を同時に学習できるところにある.

そこで,バッチごとにSGDでELBOを損失関数として変分ベイズ推定を行うStochastic Gradient-based Variational Infernce(SGDV)を,VAEのネットワークの学習に用いる.

以下の,各サンプル$\bm{x}$におけるELBOの,データセット$\mathcal{D}$全体における和を用いてSGDによる最適化を行う:

\[
\mathcal{L}_{\bm{\theta},\bm{\phi}}(\mathcal{D})
= \sum_{\bm{x} \in \mathcal{D}}\mathcal{L}_{\bm{\theta},\bm{\phi}}(\bm{x})\tag{4.5}
\]

4.3 再パラメータ化トリック

再パラメータ化トリック(Reparametrization Trick)
図4 再パラメータ化トリック.

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

そこで [Kingma and Welling, 2014], [Rezende et al., 2014] では,再パラメータ化トリック(Reparametrization Trick)が提案された(図4).

中間層$\bm{z}$の正規分布のサンプリングを,外の別の確率変数$\epsilon$側へ逃した以下の表現に再パラメータ化することで,$\bm{z}$を確率からのサンプリングから(図4-a),決定論的変数の推定(図4-b)に変更した:

\begin{align}
\bm{\epsilon} \sim \mathcal{N}(0,\bm{I}) \tag{4.6}\\
\bm{z} = \bm{\mu} + \bm{\sigma} \odot \bm{\epsilon} \tag{4.7}
\end{align}

これらの再パラメータ化(4.6),(4.7)により,再パラメータ化された勾配を用いて$\bm{z}$の正規分布もつながり,VAE全体に出力の損失の勾配を逆伝播させて,一括にエンコーダとデコーダを学習することが可能となった(図4-b).確率的ノイズ変数$\epsilon$は,順伝搬時のみ別途サンプリングし,サンプリングした$\epsilon \sim p(\epsilon)$が$\bm{z}$で確率変数的な値を表現する役割を担当できる.

こうして,グラフィカルモデルとしてのVAEの挙動が$\bm{z}$が確率変数である場合(図4-a)と等価なまま,再パラメータ化のおかげで,SGDのみでエンコーダとデコーダをend-to-end学習できるようになった.(ただしVAEは,完全な確率モデルではないことに注意)

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を学習できる.よって,多クラス識別問題向けの画像データセットからも,CVAEを教師あり学習できるようになった.

6. VAEの まとめ

この記事では,変分オートエンコーダ(VAE)の,KingmaとWelligによる元祖の提案について,その構造と学習方法についてまとめた.VAEは,潜在変数付き深層生成モデルであり,変分ベイズ推定を利用して学習するオートエンコーダである.正規分布を事前確率とする潜在変数$\bm{z}$に,もつれほどきされた低次元表現を抽出しやすいので,各$z_i$が意味のある変化方向になりやすい(姿勢や表情や照明変化など).

VAEでは,ELBOの最大化により,変分ベイズ学習を行う.End-to-End学習することにより,中間層$\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.
  • [Rezende 2016] Rezende, D. J., S. Mohamed, and D. Wierstra. “Stochastic back- propagation and approximate inference in deep generative models”. In ICLR 2014
  • [Sohn et al., 2015] K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep conditional generative models. In, NIPS 2015.
  • [Yan 2016] Xinchen Yan, Jimei Yang, Kihyuk Sohn, and Honglak Lee. Attribute2image: Conditional image generation from visual attributes. In ECCV, 2016.

参照外部リンク