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

はじめに

ここではベイズ推論（データからベイズ予測分布を得る手続き）をある程度知っている前提で、ベイズ推論を具体的に実行する手立てとしての近年の変分推論について適当にまとめておきます。

ベイズ推論では観測データ $\bf X$ から未観測の潜在変数 $\bf Z $ とモデルのパラメータ $\bf w$ の事後分布

$$ p({\bf Z, w}, \mid {\bf X}) $$

を求め、新たなデータ $\bf x$ に関する統計モデル $p({\bf x}\mid {\bf Z, w})$ の期待値（すなわち、事後分布を重みとした統計モデル分布の加算平均分布）を得ます。ここで、統計モデルを作ったからには何かしらの人手で決めなければならないハイパーパラメータが出てきてしまいます。それは階層モデルを考えようが、その階層の最も端っこの方で仮定された事前分布で生じてしまうことに注意してください。

このようなハイパーパラメータをこの記事では $\theta$ と書くことにします。この $\theta$ に関しては決め打ち、あるいは何らかの方法によって点推定するものとして、ベイズの枠組みでは扱いません（当然、 $\theta$ に事前分布を用意することはできますが、そうしたらば、新たにできた事前分布に対して、結局何らかのパラメータが潜んでしまうのです。ここではそのように、考えに考え抜いた先で出てくる機械的なパラメータをハイパーパラメータと呼んでいるということです）。

そうすると、変数の種別としては、観測データ、ベイズ推論の枠組みで事後分布を求める変数、ハイパーパラメータに分けられます。事後分布を求めようという手続きの中では、事後分布を求める変数、すなわち潜在変数 $\bf Z$ と 統計モデルのパラメータ $\bf w$ を明確に区別する意義はあまりないので、ここではまとめて潜在変数 $\bf Z$ としてしまいます。

（もちろん、ベイズ推論のアルゴリズムを回す手続きの外で、実データと向き合ったモデリングの段階では $\bf Z$ と $\bf w$ に信念としての区別はあったことでしょう。$\bf Z$ にはデータの背後にある現象や物理的実態を仮定しており、$\bf w$ はよく知られた確率分布や数理モデルの係数なりパラメータなり、数学的都合を押し込んだものであるケースが多いかもしれません。それはモデルを立てて、分析結果と向き合うときに再び思い起こすことにします。一度、事後分布を求めるという段階に入ると、兎にも角にもデータ $\bf X$ から上手い分布を探ろうという姿勢になり、区別する意義は大きくありません）

変分推論

モデルの書き方

さて、潜在変数 $\bf z$ と観測変数 $\bf x$ に関して、通常下記のような同時分布を構築します。

$$ p({\bf x, z}) = p({\bf x\mid z})p({\bf z}) $$

このとき、同時分布を条件付き分布に分解したときに現れる潜在変数の分布 $p(\bf z)$ を事前分布と呼び、$p({\bf x\mid z})$ を統計モデルと言ったりします。事前分布にはハイパーパラメータ $\theta$ が仮定されていることになります。（また近年、統計モデルの方も非常に複雑なモデルを仮定することが多くなりました。）

構築した同時分布には、事前分布と統計モデル、両者含めて、事後分布を求めることはしないハイパーパラメータが内在しているということです。それら全てをまとめて $\theta$ と記述することにし、仮定した同時分布を改めて

$$ p _ \theta({\bf x, z}) = p _ \theta({\bf x\mid z})p _ \theta({\bf z}) $$

と乱暴に書いてしまうことにします。分布を求めるつもりは無い（確率変数として扱うつもりはない）という意図を込めて下付き文字にしています（いやいや、それも事後分布を求めたい！と思うならば、そういう変数に限って $\bf z$ に含めてしまえばいい。繰り返しになるが、そうしたとて、最終的に $\theta$ に相当するハイパーパラメータがどこかで出るのだから、そういうものをここでまとめて $\theta$ と書いたのである）

立てたモデルに対する変分近似分布

さて、ここで事後分布は

$$ p _ \theta({\bf z \mid x}) = \frac{p _ \theta({\bf x\mid z})p _ \theta({\bf z})}{\int _ {\bf z} p _ \theta({\bf x\mid z})p _ \theta({\bf z})d\bf z} $$

