BoTorch AX相关介绍

BoTorch 基本概念

最常用的是Ax可以直接跳转第五章去看

1 概述（翻译自官方文档）

1.2 黑盒优化

在高层次上，贝叶斯优化 (BayesOpt) 的基本问题是最大化一些评估成本高的黑盒函数$F$。换句话说，我们无法访问函数形式$F$我们唯一的办法是在一系列测试点上评估$F$，希望在少量评估后确定一个接近最优的值。在许多设置中，函数值没有被准确地观察到且只有带噪声观察结果可以使用。

贝叶斯优化是一种自适应地选择这些测试点（或要并行评估的测试点批次）的通用方法，它允许在具有良好观察性能的区域和高不确定性区域之间对$F$进行原则性权衡评估。

1.3 贝叶斯优化

1.3.1 模型

为了在少量的评估中优化$f$，我们需要一种方法来推断我们对$f$在我们尚未评估的点上的样子的看法。 在贝叶斯优化中，这被称为代理模型。 重要的是，代理模型应该能够以点$x$处函数值 $f(x)$的后验分布的形式量化其预测的不确定性。

$f$的代理模型通常是高斯过程 (Gaussian Process,GP)，在这种情况下，任何有限点集合上的后验分布都是多元正态分布。高斯过程由均值函数 $\mu (x)$和协方差核 $k(x,x^{\prime })$ 指定，其中均值向量$(\mu (x_0),\dots ,u(x_k))$ 和协方差矩阵 $\sum$Σ其中${\sum }_{ij}=k(x_i,x_j)$可以为任何点集 $(x_1,\dots x_k)$计算。 对$f$使用高斯过程代理模型意味着我们假设 $(f(x_1),\dots ,f(x_k))$ 是多元正态且平均向量和协方差矩阵由$\mu (x)$ 和$k(x,x^{\prime })$ 确定。

BoTorch 为 GPyTorch 提供一流的支持，GPyTorch 是一个用于在 PyTorch 中实现的可扩展高斯过程和贝叶斯深度学习的包。

虽然高斯过程是一种非常成功的建模方法，但 BoTorch 对基于蒙特卡洛采样的采集功能的支持使得也可以直接使用其他模型类型。 特别是，BoTorch 对使用哪种模型没有特别的假设，只要能够在给定输入$x$的情况下从输出的后验生成样本。

1.3.2 后验

后验代表模型对某个点（或一组点）的函数值的“信念”，基于它已经训练过的数据。也就是说，输出的后验分布取决于迄今为止观察到的数据。

当使用高斯过程模型时，后验以多元高斯显式给出（由其均值和协方差矩阵完全参数化）。在其他情况下，后验可能隐含在模型中，并且不容易通过一小组参数来描述。

BoTorch 通过提供一个简单的PosteriorAPI 来抽象出后验的特定形式，该API只需要实现一种rsample()从后验采样的方法。

1.3.3 采集函数

采集函数是用于评估多个设计点之一的有用性的启发式方法，以实现最大化底层黑盒函数的目标。

其中一些采集函数在高斯后验下具有封闭形式的解决方案，但其中许多（尤其是在评估多个并行点的联合值时）没有。在后一种情况下，可以求助于使用蒙特卡洛 (MC) 采样来近似采集函数。

BoTorch 支持基于分析和（准）蒙特卡罗的采集功能。它提供了一个AcquisitionFunction从特定类型中抽象出来的 API，以便可以对相同的对象执行优化。

评估蒙特卡洛采集函数

使用蒙特卡洛采样来评估采集函数背后的想法很简单：我们不是计算后验的（难以处理的）期望，而是从后验采样并使用样本平均值作为近似值。

目标

为了在基于蒙特卡罗的采集函数的情况下提供额外的灵活性，BoTorch 提供了通过Objective模块转换模型输出的选项，该模块返回传递给采集函数的一维输出。该类MCAcquisitionFunction默认其目标为IdentityMCObjective，它仅返回模型输出的最后一个维度。因此，对于直接对黑盒函数建模的单输出高斯过程的标准用例$F$，不需要特殊的目标。

1.4 重新参数化技巧

