はじめに

Improved Variational Inference with Inverse Autoregressive Flowを読んだのでメモ．

変分推定

論文に章立ててあるのと復習も含め何度も書くことが重要かなということでおさらい．

観測変数を $\mathbf{x}$ ，潜在変数を $\mathbf{z}$ とし，同時分布を $p(\mathbf{x},\mathbf{z})$ とする．データセット $\mathbf{X}=\{\mathbf{x}^1,\dots,\mathbf{x}^N\}$ が与えられた時，以下の周辺尤度の最大化をしたいというのが目標．

$\displaystyle \log p(\mathbf{X})=\sum_i^N\log p(\mathbf{x}^{(i)})$

一般的には上のような周辺尤度は解析的に計算困難な場合が多い．そこで解決策の一つとして $q(\mathbf{z}|\mathbf{x})$ で表される推論モデルと呼ばれるものを導入し，変分下限を最適化することが考えられる．すなわちある観測 $\mathbf{x}$ に対し以下の周辺対数尤度の変分下限を最適化する．

$\displaystyle \log p(\mathbf{x})\geq\mathbb{E}_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x},\mathbf{z})-\log q(\mathbf{z}|\mathbf{x})]=\mathcal{L}(\mathbf{x};\mathbf{\theta})$

$\mathbf{\theta}$ は $p,q$ に関するモデルパラメータを表す．さらに式変形を行えば，

$\displaystyle \mathbb{E}_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x},\mathbf{z})-\log q(\mathbf{z}|\mathbf{x})] \\ \displaystyle =\mathbb{E}_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{z}|\mathbf{x})p(\mathbf{x})-\log q(\mathbf{z}|\mathbf{x})] \\ \displaystyle =\log p(\mathbf{x})-D_{KL}(q(\mathbf{z}|\mathbf{x})||p(\mathbf{z}|\mathbf{x}))$

として書き換えることができる．この式を最適化する方法は数多く存在していて，VAEのreparameterizationなどはその一例．ここで注意としては，一般的に潜在変数が多変数のモデルにおいては $q(\mathbf{z}_a,\mathbf{z}_b|\mathbf{x})=q(\mathbf{z}_a|\mathbf{x})q(\mathbf{z}_b|\mathbf{z}_a,\mathbf{x})$ のように分解できることを仮定としておく．

変分下界を最適化する上で必要な性質としては（1） $q(\mathbf{z}|\mathbf{x})$ が効率的に計算できること（2）そこから効率的にサンプリングできることの二つがあげられる．例えば， $\mathbf{z}$ が高次元の場合，効率的に計算するために並列計算できることが望ましく，実践的には $q(\mathbf{z}|\mathbf{x})$ を等方性のガウス分布などに制限することが多い．ただし， $q(\mathbf{z}|\mathbf{x})$ は真の分布 $p(\mathbf{z}|\mathbf{x})$ に一致できるような柔軟性が必要でもある．

flow

散々勉強したけど今一度．flowは前述したように計算効率と分布の柔軟性を兼ね備えた $q(\mathbf{z}|\mathbf{x})$ を得る方法で，変数変換を何度も用いることで単純な分布を複雑な形にしようというもの．つまり初期の分布 $q(\mathbf{z}_0|\mathbf{x})$ を球面ガウス分布などで考えて， $\mathbf{z}_t=\mathbf{f}_t(\mathbf{z}_{t-1},\mathbf{x})$ のように変数変換を行っていくというもの．この時変数変換された確率密度は公式によって以下の形で与えられる．

$\displaystyle \log q(\mathbf{z}_T|\mathbf{x})=\log q(\mathbf{z}_0|\mathbf{x})-\sum_t^T\log\mathrm{det}\left|\frac{d\mathbf{z}_t}{d\mathbf{z}_{t-1}}\right|$

flowの問題点の一つとして変数変換の関数 $f$ に強い制約があることがあげられる．

Inverse Autoregressive Transformations

この論文では自己回帰モデルのガウス版を変数変換に使ったflowを考える．

今自己回帰モデルにおける変数を $\mathbf{y}$ とし，ガウス分布のパラメータを $\mu_i(\mathbf{y}_{1:i-1}),\sigma_i(\mathbf{y}_{1:i-1})$ と表す．自己回帰の構造から， $j\geq i$ の時には $i,j$ に依存関係がないため $\partial[\mathbf{\mu}_i,\mathbf{\sigma}_i]/\partial\mathbf{y}_i=[0,0]$ が成り立つ．（これはヤコビアンが三角行列になることを意味している）．

