「確率分布をファーストクラスオブジェクトとして扱う」という観点で Tensorflow Probability を理解する

何の話かと言うと

Tensorflow Probability の公式ページを見ると、

「TensorFlow Probability は確率的推論と統計的分析のためのライブラリです。」

という言葉が目に飛び込んできますが、機械学習モデルを扱うライブラリーとしての Tensorflow とどういう関係にあるのかがよくわかりません。

ここでは、「確率分布をファーストクラスオブジェクトとして扱う」というプログラミング言語的な観点から、Tensorflow Probability を説明してみます。

「確率分布」というオブジェクト

数学で言うところの「変数 x」には、通常、実数や複素数などのスカラー値が入ります。一方、プログラミング言語の「変数」には、もっと多様なものを代入することができます。次の例では、変数 f に対して、「関数 is_even()」を代入しています。

def is_even(x):
  if x % 2 == 0:
    return True
  else:
    return False

f = is_even
f(10)   # True

プログラミング言語では、一般に、変数に代入したり、さまざまな演算処理をほどこすことができる対象物（オブジェクト）を「ファーストクラスオブジェクト」と言います。上記の例は、「Python は、関数をファーストクラスオブジェクトとして扱える」ことを示しています。

一方、数学の中にも、ちょっと変わり種の「オブジェクト」を格納した変数があります。数学用語では、「確率変数 X」と呼ばれるもので、この中には、「確率分布」と呼ばれるオブジェクトが格納されています。確率分布というのは、プログラミング言語的に言うと、「（ある一定の分布に従って）ランダムな値を返す関数」と考えても構いません。

たとえば、

　 $X \sim N(0, 1)$ （平均 0、分散 1 の正規分布）

と表した場合、変数 X には、次図のように、0 を中心におよそ -1 〜 1 の範囲の値がランダムに得られる確率分布が格納されています。（グラフの縦軸は、その値が得られる確率に比例すると思ってください。）

実は、Tensorflow Probability を用いると、このような確率分布をファーストクラスオブジェクトとして扱うことができます。つまり、Python の変数に確率分布が代入できるようになります。

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

x = tfd.Normal(loc=0., scale=1.)

x  # <tfp.distributions.Normal 'Normal' batch_shape=[] event_shape=[] dtype=float32>
x.sample()  # <tf.Tensor: shape=(), dtype=float32, numpy=-1.5092232>
x.sample(3).numpy()  # array([ 0.6608737 , -0.99310094,  1.223854  ], dtype=float32)

この例では、変数 x に、先ほどの正規分布 $N(0, 1)$ を代入して、sample() メソッドで、乱数を1つ取得しています。x.sample() の出力値は、Tensorflow のオブジェクト形式になっていますが、「numpy=....」の部分が実際に得られた値に対応します。その次の x.sample(3).numpy() は、乱数を3つ取得した後に、numpy() メソッドで、NumPy の Array に変換しています。

この他にも、複数の確率分布を合成して新しい確率分布を作るなど、確率分布どうしの「演算処理」ができます。こういった Tensorflow Probability の特徴は、「確率分布がファーストクラスオブジェクトである」と考えれば、スッキリと理解できるでしょう。

確率分布を出力する機械学習モデル

Tensorflow Probability を Tensorflow（もしくは Keras）と組み合わせると、「確率分布を出力する機械学習モデル」を定義することができます。たとえば、次の図は、青点で与えられた学習データをニューラルネットワークによる回帰モデルでフィッティングした結果です。

ニューラルネットワークなど、通常の機械学習モデルの出力値は、実数値になります。出力される実数値が、学習データになるべく近くなるようにモデルのパラメーターを調整するのが回帰モデルの学習処理です。

しかしながら、このような回帰モデルでは、「学習データに含まれる確率的要素」を捉えることができません。上図の学習データを見ると、グラフの中央付近はデータのばらつきが大きく、両端付近はデータのばらつきが小さくなっています。実は、この学習データは、コサインの山形のカーブに対して、正規分布のノイズを乗せて生成しており、正規分布の分散（ノイズの幅）が中央付近になるほど大きくなっているのです。