と、中々に途方もない形をしていることを思い出しましょう。今、我々は複雑な潜在変数と観測データを含み、更に $\theta$ でパラメトライズされた解析的に事後分布が求まってなさそうな姿を目の当たりにしています。 モデルを仮定できたということと、それを数式でピシッと表せることには乖離があるのです。

ここで、元の同時分布のモデルから導かれる事後分布をそれほど損なわないような、しかし、元のモデルよりは簡単に形状を求められそうな $q _ \phi (\bf z)$ なる分布を準備します。これはパラメータ $\phi$ によって、ある程度複雑な形状を取れるようにした分布です。下付き文字なのは、こいつも決め打ち、あるいは何らかの方法で点推定を行うだけのパラメータだからです。

この $\phi$ というパラメータをどのようにもたせるかが、元の分布を表現しうるか、あるいは分布を構築することを（少なくとも元のモデルから導かれる事後分布を求めるよりは）簡単にできるかが決まります。近年はニューラルネットワークによって変分近似分布を構成し、ニューラルネットワークのパラメータを $\phi$ として選ぶような方法が取られます。

変分推論

今、下記のもともとのモデル

$$ p _ \theta({\bf x, z}) = p _ \theta({\bf x\mid z})p _ \theta({\bf z}) $$

から導かれる潜在変数の事後分布

$$ p _ \theta({\bf z \mid x}) = \frac{p _ \theta({\bf x\mid z})p _ \theta({\bf z})}{\int _ {\bf z} p _ \theta({\bf x\mid z})p _ \theta({\bf z}) d\bf z} $$

を求めたいのだが、これがどうも扱いづらそうということで下記の変分近似分布

$$ q _ \phi (\bf z) $$

を代わりに使おうということでした。そこで、ある固定された $\theta$ の元で

