HELLO CYBERNETICS

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

ベイズニューラルネットワーク基本

 

 

follow us in feedly

はじめに

まず下記の記事について、通常のニューラルネットワーク、正則化を用いたニューラルネットワーク、また、ニューラルネットワークのパラメータに事前分布を構えた点推定を行うことで正則化と同等の働きをすることを理解した上で解説をします。

www.hellocybernetics.tech

今回からは、ニューラルネットワークの具体的な構造は書かず、確率モデルとして記述していくことになるため、別途

www.hellocybernetics.tech

www.hellocybernetics.tech

www.hellocybernetics.tech

なども理解していると望ましいです。

ベイズニューラルネットワーク

確率モデルとしてのベイズニューラルネットワーク

まず通常のニューラルネットワークのパラメータを $\theta$ として、入力を $x$ 、出力を $y$ とします。そのまま $y = f(x; \theta)$ という関数を考え、$f$ をニューラルネットワークとしておいて、最小二乗法等によって $\theta$ を求めるのが学習になります。

一方で確率モデルで、$x$ を入力したら $y$ が何らかの確率法則に従って出力されるような場合には、その確率分布のパラメータを $\theta$ として $x$ に条件付けられた $y$ のパラメトリックな分布

$$ p(y \mid x , \theta) $$

という条件付き分布を考えることができます。この条件付き分布とニューラルネットワークを上手に融合させることを考えると、例えば、正規分布の平均が $f(x; \theta) $ で分散が適当な $\sigma$ で表される正規分布であるとし、

$$ p(y \mid x, \theta) = {\rm Normal} (y \mid f(x ; \theta), \sigma) $$

というふうなモデルを考えることができます。ここで $f(\cdot ; \theta)$ は何らかのニューラルネットワークということになります。このモデルに対して、最尤推定を実施すると、最小二乗法と同じ結果が得られるのでした(正確には同じ最適化問題になる)。ベイズニューラルネットを考えていく場合には、常に後者のような確率モデルによるモデリングを行っていくことに注意してください。

例えば、条件付き分布 $p(y \mid x , \theta) $ として、ベルヌーイ分布を選び、そのパラメータを $f(x; \theta)$ で表現することにすれば

$$ p(y \mid x , \theta) = {\rm Bern}(y\mid f(x; \theta)) $$

と言った具合で書くことができ、最尤推定の手続きを踏んでいくとバイナリクロスエントロピーが出現します。以降、特別に強調したい場合を除き、ニューラルネットワーク $f (\cdot ; \theta)$ を一々書くことをやめ、出力 $y$ が入力 $x$ とモデルパラメータ $\theta$ に依存しているという関係性だけを表す $p(y \mid x, \theta)$ という形式で(ニューラルネットワークを含めた)モデルを書いていくことにします。

ベイズニューラルネット

ベイズニューラルネットは確率モデル $p(y \mid x, \theta)$ に対してパラメータの事前分布 $p(\theta)$ を考慮したニューラルネットワークになります。確率モデルと事前分布のペアによってベイズニューラルネットのモデリングが実施されるということです。このとき、$\theta$ と入力 $x$ は独立で、$y$ が $\theta$ と $x$ に依存するというモデルにおいて

$$ p(y\mid x, \theta)p(\theta) = p(y, \theta \mid x) = p(\theta \mid x, y)p(y \mid x) $$

なる式変形が実施できます。この左辺と右辺を見比べて、

$$ p(\theta \mid x, y) = \frac{p(y \mid x, \theta)p(\theta)}{p(y \mid x)} $$

という式を書くことができます。今、教師あり学習を想定すれば $x, y$ のデータセット $D _ x = \{x _ 1, \cdots, x _ N\}, D _ y = \{y _ 1, \cdots, y _ N \} $ は所与ですから、

$$ p(\theta \mid D _ x, D _ y) = \frac{p(D _ y \mid D _ x, \theta)p(\theta)}{p(D _ y \mid D _ x)} = \frac{p(D _ y \mid D _ x, \theta)p(\theta)}{Z _ N} $$