重新参数化技巧（参见例如^{1 2}）可用于将后验分布写为辅助随机变量$ϵ$的确定性变换.。例如，具有均值$\mu$和标准差 $\sigma$的正态分布随机变量 $X$与$\mu +\sigma ϵ$ 具有相同的分布，其中 $ϵ$ 是标准正态分布。因此，关于 $X$的期望可以使用来自$ϵ$ 的样本来近似。 在$\mu$和$\sigma$是优化问题的参数的情况下，可以使用单组“基础样本”计算目标在不同 μ 和 σ 值下的蒙特卡罗近似值。 重要的是，这种重新参数化允许通过样本反向传播梯度，这使得基于蒙特卡罗的采集函数能够相对于候选点进行自动微分。

在 BoTorch 中，基本样本是使用 MCSampler对象构建的，该对象提供了一个允许不同采样技术的接口。IIDNormalSampler使用独立的标准正态抽取，而SobolQMCNormalSampler使用准随机、低差异的“Sobol”序列作为均匀样本，然后将其转换为构造准正态样本。 Sobol 序列比 i.i.d. 分布更均匀。 均匀样本，并倾向于提高积分/期望的蒙特卡罗估计的收敛速度。 我们发现 Sobol 序列大大提高了基于蒙特卡罗的采集功能的性能，因此默认使用SobolQMCNormalSampler。

2 模型

模型在贝叶斯优化 (BO) 中起着至关重要的作用。模型用作要优化的实际底层黑盒函数的代理函数。在 BoTorch 中，一个Model将一组设计点映射到其输出在设计点上的后验概率分布。

在贝叶斯优化中，使用的模型传统上是高斯过程 (GP)，在这种情况下，后验分布是多元正态分布。虽然 BoTorch 支持许多高斯过程模型，但 BoTorch 没有假设模型是高斯过程或后验是多元正态的。除了 botorch.acquisition.analytic模块中的一些分析采集功能外，BoTorch 基于蒙特卡罗的采集功能与任何符合Model接口的模型兼容，无论是用户实现的还是提供的。

在底层，BoTorch 模型是实现轻量级Model接口的 PyTorch Modules。使用高斯过程时，GPyTorchModel为方便地包装 GPyTorch 模型提供了一个基类。

用户可以扩展Model并GPyTorchModel生成自己的模型。

2.1 术语

2.1.1 多输出和多任务（Multi-Output and Multi-Task）

一个模型Model（如在 BoTorch 对象中）可能有多个输出、多个输入，并且可以利用不同输入之间的相关性。 BoTorch 使用以下术语来区分这些模型类型：

多输出模型：具有多个输出的模型。大多数 BoTorch 模型都是多输出的。
多任务模型：使用输入/观察的逻辑分组的模型（如在底层过程中）。例如，可能有多个任务，其中每个任务具有不同的保真度。在多任务模型中，不同输出之间的关系被建模，具有跨任务的联合模型。
请注意以下事项：

多任务 (MT) 模型可能是也可能不是多输出模型。例如，如果多任务模型使用不同的任务进行建模，但只输出其中一项任务的预测，则它是单输出的。
相反，多输出 (MO) 模型可能是也可能不是多任务模型。例如，独立建模不同输出而不是建立联合模型的多输出模型不是多任务的。
如果一个模型两者兼而有之，我们将其称为多任务多输出 (MTMO) 模型。

2.1.2 噪声：同方差、固定和异方差（Noise: Homoskedastic, fixed, and heteroskedastic）

可以通过几种不同的方式处理噪声：

• Homoskedastic：噪声不作为输入提供，而是推断出来的，具有不依赖于$X$ 的恒定方差。许多模型，例如 SingleTaskGP，采用这种方法。 如果您知道您的观察结果很嘈杂，但不知道有多少嘈杂，请使用这些模型。
• fixed：噪声作为输入提供，不适合。 在像 FixedNoiseGP这样的“固定噪声”模型中，由于尚未建模，因此无法在样本外预测噪声。 如果您对观测值中的噪声进行了估计（例如，观测值可能是单个样本的平均值，在这种情况下，您将提供均值作为观测值，将均值的标准误差作为噪声估计值），或者如果您知道您的观测值，请使用这些模型 观察是无噪音的（通过零噪音水平）。
• Heteroskedastic：噪声作为输入提供，并被建模以允许预测样本外的噪声。 像HeteroskedasticSingleTaskGP这样的模型采用了这种方法。

2.1.3 标准 BoTorch 模型

