はじめに

Sylvester Normalizing Flows for Variational Inferenceを読んだのでメモ．planer flowの一般化であるSylvester normalizing flowを提案．planer flowは以前勉強した．

Normalizing flow

復習がてらnormalizing flowについて．

確率変数 $\mathbf{z}$ を関数 $f$ を使って $\mathbf{z}'$ に変数変換する．その時， $\mathbf{z}'$ の確率密度は次のように計算できる．

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

上記は変数変換の公式として知られていて調べればいくらでも出てくる．基本的には $P_0(\mathbf{z})$ も $P_1(\mathbf{z}')$ も確率密度の定義から積分が1であるため $P(\mathbf{z})d\mathbf{z}=P(\mathbf{z}')d\mathbf{z}'$ という関係が成り立つことから導かれる．

次にこの変数変換を変分推定に組み込むことを考える．初期の確率変数 $\mathbf{z}_0$ をdiagonal Gaussian $\mathcal{N}(\mathbf{z}_0|\mathbf{\mu}(\mathbf{x}),\mathbf{\sigma}^2(\mathbf{x}))$ に従う確率変数とする．この変数 $\mathbf{z}_0$ を $\mathbf{z}_K=(f_K\circ\dots\circ f_1)(\mathbf{z}_0$ のように $K$ 回変数変換することを考える．すると確率密度 $q_K(\mathbf{z}_K)$ は以下の計算により求めることができる．

$\displaystyle \log q_K(\mathbf{z}_K|\mathbf{x})=\log q_0(\mathbf{z}_0|\mathbf{x})-\sum_k\log\left|\mathrm{det}\left(\frac{\partial f_k(\mathbf{z}_{k-1};\lambda_k(\mathbf{x}))}{\partial\mathbf{z}_{k-1}}\right)\right|$

$\lambda_k$ は $k$ 番目の変換のパラメータを表す．上記の $q_K(\mathbf{z}|\mathbf{x})$ を変分事後分布とすると変分推論における下限の式は以下のように表現ができる．

$\displaystyle \mathcal{F}(\theta,\phi) \\ \displaystyle =\mathbb{E}_{q_phi}[\log q_\phi(\mathbf{z}|\mathbf{x})-\log p_\theta(\mathbf{x},\mathbf{z})] \\ \displaystyle =\mathbb{E}_{q_0}[\log q_0(\mathbf{z}_0|\mathbf{x})-\log p_\theta(\mathbf{x},\mathbf{z})] - \mathbb{E}_{q_0}\left[\sum_k\log\left|\mathrm{det}\left(\frac{\partial f_k(\mathbf{z}_{k-1};\lambda_k(\mathbf{x}))}{\partial\mathbf{z}_{k-1}}\right)\right|\right]$

通常の変分推論では変分事後分布は勾配計算のためにサンプリングを用いるため，サンプリングが簡単にできるような単純な分布である必要がある．それに対し，noramlizing flowは単純な分布から複雑な分布へ変換することで複雑な分布を変分事後分布として用いることが可能で，サンプリングも単純な分布から行うことができる．それにより変分推論の近似精度をあげることができるというもの．

先行研究ではこの変数変換の関数を以下のplanar flowと呼ばれる関数族で定義した．

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

各パラメータは $\mathbf{u},\mathbf{W}\in\mathbb{R}^D,b\in\mathbb{R}$ として定義されれ $h$ は $\mathrm{tanh}$ のような全単射の活性化関数を表す．このplanar flowは $\mathbf{u}^T\mathbf{w}\geq -1$ である限り逆変換が可能である．

またplanar flowのヤコビアンの行列式は以下の形で計算が可能．

$\displaystyle \mathrm{det}\frac{\partial\mathbf{z}'}{\partial\mathbf{z}}=\mathrm{det}(\mathbf{I}+\mathbf{u}h'(\mathbf{w}^T\mathbf{z}+b)\mathbf{w}^T)=1+\mathbf{u}^Th'(\mathbf{w}^T\mathbf{z}+b)\mathbf{w}$

ここで $h'$ は $h$ の導関数を表し，このように計算できることで行列式の計算コストが $\mathcal{O}(D)$ になる（普通は $\mathcal{O}(D^3)$ ）．

Sylvester normalizing flow

名前がかっこいい．ここでは次のようなresidual connectionを持つ隠れユニット数 $M$ の単層ニューラルネットのような変換を考える．

$\displaystyle \mathbf{z}'=\mathbf{z}+\mathbf{A}h(\mathbf{B}\mathbf{z}+\mathbf{b})$

各変数は $\mathbf{A}\in\mathbb{R}^{D\times M},\mathbf{B}\in\mathbb{R}^{M\times D},\mathbf{b}\in\mathbb{R}^M,M\leq D$ で定義されている．この変換のヤコビアンの行列式は以下のSylvester's determinant identityを使うことにより求めることが可能とのこと．

$\displaystyle \mathrm{det}(\mathbf{I}_D+\mathbf{A}\mathbf{B})=\mathrm{det}(\mathbf{I}_M+\mathbf{B}\mathbf{A})$

ここで $\mathbf{I}_D$ は $D$ 次元単位行列を表す．上記の関係を用いれば $M\lt D$ の時に $D\times D$ 行列の行列式を $M\times M$ 行列の行列式として計算ができるため，計算コストの削減ができる．このSylvester's determinant identityを使えば最初に定義した変換のヤコビアンの行列式は以下のように計算できる．

$\displaystyle \mathrm{det}\frac{\partial\mathbf{z}'}{\partial\mathbf{z}}=\mathrm{det}(\mathbf{I}_M+\mathrm{diag}(h'(\mathbf{B}\mathbf{z}+\mathbf{b}))\mathbf{B}\mathbf{A})$