しかしながら、一般的な回帰モデルでは、このようなノイズに関する情報はすっぽりと抜け落ちてしまいます。回帰モデルで何らかの予測を行おうと言う場合、このようなノイズの情報も取得できれば、「中央付近の予測値は外れる可能性が高い（±○○程度にずれる恐がある）」といった（予測値を利用する上で役に立つ）付加情報が得られることになります。

そして・・・、Tensorflow Probability を用いると、このような「ノイズの情報を含めた回帰モデル」を簡単に作ることができます。（じゃじゃーん！）

具体的には、モデルの出力値を実数値ではなく、「平均 $\mu$ 、分散 $\sigma^2$ の正規分布 $N(\mu,\sigma^2)$ 」に置き換えます。Tensorflow Probability では、確率分布もファーストクラスオブジェクトですので、こんなことが簡単にできてしまいます。（ここでは、簡単のために学習データに含まれるノイズが正規分布であるという前提にしていますが、より一般に、もっと広い範囲の確率分布を想定して学習することもできます。）

そして、入力値（横軸の値 x）に対して、モデルが出力する分布 $N(\mu,\sigma^2)$ が学習データの分布に近くなるように、モデルのパラメーターをチューニングします。もう少し詳しく言うと、モデルが出力した分布 $N(\mu,\sigma^2)$ からサンプル群を取得した際に、それが、学習データに近いサンプル群になることを目指してチューニングを行います。

具体的なコードは、こちらにありますが、ここではポイントを抜粋して説明します。

・DNN_regression_for_means_and_variances.ipynb

まず、実数値で予測する通常のモデルはこちらです。

model = models.Sequential()
model.add(layers.Dense(64, activation='relu', input_shape=(1,)))
model.add(layers.Dense(32, activation='relu'))
model.add(layers.Dense(1))
model.summary()

ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 64)                128       
_________________________________________________________________
dense_1 (Dense)              (None, 32)                2080      
_________________________________________________________________
dense_2 (Dense)              (None, 1)                 33        
=================================================================
Total params: 2,241
Trainable params: 2,241
Non-trainable params: 0
_________________________________________________________________

ReLU を活性化関数とする、ごく普通のフィードフォワードネットワークです。最後の出力層は、活性化関数を持たないノードが1つで、ここから実数値が出力されます。

Tensorflow Probability を用いて、これを正規分布 $N(\mu,\sigma^2)$ を出力するモデルに書き換えると、次のようになります。

model = models.Sequential()
model.add(layers.Dense(64, activation='relu', input_shape=(1,)))
model.add(layers.Dense(32, activation='relu'))
model.add(layers.Dense(2)) # predict loc and scale
model.add(tfp.layers.DistributionLambda(
      lambda t: tfd.Normal(loc=t[..., :1], scale=tf.math.softplus(t[..., 1:]))
      ))
