HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

ニューラルネットワークを用いた近年の変分推論2

 

 

follow us in feedly

はじめに

この記事は下記の記事の続編です。ただし、下記の記事は変分ベイズの話の基本的事項を、ニューラルネットワークを使いそうな雰囲気を醸しながら書いているに過ぎなく、既に理解している人はわざわざ読まなくてもいいと思われます。

ちなみにニューラルネットワークを使いそうな雰囲気を出しているというのは、潜在変数を含んだ統計モデルの同時分布 $p(x,z)$ に対してハイパーパラメータ $\theta$ があり、$p _ \theta(x, z)$ と明記しつつ、この $\theta$ がニューラルネットワークのパラメータになってたりする場合がある(普通は事前分布のハイパーパラメータだけであることが多い)という意識をして書いている点です。また、変分事後分布 $q(z)$ に対しても、変分パラメータ $q _ \phi (z)$ と明記し、この $\phi$ がニューラルネットワークのパラメータになっていることを意識している点もあります。この形式で、データ集合$X = \{x _ 1, ..., x _ N\}$ から変分事後分布 $q _ \phi (z \mid X)$ によって事後分布を形成する手法を、変分オートエンコーダと言います。

この記事では、変分ベイズとニューラルネットワークが融合してきた領域を扱うため、上記の記事と下記のベイズの基本は前提となります。

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

また、ボールド体を用いるのはまあまあなしんどさであるので、基本的には用いません。問題に応じてデータの次元が異なるので、当然ベクトルであるかスカラーであるかは本来書き分ける必要があるところですが、今回は誤解を招かない限り書き分けることはしません。大文字に関しては行列という意味の場合と、データ集合である場合があるので、わからなさそうな場合明記します。また、確率変数とその実現値は表記の濫用になりますが全て同一視しながら同じ表記で書いてしまいます(確率変数という関数を意識しなければならない場面は無いはずなので)。

ベイズ推論基本

ベイズ推論

ベイズ推論では観測変数 $x$ と潜在変数 $z$ に関して同時分布$p(x, z)$が

$$ p(x, z) = p(z\mid x) p(x) = p(x\mid z)p(z) $$

と書き表せることを基本原理にしていきます。これは乗法定理として常に成り立っている一般的な式ですが、観測変数や潜在変数が複数ある場合に、条件付き独立性などを付与することでモデル化を行っていきます。これは一般的に成り立っている式から、特殊な制約を加えることで、モデリングするということです。

このとき、

$$ p(z \mid x) = \frac{p(x\mid z)p(z)}{p(x)} $$

をベイズの定理と呼び、観測変数から潜在変数の分布を計算することが可能となります。既に手元に観測データ集合 $X$ があるときに、このベイズの定理を用いて事後分布 $p(z \mid X)$ を求めるのがベイズ推論の当面の目標になります。ここで現れる $p(z)$ を事前分布と呼び、設計者が何らかの知見を導入することで決定します。

$p(x)$ は周辺尤度と呼ばれ、加法定理(周辺化)によって

$$ p(x) = \int _ z p(x, z)dz = \int _ z p(x \mid z)p(z) dz $$

と書き表されます(最右辺は乗法定理を利用している)。 この周辺尤度が大きいほど、確率分布 $p$ の設計がデータによく当てはまっていることを示しているため、第二種最尤推定や変分ベイズ推論で活用されます。

ベイズ予測分布

ベイズ予測分布は、求まった事後分布 $p(z \mid X)$ を利用して、新しいデータ $x ^ *$ に対する予測分布を構築します。 ここで興味があるのが観測変数 $x$ だけであるとすれば、潜在変数を下記のように積分消去することでベイズ予測分布を作ることできます。

$$ p(x ^ * \mid X) =\mathbb E _ {p(z\mid X)} [p(x ^ * \mid z)] = \int _ z p(x ^ * \mid z) p(z \mid X) dz $$

これはモデリングする上で現れただけで、潜在変数自体を予測したりする必要はないという意図があり、実施される計算です。もしも、予測したい変数がある場合には、それに応じた予測分布を作る必要があります。

