はじめに

深層生成モデルを巡る旅をしています。
前回はFlowについて理論と各種法の簡単な紹介をしました.
今回はまた別の深層生成モデルとして 変分自己符号化器 (VAE; variational autoencoder) [1] を紹介します.
VAEはGANに比べて安定した学習ができ, Flowと異なり潜在変数を低次元に落とすことができるので, その扱いやすさや解釈性から好まれることが多いように感じます. 一方で, 生成画像がぼやけがちである, 尤度の計算ができないといった欠点もあります.
この1ヶ月くらいVAEについていろいろと調べてみましたが, GANやFlowのようなモデル乱立状態にはなっていなさそうだったので, 主要な手法を少し詳しめに紹介できればと思います.

VAEの基本

最初に, 本記事を読むのに必要な事項をまとめます.

全体像

式での説明は少々長いので, 先にMNISTの画像を扱う例で全体像を示します.

VAEはEncoderとDecoderという2つのネットワークからなり, 名前の通りautoencoderのような構造をしています. しかし, 出自は変分推論でautoencoderとは全く異なるものなので, 以降はBayes流で説明します.
潜在変数 $z$ は, 我々が観測している画像の背後にあるシードのようなもので, 「どの数字か」「筆跡」「線の太さ」といったものを示す低次元の変数です. 多様体仮説によれば, MNISTの画像は28x28で784次元ですが, 実際にはデータは高次元空間のごく一部に分布していて, 情報を落とさずに低次元の空間に埋め込むことが可能です.
詳しくは立ち入りませんが, グラフィカルモデルで表すと次のような関係になっています.

目的関数の導出

潜在変数 $z$ からデータ $x$ （通常画像など高次元）を生成する確率的なDecoderを $p (x | z)$ で表します. DecoderはCNNなどで構成され, パラメータ $θ$ を明示して $p_{θ} (x | z)$ と表すこともあります.
VAEも, 前回のFlowと同様に最尤推定で $θ$ を決めるために（対数）尤度の最大化を目指しますが, 一般にデータ $x$ の分布は複雑で $p (x)$ を直接求めることは難しいです. また, 事後分布 $p (z | x)$ を使って式変形をしようとしても, 「データ $x$ を生成するような潜在変数 $z$ を求める操作」はNNの逆変換となり難しいです.
これを解決するために, 近似事後分布 $q (z | x)$ を導入します. これは確率的なEncoderであり, パラメータ $ϕ$ を持つCNNとして $q_{ϕ} (z | x)$ と表すこともあります.
それでは式変形をしていきます.

\begin{array}{rcl} \log p (x) & = & \log p (x) \int q (z | x) d z \\ = & \int q (z | x) \log \frac{p (z | x) p (x)}{p (z | x)} d z \\ = & \int q (z | x) \log \frac{p (x, z)}{p (z | x)} d z \\ = & \int q (z | x) (\log \frac{q (z | x)}{p (z | x)} + \log \frac{p (x, z)}{q (z | x)}) d z \\ = & D_{K L} [q (z | x) | | p (z | x)] + L (x, z) \end{array}

ここで, 右辺第2項の $L (x, z)$ は 変分下限 (variational lower bound) または ELBO (evidence lower bound) と呼ばれ, 次式で定義されます.

L (x, z) := \int q (z | x) \log \frac{p (x, z)}{q (z | x)} d z

右辺第1項は事後分布を含むため計算が困難ですが, KLダイバージェンスなので次の不等式が成り立ちます（c.f. イェンセンの不等式. 等号は $q (z | x)$ の近似性能が十分なときに成立します）.

D_{K L} [q (z | x) | | p (z | x)] \geq 0

したがって, 次のように対数尤度をELBOで下から抑えることができます.

\log p (x) \geq L (x, z)

ELBOはさらに分解することができます.

