はじめに

Avoiding Latent Variable Collapse with Generative Skip Modelsを読んだのでメモ．勉強メモがわりにブログ始めてからずっと生成モデルしか勉強してない...

気持ち

VAEではdecoderが潜在変数を無視して画像を生成してしまうlatent variable collapseという現象がある．これはVAEが良いデータの表現を獲得できないことに繋がるため回避したい現象である．latent variable collapseが起こる要因は乱暴に言ってしまえば，evidence lower bound(ELBO)のKL項が最小化された結果prior $p(\mathbf{z})$ と近似事後分布 $q(\mathbf{z}|\mathbf{x})$ がほぼ一致してしまうことによる．

この論文ではこれを回避するためにskip connectionをdecoder側に導入して，理論的にも実験的にもlatent variable collapseが回避できることを示したという話．この理論的側面が気になって読んでみた．

latent variable collapse

まず必要となる式を諸々導出や定義を行いつつlatent variable collapseが起こる理由について簡単に．

VAEのdecoderは $p_\theta(\mathbf{x}|\mathbf{z})$ として表現され，以下で表されるニューラルネットによってparameterizeされた指数型分布族と考えられる．

$\displaystyle p_\theta(\mathbf{x}|\mathbf{z})=\mathrm{ExpFam}(\mathbf{x};\eta(\mathbf{z};\theta)) =\nu(\mathbf{x})\exp\left\{\eta(\mathbf{z};\theta)'\mathbf{x}-A(\eta(\mathbf{z};\theta))\right\}$

ここで $A$ はlog-normalizerを表す． $\eta(\mathbf{z};\theta)$ は $\mathbf{z}$ のhierarchical function（要は多層のニューラルネット）で，以下のように表される．

$\displaystyle \mathbf{h}^{(1)}=f_{\theta_0}(\mathbf{z})\\ \displaystyle \mathbf{h}^{(l+1)}=f_{\theta_l}(\mathbf{h}^{(l)})\\ \displaystyle \eta(\mathbf{z};\theta)=f_{\theta_L}(\mathbf{h}^{(L)})$

ただし， $\mathbf{h}^{(l)}$ は $l$ 層目の隠れ状態を表す．

次に目的関数を考える．ELBOを以下の形で与える．

$\displaystyle ELBO=\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log p_\theta(\mathbf{x}|\mathbf{z})\right]\right]-KL(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}))$

上のELBOは第一項目の自由度の高さがKL項を0にする要因となっていてlatent variable collapseを起こす． Latent variable collapseをさらに理解するためadversarial autoencodersで議論されていたaggregated posterior $q_\phi(\mathbf{z})$ を考える．aggregated posteriorは次の形で表される．

$\displaystyle q_\phi(\mathbf{x},\mathbf{z})=p(\mathbf{x})q_\phi(\mathbf{z}|\mathbf{x}) \\ \displaystyle q_\phi(\mathbf{z})=\mathbb{E}_{p(\mathbf{x})}q_\phi(\mathbf{z}|\mathbf{x})$

するとELBOは式変形により次のように表現することができる．（具体的な式変形は割愛するが，以下のELBOにaggregated posteriorの具体形を代入して整理すれば元のELBOに戻る．）

$\displaystyle ELBO=\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log p_\theta(\mathbf{x}|\mathbf{z})\right]\right]-\mathcal{I}^{VAE}_q(\mathbf{x},\mathbf{z})-KL(q_\phi(\mathbf{z})||p(\mathbf{z}))$

ここで $\mathcal{I}$ は相互情報量を表していて，その定義から $\mathcal{I}^{VAE}_q(\mathbf{x},\mathbf{z})=\mathbb{E}_{p(\mathbf{x})}\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\log q_\phi(\mathbf{z}|\mathbf{x})-\mathbb{E}_{q_\phi(\mathbf{z})}\log q_\phi(\mathbf{z})$ と計算される．KL項が0になるとういうことは上記の形で表現されるELBOにおいて

