はじめに

Masked Autoregressive Flow for Density Estimationを読んだのでメモ．

お気持ち

Inverse Autoregressive Flowと密接な関係がある，Real NVPの一般化手法Masked Autoregressive Flowを提案したぜということ．

密度推定の課題ではnormalizing flowやautoregressive modelがうまくいっているが，autoregressive modelは変数のオーダーにセンシティブな部分がありそこが欠点である．実際にFigure 1に例が載っていて，ターゲットの分布とモデルの分布の変数のオーダーを等しくしてもautoregressive modelは真の分布にフィットできず，フィットするのに適したオーダーを知るすべもないので良くない．さらに並列計算もあまり有効ではないがそれはMasked Autoencoder for Distribution Estimation(MADE)が解決した． Normalizing flowにおいてはInverse Autoregressive Flow(IAF)などのモデルは変分推定に基づいて作られているため密度推定に最適なものになっていない．

だからこの辺うまいこと統合して密度推定に適したものを作りたいというのが気持ち．

Masked Autoregressive Flow

まず，導入として自己回帰モデルをnormalizing flowの枠組みで考える．

はじめにガウス分布でparameterizedされた自己回帰モデルを考える．すると， $x_i$ の分布は

$\displaystyle p(x_i|\mathbf{x}_{1:i-1})=\mathcal{N}(x_i|\mu_i,(\exp\alpha_i)^2) \\ \displaystyle \mu_i=f_{\mu_i}(\mathbf{x}_{1:i-1}),\:\alpha_i=f_{\alpha_i}(\mathbf{x}_{1:i-1})$

と表すことができる．ただし， $f$ は平均と標準偏差の対数を計算するなんらかの関数．上記のモデルからデータを生成する際にはreparameterization trickから

$\displaystyle x_i=u_i\exp\alpha_i+\mu_i\\ \displaystyle u_i\sim\mathcal{N}(0,1)$

という計算を行えばいい．この式を考えてみれば，確率変数 $\mathbf{u}$ からデータ $\mathbf{x}$ への写像を与えていることになるため，モデルとしては $\mathbf{x}=f(\mathbf{u})$ として書くことができる．さらにこの関数は簡単に逆変換が $u_i=(x_i-\mu_i)\exp(-\alpha_i)$ のように記述でき，その逆変換のヤコビアンは自己回帰モデルの構造から三角行列になる．よって行列式は

$\displaystyle \left|\mathrm{det}\left(\frac{\partial f^{-1}}{\partial\mathbf{x}}\right)\right|=\exp\left(-\sum_i\alpha_i\right)$

として簡単に計算ができる．この辺りはInverse Autoregressive Flowと同様の話．

ここでflowとして自己回帰モデルを記述し，モデルを積み重ねて深層構造にすることで改善を行う．自己回帰モデルが $M_1,\dots,M_K$ と与えられたとき，モデル $M_k$ の乱数 $\mathbf{u}_k$ の密度を $M_k+1$ でモデル化するモデルを考える．これは最終的には乱数 $\mathbf{u}_k$ を標準ガウス分布でモデリングすることになるが，モデルを積み重ねることで柔軟性を持ったモデルを作ることができる．この論文では関数 $f_{\mu_u},f_{\alpha_i}$ をMADEを使ったアプローチで構成し，それをMasked Autoregressive Flow(MAF)とする．

IAFとの関係

今までの流れで定義されたMAFはほぼほぼIAFと同じ形をしていて，違うのはガウス分布のパラメータの計算が何を入力とするかだけ．MAFは上で書いたように $\mathbf{x}_{1:i-1}$ から計算されるのに対し，IAFは $\mathbf{u}_{1:i-1}$ から計算される．この違いは計算のコストに大きく現れ，MAFはデータ点 $\mathbf{x}$ の密度分布 $p(\mathbf{x})$ を一つのパスから計算できるが，サンプリングにはデータ点の次元数 $D$ のパスが必要になる．それに対しIAFは逆になる．このことからMAFは密度推定のためにデザインされていることが伺える．

理論的な面としては，初期の分布を $\pi$ を使って表せば，MAFにより計算される密度分布は変数変換の公式から

$\displaystyle p_x(\mathbf{x})=\pi_u(f^{-1}(\mathbf{x}))\left|\mathrm{det}\left(\frac{\partial f^{-1}}{\partial \mathbf{x}}\right)\right|$

となるのに対し，IAFは

$\displaystyle p_u(\mathbf{u})=\pi_x(f(\mathbf{u}))\left|\mathrm{det}\left(\frac{\partial f}{\partial \mathbf{u}}\right)\right|$

となる．ここから以下の関係が成り立つ．

$\displaystyle D_{KL}(\pi_x(\mathbf{x})||p_x(\mathbf{x}))=D_{KL}(p_u(\mathbf{u})||\pi_u(\mathbf{u}))$

証明はsupplementary materialにあったが簡単だったので割愛．基本的には $D_{KL}(\pi_x(\mathbf{x})||p_x(\mathbf{x}))$ を $\mathbf{u}=f(\mathbf{x})$ で変数変換することで $D_{KL}(p_u(\mathbf{u})||\pi_u(\mathbf{u}))$ の形になることを導くだけ．

この後も論文中では色々書いてあったが，要するにIAFは変分推定を使って問題を解いていて，MAFはnormalizing flowとして問題を解いているくらいの違いだと思う．そこの違いで出口が潜在変数になるかデータになるか変わってくると言ったところ．

realNVPとの関係

realNVPはNICEをベースにaffine coupling layerを使って拡張したモデルで，affine coupling layerは今まで定義された変数を使えば以下の形で定義される．

$\displaystyle \mathbf{x}_{1:d}=\mathbf{u}_{1:d},\:\mathbf{x}_{d+1:D}=\mathbf{u}_{d+1:D}\odot\exp\mathbf{\alpha}+\mathbf{\mu}$

式を見ればわかるようにIAFヤMAFとほぼほぼ同じ形をしていて，実際MAFの特別な場合とみなすことができる．realNVPでは $\mathbf{x}$ を二つに分けて一つは恒等写像されるが，IAFやMAFでは自己回帰性を満たすためのマスクにより $i\leq d$ については恒等写像で $i\gt d$ については変換が行われるため，realNVPの柔軟性を持つ一般化としてみなすことができる．ただし，realNVPはデータの生成も密度の推定もone forward passで計算できるため利点も存在する．

Conditional MAF

ここではサンプルのペア $(\mathbf{x}_n,\mathbf{y}_n)$ が与えられた時の条件付き分布 $p(\mathbf{x}|\mathbf{y})$ の推定を行うタスクについて考える．自己回帰モデルはこの問題に自然に拡張することが可能．確率における連鎖律から自己回帰モデルでは $p(\mathbf{x}|\mathbf{y})=\prod_ip(x_i|\mathbf{x}_{1:i-1},\mathbf{y})$ として条件付き分布を書くことが可能で，入力の変数を拡張するだけで記述できる．

まとめ

realNVPの一般化としてIAFを定義したというところか．密度推定を意識してモデリングされているからタスク次第で，IAFとMAFの優劣はつかない感じか．MAFはSupplementary materialに生成した画像載っていたが全然うまくいってなかったところを見るとIAFの方がモデルの形式からも生成には向いている？（IAFは生成画像が載ってないからわからないが。。。）ただそれにしてもモデルの形式上は大差ないGlowがあれだけ綺麗に画像を生成できたのは何が起きているのか気になる．