model.summary()

ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー
Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_3 (Dense)              (None, 64)                128       
_________________________________________________________________
dense_4 (Dense)              (None, 32)                2080      
_________________________________________________________________
dense_5 (Dense)              (None, 2)                 66        
_________________________________________________________________
distribution_lambda (Distrib ((None, 1), (None, 1))    0         
=================================================================
Total params: 2,274
Trainable params: 2,274
Non-trainable params: 0
_________________________________________________________________

まず、先ほどのモデルの出力層が、「活性化関数をもたない2つのノード」に置き換わっています。ここでは、これを正規分布の平均 $\mu$ 、および、標準偏差 $\sigma$ （分散の平方根）と解釈して、最後の出力層（tfp.layers.DistributionLambda）で、正規分布 $N(\mu,\sigma^2)$ に変換しています。

そして、実際の学習処理は次になります。

negloglik = lambda y, rv_y: -rv_y.log_prob(y)
model.compile(optimizer='adam', loss=negloglik)
history = model.fit(xs, ys, batch_size=len(xs), epochs=1000, verbose=0)

誤差関数 loss にちょっと特殊なもの（negloglik）が入っていますが、先ほど説明したように、「モデルが出力した分布 $N(\mu,\sigma^2)$ から得られるサンプル群と実際の学習データの差異」を表していると考えてください。これが小さくなるように、前段のニューラルネットワークに含まれるパラメーターをチューニングするわけです。

学習結果は次のようになります。この図では、得られた正規分布 $N(\mu,\sigma^2)$ の平均 $\mu$ を赤いラインで示した上で、その上下に標準偏差 $\sigma$ （の2倍）の幅のラインを示してあります。

また、モデルの出力値は、確率分布ですので、ここから新しいサンプルを取得することもできます。次の図は、学習済みのモデルから新たに取得したサンプルです。学習データに類似した広がりを持つサンプルが得られていることがわかります。

「関数」を出力する「確率変数」としてガウス過程を理解する

Tensorflow Probability のチュートリアルに、ガウス過程（Gaussian Process）を用いた回帰モデルが出てきます。「ガウス過程ってなんぞや？」という方も多いと思いますが、これは、「関数を出力する確率変数」と理解することができます。

たとえば、先ほどの（正規分布を出力する）回帰モデルでは、横軸の値 x を1つ決めると、その点の確率分布が得られて、そこからさらにサンプル値を得ることができます。そこで、すべての点 x について、1つずつサンプルを取って折れ線グラフを作ったとします。中央付近は上下に大きく変動するグラフになるものと想像できます。

しかしながら、学習データの種類によっては、このような捉え方は適切でない場合があります。

たとえば、この学習データは、「一年間の気温データを過去50年間分まとめたもの」だと考えてみます。夏の気温は、冬の気温に比べて、年による変動が激しいというわけですね。このため、中央付近のデータは上下にはげしくちらばっているわけですが、これは、「特定の年」における気温変動が激しいというわけではありません。その年が猛暑であれば、中央付近のデータは一様に高い値を示すはずですし、冷夏であれば、一様に低い値を示すはずです。さきほどのモデルでは、「ある特定の１年間の気温変化のサンプル」を得ることができないのです。

このようなデータをモデル化する際は、「１年間の気温変化全体」を１つのサンプルとしてまとめて出力するようなモデルが必要になります。「１年間の気温変化全体」というのは、数学的に言えば、日付 x を変数とする関数 f(x) ですので、「関数 f(x)」がランダムに1つ得られるモデルということになります。（当然ながら、得られるサンプル f(x) は、特定の年の気温変化に相当する滑らかな関数という想定です。）

関数がランダムに得られる・・・・というと何か気持ち悪い気もしますが、プログラミング言語的に考え直してみましょう。冒頭で述べたように、Python などのプログラミング言語では、「関数」もファーストクラスオブジェクトですので、「関数を返り値として返す関数」といったものはごく普通に作ることができます。Tensorflow Probability では、これをもう一段拡張して、「関数をサンプル値として返す確率変数」が扱えるのです。サンプルとして得られた関数 f(x) に対して、個別の x に対応する値を計算してグラフを描けば、滑らかなグラフになりますが、多数のサンプル f(x) を取得してグラフを重ね合わせると、場所によって、大きく変動する場所や、あまり変動しない場所が出てくると言う寸法です。

表題のガウス過程は、このような、「関数を出力値とする確率変数」の一種になり、Tensorflow Probability には、ガウス過程を扱うためのモジュールが事前に用意されているのです。具体的なコードの例はこちらになります。

・Gaussian_Process_Regression.ipynb

これはまだ Keras と統合されておらず、若干冗長なコードになっていますが、モデルを定義して、勾配効果法でモデルに含まれるパラメーターを最適化するという流れは、通常の Tensorflow と同じです。

ガウス過程のモデルでは、カーネルと呼ばれる特別な関数に含まれるパラメータが（主な）チューニング対象のパラメーターになります。上記のノートブックの場合、学習データと学習済みのモデルから得られたいくつかのサンプル（関数）は次のようになります。