\begin{array}{rcl} L (x, z) & = & \int q (z | x) \log \frac{p (x, z)}{q (z | x)} d z \\ = & \int q (z | x) \log p (x | z) d z - \int q (z | x) \log \frac{q (z | x)}{p (z)} d z \\ = & E_{q (z | x)} [\log p (x | z)] - D_{K L} [q (z | x) | | p (z)] \end{array}

VAEの学習では, この目的関数を最大化するようにパラメータ $θ, ϕ$ を選びます.

目的関数の意味と計算方法

いま導出した目的関数の第1項は, データ $x$ を $q_{ϕ} (z | x)$ でエンコードした $z$ を $p_{θ} (x | z)$ でデコードしたときの再構成誤差で, これがAutoencoderと呼ばれる所以です. 実際, $q (z | x)$ を正規分布とすれば, 期待値の計算をバッチ内の平均で近似して,

\begin{array}{rcl} E_{q (z | x)} [\log p (x | z)] & ≃ & \frac{1}{B} \sum_{b = 1}^{B} \log p (x_{b} | z_{b}) \\ z_{b} & \sim & N (μ (x_{b}), Σ (x_{b})) \end{array}

と求めることができます.
$μ (x_{b}), Σ (x_{b})$ はエンコーダ $q_{ϕ} (z | x)$ の出力で, それぞれ多変量正規分布の平均と分散（共分散は0）を表します. $\log p (x_{b} | z_{b})$ の部分の計算は, Bernoulli分布や正規分布を仮定して, それぞれ交差エントロピーや平均二乗誤差の最小化に置き換えて行います.

目的関数の第2項は, 事前分布の正則化を担います. すなわち, データ $x$ を $q_{ϕ} (z | x)$ でエンコードした $z$ が事前分布（標準正規分布）に従うように, KLダイバージェンスを最小化します. 正規分布同士のKLダイバージェンスは解析的に求まります.

\begin{array}{rcl} D_{K L} [q (z | x) | | p (z)] & = & D_{K L} [N (μ (x), Σ (x)) | | N (0, I)] \\ = & \frac{1}{2} (t r Σ (x) + μ (x)^{T} μ (x) - d - \log det Σ (x)) \end{array}

$d$ は潜在変数 $z$ の次元数です. 導出は[6]が詳しいです.

Reparametrization Trick

さて, 以上の定式化を踏まえてVAEの図をもう一度眺めてみましょう. .

このままでは誤差逆伝搬でVAEを学習させることができません. なぜなら, $z$ をサンプリングする操作は微分不可能であり, それより前のEncoder側に勾配を伝えることができないからです.
そこで, reparametrization trickという方法でこの問題を回避します. これは, $z$ の代わりに標準正規分布から $ϵ$ をサンプリングし, これを使って $z$ を構成するという方法です.

\begin{array}{rcl} z & = & μ (x) + ϵ ⊙ Σ (x) \\ ϵ & \sim & N (0, I) \end{array}

これで, Encoder側にも勾配が伝わるようになりました.

これまでの流れを一旦まとめます.

対数尤度を直接最大化するのは難しいため, 近似事後分布 $q (z | x)$ を導入してELBO（変分下限）の最大化に帰着
ELBOを再構成誤差と正則化項に分解し, 計算可能な形にする
誤差逆伝搬を可能にするため

様々なVAE

それでは, 様々なVAE（といっても5種類ですが）を紹介していきます.

"Vanilla" VAE [Kingma+, 2013]

元祖VAEは, ここまでで説明したVAEを3層MLPというシンプルなモデルで実装しました. 潜在変数 $z$ を2次元にして, $z$ を動かしたときにDecoderから得られた画像 $x$ を並べると, 数字が連続的に変化していく様子が見て取れます.

Auto-Encoding Variational Bayes

Conditional VAE　[Kingma+, 2014]

MNISTの例で言うと, 生成時にランダムなサンプルではなく「0」の数字だけを生成したい, といったケースがあると思います. これに応えられるように, EncoderとDecoderをラベル $y$ で条件付けたのがConditional VAE (CVAE) です.