BoTorch 提供了几个 GPyTorch 模型来涵盖大多数标准 贝叶斯优化 用例：

单任务高斯过程

这些模型对所有输出使用相同的训练数据，并假设给定输入的输出具有条件独立性。如果每个输出需要不同的训练数据，请改用ModelListGP。

• SingleTaskGP：单任务精确高斯过程，推断同方差噪声水平（无噪声观察）。
• FixedNoiseGP：与SingleTaskGP不同的单任务精确高斯过程，它使用固定的观察噪声水平。它需要噪声观察。
• HeteroskedasticSingleTaskGP：与 SingleTaskGP 和 FixedNoiseGP 不同的单任务精确高斯过程，它使用额外的内部 高斯过程模型对异方差噪声进行建模。它需要噪声观察。
• MixedSingleTaskGP：支持混合搜索空间的单任务精确高斯过程，它结合了离散和连续特征。
• SaasFullyBayesianSingleTaskGP：具有SAAS先验的完全贝叶斯单任务高斯过程。该模型适用于样本高效的高维贝叶斯优化。

单任务高斯过程模型列表

ModelListGP：一个多输出模型，其中结果被独立建模，给定任何类型的单任务高斯过程列表。 当并非所有输出都使用相同的训练数据时，应使用此模型。

2 后验概率

BoTorch 后验对象Posterior是一个抽象层，将使用的特定模型与采集功能的评估（和后续优化）分开。在最简单的情况下，后验是围绕来自torch.distributions（或gpytorch.distributions）的显式分布对象的轻量级包装器。但是，BoTorch 后验可以是任何分布（甚至是隐式分布），只要可以从该分布中采样。例如，后验可以由通过神经网络映射的一些基本分布隐式表示。

虽然分析采集函数假设后验是多元高斯分布，但基于蒙特卡洛 (MC) 的采集函数不对潜在分布做出任何假设。相反，基于蒙特卡罗的采集功能只要求后验可以通过rsample方法生成样本。只要后验实现了后验接口Posterior，就可以与基于MC的采集功能一起使用。此外，请注意，基于梯度的采集函数优化需要能够通过蒙特卡罗样本反向传播梯度。

对于基于 GPyTorch 且后验分布为多元高斯分布的高斯过程模型，应使用GPyTorchPosterior。

3 采集函数

采集函数是用于评估多个设计点之一的有用性的启发式方法，以实现最大化底层黑盒函数的目标。

BoTorch 支持基于分析和（准）蒙特卡罗的采集功能。它提供了一个从特定类型中抽象出来的通用 AcquisitionFunctionAPI，以便可以对相同的对象执行优化。

4 示例代码

1.将高斯过程模型拟合到数据

import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_model
from gpytorch.mlls import ExactMarginalLogLikelihood

train_X = torch.rand(10, 2)
Y = 1 - (train_X - 0.5).norm(dim=-1, keepdim=True)  # explicit output dimension
Y += 0.1 * torch.rand_like(Y)
train_Y = (Y - Y.mean()) / Y.std()

gp = SingleTaskGP(train_X, train_Y)
mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_model(mll);
1
2
3
4
5
6
7
8
9
10
11
12
13

2.构造采集函数

from botorch.acquisition import UpperConfidenceBound

UCB = UpperConfidenceBound(gp, beta=0.1)
1
2
3

3.优化采集函数

from botorch.optim import optimize_acqf

bounds = torch.stack([torch.zeros(2), torch.ones(2)])
candidate, acq_value = optimize_acqf(
    UCB, bounds=bounds, q=1, num_restarts=5, raw_samples=20,
)
1
2
3
4
5
6

5 使用带有 Ax 的自定义 BoTorch 模型

在本教程中，我们将解释如何在 Ax 的 botorch_modularAPI 中使用自定义 BoTorch 模型。 这使我们能够利用 Ax 的便利性来运行贝叶斯优化循环，同时保持建模的完全灵活性。获取函数和优化策略可以以几乎相同的方式进行交换。

Next cell 专门设置了一个装饰器来加速笔记本测试。 在整个教程中，您可以放心地忽略此单元格和装饰器的使用。

import os
from contextlib import contextmanager