$$ {\rm argmin _ \phi KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

を最小化するのが変分推論の目的になります。これを最小化したときには、ある固定した $\theta$ で表されるもともとのモデルでの事後分布と、変分近似分布が似ていると言えるようになります。

EMアルゴリズムの発展としての変分推論

ところで、今、固定した $\theta$ での状況を考えていますが、もともとのモデルにおいて、$\theta$ が決め打ちで良かったのかというと疑問が残ります。元のモデルの世界に戻りましょう。今手元にある観測変数 $\bf x$ に関する対数確率は

$$ \log p _ \theta({\bf x}) =\log \int _ {\bf z} p _ \theta ({\bf x, z})d{\bf z} $$

と、潜在変数 $\bf z$ の周辺化によって表現することが可能になります。この対数確率が手元のデータに関して高い値を持っていれば、もともとのモデル自体がそれなりにデータを表現していると言えます。ところで、このモデルは $\theta$ によってその形状を変えるのでした。そこでこの $\theta$ を決め打ちするのではなく、

$$ {\rm argmax} _ \theta \log p _ \theta ({\bf x}) $$

としてしまうことで、対数確率を手元のデータに対して高めるような $\theta$ を選べたことになります。ここで対数確率は、普通は「対数周辺尤度」とか「エビデンス」だとか言われます。このエビデンス最大化に基づくハイパーパラメータ $\theta$ の選定を第二種最尤推定等と言います。

なぜ変分推論の話で、急に元のモデルの話に戻ったのかというと、上記の対数周辺尤度は、変分推論で最小化しようとしていたKLダイバージェンスを使って

$$ \log p _ \theta ({\bf x}) = {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] + {\rm ELBO} $$

という形で表すことができるからです（下記記事も参考にしてください）。この式が意味するところは、なんだかよくわからない $\rm ELBO$ なるものを考えてあげれば、エビデンスの値を評価することと、変分近似分布とモデルの事後分布のKLダイバージェンスを評価することは同じようなことであるという意味になります。

www.hellocybernetics.tech

www.hellocybernetics.tech

さて、ここで出てきた $\rm KL$ は先程変分推論で $\phi$ に関して最小化しようとしたものでした。また突如現れた $\rm ELBO$ は evidence lower bound の略になりますが、これは、エビデンスの下界という意味であり、先程元のモデルを $\theta$ に関して最大化しようとした $\log p _ \theta ({\bf x})$ の下界であるということです。なぜなら、KLダイバージェンスは一般に非負の値を取りますので

$$ {\rm ELBO} = \log p _ \theta ({\bf x}) - {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

という式で見たときには、ELBOはエビデンスよりも必ず小さな値を取るからです。そして変分近似分布とモデルの事後分布が完全に一致するときに限って、ELBO＝エビデンスという状態になります。対数周辺尤度（エビデンス）の最大化は、その数式が複雑なため一般には困難です。

対数周辺尤度を直接評価はできなくても、その下界を評価することが仮に可能なのであれば、そいつをとりあえず最大化するという話で済ましてしまおうと考えることができます。とは言っても、ELBOがあまりにも対数周辺尤度から乖離しているにも関わらず、それを最大化しましたと言われても、対数周辺尤度が最大化されたかは怪しい話です。つまり、ELBOが対数周辺尤度に近いという状態が望ましいのです。

$$ {\rm ELBO} = \log p _ \theta ({\bf x}) - {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

さて、上記の数式を見ると、「元のモデルを$\theta $ を調節していい感じにしたい」ということ、「元のモデルを近似する変分近似分布を $\phi$ を調節して求めよう」という目的がたった一本の式に集約されているように見えます。

具体的には $\phi$ を調整して $\rm KL$ が小さくなれば ELBOとエビデンス（対数周辺尤度）が近づきます。そして次に $\theta$ を調節して、ELBOの最大化による対数周辺尤度の底上げを実施します。するとこのとき、ELBOと対数周辺尤度（$\theta$ に依存している）が剥離がしてしまうかもしれません。そこでもう一度 $\phi$ を調整して $\rm KL$ を最小化し直すことでELBOと対数周辺尤度の剥離を縮めます。そうした後に、更に $\theta$ の調整によって底上げを実施します…。

これを繰り返すというのはすなわち、 $\theta$ の調整による元のモデルの改善と、その元のモデルを近似する変分近似分布の形状を支配する $\phi$ の調整を順次行っていくということです。

結局のところ、我々は

$$ {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

を $\theta $ と $\phi$ の両方に関して最小化していくという作業を行うのです（それはエビデンスを最大化したいという思いそのものである）。これはちょうどEMアルゴリズムと呼ばれる潜在変数モデルの最尤推定に近い形式を取っています。

変分推論の違った見方

近似したいという動機

私は変分推論とは対数周辺尤度を最大化するという話から出発して、Jensenの不等式によってその下界、ELBOを最大化するという話が「自然な流れ」だと思っていました。EMアルゴリズムの拡張的な立場で、「学びやすい」と思いますし、そしてELBOの最大化問題は既に述べたとおり

$$ \mathcal L (\phi, \theta) = {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

を最小化するという問題に落ち着き、扱いやすい気がするからです。しかし思えば、ベイズ推論では、我々は自分で立てた同時分布により導かれる、事後分布を推論したかったのです。そうすれば、その事後分布を用いてベイズ予測分布を構成できるという流れです。そして事後分布を求めるのが難しそうなので他の求めるのが楽そうな分布で、事後分布を代替したかったというのが変分推論のモチベーションです。

そうであれば、なぜに対数周辺尤度を最大化するという出発点が必要なのでしょうか（それは、モデルのハイパーパラメータを推定する際には妥当な話なのかもしれないが…、ハイパーパラメータはハイパーパラメータらしく、一旦定数で決め打ちだと思って話を進めたほうが良い気がする）。事後分布を近似したいというモチベーションの中では、何か小難しい話を間に挟まずに、いきなり

$$ \mathcal L (\phi) = {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] $$

を最小化したいと思っても良いような気がします（すなわち、自分で立てた複雑なモデルから得られる事後分布を、計算しやすい分布 $q$ で近似したいというだけの話）。ここで、$\mathcal L(\phi)$ という表記は意図的です。ここで調整しようと思える必然性があるのは、$\phi$の方だけだからです。なぜなら、既に立て終えたが数式として取り扱いづらい $p _ \theta({\bf z \mid x})$ を $q _ \phi ({\bf z}) $ のパラメータ $\phi$ を調整して近似したいという気持ちだからです。

ひとまずこれを出発点にしてみます。

近似のための損失関数の分解

さて、ここでKLダイバージェンスを書き下して式変形してみます。

$$ \begin{align} \mathcal L (\phi) &= {\rm KL} \left[q _ \phi ({\bf z}) || p _ \theta({\bf z \mid x}) \right] \\ & = \int q _ \phi ({\bf z}) \log \frac{q _ \phi ({\bf z})}{p _ \theta({\bf z \mid x})} {\rm d}{\bf z} \\ & = \int q _ \phi ({\bf z}) \log \frac{q _ \phi ({\bf z})}{\frac{p _ \theta({\bf x \mid z})p _\theta ({\bf z})}{p _ \theta({\bf x})} } {\rm d}{\bf z} \\ & = \int q _ \phi ({\bf z}) \log q _ \phi ({\bf z}) {\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x \mid z}){\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf z}){\rm d}{\bf z} + \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x}){\rm d}{\bf z} \\ & = \int q _ \phi ({\bf z}) \log q _ \phi ({\bf z}) {\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x \mid z}){\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf z}){\rm d}{\bf z} + \log p _ \theta ({\bf x}) \end{align} $$

