Thompson samplingによる離散バッチベイズ最適化

離散バッチベイズ最適化の需要と困難

材料科学のベイズ最適化において、1サイクルの実験で1条件のみを試すことは効率が悪く、1サイクルにつき複数条件を同時に実験したいという要望が多い。
このようなベイズ最適化をバッチベイズ最適化と呼び、BoTrochではq～～という獲得関数を利用することで対応することができる。
predictorが連続変数のときは獲得関数の最適化が容易であり、optimize_acqf関数により効率良く最適化できるが、材料科学においてはカテゴリ変数や離散値 (分子記述子など) がpredictorにある場合が多く、獲得関数の最適化が困難となる。
特にバッチベイズ最適化においては、考えられる条件の組合せ$N$通りからバッチ数$q$個を取り出した場合の全ての獲得関数を計算すると、$_NC_q$通りの計算が必要となり、容易に組合せ爆発を起こす。

簡易な対応方法

ベイズ最適化ライブラリBoTorchのDocumentでは、連続変数の最適化において、バッチ数が大きくなると逐次最適化法 (Fantasyモデルを用いた方法) が望ましいとの記載があり、この方法を離散最適化でも適用することで、離散バッチベイズ最適化に対応できると考えられる。
逐次最適化法は、明治大学 金子先生も紹介 しており、実装については金子先生のGithubを参照いただきたい。
GPや獲得関数の計算の部分をGPyTorch BoTorchに変更することで、さらに柔軟な対応が可能となる。

Thompson samplingによる対応

ガウス過程回帰から関数をバッチ数だけサンプリングし、各関数における最適値を次の実験候補とする方法。
この方法は、全条件の数 ($N$) に依存する $N \times N$ 行列の逆行列を計算するため、探索空間が広いとき計算困難とデメリットはあるものの、対応方法が少ない離散バッチベイズ最適化に対する解法の1つとして紹介する。
ベイズ最適化ライブラリPhysboでは、この逆行列計算のコストを低減しており、'BoTorch'で実装するよりも適用範囲が広がる。

実装 (全条件組合せ数 N=10000)

インポート

import sys
import physbo
import itertools
import numpy as np
import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.test_functions import Branin
import matplotlib.pyplot as plt
import seaborn as sns
import japanize_matplotlib

sys.path.append("src/")
import visualize

sns.set(rc={
    "savefig.bbox": "tight",
    "legend.frameon": False,
    "figure.figsize": (5, 5)
})
japanize_matplotlib.japanize()

visualize.py

import torch
from botorch.test_functions import Branin
import matplotlib.pyplot as plt

def plot_Branin(
    fig=None, 
    axes=None, 
    func=Branin(negate=True), 
    ncol=1, 
    nrow=1, 
    i=0
):
    
    if (fig is None) or (axes is None):
        fig = plt.figure()
        axes = []
    
    x1 = torch.linspace(-5, 10, 100)
    x2 = torch.linspace(0, 15, 100)
    X1, X2 = torch.meshgrid(x1, x2, indexing="ij")
    X = torch.stack([X1.reshape(-1), X2.reshape(-1)], dim=-1)
    
    Z = func(X.double()).reshape(100, 100).detach().numpy()
    
    optima = torch.tensor([
        [-torch.pi, 12.275],
        [torch.pi, 2.275],
        [9.42478, 2.475]
    ])
    
    axes.append(fig.add_subplot(nrow, ncol, i+1))
    contours = axes[i].contour(X1.numpy(), X2.numpy(), Z, levels=20, cmap="viridis")
    axes[i].clabel(contours, inline=True, fontsize=8)
    axes[i].set_aspect("equal", "box")
    
    return fig, axes

データの生成

func = Branin(negate=True)

bounds = torch.tensor([
    [-5, 0],
    [10, 15],
], dtype=torch.float64)

choices = torch.tensor(list(itertools.product(
    torch.linspace(bounds[0, 0], bounds[1, 0], 100), 
    torch.linspace(bounds[0, 1], bounds[1, 1], 100)
)))
y = func(choices)

train_id = np.random.choice(range(choices.shape[0]), 20, replace=False)
train_X = choices[train_id]
train_y = y[train_id].unsqueeze(-1)

最適化対象の関数の形状は以下の通り。

plot_Branin()
plt.show()

Physboによる実装

policy = physbo.search.discrete.policy(
    test_X=choices.numpy(), 
    initial_data=[train_id, y.numpy()[train_id]]
)

actions = policy.bayes_search(
    max_num_probes=5,
    num_search_each_probe=10, 
    simulator=None, 
    score="TS", 
    interval=0, 
    num_rand_basis=5000
)

fig = plt.figure()
ax = []

visualize.plot_Branin(fig, ax)
ax[0].scatter(*choices[actions].T, marker="*", s=200, c="r")
ax[0].scatter(*train_X.T, c="k")
plt.show()

BoTorch による実装

GPの学習

model = SingleTaskGP(train_X, train_y)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll(mll)

全候補条件に対するThompson sampling獲得関数の計算

model.eval()
with torch.no_grad():
    posterior = model.posterior(choices)
    samples = posterior.rsample(sample_shape=torch.Size([10])) 

idx = list(range(choices.shape[0]))
candidates = []

for batch in samples:
    batch = batch[idx, :]
    best_idx = batch.argmax()
    best_choice = choices[best_idx]
    idx.pop(best_idx)
    candidates.append(best_choice)

candidates = np.array(candidates)

fig = plt.figure()
ax = []

visualize.plot_Branin(fig, ax)
ax[0].scatter(*candidates.T, marker="*", s=200, c="r")
ax[0].scatter(*train_X.T, c="k")
plt.show()

サンプリングした各関数の予測値の分布

ncol = 5
nrow = 2

fig = plt.figure(figsize=4 * np.array([ncol * 4/3, nrow]))
axes = []
for i, sample in enumerate(samples):
    axes.append(fig.add_subplot(nrow, ncol, i+1))
    mappable = axes[i].scatter(*choices.T, c=sample, cmap="viridis")
    plt.colorbar(mappable, ax=axes[i])
    best_idx = sample.argmax()
    axes[i].scatter(*choices[best_idx].T, marker="*", c="r", s=500)
    axes[i].set_aspect("equal", "box")
plt.tight_layout()
plt.show()

2つのライブラリによる計算結果の比較

2つのライブラリによる計算結果を比較すると、Physboは最適条件付近の提案数が多いが、条件の多様性に乏しいような…。
全条件数$N$が100,000になると、Physboでは計算できるが、BoTorchはメモリオーバーフロー (使用した計算機のメモリは32Gb) により計算できなかった。
$N$が1,000,000になると、Physboでも計算不可。