という式を作ることができます。ここで $Z _ N$ は $\theta$ を含まないので1つの文字でまとめてしまいました。これはデータに条件付けられたモデルのパラメータ $\theta$ の分布を記述していることになります。最尤推定では上記右辺の分子の第一因子 $p(D _ y \mid D _ x, \theta)$ を最大化(あるいは負の対数尤度を最小化)することによって点推定値 $\hat \theta$ を求めていたのでした。このことを上記の条件付き分布全体を考慮して表現するならば、$p(\theta) = 1$ という一様分布を仮定した条件付き分布のモードを推定値に選んでいることになります。

ベイズニューラルネットでは $\theta$ の点推定ではなく分布による推論を実施していることが大きな違いです。また、通常のニューラルネットワークでは点推定値 $\hat \theta = \arg\max _ \theta p(D _ y \mid D _ x, \theta)$ を用いて

$$ \hat y = f(\hat x ; \hat \theta) $$

と新しい入力 $\hat x$ に対する予測を行うのでした。一方でベイズニューラルネットの場合は新しいデータに対する予測 $\hat y$ に関しても分布で表現することにします。

$$ p(\hat y \mid \hat x) = \int _ \Theta p(y \mid x, \theta) p(\theta \mid D _ x, D _ y) d\theta $$

と表され、ベイズ予測分布と呼ばれます。これはパラメータ $\theta$ を持つ確率モデル $p(y\mid x, \theta)$ の分布 $p(\theta \mid D _ x, D _ y) $ による期待値という形式になっています。ベイズ予測分布は所与となっているデータ $D _ y, D _ x$ に依存するのですから、そのことを明記して、ベイズ予測分布を

$$ p(\hat y \mid \hat x, D _ x, D _ y) = \int _ \Theta p(y \mid x, \theta) p(\theta \mid D _ x, D _ y) d\theta $$

と表記することもあります。

今、ベイズニューラルネットではパラメータの点推定を実施する代わりに条件付き分布によってパラメータの情報を表現し、その分布による確率モデルの期待値で予測分布を構成するという話をしました。これはむしろ話を逆で、新しいデータを予測しようという場合に、得られる確率モデル $p(y \mid x, \theta)$ 自体が所与のデータ依存しているのだから、それは一種の確率変数であるはずで、では確率モデルという確率変数の一実現値を決め打ちで使うより、期待値により得られる分布を新しいデータの予測分布として使おう決心したという手順で話をしてみましょう。すると、ではどんな分布で確率モデルの期待値を取ろうか?ということが問題になります。

ここで $\theta$ の真の分布(ないし真の値に相当するデルタ分布)を $\psi(\theta)$ として、それを知っているのであれば、無論

$$ p(\hat y \mid \hat x) = \int _ \Theta p(y \mid x, \theta) \psi(\theta) d\theta $$

と構成するのが良さそうです。しかし、当然そんな真の $\psi(\theta)$ 等知らないのですから、期待値を取るための分布をデータに基づいて決定する必要があります。それはデータに条件付けられた $\theta$ の分布を考えればよく、条件付き分布は常に同時分布を介して

$$ p(\theta \mid D _ x, D _ y) = \frac{p(D _ y \mid D _ x, \theta)p(\theta)}{p(D _ y \mid D _ x)} $$

と構成できるのでした。すると通常で言うところの尤度関数 $p(D _ y \mid D _ x, \theta)$ と事前分布 $p(\theta)$ を構える必要があり、これらのペアをベイズモデルと呼ぶことにするのです。そして、ベイズモデルに基づいて、$\theta$ の条件付き分布を求めることが、ベイズニューラルネットの学習であり、この条件付き分布のことを事後分布と呼びます。

事後分布の推論

さて、実際にニューラルネットワークの学習をすべく、

$$ \begin {align} p(\theta \mid D _ x, D _ y) &= \frac{p(D _ y \mid D _ x, \theta)p(\theta)}{Z _ N} \\ Z _ N &= p(D _ y \mid D _ x) \end{align} $$

を求めるのは容易ではありません。なぜなら、$Z _ N = p (D _ y \mid D _ x)$ は確率分布の周辺化公式を用いて

$$ Z _ N = \int _ \Theta p(D _ y \mid D _ x, \theta)p(\theta) d \theta $$

