Normalizing Flow入門第1回変分推論

Tatsuya Yatagawa

2021-01-01

こんにちはtatsyです。今回は最近興味をもっているNormalizing Flowの技術について紹介したいと思います。

Normalizing Flowという言葉はRezendeらによって2015年に発表された以下の論文で一般的に使われるようになった言葉です。

Rezende et al. 2015, “Variational Inference with Normalizing Flows”, ICLR 2015. [arXiv]

この論文のタイトルにある通り、Normalizing Flowという技術はVariational Inference、すなわち変分推論の技術の一つです。そこで、第1回では、変分法と変分推論について見ていくこととします。

変分法

変分法というと、大学の教育課程では、おそらく解析力学のような分野で最初に登場する言葉かと思います。例えば、坂道を転がる球が最短の時間で転がり落ちる坂道の曲がり具合はどのようなものかを求める最速降下曲線の問題を解く際には、この変分法の考え方が使われていたのでした (物理かぎのしっぽ - 変分法)。

変分、というのは、関数を変数とする関数、つまり関数の関数を関数によって微分したもの、というふうに例えられることが多いです。有名な「パターン認識と機械学習」によると、関数の微分と変分の違いは次のように説明されています。

微分との対比による変分の説明

関数$f(x)$の変数$x$を微少量$\epsilon$だけ変化させた場合、その時の関数の値は、テイラー展開を用いて次のようにかけます。

$$ f(x + \epsilon) = f(x) + \epsilon \frac{df}{dx}(x) + O(\epsilon^2) $$

さらに、関数$f$が多変数関数$f(x_1, x_2, \ldots, x_D)$であるとし、各々の$x_d$の値が微少量$\epsilon_d (d=1, 2, \ldots, D)$だけ変化したとすると、その時の関数の値は次のようにかけます。

$$ f(x_1 + \epsilon_1, \ldots, x_D + \epsilon_D) = f(x_1, \ldots, x_D) + \sum_{d=1}^D \epsilon_d \frac{df}{dx_d}(x_1, \ldots, x_D) + O(\epsilon^2) $$

多変数関数というのは見方を変えるとベクトルを変数とする関数と考えることができる、というのがポイントです。

ところで、関数、というのは、とある入力に対して単一の値を返すものなので、別の見方をすると無限次元のベクトルであると考えることができます。そこで、$y$を関数として関数の関数$F(y)$を考え、関数$y$の任意の関数$\eta(x)$と微量変化量$\epsilon$により微小変化させた場合を考えてみます。

すると、上記のベクトル値関数の場合と同様に次のような式が得られると考えられます。

$$ F(y(x) + \epsilon \eta(x)) = F(y) + \epsilon \int \frac{\delta F}{\delta y(x)} \eta(x) dx + O(\epsilon) $$

上記の式は、ベクトル関数の時に$x_d$であったものが$y(x)$に、$\epsilon_d$であったものが$\epsilon \eta(x)$に置き換わったものと考えられます。この式に現れる$\frac{\delta F}{\delta y}$が汎関数$F$の関数$y$による変分になります。

変分を用いた最適化

通常の関数の最適化では、変数$x$を動かした時に$f(x)$が極大・極小の値をとる場所を探すわけですが、変分を用いた汎関数の最適化では、関数$y$を変化させて、汎関数$F(y)$が極大・極小となる場所を探します。

一般的な変分の解説では、汎関数$F$をラグランジアンとして$L(x, y, y’)$を用いて、

$$ F(y) = \int_{x_1}^{x_2} L(x, y, y’) dx $$

のように定義して、オイラー・ラグランジュの式が得られることを利用します。

以下で説明する、機械学習での「変分推論」での変分とは、例えば確率密度「関数」$p(x)$の関数として定義されるエントロピー

$$ H(p(x)) = -p(x) \log p(x) $$

を最大化する密度関数を最適化により求める、という考えに従っています。

変分推論

変分推論は、ベイズ推論に基づいた、確率モデルの一種です。ベイズ推論においては、観測できる値$x$を用いて、未知の値$y$を確率密度関数として表現します。この時、$x$を観測変数、$y$を潜在変数と呼びます。

ベイズ推論による確率モデル化では、確率変数$x$, $y$の同時密度関数$p(x, y)$を何らかの方法で決めてあげて、それを用いて$x$の確率密度分布を定義します。

既知の$p(x, y)$を用いると、$x$に対する$y$の事後分布$p(y|x)$は、その定義から次のように書けます。