from ax.utils.testing.mock import fast_botorch_optimize_context_manager
import plotly.io as pio
# Axe 使用 Plotly 生成交互式绘图。 这些非常适合查看和分析，
# 虽然它们也会导致文件很大，这对于 GH 中的文件来说并不理想。
# 将默认值更改为 `png` 会剥离交互式组件以解决此问题。
pio.renderers.default = "png"

SMOKE_TEST = os.environ.get("SMOKE_TEST")
NUM_EVALS = 10 if SMOKE_TEST else 30


@contextmanager
def dummy_context_manager():
    yield


if SMOKE_TEST:
    fast_smoke_test = fast_botorch_optimize_context_manager
else:
    fast_smoke_test = dummy_context_manager
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

5.1 实现自定义模型

在本教程中，我们实现了一个非常简单的 gpytorch Exact GP 模型，它使用 RBF 内核（带有 ARD）并推断（同方差）噪声水平。

模型定义很简单——这里我们实现了一个 gpytorch ExactGP，它也继承自 GPyTorchModel——这添加了 botorch 在其各种模块中期望的所有 api 调用。

注意：botorch 还允许实现其他自定义模型，只要它们遵循最小模型 API。

from botorch.models.gpytorch import GPyTorchModel
from botorch.utils.datasets import SupervisedDataset
from gpytorch.distributions import MultivariateNormal
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.means import ConstantMean
from gpytorch.models import ExactGP


class SimpleCustomGP(ExactGP, GPyTorchModel):

    _num_outputs = 1  # 通知 GPyTorchModel API

    def __init__(self, train_X, train_Y):
        super().__init__(train_X, train_Y.squeeze(-1), GaussianLikelihood())
        self.mean_module = ConstantMean()
        self.covar_module = ScaleKernel(
            base_kernel=RBFKernel(ard_num_dims=train_X.shape[-1]),
        )
        self.to(train_X) 

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

5.2 在 Ax 中实例化BoTorchModel

Ax 中的BoTorchModel封装了代理（在 BoTorch 中通常称为模型Model）和采集函数。 在这里，我们将只指定自定义代理，让 Ax 选择默认的获取函数。

大多数模型都应与 Ax 中的基本代理Surrogate一起使用，但 BoTorchModelListGP除外，它与ListSurrogate一起使用。 请注意，模型Model（例如SimpleCustomGP）必须实现construct_inputs，因为它用于构造从实验数据实例化模型实例所需的输入。

如果模型Model需要一组无法使用construct_inputs方法构造的复杂参数，则可以初始化Model并通过Surrogate.from_botorch(model=model, mll_class=) 提供它，替换Surrogate(...)以下。

from ax.models.torch.botorch_modular.model import BoTorchModel
from ax.models.torch.botorch_modular.surrogate import Surrogate


