はじめに

Variational Inference with Normalizing Flowsを読んだのでメモ．OpenAIからでたGlowという生成モデルの論文が話題になっていて面白そうだったので関連知識を一から勉強してみる．まず手始めにnormalizing flowを使った変分推論について．

Variational Inference

この論文で問題としているのは生成モデルにおいて変分推定を適用するときに，近似事後分布の表現力が低いというもの．より良い近似を使えば強い生成モデルが作れるはずなのでその方法としてNormalizing Flowを使った変分推定を提案しますとのこと．

変分推定は以下のようにJensenの不等式を使って変分下限（または変分下界）を求め，それを最大化しようというもの．

$\displaystyle \log p_\theta \\ \displaystyle = \log\int p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})d\mathbf{z} \\ \displaystyle \geq -\mathbb{D}_{KL}[ q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z})]+\mathbb{E}_q[\log p_\theta(\mathbf{x}|\mathbf{z})]$

ここで $\mathbf{x}$ は観測値で $\mathbf{z}$ は潜在変数． $\theta$ はモデルのパラメータを表し， $q_\phi(\mathbf{z}|\mathbf{x})$ は近似事後分布を表す．近似事後分布は因子分解可能であることなどを仮定して解析的に扱いやすくしているため，元の分布を完全に表現することは不可能である．一応，近似事後分布にすることでの表現力の低下以外の問題として第2項目の勾配の計算コストが高いことも問題としてあげられる(VAEのreparameterizationはそこをうまく扱ったもの)．

Deep Latent Gaussian Models（DLGM）

Deep Neural Networksにおいて各層が潜在変数 $z_l$ を出力するとする．ここで $z_l$ は $l$ 番目の層に対応する潜在変数．このとき各々の潜在変数は自分の上位の層以外とは独立と仮定している．層の数を $L$ とするとDLGMにおける同時分布は各潜在変数の独立性より

$\displaystyle p(\mathbf{x},\mathbf{z}_1,...,\mathbf{z}_L)=p(\mathbf{x}|f_0(\mathbf{z}_1))\prod_lp(\mathbf{z}_l|f_l(\mathbf{z}_{l+1}))$

と表される．各々の潜在変数は事前分布として標準ガウス分布 $p(\mathbf{z}_l)=\mathcal{N}(\mathbf{0},\mathbf{I})$ が仮定されている．DLGMはDeepな分表現力は高いが，結局は分散が対角成分しかないような単純化されたガウス分布を用いるため真の分布を表現するには無理がある．

Normalizing flow

Normalizing flowは近似事後分布の作り方の一つで表現力の高い近似事後分布を与えることができる．イメージとしては確率分布の変数変換を繰り返し用いることで単純な分布を複雑な分布へと変換できるというもの．

まず，確率変数の変数変換として基本的な規則について．逆変換を持つ連続な写像 $f:\mathbb{R}^d\rightarrow\mathbb{R}^d,\:f^{-1}=g$ を考える．分布 $q(\mathbf{z})$ を持つ確率変数 $\mathbf{z}$ に対し変換を行った確率変数 $\mathbf{z}'=f(\mathbf{z})$ を考える．このとき， $\mathbf{z}'$ に対する確率密度を $q'(\mathbf{z}')$ とすると