上記のようなモデルからサンプリングを行うには， $\mathbf{\epsilon}\sim\mathcal{N}(0,\mathbf{I})$ とすればreparameterizationにより， $y_i=\mu_i(\mathbf{y}_{1:i-1})+\sigma_i(\mathbf{y}_{1:i-1})\epsilon_i$ としてサンプリングすることができる．ただし，基本的に変分推論においては事後分布からのサンプリングが必要となるため上記のサンプリングが効率に行えたところであまり嬉しさはない．ここで嬉しいのは逆変換が $\epsilon_i=\frac{y_i-\mu_i(\mathbf{y}_{1:i-1})}{\sigma_i(\mathbf{y}_{1:i-1})}$ として記述が可能なことで，このように逆変換が簡単にかけることはnormalizing flowの文脈においてはとても嬉しい．ここではこのような自己回帰モデルの逆変換をinverse autoregressive flowと呼ぶ．

上記二つの性質から，reparameterizationをすれば分布の計算で必要なヤコビアンの行列式の対数は

$\displaystyle \log\mathrm{det}\left|\frac{d\mathbf{\epsilon}}{d\mathbf{y}}\right|=\sum_i^D-\log\sigma_i(\mathbf{y})$

として計算ができることがわかる．

Inverse Autoregressive Flow

ここではモデルの全体像を見る．基本的にはinverse autoregressive transformationをnormalizing flowに適用する．処理の流れは論文中のAlgorithm 1とFigure 2によくまとまっている．まず標準正規分布から $\mathbf{\epsilon}$ をサンプリングし，初期の潜在変数 $\mathbf{z}_0$ を

$\displaystyle \mathbf{z}_0=\mathbf{\mathbf{\mu}}_0+\mathbf{\sigma}_0\odot\mathbf{\epsilon}$

と計算する．ただし， $\mathbf{\mu}_0,\mathbf{\sigma}_0$ はEncoderNNの出力として得られる．その後 $t>0$ の潜在変数を

$\displaystyle \mathbf{z}_t=\mathbf{\mu}_t+\mathbf{\sigma}_t\odot\mathbf{z}_{t-1}$

として計算していく．ここでの $\mathbf{\mu}_t,\mathbf{\sigma}_t$ はEncoderNNとは別なAutoregressiveNNから計算される．ここでは $\frac{d\mathbf{\mu}_t}{d\mathbf{z}_{t-1}},\frac{d\mathbf{\sigma}_t}{d\mathbf{z}_{t-1}}$ が対角成分0の三角行列となるようにすることでニューラルネットに自己回帰としての性質を持たせる．結果として $\frac{d\mathbf{z}_t}{d\mathbf{z}_{t-1}}$ は対角成分に $\mathbf{\sigma}_t$ を持つ三角行列となるため， $t$ 番目のflowにおける行列式は $\prod_i^D\sigma_{t,i}$ として計算ができる．以上から，flowにより得られる確率密度は $q(\mathbf{z}_0|\mathbf{x})$ が標準ガウス分布であることに注意すれば

$\displaystyle \log q(\mathbf{z}_T|\mathbf{x})=-\sum_i^D\left(\frac{1}{2}\epsilon_i^2+\frac{1}{2}\log2\pi+\sum_t^T\log\sigma_{t,i}\right)$

として計算される．

ちなみに実験的な部分として，計算の安定性からautoregressiveNNの出力を $\mathbf{m}_t,\mathbf{s}_t$ とし， $\mathbf{\sigma}_t=\mathrm{sigmoid}(\mathbf{s}_t)$ として分散を計算する．さらに潜在変数は $\mathbf{z}_t=\mathbf{\sigma}_t\odot\mathbf{z}_{t-1}+(1-\mathbf{\sigma}_t)\odot\mathbf{m}_t$ として計算するといいらしい．あと自己回帰の要請を満たすためにautoregressiveNNはMADEなどを使ったとのこと．

個人的に悩んだところとして目的関数である変分下限の $p(\mathbf{z}|\mathbf{x})$ がどこから得られるのかわからなかった．例えばオリジナルのVAEだとそもそも $p(\mathbf{z})$ になっていて $\mathbf{x}$ で条件付けられてなく，標準正規分布としている．この論文では $p(\mathbf{z}|\mathbf{x})$ もモデルから得られて， $q(\mathbf{z}|\mathbf{x})$ と $p(\mathbf{z}|\mathbf{x})$ の違いはどっち向きに流すかだけ．すなわち $\mathbf{z}\rightarrow\mathbf{x}$ と $\mathbf{x}\rightarrow\mathbf{z}$ のどっちの処理から得られた潜在変数かという感じ．Figure 3を見ればなんとなくわかる．