とできますが、これは、モデル $p(D _ y \mid D _ x, \theta)p(\theta)$ を $\theta$ の取りうる値全域で積分した周辺分布になってなければならず、到底計算できる代物ではないからです。実応用では、変分ベイズ推論あるいはMCMCが選択肢として上がりますが、ニューラルネットワークのようにパラメータ $\theta$ の次元が膨大である場合には、現状、変分ベイズ推論による方法が現実的です。

変分ベイズ推論では代理事後分布 $q(\theta ; \lambda)$ を仮定し、ニューラルネットワークの"パラメータ $\theta$ の事後分布" をパラメトリックな分布関数 $q(\theta; \lambda)$で近似することを目指します。このとき $\lambda $ を変分パラメータと呼びます。事後分布 $p(\theta \mid D _ y, D _ x)$ に似た代理事後分布 $q(\theta ;\hat \lambda)$ が手に入れば、この分布を用いて確率モデルの期待値を計算することに使えます。すなわち予測分布を

$$ p(y \mid x, D _ y, D _ x) = \int _ \Theta p(y \mid x, \theta) q(\theta ; \hat \lambda) d \theta $$

と構成することになります。分布の似てなさ具合は

$$ {\rm KL} (q(\theta ; \lambda) || p(\theta \mid D _ y, D _ x)) = \int q(\theta ; \lambda) \log \frac { q(\theta; \lambda) } {p(\theta \mid D _ y, D _ x)} d\theta $$

で測ることができます。これは確率密度比の対数 $\log (q(\theta; \lambda)/p(\theta\mid D _ y, D _ x))$ に関して $q$ で期待値を取っている形になっています。従って、例えば、現在の $\lambda$ を用いて $\theta \sim q (\theta ; \lambda)$ をいくつかサンプリングしてみて、その値で上記の${\rm KL}(q || p)$ を評価し、値を勾配法により減少させていく等の方法が考えられます。

概ね、ベイズニューラルネットに対する変分推論法は、ニューラルネットワークのパラメータ $\theta$ の事後分布を近似する代理事後分布 $q(\theta; \lambda)$ を構成し $\lambda$ に関する最適化問題を解く という形式に落ち着きます。

このとき、$\lambda$ が求まることによって $\theta$ の"分布"が得られたことになり、元々、パラメータを分布で表現するという話や、予測分布を構成したいという話は(事後分布の近似により情報が落ちている可能性はあれど)ちゃんと実施できていることに注意しましょう(ときおり、分布の推論が"最適化"問題になってしまったことで、あれ、もとのニューラルネットワークと何が違うんだっけ?となる人がいる)。

応用上の参考

ベイズニューラルネットを簡単に取り扱えるフレームワークとしてはPyro on PyTorchとTensorFlow Probability on TFがあります。 それぞれライブラリの構成が大きく異なっており、Pyroは確率モデルを構成する上で便利なモジュール群が整備されており、確率モデルとしてPyTorchの nn.Module を取り扱えるMixInであるPyroModuleクラスが準備されています。PyroModuleクラスではパラメータを確率変数として扱い事前分布を設定するAPI等が提供されており、また、パラメータは確率分布からのサンプリングにより毎回異なる値を持つようにできます。

docs.pyro.ai

一方でTensorFlow Probabilityはtf.keras.Layersと同等の構成を持ち、ユーザー側はあたかも普通にkerasのLayerを利用しているかのように自然に確率的な層を導入できます。TFPのlayerの内部には代理事後分布と変分パラメータがあり、普通にNNの学習を行うコードを実行すると変分パラメータが最適化されます。これは、変分ベイズ推論においては変分パラメータの最適化における目的関数が、代理分布からのサンプリングで得られるパラメータを利用したフィードフォワードによる損失の評価項と、サンプリングされたパラメータが事前分布から見てどの程度の確率を持ちうるかの項に分解されるため、あたかも前者を普通のNNの損失関数、後者を正則化項の如く扱えば、従来のAPIにそのまま合致させられるからです。

従ってTFPのlayers APIでは代理分布と事前分布、そして事前分布との隔たりを測るダイバージェンスの設定が可能となっています。

www.tensorflow.org