高次元ベイズ最適化は1行追加で精度が上がるかもしれない

[株式会社Enginee](https://www.enginee.co.jp)の@__Kat__です。最近日傘買いました。



### TL;DR

高次元でのベイズ最適化は多くの手法があるが、ハイパーパラメータ一つ変更するだけで十分精度が出るかもしれない

:::note info
ガウス過程ベースのベイズ最適化の話です。OptunaでTPEとか使ってる場合は関係ないです。
:::


2024/09/19追記
---
BoTorchの[v0.12](https://github.com/pytorch/botorch/releases/tag/v0.12.0)でデフォルトの事前分布として追加されたようです

## ベイズ最適化
ベイズ最適化は基本的に
1. 関数のモデリング
1. 獲得関数を最適化して次の探索点を決定する
を繰り返します。

モデリングはガウス過程やTPE(Optuna)が有名ですが、決定木やBNNなどの例もあります。


### ガウス過程
学習データ$(X_{1,n},y_{1:n})$があったとき、ガウス過程の事後分布は次式で表現できます。
```math
\boldsymbol{f}^*| X_{1:n},\mathbf{y}_{1:n},X^* \sim N(\boldsymbol \mu_n(X^*),V_n(X^*))
```
```math
\displaylines{
 \boldsymbol \mu_n(X^*) &= V(X^*,X_{1:n})(V(X_{1:n},X_{1:n})+\sigma^2I)^{-1}(\mathbf{y}_{1:n}) - \boldsymbol \mu(X_{1:n})) + \boldsymbol \mu(X^*)\\
    V_n(X^*) &= K(X^*,X^*) - K(X^*,X_{1:n})(K(X_{1:n},X_{1:n})+\sigma^2I)^{-1}K(X_{1:n},X^*)
}
```


https://qiita.com/nouyakko/items/e32275b0a82f9145f37f

獲得関数は最も一般的なものはExpected Improvementであり、とくにガウス過程を用いた場合は以下のように解析的に計算できます。

```math
\begin{align}
a_{EI}(\mathbf{x}) &= \mathbb{E}[(y-y^*)^+| x_{1:n},y_{1:n}]\\
&= \left ( \mu(\mathbf{x}) - y^*\right )\Phi(Z) + \sigma(\mathbf{x})\phi(Z)
\end{align}

```

ガウス過程のハイパーパラメータは、例えばRBFカーネルなら以下の$\theta_0,\theta_1,\theta_2$となります。
ARDを利用する場合は次元毎に$\theta_1$(lengthscale)が設定されるため、ハイパーパラメータもその分増えます。

```math
k(x,x') = \theta_0 exp(-  \frac{||x-x'||^2}{\theta_1}) + \theta_2
```

このハイパーパラメータは手動で決める必要があるように見えますが、周辺尤度が計算できるため勾配法を用いて最適化可能です。一般的にハイパーパラメータはMLE, MAP, MCMCのいずれかで決定されます。BoTorchはデフォルトではMAP推定が採用されています。


https://arxiv.org/abs/1807.02811

### 高次元
ベイズ最適化は一般的に20次元程度までしか使えない、などと言われます。次元が大きくなるとガウス過程のフィッティングの精度が下がると同時に、獲得関数の最適化がうまく進まないからです。

　そのため高次元でも最適化が上手く進むような研究がされているのですがベストプラクティスがあるわけではなく、多くの手法が登場した結果高次元に特化したサーベイ論文も最近出ているほどです。[^1]
 
図を見ると分かる通り、魔境と化しています。
<img width="500" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1793316/91f2ca23-049f-1f98-7c70-1a7772e15be8.png">

このサーベイ論文では、手法を
1. Variable selection
1. Additive Models
1. Trust regions
1. Linear embeddings
1. Non-Linear embeddings
1. Gradient information
1. Structured Space
1. Other

というふうに分類しています。

> この中のAdditive Modesl, Trsut Region, Embeddingsといった手法は松井先生の資料を参考にしていただけるとよく分かるかと思います。[^5]

そして、Otherで取り上げられている2つの論文では、**通常のガウス過程、もしくは軽微な修正で十分に性能が出ると報告されています。**[^2][^3]

## 本題
今回はこの内の[Vanilla Bayesian Optimization Performs Great in High Dimensions](https://arxiv.org/abs/2402.02229)を実装してみます。
この論文はすでに検証記事もあるので、こちらを参考にしてもいいです[^4]。


この論文では、学習時のlengthscaleの事前分布に対数正規分布をおいてMAP推定しろといっています。
(Dは次元の数。)
$$
l \sim LogNormal(\mu_0 + \frac{1}{2}\log{D}, \sigma_0)
$$

これをBoTorchで実装します。
BoTorchでは何も指定しない場合、MaternKernel + ARDが用いられ、lengthscaleの事前分布にGamma分布をおいてMAP推定しています。

そのためこれをLogNormalに変えればいいです。

```python
gp = SingleTaskGP(
    train_X=data_X,
    train_Y=data_y,
    train_Yvar=1e-5 * torch.ones_like(data_y),
    input_transform=Normalize(d=dim),
    outcome_transform=Standardize(m=1),
    # ---- 追加 --- 
    covar_module=ScaleKernel(
        base_kernel=MaternKernel(
            nu=2.5,
            ard_num_dims=dim,
            # -----   本質的な変更点はここだけ   ------
            lengthscale_prior=LogNormalPrior(mu_0 + np.log(dim)/ 2, sigma_0),
        ),
        outputscale_prior=GammaPrior(2.0, 0.15),
    )
)

mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_mll(mll)
```

## 検証
検証ブログと同様に、いかの関数を設定する($\boldsymbol{r}$は適当な一様乱数)
$$
f_\boldsymbol{r}(\boldsymbol{x}) = \sum(x_i - r_i)^2
$$

- 次元数を64
- 初期の学習データ数を500
としたとき、ランダムにサンプリングした100個の点を予測させます。

### 通常のガウス分布によるフィッティング
横軸が真の値、縦軸が予測値
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1793316/b1403bbb-c87c-4765-ce48-af9519f16977.png)

### 対数正規分布を利用したガウス分布によるフィッティング
横軸が真の値、縦軸が予測値

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1793316/bdb5d5bf-bc08-d5e2-5bf7-cc97e45eb9c6.png)


ずいぶんとフィッティングの精度が上がりました。

この流れでベイズ最適化もしました

### 通常のガウス分布によるベイズ最適化
横軸がiteration、縦軸がその地点での最大値、95%CI

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1793316/69be18ef-27fa-6bc4-0ba3-e225f4c3b950.png)