$\displaystyle \mathcal{I}^{VAE}_q(\mathbf{x},\mathbf{z})=KL(q_\phi(\mathbf{z})||p(\mathbf{z}))=0$

となることを意味する．相互情報量が0であるということは $\mathbf{x},\mathbf{z}$ が独立であることを意味していて，独立になるということはlatent variable collapseと同義．

avoid latent variable collapse

論文のメインの主張であるgenerative skip modelsについて．generative skip modelsの $\eta(\mathbf{z},\theta)$ は以下のように表現できる．

$\displaystyle \mathbf{h}^{(1)}=f_{\theta_0}(\mathbf{z})\\ \displaystyle \mathbf{h}^{(l+1)}=gw_l\left(f_{\theta_l}(\mathbf{h}^{(l)}),\mathbf{z}\right) \\ \displaystyle \eta(\mathbf{z};\theta)=gw_L\left(f_{\theta_L}(\mathbf{h}^{(L)}), \mathbf{z}\right)$

各層において隠れ状態と潜在変数を入力とする関数 $gw_l$ を導入した形で，ResNetやU-Netにあるいわゆるスキップコネクションと同じものと考えればいい．この論文では $gw_l$ を以下の形で表す．

$\displaystyle gw_l\left(f_{\theta_l}(\mathbf{h}^{(l)},\mathbf{z}\right)=\sigma\left(W_l^{(h)}f_{\theta_l}(\mathbf{h}^{(l)})+W_l^{(z)}\mathbf{z}\right)$

ここで $\sigma$ はシグモイドやReLUといった非線形関数で $W$ は学習される重みパラメータ．

上記のようなモデルを与えるとgenerative skip modelにおける相互情報量 $\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})$ は次のような関係を持つ．

$\displaystyle \mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})\geq\mathcal{I}^{VAE}_p(\mathbf{x},\mathbf{z})$

すなわち提案手法の方が相互情報量が大きくなるため潜在変数との結びつきが強くなる．証明としては， $\mathbf{x}$ は層を重ねるごとに $\mathbf{z}$ との関係性が薄れていくものである，すなわち， $\mathcal{I}^{VAE}_p(\mathbf{x},\mathbf{h}^{(1)})\geq\mathcal{I}^{VAE}_p(\mathbf{x},\mathbf{z})$ が成り立つことを使う．SKIP-VAEでは全ての層に $\mathbf{z}$ が入力されるため逆に層を重ねるごとに $\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})\geq\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{h}^{(l)})$ と依存関係が増していく．よって $\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})\geq\mathcal{I}^{VAE}_p(\mathbf{x},\mathbf{z})$ が成り立つ．

この関係性を使えばELBOの最大化は次のような制約付きの形で表すことができる．

$\displaystyle \theta^\ast,\phi^\ast=\mathrm{arg}\max_{\theta, \phi}ELBO\:\:\:\mathrm{s.t}\:\:\:\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})\geq\delta=\mathcal{I}^{VAE}(\mathbf{x},\mathbf{z})$

よってskip-vaeにおける目的関数は

$\displaystyle \mathcal{L}=ELBO+\lambda\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z}) \\ =\mathcal{H}_p(\mathbf{x})-\lambda\mathcal{I}^{SKIP-VAE}_p(\mathbf{x},\mathbf{z})+KL(q_\phi(\mathbf{x},\mathbf{z})||p_\theta(\mathbf{x},\mathbf{z}))$

となる．ただし $\lambda$ はラグランジュ乗数で $\mathcal{I}_p(\mathbf{x})$ はデータ分布のエントロピーを表す．また最後の式の導出はaggregated posteriorで導入した関係性を使った．

まとめ

論文では目的関数導出後，4章で実験において目的関数をどう計算していくか具体的に説明していたが実装予定はないので割愛．アブストに理論的にlatent variable collapseを防ぐと書いてある割に相互情報量の不等式が雰囲気で導入されててなんだか腑に落ちない（自分の理解不足または英語力不足かも知れないけど）．とは言え感覚としてはわかるので手法としては実装も容易だし優秀な気がする．