概要
ハイパラ探索や材料工学などの分野で利用されているベイズ最適化ですが，一般的には説明変数が高次元（大体20次元以上くらい？）になってくると，パフォーマンスを発揮できないことが多いとされています．
高次元ベイズ最適化のためのアプローチとしては，
- 高次元を低次元に埋め込む (REMBO)
- 探索する次元を制限する (LINEBO)
などがありますが，今回は下記論文で提案されたSAASBO(スパース軸整列部分空間BO)を試してみようと思います．
High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspaces
目的変数に対して重要度が階層的になっていると仮定します
(例)
$\lbrace x_3, x_{52} \rbrace$は重要な特徴
$\lbrace x_7, x_{14}, x_{31}, x_{72} \rbrace$は中程度に重要な特徴
その他はわずかに寄与
上記の仮定から，サンプル効率を上げるためにスパース性を誘導するような事前分布を導入し，適度に簡素なモデルを用いてベイズ最適化を行うことがSAASBOのアイデアです．
SAAS事前分布
カーネルのハイパラ（特にlength scale）に対し階層的かつスパース性を誘導するような事前分布としてSAAS(Sparse Axis-Aligned Subspaces) を導入します
モデルの詳細は以下のようになっています
- [observations] : $\mathbf{y}\sim \mathcal{N}(\mathbf{f}, \sigma^2\mathbb{1}_N)$
- [function values] : $\mathbf{f} \sim \mathcal{N}(\boldsymbol{0}, K_{\mathbf{xx}}^{\psi}) \quad \psi=\lbrace \rho_{1:d}, \sigma_k^2 \rbrace$
- [kernel variance] : $\sigma^2_k \sim \mathcal{LN}(0, 10^2)$
- [global shrinkage] : $\tau \sim \mathcal{HC}(\alpha)$
- [length scales] : $\rho_i \sim \mathcal{HC}(\tau) \quad i=1, \ldots, D$
本論文中の実験では，カーネル関数にはRBFカーネルを用いていますが，Matern-5/2カーネルなども選択可能です．また，$\mathcal{LN}$, $\mathcal{HC}$はそれぞれ対数正規分布，ハーフ・コーシー分布を表しています．
カーネルのlength scale $\rho_i$の事前分布にパラメータ$\tau >0$のハーフ・コーシー分布を採用しており，$\tau$の事前分布にもパラメータ$\alpha$(デフォルトで0.1)のハーフ・コーシー分布を採用しています．
つまり
\begin{aligned}
p(\tau | \alpha) &\propto (\alpha^2 + \tau^2)^{-1} \mathbb{1}(\tau > 0)
\\
p(\rho_i | \tau) &\propto (\tau^2 + \rho_i^2)^{-1} \mathbb{1}(\rho_i > 0)
\end{aligned}
です．
ハーフ・コーシー分布は以下のような形をしています．
ハーフ・コーシー事前分布によって$\tau$は$0$付近に集中し，$\tau$が$0$付近においては$\rho_i$も$0$付近に集中します．
つまり，事前分布においては，length scale $\rho_i$はほとんどの次元で$0$付近に集中し，「オフ」の状態になります（周りに影響を与えない）
このように大域的なスパース性が$\tau$によって制御されています．
また，ハーフ・コーシー分布はheavy tail (裾の思い分布) であるため，観測データ$\mathbf{y}$に十分な情報があれば，事後分布としては$\tau$の値は大きくなりやすく，スパース性が減少します．
→ より多くの$\rho_i$が$0$付近から脱出
事前分布ではスパース性を誘導し，より多くのデータが蓄積されるにつれて, より多くの$\rho_i$が$0$付近を脱出(「オン」の状態)するように設計されています.
→ 探索初期でもオーバーフィットしにくく，効率よく探索できる．
探索候補点の取得
上記で定義したモデルは潜在空間の次元が多く，かつ非線形です．
そこで，ここではNo-U-Turn (NUT) Sampler [Hoffman and Gelman, 2014]を用います.
NUT Samplerを用いて，$L$個の事後分布からのサンプル$\lbrace \psi_l \rbrace_{l=1}^L$を取得します．
各$\psi_l$についての推論は特に変わらず，以下のようになります．
\begin{aligned}
\mu_{\mathbf{f}}\left(\mathbf{x}^*\right) & =k_{* \mathbf{X}}^{\psi^{\mathrm{T}}}\left(K_{\mathbf{X} \mathbf{X}}^\psi+\sigma^2 \mathbb{1}_N\right)^{-1} \mathbf{y} \\
\sigma_{\mathbf{f}}\left(\mathbf{x}^*\right)^2 & =k_{* *}^\psi-k_{* \mathbf{X}}^\psi{ }^{\mathrm{T}}\left(K_{\mathbf{X} \mathbf{X}}^\psi+\sigma^2 \mathbb{1}_N\right)^{-1} k_{* \mathbf{X}}^\psi
\end{aligned}
獲得関数は以下のように，各$\psi_l$についてのEI(Expected Improvement) の平均で定義します．
\mathrm{EI}\left(\mathbf{x} \mid \mathrm{y}_{\min },\left\{\psi_{l}\right\}\right) \equiv \frac{1}{L} \sum_{l=1}^{L} \mathrm{EI}\left(\mathbf{x} \mid \mathrm{y}_{\min }, \psi_{l}\right)
獲得関数を最大化する$\mathbf{x}_{\rm next}$を次の探索点とします
\mathbf{x}_{\text {next }}=\operatorname{argmax}_{\mathbf{x}} \operatorname{EI}\left(\mathbf{x} \mid y_{\min },\left\{\psi_{l}\right\}\right)
BoTorchでSAASBOを試す
BoTorchではSAASBOのためのSaasFullyBayesianSingleTaskGPが実装されています．
今回はSAASBOとナイーブなBOの比較を行ってみようと思います．
SAASBOの設定に合わせて，重要な説明変数とそうでない説明変数が存在している条件で実験を行います．
目的関数にはBranin function (以降Braninと表記)を用います．
Braninの最小値をベイズ最適化を用いて探索します．
Braninの説明変数は本来2次元であるため，この2次元を「重要な説明変数」とし，目的変数に寄与しない28次元を追加した30次元のベイズ最適化を行います．
from botorch.test_functions import Branin
branin = Branin()
def branin_d30(x):
"""
30次元の入力のうち，1,2次元目を用いてbraninの目的変数の値を取得
※ xは0-1正規化されているとする
"""
lb, ub = branin.bounds
return branin(lb + (ub - lb) * x[..., :2]).unsqueeze(-1)
以下でEIによるナイーブなベイズ最適化とSAASBOのoptimizerをそれぞれ実装しています．
(BoOptimzer, SaasBoOptimzer)
import torch
from botorch.models.transforms.outcome import Standardize
from botorch import fit_gpytorch_model
from botorch import fit_fully_bayesian_model_nuts
from botorch.models import SingleTaskGP
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.acquisition.monte_carlo import qExpectedImprovement
from botorch.optim import optimize_acqf
class BoOptimzer:
"""
ナイーブなBOの実行クラス
"""
def __init__(self):
self.model = None
self.sampler = SobolQMCNormalSampler(sample_shape=torch.Size([128]))
def fit_gp(self, X, y):
"""
GPのfitting
"""
self.model = SingleTaskGP(X, y, outcome_transform=Standardize(m=1))
mll = ExactMarginalLogLikelihood(self.model.likelihood, self.model)
fit_gpytorch_model(mll)
def get_candidates(self, X, y, direction='minimize'):
"""
探索候補点の取得
"""
if direction == 'minimize':
# 最小化タスクの場合は目的変数の符号反転
y = -y
self.fit_gp(X, y)
# 獲得関数定義(EI)
acq_func = qExpectedImprovement(model=self.model, sampler=self.sampler, best_f=y.max())
# 候補点取得
X_dim = X.size()[1]
bounds = torch.zeros(2, X_dim)
bounds[1] = 1
candidates, acq_values = optimize_acqf(acq_func, bounds=bounds, q=1, num_restarts=10, raw_samples=100)
return candidates
class SaasBoOptimzer:
"""
SAASBOの実行クラス
"""
def __init__(self):
self.model = None
self.sampler = SobolQMCNormalSampler(sample_shape=torch.Size([128]))
def fit_gp(self, X, y):
"""
GPのfitting
"""
self.model = SaasFullyBayesianSingleTaskGP(X,
y,
train_Yvar=torch.full_like(y, 1e-6),
outcome_transform=Standardize(m=1))
fit_fully_bayesian_model_nuts(self.model,
warmup_steps=256,
num_samples=128,
thinning=16,
disable_progbar=True)
def get_candidates(self, X, y, direction='minimize'):
"""
探索候補点の取得
"""
if direction == 'minimize':
# 最小化タスクの場合は目的変数の符号反転
y = -y
self.fit_gp(X, y)
# 獲得関数定義(EI)
acq_func = qExpectedImprovement(model=self.model, sampler=self.sampler, best_f=y.max())
# 候補点取得
X_dim = X.size()[1]
bounds = torch.zeros(2, X_dim)
bounds[1] = 1
candidates, acq_values = optimize_acqf(acq_func, bounds=bounds, q=1, num_restarts=10, raw_samples=100)
return candidates
今回，獲得関数にはEIのモンテカルロ獲得関数であるqExpectedImprovementを用いていますが，BoTorchのモンテカルロ獲得関数は基本的に最大化タスクを想定された設計になっているため，目的変数の符号を反転することで最小化タスクに対処しています．
(BoTorchのMC獲得関数で最小化タスク解きたいとき)
ナイーブなベイズ最適化ではfit_gpytorch_modelによってモデルのfittingを行いますが，SaasFullyBayesianSingleTaskGPを用いる場合はfit_fully_bayesian_model_nutsを用いてNUT Samplerによるfittingを行なっています．
シミュレーションのメインループは以下になります．
初期観測データとして同じ点を10点与え，以降はそれぞれの探索戦略に従って1点ずつ観測データを取得していきます．
from torch.quasirandom import SobolEngine
# 初期観測データを10点生成
Xs = SobolEngine(dimension=30, scramble=True, seed=11).draw(10)
ys = branin_d30(Xs)
print(f"Generate init date. best: {ys.min()}")
Xs_bo, Xs_saasbo = Xs, Xs
ys_bo, ys_saasbo = ys, ys
bests_bo = [ys.min()]
bests_saasbo = [ys.min()]
# ベイズ最適化による探索開始
for i in range(25):
# 候補点取得
X_cs_bo = BoOptimzer().get_candidates(Xs_bo, ys_bo)
X_cs_saasbo = SaasBoOptimzer().get_candidates(Xs_saasbo, ys_saasbo)
# 観測
y_cs_bo = branin_d30(X_cs_bo)
y_cs_saasbo = branin_d30(X_cs_saasbo)
# 観測データに追加
Xs_bo = torch.cat((Xs_bo, X_cs_bo))
ys_bo = torch.cat((ys_bo, y_cs_bo))
Xs_saasbo = torch.cat((Xs_saasbo, X_cs_saasbo))
ys_saasbo = torch.cat((ys_saasbo, y_cs_saasbo))
# best更新
bests_bo.append(ys_bo.min())
bests_saasbo.append(ys_saasbo.min())
print(f"Iter:{i}")
print(f"Bo. best: {ys_bo.min()}")
print(f"SaasBo. best: {ys_saasbo.min()}")
シミュレーション結果は以下のようになりました．
SAASBOに合わせた実験設定ではありますが，ナイーブなベイズ最適化よりもSAASBOの方が効率よく探索できていることが確認できました．
一方で，一回の探索にかかる時間はSAASBOの方がGPのfittingに時間がかかる分長かったです．(NUT Samplerは$O(N^3)$)