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

Sansan Advent Calendar 2019の10日目の記事です．

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

環境

• Python 3.6.5
• GPyTorch 0.3.6
• PyTorch 1.3.0
• PyTorch Geometric 1.3.2

GPyTorchとは

GPyTorchとは，ガウス過程周りの種々の手法を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から引用）．

コード

コードを踏まえて，使い方を解説したいと思います．コード全体はこちら．
以下は各種インポート．

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 として定義します．

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 を定義します．

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 を構築します．

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 を連結させます

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 を定義します．

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を定義しています．

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

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の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.