導出は普通に微分するだけ．嬉しいこととしては $\mathbf{BA}$ のように計算している点で，さっきも書いたように行列のサイズが落ちるため計算コストが下がる．

注意としてこの変換自体は基本的には逆変換を持たない．逆変換を持つためには $\mathbf{A},\mathbf{B}$ が次のようなQR分解の形で表される必要がある．

$\displaystyle \mathbf{z}'=\mathbf{z}+\mathbf{QR}h(\mathbf{\tilde{R}Q}^T\mathbf{z}+\mathbf{b})=\phi(\mathbf{z})$

ただし， $\mathbf{R},\mathbf{\tilde{R}}$ は $M\times M$ 上三角行列で $Q$ は $D\times M$ 直交行列となる．この時ヤコビアンの行列式は次の形で表される．

$\displaystyle \mathrm{det}\frac{\partial\mathbf{z}'}{\partial\mathbf{z}}=\mathrm{det}(\mathbf{I}_M+\mathrm{diag}(h'(\mathbf{\tilde{R}Q}^T\mathbf{z}+\mathbf{b}))\mathbf{\tilde{R}Q}^T\mathbf{QR})=\mathrm{det}(\mathbf{I}_M+\mathrm{diag}(h'(\mathbf{\tilde{R}Q}^T\mathbf{z}+\mathbf{b}))\mathbf{\tilde{R}}\mathbf{R})$

最後は $Q$ が直交行列であることを利用した．ここで $\mathbf{\tilde{R}R}$ は上三角行列であるため行列式の計算オーダーは $\mathcal{O}(M)$ となる．

この変換が逆変換を持つための十分条件は， $\mathbf{R},\mathbf{\tilde{R}}$ が上三角行列で $h:\mathbb{R}\rightarrow\mathbb{R}$ が連続な関数である時， $\mathbf{R},\mathbf{\tilde{R}}$ が $r_{ii}\tilde{r}_{ii}\gt -1/||h'||_\infty$ かつ $\mathbb{\tilde{R}}$ が逆行列を持つという条件で与えられる．以下 $\mathbf{R},\mathbf{\tilde{R}}$ がともに対角行列であるときの証明．

$\mathbf{Q}$ は直交行列であるため， $\mathbf{Q}$ が張る部分空間は $\mathcal{W}=\mathrm{span}\{\mathbf{q}_q,\dots,\mathbf{q}_M\}$ として表すことができ，その直交補空間を $\mathcal{W}^\perp$ で定義する．すると $\mathbf{z}=\mathbf{z}_{||}+\mathbf{z}_perp$ として分解することができる．ただし， $\mathbf{z}_{||}\in\mathcal{W},\mathbf{z}_\perp\in\mathcal{W}^\perp$ ．当然， $\mathbf{QR}h(\mathbf{\tilde{R}Q}^T\mathbf{z}+\mathbf{b})\in\mathcal{W}$ となる．したがって関数 $\phi$ は $\mathbf{z}_{||}$ 上のみで振る舞うことになる．したがって， $\mathbf{z}_{||}$ 上での $\phi$ の振る舞いのみを考えればよく， $\phi(\mathbf{z})$ に左から $\mathbf{Q}$ をかければ

$\displaystyle \mathbf{Q}^T\mathbf{z}' = \mathbf{Q}^T\mathbf{z}+\mathbf{R}h(\mathbf{\tilde{R}}\mathbf{Q}^T\mathbf{z}+\mathbf{b})=(f_1(v_1),\dots,f_M(v_M))^T$

としてかける．ここでベクトル $\mathbf{v}$ は $\mathbf{Q}$ が張る部分空間上での $\mathbf{z}_{||}$ の座標を表す．ここでは $\mathbf{R}$ は対角行列であることを仮定しているため， $\mathbf{v}$ の各次元は独立であり， $f_i(v)=v+r_{ii}h(\tilde{r}_{ii}v+b_i)$ という変換としてかける．よって各次元ごとに考えることができ， $||h'||_\infty r_{ii}\tilde{r}_{ii}\gt -1$ であることから， $f_i'(v)\gt 0$ であることが言える．導関数が正の1次元の実関数は逆変換を持つため $f_i$ は逆変換を持つ．よって $\phi(\mathbf{z})$ は逆変換を持つことが言え，その変換は

$\displaystyle \phi^{-1}(\mathbf{z}')=\mathbf{z}_\perp'+f^{-1}(\mathbf{z}_{||}')=\mathbf{z}$

としてかける．論文では当然 $\mathbf{R}$ が三角行列の時の証明も行われているがかくのが面倒なのでここでは割愛．

実践的な面での課題ではパラメータの学習時にどのようにして $\mathbf{Q}$ を直交行列として保ち続けるかという問題がある．ここでは $\mathbf{Q}$ を直交行列として保ち続けるために，Orthogonal Sylvester flows, Householder Sylvester flows, Triangular Sylvester flowsの三つのflowを提案する．

Orthogonal Sylvester flows（O-SNF）

ここでは $||\mathbf{Q}^{(0)T}\mathbf{Q}^{(0)}-\mathbf{I}||_2\lt 1$ を満たしながら以下の手順を用いることで $\mathbf{Q}$ の直交性を保つ．ただし，行列の2-normは $||\mathbf{X}||_2=\lambda_{\max}(\mathbf{X})$ で計算され， $\lambda_\max(\mathbf{X})$ は $\mathbf{X}$ の最大特異値を表す．

$\displaystyle \mathbf{Q}^{(k+1)}=\mathbf{Q}^{(k)}\left(\mathbf{I}+\frac{1}{2}\left(\mathbf{I}-\mathbf{Q}^{(k)T}\mathbf{Q}^{(k)}\right)\right)$

ただし実験の時には $||\mathbf{Q}^{(k)T}\mathbf{Q}^{(k)}-\mathbf{I}||_F\lt \epsilon$ を満たすように行ったらしい． $||\cdot||_F$ はフロベニウスノルム．この手順は微分可能なためbackpropに組み込めるので嬉しいという点も．

Householder Sylvester flows（H-NSF）

ここでは次のHouseholder変換の積によって直交行列を作る．ハウスホルダー変換は超平面上での反射を表していて， $\mathbf{v}\in\mathbb{R}^D$ とすると $\mathbf{v}$ に直交する超平面上での反射（つまりハウスホルダー変換）は以下で与えられる．

$\displaystyle H(\mathbf{z})=\mathbf{z}-\frac{\mathbf{vv}^T}{||\mathbf{v}||^2}\mathbf{z}$

これは計算コストは低くパラメータも $D$ しか持たないため嬉しい．ちなみに $M\times M$ 直交行列は $M-1$ 回のハウスホルダー変換の積で記述できることが知られているらしい．注意としてはハウスホルダー変換を使う場合は，ハウスホルダー変換の結果が正方行列であるため $M=D$ とする必要がある．

Triangular Sylvester flows（T-SNF）

これは $\mathbf{Q}$ を単位行列または置換行列として，flowの変換の過程で単位行列と置換行列が交互に現れるようなflow．つまり， $\mathbf{\tilde{R}}$ と $\mathbf{R}$ が上三角と下三角として交互にflowに現れるようにすることと等しい．すなわちT-SNFでは $\mathbf{Q}$ は固定．