変分推論

実際には事後分布 $p(z \mid X)$ を作ること自体が困難を極める場合が多々あります。そこで、導出される事後分布よりも求めるのが簡単であるが、その表現力を失わない程度には複雑な確率分布 $q(z)$ によって事後分布を代替します。この分布を変分近似分布とか変分事後分布などと呼びます。

本来表されるべき事後分布 $p(z \mid X)$ と $q(Z)$ がそれほど違ったものにはならないように、下記のKLダイバージェンスを最小化することで変分近似分布を求めます。

$$ {\rm KL} [ q(z) : p(z \mid X) ] = \int _ z q(z) \log \frac {q(z)} {p(z \mid X)}dz $$

ここで、この式は

$$ {\rm KL} [ q(z) : p(z \mid X) ] = - \mathcal L (X) + \log p(X) $$

と評価することができます。 $\mathcal L (X)$ はデータ集合 $X$ の対数周辺尤度 $\log p(X)$ の下界を表す値を返します。仮に $\mathcal L = 0$ のとき ${\rm KL} [ q(z) : p(z \mid X) ] = \log p(X)$ であり、これは変分近似分布が元の事後分布と完全に一致しており、既に手元にあるデータに対する対数周辺尤度だけが値を持ち、完全に最適化が上手く実施された場合を指します。

通常はそうはならないため、KLダイバージェンスを最小化する、あるいは $\mathcal L$ (ELBO)を最大化するというのが基本方針になります。

変分推論のアルゴリズム

対数周辺尤度 $\log p(X)$ は、 $p$ 自体のデータへの当てはまりを評価することに使うことができます。従って、 $\log p(X)$ をモデル $p$ に対して最大化するということも意味のある行為です。 変分推論では $\log p(X)$ の下界を最大化することが1つの目標になるため、下界を最大化、そして $\log p(X)$ 自体を最大化…という手続きを繰り返すアルゴリズムが一般的には利用されます。

仮に、元のモデル $p(x) =\int _z p _ \theta (x, z)p(z) dz$ とモデル化しているのであれば、対数周辺尤度を最大化するために $\theta $ を最適化し、このようにして定められた $\theta$ を持つ事後分布 $p _ \theta (z \mid X)$ に対して、$KL [q _ \phi (z) : p _ \theta(z \mid X)]$ を最小化することで変分事後分布の変分パラメータ $\phi$ を最適化する…というのを繰り返し実施することでアルゴリズムを構成します。

上記は実際にアルゴリズムを実行するという意識での記述になりますが、上記で出てきたELBO $\mathcal L$ に関して下記のように記述できることが知られています。

$$ \mathcal L = \mathbb E _ q [\log p (x, z)] + \mathcal H [q] $$

ここで $\mathcal H[q]$ は確率分布 $q$ のエントロピーと呼ばれる量を返します。エントロピーは確率分布が広がっているほど大きく、ここの項が、変分近似分布が局在しにくくなるような効果をもたらします。 第一項は、モデルの対数同時分布を変分事後分布で期待値演算したものとなっており、観測変数と潜在変数の尤度を確率分布 $q$ で評価することで、元のモデルに対する当てはまりの良さを上げる効果をもたらします。これら2つの最大化が変分推論アルゴリズムの目的となっていました。

Normalizing Flows

確率分布の変換規則

Normalizing Flowsは日本語では正規化流と呼ばれ、確率分布を柔軟に変更することで生成モデルの表現力を高めたりすることに使われます。例えば潜在変数 $z$ が面白くない分布 $q(z)$ により構築された場合に、生成モデル $p(x \mid z)$ というのは、あまり多様性の無い分布になってしまっている可能性があります(なぜなら $z$ の発生がそもそも面白くないから)。

ここで、何らかの方法によって $q (z) \mapsto q' (z')$ と面白い分布に変更ができたら嬉しいとなるわけです。確率分布は、確率変数を $f : z \mapsto z'$ と変更した場合には、

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

