Tensorflowで実装されたBayesian Layersを試してみる

はじめに

先日、 Bayesian Layers: A Module for Neural Network Uncertainty というのを見かけまして、Tensorflow上で確率的な振る舞いをするLayerの話でした。
見てみると、こんな感じで、

in_x = tf.keras.layers.Input(shape=(2, ), name="in_x")
x = layers.DenseReparameterization(10, activation="tanh")(in_x)
x = layers.DenseReparameterization(1, activation="sigmoid")(x)
probabilistic_model = tf.keras.Model(in_x, x)

KerasのLayerのように使うことができるようだったので、実際に試してみました。

Version

• Python: 3.6.4
• tensorflow: 1.12.0
• tensorflow-probability: 0.5.0
• tensor2tensor: 1.11.0

サンプルデータの準備

今回は、平面上に2種類の点があって、それらを識別するタスクをやってみることにします。

# 今回使うライブラリたち
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import tensorflow as tf

# jupyterlab 上のおまじない達
%matplotlib inline
sns.set_style('whitegrid')
%reload_ext autoreload
%autoreload 2

# データ生成
N = 200
np.random.seed(999)
X = np.random.random((N, 2)) * 2 - 1  # -1 ~ 1
X_ = X + np.random.normal(scale=0.1, size=(N,2))  # 適当にノイズを入れて
Y = (X_[:, 0] ** 2 < X_[:, 1]) + 0  # x^2 < y を 1とする

# plot
def plot_data(x, y):
    plt.figure(figsize=(5,5))
    plt.scatter(x[y==0, 0], x[y==0, 1], c="blue")
    plt.scatter(x[y==1, 0], x[y==1, 1], c="red")
    plt.plot(np.linspace(-1, 1), np.linspace(-1, 1)**2)

plot_data(X, Y)    

上記のような感じで生成して、こんな感じの分布になっています。赤い点が "1" で、青い点が "0"です。

普通のNNで

Modelの生成と学習

まずは、普通のNNでやってみることにします。

# Create model
in_x = tf.keras.layers.Input(shape=(2, ))
x = tf.keras.layers.Dense(10, activation="tanh")(in_x)
x = tf.keras.layers.Dense(1, activation="sigmoid")(x)
deterministic_model = tf.keras.Model(in_x, x)

# Training
deterministic_model.compile(optimizer="adam", loss=tf.keras.losses.binary_crossentropy)
history = deterministic_model.fit(X, Y, epochs=2000, verbose=False)

# Plot Losses
plt.plot(history.history["loss"])

上記のような感じで、モデルを作り学習させます。
Lossはちょっと、頼りない感じはありますが、まあ、順調に学習していたようです。

ModelのOutputの確認

以下のように 0.5 を境に0, 1と離散化してPlotしてみると、

pred_y = (deterministic_model.predict(X)[:, 0] > 0.5) + 0
plot_data(X, pred_y)

となり、まあそれほど文句のない感じ。

---

-1~1の範囲の予測値をヒートマップにしてみると、

z = np.linspace(-1, 1)
xx, yy = np.meshgrid(z, -z)
mesh_x = np.concatenate([xx.reshape(-1, 1), yy.reshape(-1, 1)], axis=1)
mesh_pred_y = deterministic_model.predict(mesh_x)
plt.figure(figsize=(5,5))
sns.heatmap(mesh_pred_y.reshape(-1, 50))

こんな感じになりました。
理想的な y=x^2 という形状ではないですが、例えば学習データの左上付近は隙間があるのでぼんやりしてしまうのはしょうがないかなという感じです。

確率的なNNで

Modelの生成

さて、ここからが本題です。
まず、モデルですが、以下のように定義できました。

import tensor2tensor.layers.bayes as layers

in_x = tf.keras.layers.Input(shape=(2, ), name="in_x")
x = layers.DenseReparameterization(10, activation="tanh")(in_x)
x = layers.DenseReparameterization(1, activation="sigmoid")(x)
probabilistic_model = tf.keras.Model(in_x, x)

tf.keras.layers.Dense が tensor2tensor.layers.bayes.layers.DenseReparameterization に置き換わっただけです。Edwardなどを使うと、事前分布と事後分布を別々に作ったりしないといけなかったことを考えるとかなり楽に定義できているんじゃないでしょうか。

DenseReparameterization は ソースコードを見ると weight や bias を 正規分布 で表現し、reparameterization とかうまく使ってパラメータを学習してくれるっぽいです。
このLayer自体がlossを持っていて、それも一緒に最小化することで正則化的な働きがあるっぽい(用語にちょっとついていけてない...)。

Modelの学習

このまま probabilistic_model.compile(...) とやってみましたが、さすがにエラーになりました。現状ではここからは、 tensorflow的なお作法で進めないといけないようです。