### 対数正規分布を利用したガウス分布によるベイズ最適化
横軸がiteration、縦軸がその地点での最大値、95%CI

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1793316/aeffd66f-372d-bbc4-9bcb-b788c027ab6f.png)


<details><summary>コードの一部</summary>

```python
# ベイズ最適化
dim = 64
r = torch.rand(dim,dtype=torch.float64).to(device)

initial_points_num = 500
iter_num=10
func = f
sample_num =100
sample_point_num = 100
bounds = torch.stack([torch.zeros(dim), torch.ones(dim)]).to(torch.double).to(device)

batch = 10

def f(X):
  """
  function f = \sum (x_i-r_i)^2
  """
  d = X.shape[-1]
  return torch.sum((X-r)**2,dim=1)

  
all_best = []
for _ in tqdm_notebook(batch):
  initial_X = torch.rand([initial_points_num,dim],dtype=torch.float64).to(device)
  initial_y = func(initial_X).unsqueeze(1)

  train_X = initial_X
  train_y = initial_y
  best = [train_y.max().item()]

  for i in tqdm_notebook(range(iter_num)):
    gp = SingleTaskGP(
      train_X=train_X,
      train_Y=train_y,
      input_transform=Normalize(d=dim),
      outcome_transform=Standardize(m=1),
    )
    # gp.covar_module.base_kernel.lengthscale_prior = LogNormalPrior(mu_0 +np.log(n_dim)/2, sigma_0)

    mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
    fit_gpytorch_mll(mll)
    # EI -> logEIが推奨されている
    logEI = LogExpectedImprovement(model=gp, best_f=train_y.max())
    candidate, acq_value = optimize_acqf(
      logEI, bounds=bounds, q=1, num_restarts=5, raw_samples=50,
    )

    y = func(candidate).unsqueeze(0)

    train_X = torch.cat([train_X, candidate], dim=0)
    train_y = torch.cat([train_y, y], dim=0)

    best.append(train_y.max().item())
  all_best.append(best)



plt.errorbar(range(iter_num+1), np.mean(all_best,axis=0), yerr=1.96 * np.std(all_best,axis=0) / np.sqrt(batch))
plt.show()
```

```python
# 事後分布からのサンプル部分
pred = gp.posterior(X).sample(torch.Size([sample_num]))
pred= pred.squeeze(-1)
mean = pred.mean(0)
ci = 1.96 * pred.std(0) / np.sqrt(sample_num)

values = torch.stack([y,mean,ci],dim=1).detach().numpy()
values = values[np.argsort(values[:,0])]

## plot y = x
plt.plot(values[:,0],values[:,0],linestyle="dashed",label='y=x', )
plt.errorbar(values[:,0],values[:,1], yerr=values[:,2],label="", capsize=3, fmt='.', markersize=10, ecolor="black", elinewidth=0.5, markeredgecolor="black", color="w")

plt.legend()
plt.show()
```
</details>

---

Engineeでは機械学習エンジニアを募集してるので興味のある方は是非お問い合わせください！

https://www.enginee.co.jp/

[^1]: [A survey and benchmark of high-dimensional Bayesian optimization of discrete sequences](https://arxiv.org/abs/2406.04739)
[^2]: [Vanilla Bayesian Optimization Performs Great in High Dimensions](https://arxiv.org/abs/2402.02229)
[^3]: [Standard Gaussian Process Can Be Excellent for High-Dimensional Bayesian Optimization](https://arxiv.org/abs/2402.02746v3)
[^4]: [Do Gaussian Processes scale well with dimension?](https://www.miguelgondu.com/blogposts/2024-03-16/when-does-vanilla-gpr-fail)
[^5]: [機械学習による適応的実験計画](https://speakerdeck.com/kmatsui/ji-jie-xue-xi-niyorushi-ying-de-shi-yan-ji-hua?slide=180)