はじめに

ニューラルネットワークの予測の不確実性を算出する手法を検証します。
Jupyter notebookは下記にあります。

概要

連続値を予測する回帰のためのニューラルネットワークを構築
deep Gaussian process で予測の不確実性を算出

deep Gaussian process (深層ガウス過程)

ニューラルネットワークの各層にガウス過程を用いたものが、深層ガウス過程になります。
3層のニューラルネットワークを考え、入力層を$X^1$、中間層を$X^2$、出力層を$Y$とします。
中間層と出力層は下記のようにガウス過程で定義されます。
$$
f^{X^2} \sim GP(0,k^1(X^1,X^1)) \
f^Y \sim GP(0,k^2(X^2,X^2))
$$

深層ガウス過程では、学習のためにinducing points（誘導点）を導入します。
3層のニューラルネットワークの場合は、入力層の誘導点$Z_1$と、中間層の誘導点$Z_2$を導入します。
誘導点は変分推論で最適化する変数で、観測データを用いて、誘導点、平均関数、カーネル関数の最適化を行います。
この最適化は、大規模な観測データを少数の誘導点に圧縮していると捉えることができます。
新たなデータが与えられた場合は、学習結果を用いて事後分布を計算します。

実装

1. ライブラリのインポート

必要なライブラリをインポートします。

import sys
import os

import matplotlib.pyplot as plt
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

!pip install gpytorch
import gpytorch
from gpytorch.means import ConstantMean, LinearMean
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.variational import VariationalStrategy, CholeskyVariationalDistribution
from gpytorch.distributions import MultivariateNormal
from gpytorch.models import ApproximateGP, GP
from gpytorch.mlls import VariationalELBO, AddedLossTerm
from gpytorch.likelihoods import GaussianLikelihood

from gpytorch.models.deep_gps import DeepGPLayer, DeepGP
from gpytorch.mlls import DeepApproximateMLL

2. 実行環境の確認

使用するライブラリのバージョンや、GPU環境を確認します。
Google Colaboratoryで実行した際の例になります。

print('Python:', sys.version)
print('PyTorch:', torch.__version__)
!nvidia-smi

実行結果

Python: 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0]
PyTorch: 2.0.1+cu118
Mon Sep 25 05:59:05 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P8     9W /  70W |      0MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

3. データセットの作成

線形回帰用のデータセットを作成します。
ただし、学習には[-1,1]の範囲のデータは使用しません。

def make_dataset(seed, plot=0, batch_size=64):
    np.random.seed(seed)

    x_true = np.linspace(-4, 4, 100)
    y_true = np.sin(x_true)

    x = np.concatenate([np.random.uniform(-4, -1, 100), np.random.uniform(1, 4, 100)])
    y = np.sin(x)

    # データをPyTorchのテンソルに変換
    x = torch.from_numpy(x).float().view(-1, 1)
    y = torch.from_numpy(y).float()
    x_true = torch.from_numpy(x_true).float().view(-1, 1)

    # グラフを描画
    if plot == 1:
        plt.plot(x_true, y_true)
        plt.scatter(x, y)
        plt.xlabel('x')
        plt.ylabel('y')
        plt.show()

    dataset = torch.utils.data.TensorDataset(x, y)
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader

_ = make_dataset(0, plot=1)

ニューラルネットワークの定義

GPyTorchのサンプルコードを参考にGP layerと3層のニューラルネットワークを定義します。
隠れ層のユニット数は10としてあります。

class ToyDeepGPHiddenLayer(DeepGPLayer):
    def __init__(self, input_dims, output_dims, num_inducing=128, mean_type='constant'):
        if output_dims is None:
            inducing_points = torch.randn(num_inducing, input_dims)
            batch_shape = torch.Size([])
        else:
            inducing_points = torch.randn(output_dims, num_inducing, input_dims)
            batch_shape = torch.Size([output_dims])

        variational_distribution = CholeskyVariationalDistribution(
            num_inducing_points=num_inducing,
            batch_shape=batch_shape
        )

        variational_strategy = VariationalStrategy(
            self,
            inducing_points,
            variational_distribution,
            learn_inducing_locations=True
        )

        super(ToyDeepGPHiddenLayer, self).__init__(variational_strategy, input_dims, output_dims)

        if mean_type == 'constant':
            self.mean_module = ConstantMean(batch_shape=batch_shape)
        else:
            self.mean_module = LinearMean(input_dims)
        self.covar_module = ScaleKernel(
            RBFKernel(batch_shape=batch_shape, ard_num_dims=input_dims),
            batch_shape=batch_shape, ard_num_dims=None
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)