と変更されることが知られています。ここで $f$ は全単射である必要があります。全単射であれば $g = f ^ {-1}$ が存在し、この逆写像を通じて $z'$ から $z$ を復元したり、あるいは面白い方の分布から、つまらない方の分布を知ることができるのです。Normalizing Flowでは、 $f$ の性質に全単射かつ微分可能であることを要請し、ニューラルネットワークの誤差逆伝播法をそのまま適用できるように確率分布の変換規則をモデルに組み込むのです。

実際の応用では簡単ないくつかの微分可能な全単射 $f _ k$ を準備し、複雑な全単射 $f = f _ K \circ \cdots \circ f _ 1$ を用意することで、つまらない潜在変数 $z _ 0 \sim q _ 0(z _ 0)$ から

$$ z _ K = f(z _ 0) $$

を作り出します。このときの確率分布の変換規則は単に

$$ q _ K (z _ K) = q _ 0 (z _ 0) \prod _ {k = 1} ^ K \left| {\rm det} \frac{\partial f _ k}{\partial z _ {k - 1}} \right| ^ {-1} $$

と繰り返し変換を適用するだけです。

変分推論への組み込み

学習の中では、この確率の対数が必要な場合が多く、そのような場合は単に $\log$ を取ればよく

$$ \begin{align} \log q _ K (z _ K) &= \log q _ 0 (z _ 0) + \sum _ {k = 1} ^ K \log \left| {\rm det} \frac{\partial f _ k}{\partial z _ {k - 1}} \right| ^ {-1} \\ &= \log q _ 0 (z _ 0) - \sum _ {k = 1} ^ K \log \left| {\rm det} \frac{\partial f _ k}{\partial z _ {k - 1}} \right| \end{align} $$

と表記することが可能であり、複雑な変更を伴うことなくニューラルネットワークと変分推論に組み込むことが可能です。 具体的には、変分推論で最適化すべきELBOが下記のように書きなおされます。

$$ \begin{align} \mathcal L &= \mathbb E _ q [\log p (x, z)] + \mathcal H [q] \\ &= \mathbb E _ q [\log p (x, z)] - E _ q [q (z)] \\ &= \mathbb E _ {q _ 0} [\log p (x, z _ K)] - E _ {q _ 0} [q _ K(z _ K)]\\ &= \mathbb E _ {q _ 0} [\log p (x, z _ K)] - E _ {q _ 0} [q _ 0(z _ 0)] - E _ {q _ 0} \left[ \sum _ {k = 1} ^ K \log \left| {\rm det} \frac{\partial f _ k}{\partial z _ {k - 1}} \right| ^ {-1} \right] \\ &= \mathbb E _ {q _ 0} [\log p (x, z _ K)] - E _ {q _ 0} [q _ 0(z _ 0)] + E _ {q _ 0} \left[ \sum _ {k = 1} ^ K \log \left| {\rm det} \frac{\partial f _ k}{\partial z _ {k - 1}} \right| \right]
\end{align} $$

いずれもつまらない変分近似分布 $q _ 0 (z _ 0)$ からのサンプリングで期待値計算を行えばよく、モンテカルロ計算の実行が容易にできます。また、 $f _ k$ は自分で与えた全単射の関数であるので、これも計算でき、そのヤコビアンの計算も実施できるため(できるような関数を選ぶ)問題ありません。

実際にここで最適化したいのは変分近似分布 $q _ \phi (z)$ における変分パラメータ $\phi$ と元のモデル $p _ \theta (x, z)$ におけるハイパーパラメータ $\theta$ になりますので、ELBO $\mathcal L$ をこれらの関数 $L (\theta, \phi)$ と見て勾配法などを用いて最適化します。また、このとき、Normalizing Flowを構成する全単射 $f$ が変分パラメータ $\phi$ となって現れることとなります。($q _ 0 (z) $ は正規分布など簡単な分布で固定してしまう。Normalizing Flowの方で豊かな表現を獲得できれば良い)

したがって、最適化すべきパラメータを陽に書き下せば