ただただひたすらにKLダイバージェンスを１つ１つの項に分解していきました。いわずもがな、最後の変形で出てきた積分のついていない項が「エビデンス」です。$\bf z$ による積分を「 $\log _ \theta p({\bf x})$ 」を外に出してから実行しただけの話です（積分は $1$ になる。なぜなら$q _ \phi({\bf z})$ は確率分布だから！）。

ただ、エビデンスのことは一旦忘れてしまって、とにかくこのKLダイバージェンスを最小化したいという話に集中しましょう。まずここではエビデンスは既に得られているデータに基づく定数だと思って無視します。

期待値のサンプリングによる近似

さて、分解されたKLダイバージェンスの式でまず最初に気づくのは、こんな積分計算できるのだろうか？ということです。しかし、上記の積分は気づけば何らかの関数を $q _ \phi ({\bf z})$ という確率分布を使って期待値計算しているにすぎません。例えば

$$ - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x \mid z}){\rm d}{\bf z} = - \mathbb E _ {q _ \phi ({\bf z})} \left[\log p _ \theta ({\bf x \mid z}) \right] $$

という項を見てみましょう。さて期待値（＝積分計算）はできませんが、一度これは期待値なのだということを認識すれば、適当な数のサンプリングを得てしまえば、近似できそうだと思えます。要するに

$$ \begin {align} {\bf z}& _ i \sim q _ \phi ({\bf z} _ i ) \\ - \mathbb E _ {q _ \phi ({\bf z})} \left[\log p _ \theta ({\bf x \mid z}) \right] &\simeq- \frac{1}{N} \sum _ {i=1} ^ N \log p _ \theta ({\bf x} _ i \mid {\bf z} _ i) \end{align} $$

という値で期待値を代替してしまうのです。このときまんべんなく色々な $\bf z$ をサンプリングできれば、しっかりと期待値を表現できそうです。そしてよくよく見てみると、この式は「対数尤度関数」の符号を変えたものに見えます。勝手に $\bf z$ が変分近似分布 $q$ から生成されると思ったときに、モデルの尤度関数がどうなっているのかを評価しているとみなせるわけです。

他の項はどうでしょうか、例えば下記

$$ \int q _ \phi ({\bf z}) \log q _ \phi ({\bf z}) {\rm d}{\bf z} = \sum _ {i=1} ^ N \log q _ \phi ({\bf z}) $$

こいつは、${\bf z} _ i \sim q _ \phi$ で $\log q _ phi$ の期待値を近似しています。このような具合いに $q _ \phi$ は今、適当に $\phi$ でパラメトライズした分布であるので、他の項に関してもどうにかこうにか、サンプリングで期待値を見積もることができます。

今、期待値をサンプリングで見積もるという前提で式を見直してみましょう。最大化したい目的関数 $\mathcal L(\phi)$ のサンプリングによる近似式は