class DeepGPModel(DeepGP):
    def __init__(self, input_dim=1):
        super().__init__()
        num_hidden_dims = 10

        self.input_layer = ToyDeepGPHiddenLayer(
            input_dims=input_dim,
            output_dims=num_hidden_dims,
            mean_type='linear',
        )

        self.hidden_layer = ToyDeepGPHiddenLayer(
            input_dims=num_hidden_dims,
            output_dims=num_hidden_dims,
            mean_type='linear',
        )

        self.last_layer = ToyDeepGPHiddenLayer(
            input_dims=num_hidden_dims,
            output_dims=None,
            mean_type='constant',
        )

        self.likelihood = GaussianLikelihood()

    def forward(self, inputs):
        h = self.input_layer(inputs)
        h = self.hidden_layer(h)
        output = self.last_layer(h)
        return output

    def predict(self, x):
        preds = self.likelihood(self.forward(x))
        return preds.mean, preds.variance

4. 学習

deep Gaussian processの学習を行います。

batch_size=128
data_loader = make_dataset(0, batch_size=batch_size)

model = DeepGPModel()
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

model = model.to(device)

num_epochs = 100

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
mll = DeepApproximateMLL(VariationalELBO(model.likelihood, model, len(data_loader.dataset), beta=0.01))

for i in range(num_epochs):
    for x_batch, y_batch in data_loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        optimizer.zero_grad()
        output = model(x_batch)
        loss = -mll(output, y_batch)
        loss.backward()
        optimizer.step()

5. 予測

予測結果と標準偏差を描画します。

x_true = np.linspace(-4, 4, 100)
y_true = np.sin(x_true)
x_true = torch.from_numpy(x_true).float().view(-1, 1)

model.eval()
x_true = x_true.to(device)
y_preds = []
with torch.no_grad():
    y_mean, y_var = model.predict(x_true)

x_true = x_true.to('cpu')
y_mean = y_mean.to('cpu').mean(0)
y_var = y_var.to('cpu').mean(0)
y_std = np.sqrt(y_var)


# グラフの描画
plt.figure(figsize=(8,4))
plt.subplot(121)
plt.plot(x_true, y_true, label='True Function')
plt.plot(x_true, y_mean, label='Mean Prediction')
plt.fill_between(x_true.flatten(), y_mean.flatten() - y_std.flatten(), y_mean.flatten() + y_std.flatten(), alpha=0.3, label='Uncertainty')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()


plt.subplot(122)
plt.plot(x_true, y_std)
plt.xlabel('x')
plt.ylabel('Std')

plt.tight_layout()
plt.show()

おわりに

今回の結果

予測の不確実性は、x最小値および最大値付近とデータが含まれない[-1,1]の範囲で大きくなっています。
データ数が少なく、予測が不確実と考えられる領域と、予測の標準偏差が大きい領域が一致しているため、想定通り予測の不確実性が算出できていると考えられます。
ただし、不確実性の差は非常に小さいため、今回の設定での不確実性の算出は他の手法の方が優れているかもしれません。

次にやること

予測の不確実性を算出する他の手法も検証したいと思います。

参考資料

須山敦志, ベイズ深層学習, 2019.
A. C. Damianou and N. D. Lawrence, Deep Gaussian Processes, AISTATS, 2013.
H. Salimbeni and M. P. Deisenroth, Doubly Stochastic Variational Inference
for Deep Gaussian Processes, NeurIPS, 2017.
GPyTorch https://docs.gpytorch.ai/en/stable/examples/05_Deep_Gaussian_Processes/Deep_Gaussian_Processes.html