$$ \mathcal L (\theta, \phi) = \mathbb E _ {q _ 0} [\log p _ \theta(x, z _ K)] - E _ {q _ 0} [q _ 0(z _ 0)] + E _ {q _ 0} \left[ \sum _ {k = 1} ^ K \log \left| {\rm det} \frac{\partial f _ {\phi _ k}}{\partial z _ {k - 1}} \right| \right]
$$

となり、Normalizing Flowを用いた変分推論は上記の最適化問題に帰着します。

具体的な変換の例

Normalizing Flowの変換では下記のようなものが使われます。

planner flow

$$ \mathbf{f(z)} = \mathbf{z} + \mathbf{u}\tanh(\mathbf{w}^T\mathbf{z}+b) $$

radius flow

$$ \mathbf{f(z)} = \mathbf{z} + \frac{\beta} {\alpha + |{\bf z - z'}|} {\bf(z - z')} $$

inverse autoregressive flow

$$ \mathbf{f(z)} = \mu _ t + \sigma _ t\odot\mathbf{z} $$

ここで $\mu _ t$ と $\sigma _ t$ は自己回帰ネットワークによって記述されています。また、数値的に安定する下記のバージョンもあります。

$$ \mathbf{f(z)} = \sigma_t\odot\mathbf{z} + (1-\sigma_t)\odot\mu_t $$

flow VAE

VAE 基本

VAEでは生成モデル $p _ \theta (x \mid z)$ に対してデータ集合 $X = \{x _ 1, x _ 2, ..., x _ N \}$ に対する事後分布 $p _ \theta (z \mid X)$ を考える際に、通常は近似分布を

$$ q _ \phi (z) = q _ (z ; \phi) $$

として、 ${\rm KL}(q : p)$ を最小化するところを、下記のように

$$ \begin{align} q _ \phi (z _ i) = q (z _ i ; \phi) = q(z _ i ; \phi _ i) \\ \phi _ i = {\rm NN} _ \psi (x _ i) \end{align} $$

としてしまいます。元々 $q$ の構造を司っていたのは変分分布のパラメータ $\phi$ であったところを、変分分布のパラメータに相当する部分をニューラルネットワークの出力で肩代わりしてしまったのです。そして、ニューラルネットワークのパラメータを新たに変分パラメータ $\psi$ だと考え最適化を実行します。さて、ある $x _ i$ でNNの学習をしたとしましょう。すると他の $x _ j$ を入力した時のNNの出力も変えてしまうことになります。だとすれば $x _ i$ に対応する $z _ i$ の学習を実施したことで、 $x _ j$ に対応する $z _ j$ の学習も実施されたと考えることができます(いや、本当に都合のいい方向に学習が進むかはわからない。が、似たデータを入れているのだとすれば、そうなっているだろうと信じているのである。これは償却推論と呼ばれる)。

あとはこれを変分近似分布だと思って、変分推論をするだけです。VAEの応用上、出来上がった生成モデル $p(x \mid z)$ の方に強い興味があり、 $z$ を色々調整することでバリエーション豊かに $x$ を生成できることを期待します(多くの場合は $x$ は画像である)。

flow VAE

ここで多くのケースでは変分近似分布は

$$ q (z _ i ; \phi) = \mathcal N(z _ i ; \mu _ i, \sigma _ i) $$

と設定され、そして $\mu _ i , \sigma _ i$ をニューラルネットワークに出力させるようにします。ところがそうなると、潜在変数$z _ i$ は正規分布からしか出てこず、言い換えれば、 学習時は$x$ をなんとか正規分布から生起する $z$ で表現してやらねばならないということになってしまいます。そして、そのように学習された生成モデルは、相応の $z$ に対応した $x$ しか作れないわけですから、もしかしたら得られる生成データのバリエーションは乏しくなっているかもしれません。

そこでNormalizing Flowを利用することができます。Normalizing Flowを積み重ねた関数 $f $ によって

$$ z _ {flow} = f(z _ i) $$

と作ってしまえば、その複雑な分布から生成される変数 $z _ {flow}$ を入力として、新しい生成データを得ることができるようになり、生成モデルの表現力の方も上がっていることが期待できるというわけです。