ax_model = BoTorchModel(
    surrogate=Surrogate(
        # 使用model类
        botorch_model_class=SimpleCustomGP,
        # 可选的，用于优化模型参数的 MLL 类
        # mll_class=ExactMarginalLogLikelihood,
        # 可选，模型构造函数的关键字参数字典
        # model_options={}
    ),
    # 可选，要使用的采集函数类-见自定义采集教程
    # botorch_acqf_class=qExpectedImprovement,
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

5.3 与ModelBridge结合

Ax 中的Model需要ModelBridge与Experiment交互。ModelBridge接受Experiment提供的输入并将它们转换为Model预期的输入。 对于BoTorchModel，我们使用TorchModelBridge。 用法如下：

from ax.modelbridge import TorchModelBridge
model_bridge = TorchModelBridge(
    experiment: Experiment,
    search_space: SearchSpace,
    data: Data,
    model: TorchModel,
    transforms: List[Type[Transform]],
    # And additional optional arguments.
)
# To generate a trial
trial = model_bridge.gen(1)
1
2
3
4
5
6
7
8
9
10
11

对于 Modular BoTorch 接口，我们可以将BoTorchModel和TorchModelBridge的创建合并为一个步骤，如下所示：

from ax.modelbridge.registry import Models
model_bridge = Models.BOTORCH_MODULAR(
    experiment=experiment,
    data=data,
    surrogate=Surrogate(SimpleCustomGP),  # Optional, will use default if unspecified
    # Optional, will use default if unspecified
    # botorch_acqf_class=qNoisyExpectedImprovement,  
)
# To generate a trial
trial = model_bridge.gen(1)
1
2
3
4
5
6
7
8
9
10

6 使用 Ax 中的自定义模型优化 Branin 函数（常用，自动优化）

我们将通过 Service API（更简单、更易于使用）和 Developer API（高级、更可定制）来演示这一点。

6.1 使用Service API 进行优化

为了自定义在 Service API 中创建候选者的方式，我们需要构造一个新的GenerationStrategy并将其传递给AxClient。

from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models


gs = GenerationStrategy(
    steps=[
        # Quasi-random initialization step
        GenerationStep(
            model=Models.SOBOL,
            num_trials=5,  # How many trials should be produced from this generation step
        ),
        # Bayesian optimization step using the custom acquisition function
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,  # No limitation on how many trials should be produced from this step
            # For `BOTORCH_MODULAR`, we pass in kwargs to specify what surrogate or acquisition function to use.
            model_kwargs={
                "surrogate": Surrogate(SimpleCustomGP),
            },
        ),
    ]
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

6.2 设置实验

为了使用我们刚刚创建的GenerationStrategy，我们将把它传递给AxClient。

import torch
from ax.service.ax_client import AxClient
from ax.service.utils.instantiation import ObjectiveProperties
from botorch.test_functions import Branin


# Initialize the client - AxClient offers a convenient API to control the experiment
ax_client = AxClient(generation_strategy=gs)
# Setup the experiment
ax_client.create_experiment(
    name="branin_test_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            # It is crucial to use floats for the bounds, i.e., 0.0 rather than 0.
            # Otherwise, the parameter would be inferred as an integer range.
            "bounds": [-5.0, 10.0],
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 15.0],
        },
    ],
    objectives={
        "branin": ObjectiveProperties(minimize=True),
    },
)
# Setup a function to evaluate the trials
branin = Branin()


def evaluate(parameters):
    x = torch.tensor([[parameters.get(f"x{i+1}") for i in range(2)]])
    # The GaussianLikelihood used by our model infers an observation noise level,
    # so we pass an sem value of NaN to indicate that observation noise is unknown
    return {"branin": (branin(x).item(), float("nan"))}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

6.3 运行BO循环

with fast_smoke_test():
    for i in range(NUM_EVALS):
        parameters, trial_index = ax_client.get_next_trial()
        # Local evaluation here can be replaced with deployment to external system.
        # evaluate 可以表示模型训练完成后的返回评估值的函数
        ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameters))
1
2
3
4
5
6

[INFO 05-10 20:27:10] ax.service.ax_client: Generated new trial 0 with parameters {'x1': -0.175425, 'x2': 9.532596}.
[INFO 05-10 20:27:10] ax.service.ax_client: Completed trial 0 with data: {'branin': (30.013487, nan)}.
[INFO 05-10 20:27:10] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 7.182822, 'x2': 10.563368}.
[INFO 05-10 20:27:10] ax.service.ax_client: Completed trial 1 with data: {'branin': (103.023666, nan)}.
[INFO 05-10 20:27:10] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 0.337159, 'x2': 9.081922}.
[INFO 05-10 20:27:10] ax.service.ax_client: Completed trial 2 with data: {'branin': (32.049171, nan)}.
1
2
3
4
5
6

6.4 查看评估的试验

ax_client.get_trials_data_frame()
1

parameters, values = ax_client.get_best_parameters()
print(f"Best parameters: {parameters}")
print(f"Corresponding mean: {values[0]}, covariance: {values[1]}")
1
2
3

Best parameters: {'x1': 9.53311206835773, 'x2': 2.615867701958974}
Corresponding mean: {'branin': 0.39199786406546977}, covariance: {'branin': {'branin': 0.04835207801718367}}
1
2

6.5 绘制响应面和优化进度

from ax.utils.notebook.plotting import render

render(ax_client.get_contour_plot())

1
2
3
4

best_parameters, values = ax_client.get_best_parameters()
best_parameters, values[0]
1
2

({'x1': 9.53311206835773, 'x2': 2.615867701958974},
 {'branin': 0.39199786406546977})
1
2

render(ax_client.get_optimization_trace(objective_optimum=0.397887))
1

---

1. D. P. Kingma, M. Welling. Auto-Encoding Variational Bayes. ICLR, 2013 ↩︎

2. D. J. Rezende, S. Mohamed, D. Wierstra. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML, 2014. ↩︎