$$ \begin{equation} p(y|x) = \frac{p(x, y)}{\int p(x, y) dy} \label{eq:cond-prob} \end{equation} $$

なお、分子に確率密度に対する連鎖率$p(x,y) = p(x|y) p(y)$を用いると、よく見るベイズの定理の式

$$ p(y|x) = \frac{p(x|y) p(y)}{p(x)} $$

が得られます。いずれにしても、$p(x,y)$に何らかのモデルを適用することで、上記の式から$y$の事後分布をモデル化するのがベイズ推論の目的です。この辺りの説明は「ベイズ推論による機械学習入門」が参考になると思いますので、より詳しく知りたい方はこちらの書籍をご参照ください。

変分推論による確率モデル化

上記の通り、ベイズ推論においては、同時分布$p(x, y)$を何らかの方法で決める必要があるのですが、多くの場合は、ガウス分布のような単純な形ではなく、もっと複雑な形でモデル化を行う必要があります。

そのため、事後分布の式\eqref{eq:cond-prob}の分母に現れる積分を簡単な形式で書き表すことができず、故に事後分布$p(y|x)$がうまく計算できない可能性があります。

ここでいう「簡単に計算できる」というのは、「解析的」とか「閉形式 (closed form)」などの言葉で表されるもので、すごく平たくいうと、微分や積分を使わずに、四則演算や指数、対数などの初等関数を用いて書き表せることを言います。

計算機科学的には、この閉形式で関数が書き表せれば、計算精度等に依存せず、常に一定の計算時間でその値を評価できる、という利点がありますが、これが積分計算をしなければならないとなると、一般に計算精度を上げるためには、より多くの計算量が必要となり、厳密なモデル化が難しくなります。

そこで、変分推論においては、目的となる事後分布$p(y|x)$に対して、それを近似する別の分布$q(y)$との差を最小化する問題を考えます。ここで2つの分布の間の距離はKullback-Leibler距離$D_{KL}(q||p)$を用いて以下のように書きます。

$$ D_{KL}(q(y)||p(y|x)) = - \int q(y) \log \frac{p(y|x)}{q(y)} dy $$

この右辺は分布$q(y)$の定義域全体での積分が1になることを利用すると、次のように式変形できます。

$$ \begin{aligned} D_{KL}(q(y))||p(y|x)) &= -\int q(y) \log \frac{p(x, y)}{q(y) p(x)} dy \\ &= -\int q(y) \log \frac{p(x, y)}{q(y)} dy + p(x) \int q(y) dy \\ &= -\int q(y) \log \frac{p(x, y)}{q(x)} dy + \log p(x) \\ &= -\mathcal{L}(q(y)) + \log p(x) \end{aligned} $$

今、$\log p(x)$は$q(y)$とは無関係なので、Kullback-Leibler距離を最小化する問題は $\mathcal{L}(q(y))$を$q(y)$で最大化する問題と同値であることが分かります。機械学習の教科書などでは、この$q(y)$をベクトル$y$の次元ごとに$q(y_i)$と分解して、平均場近似する方向で話が進みますが、ここではその話は割愛します。

変分推論の課題

Normalizing Flowで考慮する変分推論においては、上記の$\mathcal{L}(q(y))$を数値的に最大化することを考えます。ここで$q(y)$はいくつかのパラメータ$\theta$によってパラメータ化されているとします。特にニューラルネットを用いる場合には$\theta$はネットワークの重みパラメータになります。

すると解くべき問題は、次のように書けます。

$$ \max_{\theta} \int q(y; \theta) \left\{ \log p(x, y) - \log q(y; \theta) \right\} dy $$

この式は、確率密度$q(y; \theta)$から得たサンプル$y$を用いて$\log p(x, y) - \log q(y)$の平均を取ったものとみなせます。したがって、上記の最小化問題は、次のように書き直せます。

$$ \begin{equation} \max_{\theta} \mathbb{E}_{y \sim q(y; \theta)} [\log p(x, y) - \log q(y)] \label{eq:lower-bound} \end{equation} $$

これは、$\theta$に対する最小化問題なので、離散的には$\theta$での微分を考えて、最急降下法などを適用すれば良いのですが、そもそも、ここまでの仮定では$q(y; \theta)$という確率密度の値を求めることはできても、その確率に従うサンプル$y$を直接得ることはできないので、微分計算には工夫が必要になります。

例えば、$q(y; \theta)$が直接サンプルできるような形でなかったとしても、マルコフ連鎖モンテカルロ法などを用いて、サンプルを得ることができるかもしれません。そうすれば、近似的に被最適化関数を評価できるので、これに対して$\theta$の微分を求めることで、最小化を行うことができます。