$\displaystyle q'(\mathbf{z}')=q(\mathbf{z})\left|\mathrm{det}\frac{\partial f^{-1}(\mathbf{z}')}{\partial\mathbf{z}'}\right|$

と与えられる．詳しいことは「確率密度　変数変換」などでググるとたくさん出てくるので割愛．一応こちらのブログはnormalizing flowの内容含めわかりやすい．

ここで関数 $f$ が $\mathbf{z}_K=f_K\circ\dots\circ f_2\circ f_1(\mathbf{z}_0)$ のように $K$ 個の関数の合成関数になっていると考える．すると上記の変数変換を繰り返し用いれば，

$\displaystyle q_K(\mathbf{z}_K)=q_{K-1}(\mathbf{z}_{K-1})\left|\mathrm{det}\frac{f^{-1}_K(\mathbf{z}_K)}{\mathbf{z}_K}\right|=\dots=q_0(\mathbf{z}_0)\prod_k\left|\mathrm{det}\frac{f^{-1}_{k}(\mathbf{z}_{k})}{\mathbf{z}_k}\right|=q_0(\mathbf{z}_0)\prod_k\left|\mathrm{det}\frac{f_{k}(\mathbf{z}_{k})}{\mathbf{z}_k}\right|^{-1}$

となる．最後の等式は逆関数定理から成り立つ．さらにこれの対数をとると

$\displaystyle \log q_K(\mathbf{Z}_K)=\log q_0(\mathbf{z}_0)-\sum_k\log\mathrm{det}\left|\frac{f_k}{\mathbf{z}_k}\right|$

となる．ただし関数 $f$ の引数を省略したことに注意．また論文では $\log$ ではなく $\ln$ となっていたが，本質的には変わらないはずなのでここでは $\log$ を使って表現する．ここで注目すべきは $q_K$ を計算するのに $q_0$ しか必要がないこと．これによって適当な関数 $h(\mathbf{z})$ の $q_K$ に関する期待値の計算も $\mathbb{E}_{q_K}[ h(\mathbf{z})]=\mathbb{E}_{q_0}[h(f_K\circ\dots\circ f_1(\mathbf{z}_0))]$ のように行うことが可能．

初期の分布 $q_0$ は従来通りindependent Gaussianなど使っても，DLGMと違い各潜在変数は何か分布を仮定しているわけではないので表現力が高く，かつ $q_0$ を使って計算できるため解析的に扱いやすいという利点がある．

上記では有限のフローを扱ったが，有限があれば無限も考えたくなるのが性というもの．ただし，以下の内容は前提知識が足りず理解しきれていないのでほぼ論文の和訳メモになっていることに注意．

まずLangevin Flow（ランジュバンフロー）について．Langevin Flowでは以下のランジュバン方程式（Langevin stochastic differential equation）によって確率変数 $\mathbf{z}$ が記述される．

$\displaystyle d\mathbf{z}(t)=\mathbf{F}(\mathbf{z}(t),t)dt+\mathbf{G}(\mathbf{z}(t),t)d\mathbf{\xi}(t)$

ここで $d\mathbf{\xi}(t)$ は $\mathbb{E}[\mathbf{\xi}_i(t)]=0$ ， $\mathbb{E}[\mathbf{\xi}_i(t)\mathbf{\xi}_j(t')]=\delta_{i,j}\delta(t-t')$ で与えられるWiener過程， $\mathbf{F}$ はドリフトベクトルで $\mathbf{D}=\mathbf{G}\mathbf{G}'$ は拡散行列を表す．ここで初期の分布が $q_0(\mathbf{z})$ である確率変数 $\mathbf{z}$ が上記の方程式により変化していく場合，密度分布の変化の規則はFokker-Planck方程式で与えられる．すなわち，

$\displaystyle \frac{\partial}{\partial t}q_t(\mathbf{z})=-\sum_i\frac{\partial}{\partial z_i}[F_i(\mathbf{z},t)q_t]+\frac{1}{2}\sum_{i,j}\frac{\partial^2}{\partial z_i\partial z_j}[D_{ij}(\mathbf{z},t)q_t]$

として変化を記述できる．機械学習の文脈では $F(\mathbf{z},t)=-\nabla_z\mathcal{L}(\mathbf{z})$ ， $G(\mathbf{z},t)=\sqrt{2}\delta_{ij}$ としたLangevin flowを用いる．ただし， $\mathcal{L}(\mathbf{z})$ は正規化されていないモデルの対数密度関数．重要なこととして，定常状態における解 $q_\infty(\mathbf{z})$ はボルツマン分布になることが知られている．つまり，適当な分布 $q_0(\mathbf{z})$ を持つ確率変数 $\mathbf{z}_0$ をLangevin flowで変換していくと最終的には $q_\infty(\mathbf{z})\propto\exp(-\mathcal{L}(\mathbf{z}))$ となる．ボルツマンマシンが任意の分布を近似可能だったことを思い出せばnormalizing flowも任意の分布を近似可能と言える（のかな？）．

Inference with Normalizing Flows

ここでは有限のnormalizing flowを用いた推定について考える．基本的には変数変換における関数 $f$ は逆変換可能なものかつヤコビアンンが効率的に計算できるものでなければならない．そこで以下のような変換を考える．

$\displaystyle f(\mathbf{z})=\mathbf{z}+\mathbf{u}h(\mathbf{w}'\mathbf{z}+b)$

ここで， $\lambda=\{\mathbf{w}\in\mathbb{R}^D,\mathbf{u}\in\mathbb{R}^D,b\in\mathbb{R}\}$ は学習パラメータで， $h$ はsmoothでelement-wiseな非線形関数．上記の変換はヤコビアンの項を $\mathcal{O}(D)$ 時間で計算ができる．具体的には非線形関数の導関数を $h'$ とすると，右辺第二項目の $\mathbf{z}$ に関する微分は

$\displaystyle \psi(\mathbf{z})=h'(\mathbf{w}'\mathbf{z}+b)\mathbf{w}$

とかけるので，行列式は

$\displaystyle \mathrm{det}\left|\frac{\partial f}{\partial\mathbf{z}}\right| =|\mathrm{det}(\mathbf{I}+\mathbf{u}\psi(\mathbf{z})'|=|1+\mathbf{u}'\psi(\mathbf{z})|$

と計算可能．よって計算コストのオーダーが $\mathcal{O}(D)$ になる．

上記で定義された変換 $f$ を，元の有限なnormalizing flowの式に当てはめれば，

$\displaystyle \log q_K(\mathbf{z}_K)=\log q_0(\mathbf{z}_0)-\sum_k\log|1+\mathbf{u}_k'\psi_k(\mathbf{z}_k)|$

が導かれる．このflowは $\mathbf{w}'\mathbf{z}+b=0$ で与えられる超平面に垂直な方向での収縮と拡大を与えており，それにより初期密度 $q_0$ を変化させると考えられる．そのためこれをplanar flowと定義する．

また別なアプローチとして，初期密度を確率変数 $\mathbf{z}_0$ 周りで変化させるような変換群を考える．そのような変換群は

$\displaystyle f(\mathbf{z})=\mathbf{z}+\beta h(\alpha,r)(\mathbf{z}-\mathbf{z}_0) \\ \displaystyle \mathrm{det}\left|\frac{\partial f}{\partial\mathbf{z}}\right|=[1+\beta h(\alpha,r)]^{(d-1)}[1+\beta h(\alpha,r)+h'(\alpha,r)r)]$

として考えられる．ただし， $r=|\mathbf{z}-\mathbf{z}_0|$ ， $h(\alpha,r)=1/(\alpha+r)$ とし，パラメータとして $\lambda=\{\mathbf{z}_0\in\mathbb{R}^D,\alpha\in\mathbb{R},\beta\in\mathbb{R}\}$ ．この群による変換もまた線形時間で行列式を計算可能である．このflowは放射状の(radial)収縮と拡大を与えるのでradial flowと定義する．planer flowとradial flowで球面ガウス分布を変換した例が論文のFigure 1．より複雑な分布になっていることがわかる．

ここで注意としてplaner flow，radial flowの全てが逆変換可能な関数になっているわけではなく，一定の条件下のみで逆変換可能となる．詳細なことはAppendixに．

ここで変分推定における目的関数にnormalizing flowを適用すると，

$\displaystyle -\log p_\theta(\mathbf{x}) \\ \displaystyle \geq \mathbb{D}_{KL}[ q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z})]-\mathbb{E}_q[\log p_\theta(\mathbf{x}|\mathbf{z})] \\ \displaystyle = \mathbb{E}_{q_\phi(z|x)}[\log q_\phi(\mathbf{z}|\mathbf{x})-\log p(\mathbf{z})-\log p(\mathbf{x}|\mathbf{z})] \\ \displaystyle =\mathbb{E}_{q_0(z_0)}[\log q_K(\mathbf{z}_K)-\log p(\mathbf{x},\mathbf{z}_K)]\\ \displaystyle =\mathbb{E}_{q_0(z_0)}[\log q_0(\mathbf{z}_0)] -\mathbb{E}_{q_0(z_0)}[\log p(\mathbf{x},\mathbf{z}_K)]-\mathbb{E}_{q_0(z_0)}\left[\sum_k\log |1+\mathbf{u}'_k\psi_k(\mathbf{z}_k)|\right]$

となる．4行目はnormalizing flowでは $q_K$ の期待値が $q_0$ の期待値でかけることを利用しており，最後の行はplaner flowを代入した．

この論文ではVAEと同様にニューラルネットを使ったモデルを使っている．encoderとdecoderの間で関数 $f$ を適用していて， $k$ の数を変えて実験を行なっている．NICEと比較した結果の可視化が論文のFigure 3(a)(b)(c)．提案手法の方が真の分布をうまく表現できていることが分かる．

ちなみに学習に関して最終的な変分下限の式には $p(\mathbf{x},\mathbf{z}_K)$ という式が現れているが，目的関数としては $p(\mathbf{x}|\mathbf{z}_K),p(\mathbf{z}_K)$ で分かれるはず． $p(\mathbf{z}_K)$ はtable 1の分布を使ったと書いてあって， $p(\mathbf{x}|\mathbf{z}_K),p(\mathbf{z}_K)$ どちらも勾配は元のVAEと同様モンテカルロ推定で計算する．あと目的関数の $p(\mathbf{x},\mathbf{z}_K)$ 項には温度パラメータ $\beta$ を係数としてかけて $\beta_t=\min(1,0.01+t/10000)$ としてアニーリングしていくらしい．