y_placeholder = tf.placeholder(tf.float32, shape=(None, 1), name="y_true")
loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(y_placeholder, probabilistic_model.output))
kl = sum([tf.reduce_mean(x) for x in probabilistic_model.losses])
loss = loss + kl * 0.01
train_op = tf.train.AdamOptimizer(0.05).minimize(loss)

まず、正解データを注入するための y_placeholder で普通に binary_crossentropy を計算して loss とします。
次に、Modelの中の losses から各Layerのlossなどを集めます。
それぞれ次元が違うので単純に加算できないため、reduce_mean() が良いのか、reduce_sum() が良いのかよくわからなかったですが、とにかくreduceして、lossに足します。
あとはそのlossをベース training していきます。

sess = tf.Session()
sess.run(tf.global_variables_initializer())

losses = []
for _ in range(30000):
    l, _ = sess.run([loss, train_op], feed_dict={probabilistic_model.input: X, y_placeholder: Y.reshape(-1, 1)})
    losses.append(l)

plt.plot(losses[10:])

lossは一応減少していきますが、収束は結構時間がかかる感じです。

ModelのOutputの確認

先程と同じように 0.5 で 0, 1に離散化して、plotしてみます。

pred_y = (sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: X})[:, 0] > 0.5) + 0
plot_data(X, pred_y)

先ほどと同様、結構いい感じになっています。

確率Modelならではの振る舞い

先ほどと何が変わったかというと、同じ入力でも結果が異なることがある、ということです。
確率的なLayerの部分では、実行毎にweight, biasの値がサンプリングされるようです。

> # session.run() する度に同じ入力でも違う結果になる
> print(sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: [[0, 0.1]]}))
[[0.8825446]]

> print(sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: [[0, 0.1]]}))
[[4.8771195e-11]]

> print(sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: [[0, 0.1]]}))
[[1.]]

> # 同一のsession.run() だと同じ値になるようだ
> print(sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: [[0, 0.1], [0, 0.1], [0, 0.1]]}))
[[1.2230051e-19]
 [1.2230051e-19]
 [1.2230051e-19]]

例えば、ある点における予測を1000回行って、分布をみてみると、

zz = []
for _ in range(1000):
    z = sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: [[0.1, 0.15]]})
    zz.append(z)
sns.distplot(zz)

こんな感じになりました。

平均と標準偏差のPlot

-1~1の範囲で、1000回ずつ予測して、平均と標準偏差をPlotしてみます。

z = np.linspace(-1, 1)  # -1 ~ 1, 50 items
xx, yy = np.meshgrid(z, -z)  # 50 * 50
mesh_x = np.concatenate([xx.reshape(-1, 1), yy.reshape(-1, 1)], axis=1)

yy = []
for _ in range(1000):
    mesh_pred_y = sess.run(probabilistic_model.output, feed_dict={probabilistic_model.input: mesh_x})
    yy.append(mesh_pred_y)

pred_y_mean = np.mean(np.array(yy), axis=0)  # mean([1000, 2500, 1], axis=0) -> (2500, 1)
pred_y_std = np.std(np.array(yy), axis=0)    # std([1000, 2500, 1], axis=0) -> (2500, 1)

# 平均のPlot
plt.figure(figsize=(5,5))
sns.heatmap(pred_y_mean.reshape(-1, 50))

# 標準偏差のPlot
plt.figure(figsize=(5,5))
sns.heatmap(pred_y_std.reshape(-1, 50))

平均のplot

なんとなく、普通のNNに近い感じになっている気がします。原点付近のでっぱりがちょっと特徴的？

標準偏差のplot

このPlotが一番興味深いです。標準偏差が大きいところはだいたい境界面を表していることがわかります。確率的に確かな部分と曖昧な部分が表現できているところがこのモデルの良いところと言えそうです。

さいごに

今回の DenseReparameterization は tensor2tensor というpipに入っているのですが、これは結構余計なものをrequireしているのでちょっと嫌な感じです。まあ、いずれ、tensorflow本体に格納されていくんじゃないでしょうか。

試してはいませんが、最初に紹介した論文には他にも 以下のようなものが紹介されていました。

• Conv2DReparameterization: 2D畳み込み用
• GaussianProcess: よくわからないですが、Denseと違ってweightをチューニングするんじゃなくて、入力と出力の関係を直接(GaussianProcessで?)学習するみたいな雰囲気があります。試してみたかったのですが、今回のversionではまだ実装されていませんでした(github上には実装されていたけど)
• 他にも MixtureofLogistic とか RealNVP とかありましたが、前者はまだしも後者は何なのかわからなかったです...

確率Modelは「不確かさ」が表現されるのが結構便利だと思っています。
強化学習とかとも相性良さそうですし。
今後の展開が楽しみです。