Normalizing Flowでは$q(y; \theta)$が具体的な数式の形では与えられないような場合を考えていて、これをニューラルネットを用いて表すことを試みます。

Normalizing Flow

Normalizing Flowは、解析的に書くことが難しい確率密度分布$p_Y(y)$を、解析的に書ける確率密度$p_Z(z)$から変数変換により定義することを目指します。このように複雑な確率密度を簡単な確率密度に書き直すことを"normalize"と言っています。

言い換えると、Normalizing Flowは確率変数を$y$から$z$へと変換する写像$f$で、$z = f(y)$と変数変換を行います。

この$f$を用いると、確率密度$p_Y$の累積密度分布$P_Y(y)$は次のように書けます。

$$ P_Y(y) = \int_{-\infty}^y p_Y(y’) dy’ = \int_{-\infty}^{f(y)} p_Z(z) dz $$

この両辺を$y$で微分すると、以下の等式が得られます。

$$ p_Y(y) = p_Z(f(y)) \left| \frac{\partial f}{\partial y} \right| = p_Z(z) \left| \frac{\partial f^{-1}}{\partial z} \right|^{-1} $$

この式に従えば、変換関数$f$を、その微分ないしヤコビアンがうまく決まるように定めることで、評価が容易な確率密度$p_Z(z)$から得たサンプルの変換により、変換後のサンプル$y = f^{-1}(z)$とその$y$の確率密度$p_Y(y)$の両方を計算することができます。

そこで、変分推論の式\eqref{eq:lower-bound}の$q(y; \theta)$を$p_Y(y)$と書き直して、上記の変数変換が$\theta$でパラメータ化された$y = g(z; \theta) = f^{-1}(z; \theta)$であるとすると、以下の式が得られます。

$$ \begin{aligned} & \mathbb{E}_{y \sim p_Y(y)} \left[ \log p(x, y) - \log p_Y(y) \right] \\ & = \int p_Y(y) \left\{ \log p(x, y) - \log p_Y(y; \theta) \right\} dy \\ & = \int p_Z(z) \left| \frac{\partial f}{\partial y} \right| \left\{ \log p(x, g(z)) - \log p_Z(z) \right\} \left| \frac{\partial y}{\partial z} \right| dz \\ & = \int p_Z(z) \left\{ \log p(x, g(z; \theta)) - \log p_Z(z) \right\} dz \\ & = \mathbb{E}_{z \sim p_Z(z)} [ \log p(x, g(z; \theta)) - \log p_Z(z) ] \end{aligned} $$

この式を用いると、$z$は例えばガウス分布のようなサンプルが容易な確率密度分布から得られ、さらに、平均の値も$g(z)$が容易に評価できる限りにおいては、容易に計算可能です。

この平均を$\theta$で微分する場合には、$f(z; \theta)$の$\theta$に関する微分が出てきますが、$g(z; \theta)$がニューラルネットで表されていれば、このような勾配を自動微分で評価できるので、最急降下法などに用いる勾配の値も同様に計算可能です。

まとめ

以上によりNormalizing Flowが「理論の上では」変分推論で有効に働くことがわかりました。

ですが、ここで一つ課題が残ります。上記の議論では変数変換を行う関数$g(z; \theta)$がニューラルネットを用いて表されていて、そのパラメータ$\theta$における微分を計算する必要がありました。

パラメータによるネットワークの勾配自体は、前述の通り自動微分により計算可能なのですが、最急降下法をはじめとする微分量を用いる最適化計算においては、この勾配の計算を何度も繰り返す必要があり、ネットワークが複雑になると、非常に多くの計算量が必要になってしまいます。

しかも、ネットワークが単純になると、変数変換自体も単純になるため、$p_Z(z)$として選んだ単純な分布 (例えばガウス分布など)から大きく離れた分布を表現することが難しくなります。

これらの問題を解決するのがNormalizing Flowの中心的な課題になっていて、第2回以降では、この部分について詳しく見ていこうと思います。

参考文献

[1] Redende et al. 2015, “Variational Inference with Normalizing Flows”, ICLR 2015. [arXiv]
[2] Kobyzev et al. 2020, “Normalizing Flows: An Introduction and Review of Current Methods”, PAMI. [arXiv]
[3] ビショップ著, 『パターン認識と機械学習』, 丸善 [link]
[4] 須山著, 『ベイズ推論による機械学習入門』, 講談社 [link]