GPyTorchでニューラルネットとガウス過程を結合する

[Sansan Advent Calendar 2019](https://adventar.org/calendars/3924)の10日目の記事です．

本記事では[GPyTorch](https://gpytorch.ai/)とともに，1つガウス過程周りで気になる技術があったので，GPyTorchの実装を含めて紹介したいと思います．

## 環境
* Python 3.6.5
* GPyTorch 0.3.6
* PyTorch 1.3.0
* PyTorch Geometric 1.3.2



## GPyTorchとは
[GPyTorch](https://gpytorch.ai/)とは，ガウス過程周りの種々の手法をPyTorchバックエンドで動かせるライブラリです． これまでに GPyと呼ばれるガウス過程専用のライブラリは存在していたのですが，PyTorchの自動微分＋GPUによるスケーラブルな行列演算で，ガウス過程の計算を効率的に行えるのが特徴です．さらに，GPyTorchの元論文[1]で提案されているBlackBox Matrix-Matrix Multiplication (BBMM) により，ガウス過程の大きなボトルネックとなっている $\mathcal{O}(n^3)$ の計算量を実質的に$\mathcal{O}(n^2)$ に減らしているのも特徴の1つと言えるでしょう．


## Stochastic Variational Deep Kernel Learning
Stochastic Variational Deep Kernel Learning (SVDKL) とは，ガウス過程における確率的変分推論をニューラルネットと組み合わせることで，ニューラルネットの特徴量獲得性能をうまく使いつつ，ガウス過程による不確実性を伴う推論を行える手法です[2]．以下の図のように，入力から特徴量を取得するまでをニューラルネット部分で担い，その後ガウス過程の層を追加します（図は[1] Figure 1から引用）．


<img width="599" alt="Screen Shot 2019-12-08 at 20 50 37" src="https://user-images.githubusercontent.com/24391425/70388964-71cdff00-19fc-11ea-90e9-fa8769f74ecd.png">


## コード
コードを踏まえて，使い方を解説したいと思います．コード全体は[こちら](https://gist.github.com/wataruhashimoto52/08e8a872cbb9fd575cf7826be70ecd21)．
以下は各種インポート．

```python

import gpytorch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric.transforms as T
from gpytorch.models import ApproximateGP
from gpytorch.variational import (CholeskyVariationalDistribution,
                                  WhitenedVariationalStrategy)

from torch.nn import GRU, Linear, ReLU, Sequential
from torch_geometric.data import DataLoader
from torch_geometric.datasets import QM9
from torch_geometric.nn import NNConv, Set2Set
from torch_geometric.utils import remove_self_loops
``` 

まず最初に，データから特徴量を取得するNN部分を，`FeatureExtractor` として定義します．この部分ではConvolutional Neural Networksを使ったり，以下に示すように，分子のようなグラフ構造に対するニューラルネットワークを使うこともできます．今回は試しに，PyTorch Geometricを用いたNeural Message Passing [3] をグラフからの特徴量獲得層 `NeuralMessagePassing` として定義します．

```python

class NeuralMessagePassing(nn.Module):
    def __init__(self, num_features, output_features, dim):
        super(NeuralMessagePassing, self).__init__()
        self.num_features = num_features
        self.output_features = output_features
        self.dim = dim
        self.lin0 = torch.nn.Linear(self.num_features, self.dim)

        seq_net = Sequential(Linear(6, 128), ReLU(),
                             Linear(128, self.dim * self.dim))
        self.conv = NNConv(self.dim, self.dim, seq_net, aggr='mean')
        self.gru = GRU(self.dim, self.dim)

        self.set2set = Set2Set(self.dim, processing_steps=3)
        self.lin1 = torch.nn.Linear(2 * self.dim, self.dim)
        self.lin2 = torch.nn.Linear(self.dim, self.output_features)


    def forward(self, data):
        out = F.relu(self.lin0(data.x))
        h = out.unsqueeze(0)

        for i in range(2):
            m = F.relu(self.conv(out, data.edge_index, data.edge_attr))
            out, h = self.gru(m.unsqueeze(0), h)
            out = out.squeeze(0)

        out = self.set2set(out, data.batch)
        out = F.relu(self.lin1(out))
        out = self.lin2(out)

        return out


class FeatureExtractor(nn.Sequential):
    def __init__(self, num_features, output_features, dim):
        super(FeatureExtractor, self).__init__()
        self.num_features = num_features
        self.output_features = output_features
        self.dim = dim
        self.add_module("neural_message_passing", NeuralMessagePassing(
            self.num_features, self.output_features, self.dim))

```

そして，特徴量取得層である `feature_extractor` を定義します． 

```python
num_features = dataset.num_features
output_features = 25
num_inducing_points = 500
dim = 64
feature_extractor = FeatureExtractor(num_features=num_features,
                                     output_features=output_features,
                                     dim=dim)                              
```

次に，NN部分で得た特徴量をガウス過程に投げるためのクラス `GPRegressionLayer` を構築します．

```python

class GPRegressionLayer(ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = CholeskyVariationalDistribution(
            inducing_points.size(0))
        variational_strategy = WhitenedVariationalStrategy(
            self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(GPRegressionLayer, self).__init__(variational_strategy)
        self.mean_func = gpytorch.means.ConstantMean()
        self.covar_func = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_func(x)
        covar_x = self.covar_func(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

```

`CholeskyVariationalDistribution`は変分パラメータの共分散行列をコレスキー分解することで誘導点をreparametrizeするためのクラス， `WhitenedVariationalStrategy` は変分パラメータ最適化の際にwhiteningと呼ばれる処理を行うためのクラスです．ここで誘導点とは，ガウス過程を確率的変分推論で近似する手法でに必要な仮想的に入力点を代表するような点の集合を意味します（ガウス過程における確率的変分推論では，誘導点の配置の最適化を行っています）．

また，今回はカーネル関数に `RBFKernel` を使っていますが，`MaternKernel` など他のカーネルも存在するので，適宜問題に合わせて使っていただければと思います．

最後に，`FeatureExtractor` と `GPRegressionLayer` を連結させます

```python

class SVDKL(gpytorch.Module):
    def __init__(self, inducing_points, feature_extractor):
        super(SVDKL, self).__init__()
        self.feature_extractor = feature_extractor
        self.inducing_points = inducing_points
        self.gp_layer = GPRegressionLayer(self.inducing_points)

    def forward(self, x):
        features = self.feature_extractor(x)
        res = self.gp_layer(features)
        return res
```

`feature_extractor` で特徴量を取得した後，`gp_layer` でガウス過程に投げます．

最後に，`SVDKL` を定義します．

```python
num_inducing_points = 200
inducing_loader = DataLoader(train_dataset[:num_inducing_points],
                                 batch_size=num_inducing_points)
inducing_points = list(inducing_loader)[0]

inducing_points = feature_extractor(inducing_points)
model = SVDKL(inducing_points=inducing_points,
              feature_extractor=feature_extractor,
              output_features=output_features)
```

`inducing_points` は初期化用の誘導点です．誘導点はガウス過程に投げる前段階で作る必要があるので，`num_inducing_points` 個分のデータを `feature_extractor` に通します．
普通なら素直に`Tensor`を渡してもよいのですが，PyTorch Geometric の`Set2Set` が`torch_geometric.data`で定義されている`Batch`型を要求するため，わざとこのように`inducing_points` 専用のDataLoaderを定義しています．

ここまでで各種モデルの定義が終ったので，学習部分です． 

```python
likelihood = gpytorch.likelihoods.GaussianLikelihood().cuda()
mll = gpytorch.mlls.VariationalELBO(likelihood,
                                    model.gp_layer,
                                    num_data=len(train_dataset),
                                    combine_terms=False)

i = 0
max_iter = 6
with gpytorch.settings.max_cg_iterations(50):
    for _ in range(max_iter):
        for data in train_loader:
            data = data.to(device)
            optimizer.zero_grad()
            output = model(data)
            log_lik, kl_div, log_prior = mll(output, data.y)
            loss = -(log_lik - kl_div + log_prior)
            if i % 10 == 0:
                print(loss.item())
            optimizer.step()
            loss.backward()
```
`gpytorch.settings.max_cg_iterations` では，BBMMの行列演算に必要なConjugate Gradientsに必要な試行回数を指定しています．


## 結果
QM9データセットに対して学習を行い，テストデータ20000件に対しMAEが1.19でした．同様の条件で行ったNeural Message Passingのみの学習だとMAEは0.80だったので，このような複雑なニューラルネットに対しては学習に不安定な部分がありそうです...あるいは自分の実装に問題があるかもしれません．
現時点で[公式github](https://github.com/cornellius-gp/gpytorch)のexamplesを見たり試した限りでは，シンプルなニューラルネットに対しては動くようですが試行によっては学習があまり進まない場合もあったので，今後の発展に期待です．


## 参考文献
[1]: Jacob R. Gardner, Geoff Pleiss, David Bindel, Kilian Q. Weinberger, Andrew Gordon Wilson, "GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration." in NIPS, 2018.
[2]: Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, Eric P. Xing, "Stochastic Variational Deep Kernel Learning." in NIPS 2016.
[3]: Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, George E. Dahl, "Neural Message Passing for Quantum Chemistry." in ICML 2017.