Variational Gaussian Process と Deep Neural Netの回帰の比較

はじめに

この記事は大したこと書いてないです。

一般論としてガウス過程は入力空間で近いデータ点は出力空間でも近い値を取るという性質を持つため、直感的にはなめらかな（非線形を含む）関数の近似が強そうだと言えます。 また、確率モデルであるためハイパーパラメータをデータから推定でき、データがあれば一先ず（最尤推定の意味で）局所的な最適解は得られるというすぐれものです。

一方でニューラルネットワークは、一般線形モデルで現れる基底関数自体を中間層（非線形関数）でパラメタライズし、最後の出力層で線形回帰を実施することができます。当然非線形な関数を近似できるだけでなく、どうやら滑らかでない（あるいは不連続な）関数の近似では従来手法よりも理論的に優れることが述べられているようです（ただしハイパーパラメータは多数）。

今回は適当なデータでこれを確認してみるお遊びの時間です。

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import matplotlib.pyplot as plt

plt.style.use("seaborn")

tfd = tfp.distributions
tfkl = tf.keras.layers
tfpl = tfp.layers
tfpk = tfp.positive_semidefinite_kernels

というコードが頭にある前提で、コード付きで見ていきます。

連続で滑らかな関数の近似

データ

$$ y = w x (1 + \sin (x)) + e $$

という形式の関数を用います。 $e$ はノイズで、$x$ の値に応じて変換する分散を持つ正規分布から発生させます。

w0 = 0.125
b0 = 5.
x_range = [-20, 60]

def load_dataset(n=1000, n_tst=150):
  np.random.seed(43)
  def s(x):
    g = (x - x_range[0]) / (x_range[1] - x_range[0])
    return 3 * (0.25 + g**2.)
  x = (x_range[1] - x_range[0]) * np.random.rand(n) + x_range[0]
  eps = np.random.randn(n) * s(x)
  y = (w0 * x * (1. + np.sin(x)) + b0) + eps
  x = x[..., np.newaxis]
  return y, x

y, x = load_dataset()

Gaussian Process

tfkl モジュールを利用してガウシアンプロセスを書くには、カーネル関数を tfkl.Layers によって実装する必要があります。 このとき call は実装上ダミーで必要なメソッドです。ガウシアンプロセスで実際に利用されるのは @property によって記述されたカーネル関数のインスタンスのみです。

class RBFKernelFn(tfkl.Layer):
  def __init__(self, **kwargs):
    super(RBFKernelFn, self).__init__(**kwargs)
    dtype = kwargs.get('dtype', None)

    self._amplitude = self.add_variable(
            initializer=tf.constant_initializer(0),
            dtype=dtype,
            name='amplitude')
    
    self._length_scale = self.add_variable(
            initializer=tf.constant_initializer(0),
            dtype=dtype,
            name='length_scale')

  def call(self, x):
    return x

  @property
  def kernel(self):
    return tfpk.ExponentiatedQuadratic(
      amplitude=tf.nn.softplus(0.1 * self._amplitude),
      length_scale=tf.nn.softplus(5. * self._length_scale)
    )

このカーネルを利用して、下記のように tf.keras.Sequential クラスでの実装が可能になります。

ここで num_inducing_points は潜在的な入力データ点を何個準備するかを指定しています。ガウシアンプロセスはデータ点に対して3乗で計算量が増大します。 そこで全てのデータ点を使う代わりに、架空の入力点を仮定して、グラム行列を低ランク近似します。変分推論の美しい応用です。ビッグデータ相手にはこの手の工夫が必須です。

通常はsparse gaussian process と呼ばれる手法であり、変分推論はその具体的な実行方法という具合です（極端な話、データを間引くことでもsparseなガウシアンプロセスは作れる）。

しかしTFPではそんな難しい手法も数行で実行できてしまいます。

tf.keras.backend.set_floatx('float64')

num_inducing_points = 32
model = tf.keras.Sequential([
    tfkl.InputLayer(input_shape=[1]),
    tfpl.VariationalGaussianProcess(
        num_inducing_points=num_inducing_points,
        kernel_provider=RBFKernelFn(),
        event_shape=[1],
        inducing_index_points_initializer=tf.constant_initializer(
            np.linspace(*x_range, num=num_inducing_points,
                        dtype=x.dtype)[..., np.newaxis]),
        unconstrained_observation_noise_variance_initializer=(
            tf.constant_initializer(np.array(0.54).astype(x.dtype))),
    ),
])

損失関数はmodelがtfd.Distributionクラスになっており、メソッドとして variational_loss を有するため、これを利用します。 無名関数で第一引数が分布、第二引数が確率変数の実現値になる形にしておき、戻り値を variational_loss として、あとはKerasの通常のフォーマットに従って学習します。

batch_size = 256
loss = lambda y, rv_y: rv_y.variational_loss(
    y, kl_weight=np.array(batch_size, x.dtype) / x.shape[0])
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=loss)
model.fit(x, y, batch_size=batch_size, epochs=1000, verbose=False)

yhat = model(x_tst)

tfd.Distributionクラスが最終層に来ている場合、modelはTensorではなくDistribution型になります。ここから Tensorを得るにはsample メソッドを利用します。 試しにガウシアンプロセスによって7つの関数を表示してみます。

plt.figure(figsize=[12, 5])  # inches
plt.plot(x, y, 'b.', label='observed');

num_samples = 7
for i in range(num_samples):
  sample_ = yhat.sample().numpy()
  plt.plot(x_tst,
           sample_[..., 0].T,
           'r',
           linewidth=0.9,
           label='ensemble means' if i == 0 else None)

見事に線形に増加していく様子と、正弦波で揺れる部分を表しています。また確率モデルであるため、予測のバラつきも見ることができます（もっとたくさんサンプリングして、区間を見ると良い）。

ニューラルネット

同じデータを下記のニューラルネットに食わせます。

model_deepnn = tf.keras.Sequential([
    tfkl.Dense(100, activation=tf.nn.relu),
    tfkl.Dense(100, activation=tf.nn.relu),
    tfkl.Dense(100, activation=tf.nn.relu),
    tfkl.Dense(100, activation=tf.nn.relu),
    tfkl.Dense(100, activation=tf.nn.relu),
    tfkl.Dense(1)
])
loss_mse = tf.keras.losses.MSE
model_deepnn.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=loss_mse)
model_deepnn.fit(x, y, batch_size=batch_size, epochs=100, verbose=False,
                 workers=-1)

色々調整はしたほうが良いのでしょうけども、一先ず適当に実施。

線形に増加している様子は得られましたが、何故か正弦波で揺れている成分までもが誤差という扱いに…。

不連続な関数

データ

x_range = [-20, 60]

def load_dataset(n=1000, n_tst=150):
  np.random.seed(43)
  x = np.linspace(x_range[0], x_range[1], n)
  eps = np.random.randn(n)
  data_list = [-10] * 200 + [10] * 300 + [4] * 100 + [-7] * 100 + [0] * 300
  y = np.array(data_list, dtype=np.float32) + eps
  x = x[..., np.newaxis]
  return y, x

y, x = load_dataset()

横線を適当に引いてノイズを載せたデータです。

Gaussian Process

モデルは先ほどと全く同じコードで、データだけが変わっています。 GPの予測の平均だけを図で載せます。

不連続な変化についていけません。なぜか値がブレる現象。

ニューラルネット

割とぴしっとついていきます。