L (x, z, y) = E_{q (z | x, y)} [\log p (x | z, y)] - D_{K L} [q (z | x, y) | | p (z | y)]

潜在変数とラベルをDecoderに入力することで, 「2」「3」のようにラベルを指定した生成ができるようになります. このとき, ラベルの情報は潜在変数から取り除かれ, 筆跡や線の太さを表すようになります.

Semi-Supervised Learning with Deep Generative Models

β-VAE [Higgins+, 2017]

少し時間が空いてしまいますが, 画像特徴の"disentanglement"に注目したのがβ-VAEです.
disentangleは「もつれを解く」という意味で, 画像の特徴が潜在空間上で分離されていることを表します. 具体的には, 顔の画像に対して潜在変数の1次元目が顔の向きを, 2次元目が表情を…というように各次元がそれぞれ独立な特徴を担当している状態を想像してください. これができると, 生成された顔画像の向きを調節したければ潜在変数の1次元目を調節すれば良い, ということがわかって生成モデルとして扱いやすくなります.
β-VAEは, 目的関数の正則化項に係数 $β (> 1)$ を導入するというシンプルな変更により, 特徴のdisentanglementを実現しました.

L (x, z) = E_{q (z | x)} [\log p (x | z)] - β D_{K L} [q (z | x) | | p (z)]

係数 $β (> 1)$ で正則化項の寄与を大きくすることで, 潜在変数が標準正規分布に従うようにするための制約を強めています.
下図は, 64x64のCelebAで訓練したCNNベースのVAEで, 潜在変数の1次元のみを変化させたときの生成画像を並べたものです. VAEでは向き（上段）や表情（下段）以外の要素（眼鏡や髪の色）も変わっているのに対し, β-VAEではその要素以外は保存されています.

beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework

しかし, 係数 $β$ を大きくしすぎると再構成誤差が大きくなることが報告されていて, そこのトレードオフが課題です.
なお, 特徴のdisentanglementについては[8]が詳しいです.

Vector Quantised-VAE [van den Oord+, 2017]

「自然界の様々な事物は離散的に表現され得る」（言語は離散的ですし, 画像も離散的な言語によって記述され得ます. 生物の潜在変数とも言えるDNAも離散的ですね！）という発想から, 潜在変数を離散値にしたのがVector Quantised-VAE (VQ-VAE) です.

VQ-VAEでは, 埋め込み $z_{e}$ を丸めるためにコードブック $E = {e_{k} \in R^{D} | k = 1, . . ., K}$ を用意します. 上図のように, Encoderから出力された特徴マップ $z_{e} \in R^{H \times W \times D}$ に対して, 各 $z_{e} (i, j, :)$ をコードブックの最近傍コード $e_{k}$ で置き換え, $z_{q} \in R^{H \times W \times D}$ とします. また, 最近傍コード $e_{k}$ のインデックス $k$ を並べたものを, 離散化された潜在変数 $z \in Z^{H \times W}$ とします. Decoderは $z_{q}$ から画像 $x$ を再構成します. このとき, 近似事後分布は $K$ 個のコードから最近傍コードを選ぶone-hotな分布となります.