$$ \mathcal L (\phi) \simeq \sum _ {i=1} ^ N \log q _ \phi ({\bf z} _ i) - \sum _ {i=1} ^ N \log p _ \theta ({\bf x} _ i \mid {\bf z} _ i) - \sum _ {i=1} ^ N \log p _ \theta ({\bf z} _ i) + \rm Const $$

と表せるようになります。ここで ${\bf z} _ i \sim q _ \phi({\bf z})$ というふうにサンプリングをしています。 これでひとまず、この近似値をサンプリングしながら目的関数関数を最大化することで変分推論が実行できるというわけです。

事前分布と変分近似分布のKLダイバージェンス

さて、具体的な計算の実行方法を知ったところで、せっかく丁寧に分解した事後分布と変分近似分布とのKLダイバージェンスを少し弄くってしまいます。サンプリングによる近似前の式に戻りましょう。

$$ \mathcal L (\phi) = \int q _ \phi ({\bf z}) \log q _ \phi ({\bf z}) {\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x \mid z}){\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf z}){\rm d}{\bf z} + \log p _ \theta ({\bf x}) $$

の第１項と第３項に注目してください。

$$ \int q _ \phi ({\bf z}) \log q _ \phi ({\bf z}) {\rm d}{\bf z} - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf z}){\rm d}{\bf z} $$

という部分です。こいつらを１つの項にまとめてしまうことにすれば

$$ \int q _ \phi ({\bf z}) \log \frac{q _ \phi ({\bf z})}{p _ \theta ({\bf z})} {\rm d}{\bf z} = {\rm KL}\left[q _ \phi ({\bf z}) : p _ \theta ({\bf z}) \right] $$

となっているのです。すると

$$ \mathcal L (\phi) = {\rm KL}\left[q _ \phi ({\bf z}) : p _ \theta ({\bf z}) \right] - \int q _ \phi ({\bf z}) \log p _ \theta ({\bf x \mid z}){\rm d}{\bf z} + \rm Const $$

という形式を取っているということです。

今 $\mathcal L$ は最大化したい目的関数でした。一方でKLダイバージェンスは非負の値を取ります。したがって、${\rm KL}\left[q _ \phi ({\bf z}) : p _ \theta ({\bf z}) \right]$ は、なるべく小さな値を取っておいてほしいという願いがあるのです。これはすなわち、勝手に作った変分近似分布 $q _ \phi$ が、もともと立てたモデルの事前分布 $p _ \theta$ に近いものであるべきという正則化の効力を持ちます。

意味としてこのように理解しようがしまいが、サンプリングによって値を評価し、勾配法などで推論を実行することには変わりません。しかし、この形式を一度見てみると次のような発想が浮かびます。

第１項は正則化項、第２項は尤度関数みたいなもの…、正則化項は他のものを考えても良いのではないか？という発想です。

最後に

長くなったので一旦ここで切ります。

お茶を濁していましたが、ニューラルネットワークを使った変分推論では $\phi$ と $\theta$ がそれぞれニューラルネットワークのパラメータです。あたかもこいつらが分布の形を決定するハイパーパラメータの如く振る舞います。そして、固定された $\theta$ に対して $\mathcal L (\phi)$ を最大化することで、元のモデルの事後分布 $p _ \theta ({\bf z \ mid x})$を近似する変分近似分布 $q _ \phi ({\bf z})$ を求めよう。という流れに対して、"交互"に $\phi$ と $\theta$ を最適化するEMアルゴリズムに類似した方法がありました。これに対してニューラルネットワークを使った推論では、"同時"に $\theta$ の方も最適化してしまうことになります。

既に想像がついている通り、この方法を展開していくと変分自己符号化器の話に繋がります。しかし、その話をするまでにはローカルな変分パラメータ、グローバルな変分パラメータ、そして償却推論という考え方を知る必要があります（知らなくても変分自己符号化器は使えるが…）。

参考図書は下記です。

ベイズ深層学習 (機械学習プロフェッショナルシリーズ)

機械学習スタートアップシリーズ ベイズ推論による機械学習入門 (KS情報科学専門書)

実はPyroのexampleを理解できれば、記事を読む必要はないかもしれない。

pyro.ai