q (z (i, j) = k | x) = {\begin{cases} 1 & (if k = {a r g m i n}_{l} ∥ z_{e} (i, j, :) - e_{l} ∥_{2}) \\ 0 & (otherwise) \end{cases}

式で書くとややこしいですが, 図の水色の部分のように最近傍コードの番号を選ぶだけです.
ちなみに, ImageNetの128x128の画像の実験では, 特徴マップは $z_{e}, z_{q} \in R^{32 \times 32 \times 1}$ で, コードブックは $K = 512$ 個のコードからなります.

VQ-VAEの学習中には, 事前分布 $p (z)$ は一様分布で固定します. 学習が終わったら, 事前分布を自己回帰モデルの一種であるPixelCNNで学習させます.

次に, 目的関数についてです. VQ-VAEは近似事後分布がone-hotで, かつ事前分布が一様分布であるため, 学習中はKLダイバージェンスの項が一定となり, 無視できます（詳しくは[9]）.
最終的に目的関数は, 再構成誤差と, $z_{e}$ とコードブックの差を埋めるための項からなります.

L = E_{q (z | x)} [\log p (x | z_{q})] + ∥ sg [z_{e}] - e_{k} ∥_{2}^{2} + β ∥ z_{e} - sg [e_{k}] ∥_{2}^{2}

$sg [\cdot]$ は, 誤差逆伝播時に勾配を計算しないようにする処理です.
（最近傍コードを選ぶ操作は微分不可能ですが, これでEncoderまで勾配を伝えられる理由が納得できていません. 詳しい方がいらっしゃったら教えてください. ）

以上の工夫により, VAEで起きやすいと指摘されていたposterior collapse を解決できたそうです. posterior collapseは, PixelCNNのような強いDecoderを使うときに潜在変数を無視した生成が行われてしまうという現象です.
なお, posterior collapseに関しては, VQ-VAEチームのRazaviらが2019年にδ-VAEを提案しています.

ImageNetの128x128の画像の生成に成功しました.
――と言われても, これらの画像は鮮明ではありますが, 大域的な特徴が掴めておらず, 何を生成しているのかわからないと思います. 列ごとに異なるカテゴリで, 左から順にキツネ, クジラ, クマ, チョウ, サンゴ礁, 山, 電子レンジ, トラックだそうです.

Neural Discrete Representation Learning

また, ここには掲載しませんが, 論文では音声や映像での実験も行っています.

Vector Quantised-VAE-2 [Razavi+, 2019]

VQ-VAE-2は基本的に, VQ-VAEを階層構造にして, さらに高解像度の画像を生成できるようにしたというものです.
今度はImageNetの256x256の画像を, 大域的な特徴を保って生成できています.

また, 1024x1024のFFHQの画像が生成できるようになりました. これはStyleGANやBigGANに並んだとも言える品質ではないでしょうか.

Generating Diverse High-Fidelity Images with VQ-VAE-2

おわりに

本記事では, VAEの解説とdisentanglementや高解像度化のための工夫を紹介しました.
「VAEはGANのような鮮明な画像を生成するのが難しい」とはよく言われますが, VQ-VAE-2はそのような課題を克服できていると思います. 今後は再現実装を動かしたりしてさらに掘り下げて見たいと思います.

参考文献

[1] Diederik P. Kingma et al. "Auto-Encoding Variational Bayes." ICLR. 2014.
VAEの元論文です. 第一著者のKingmaはAdamやGlowの著者としてもお馴染みですね.

[2] Carl Doersch. "Tutorial on Variational Autoencoders." arXiv. 2016.
VAEがグラフィカルモデルから丁寧に解説されています.

[3] Variational Autoencoder徹底解説 - Qiita
こちらは日本語の丁寧な解説記事です. Chainerの実装にも触れられています.

[4] Kazuki Nitta. "Variational AutoEncoder." 2016.
こちらも日本語の丁寧な解説記事です.

[5] Lilian Weng. "From Autoencoder to Beta-VAE." 2018.
オートエンコーダの発展としてVAEを解説した記事です.

[6] 多変量正規分布の場合のKullback Leibler Divergenceの導出 - Qiita

[7] Atsushi Suyama. "線形回帰を１つ１つ改造して変分オートエンコーダ（VAE）を作る - 作って遊ぶ機械学習. ." 2018.

[8] Hideki Tsunashima. "Disentanglement Survey:Can You Explain How Much Are Generative models Disentangled?." 2019.

[9] Neural Discrete Representation Learning | Reading Notes
VQ-VAEの解説記事です.

深層生成モデルを巡る旅(2): VAE