从理论到代码：一文搞懂BoTorch/AX中的贝叶斯优化核心（采集函数、高斯过程详解）
贝叶斯优化（Bayesian Optimization, BO）作为黑盒函数优化的利器，在超参数调优、自动化实验设计等领域展现出强大潜力。而BoTorch和AX框架的组合，则为这一理论提供了工业级实现方案。本文将深入解析贝叶斯优化的核心组件——采集函数与高斯过程模型，并通过代码示例展示如何在实际项目中灵活运用这些工具。
1. 贝叶斯优化基础架构剖析
贝叶斯优化的核心思想是通过构建代理模型（Surrogate Model）来近似目标函数，再通过采集函数（Acquisition Function）指导下一步的评估点选择。这种"建模-评估-更新"的迭代过程，使其在少量评估次数下就能找到接近最优的解。
BoTorch作为PyTorch生态的贝叶斯优化库，其架构设计具有三个显著特点：
- 模块化设计：将高斯过程、采集函数、优化器等组件解耦
- 自动微分支持：基于PyTorch实现端到端的梯度计算
- 蒙特卡洛集成：通过采样处理复杂后验分布
典型的贝叶斯优化循环包含以下关键步骤：
# 伪代码展示贝叶斯优化流程 def bayesian_optimization(): # 初始化数据集 train_X, train_Y = initialize_data() for i in range(num_iterations): # 1. 训练高斯过程模型 gp = train_gp_model(train_X, train_Y) # 2. 构造采集函数 acq_func = construct_acquisition_function(gp) # 3. 优化采集函数获取下一个候选点 candidate = optimize_acquisition_function(acq_func) # 4. 评估目标函数并更新数据 new_y = evaluate_objective(candidate) train_X = torch.cat([train_X, candidate]) train_Y = torch.cat([train_Y, new_y])2. 高斯过程模型实战解析
高斯过程（Gaussian Process, GP）作为贝叶斯优化中最常用的代理模型，能够对目标函数进行概率建模。BoTorch基于GPyTorch提供了多种高斯过程变体：
| 模型类型 | 噪声假设 | 适用场景 |
|---|---|---|
SingleTaskGP | 同方差噪声 | 标准场景，噪声水平未知 |
FixedNoiseGP | 固定噪声 | 已知精确观测噪声 |
HeteroskedasticGP | 异方差噪声 | 噪声随输入变化 |
MixedSingleTaskGP | 混合空间 | 同时包含离散和连续参数 |
单任务高斯过程建模示例：
import torch from botorch.models import SingleTaskGP from gpytorch.mlls import ExactMarginalLogLikelihood # 准备训练数据 (10个2维样本) train_X = torch.rand(10, 2) train_Y = torch.sin(train_X[:, 0]) + torch.cos(train_X[:, 1]) # 初始化并训练GP模型 gp = SingleTaskGP(train_X, train_Y) mll = ExactMarginalLogLikelihood(gp.likelihood, gp) fit_gpytorch_model(mll) # 执行模型训练 # 预测新点 test_X = torch.rand(5, 2) posterior = gp.posterior(test_X) mean = posterior.mean # 预测均值 variance = posterior.variance # 预测方差在实际应用中，我们还需要考虑以下关键因素：
- 核函数选择：RBF核适合平滑函数，Matern核可控制平滑度
- 输入标准化：建议对输入进行归一化处理
- 输出标准化：对输出进行标准化可提升数值稳定性
提示：对于高维问题（>20维），考虑使用SAAS（Sparse Axis-Aligned Subspace）先验的
SaasFullyBayesianSingleTaskGP，它能自动识别重要维度。
3. 采集函数的选择与实现策略
采集函数的设计直接影响贝叶斯优化的探索-开发权衡。BoTorch提供了多种采集函数实现：
3.1 经典采集函数对比
Upper Confidence Bound (UCB)
from botorch.acquisition import UpperConfidenceBound UCB = UpperConfidenceBound(gp, beta=0.2) # beta控制探索程度Expected Improvement (EI)
from botorch.acquisition import ExpectedImprovement best_value = train_Y.max() EI = ExpectedImprovement(gp, best_f=best_value)Probability of Improvement (PI)
from botorch.acquisition import ProbabilityOfImprovement PI = ProbabilityOfImprovement(gp, best_f=best_value)
3.2 蒙特卡洛采集函数
对于复杂场景（如并行评估、组合优化），蒙特卡洛采集函数展现出独特优势：
from botorch.acquisition import qExpectedImprovement from botorch.sampling import SobolQMCNormalSampler # 使用Sobol序列采样器 sampler = SobolQMCNormalSampler(num_samples=512) qEI = qExpectedImprovement( model=gp, best_f=best_value, sampler=sampler )关键参数说明：
num_samples：影响估计精度，通常取256-1024q：批量评估的点数mc_samples：蒙特卡洛采样次数
3.3 采集函数优化技巧
采集函数优化是BO中最耗时的步骤之一，BoTorch提供了多种优化策略：
from botorch.optim import optimize_acqf # 标准优化 candidate, acq_value = optimize_acqf( acq_function=qEI, bounds=torch.tensor([[0., 0.], [1., 1.]]), q=3, # 同时优化3个点 num_restarts=10, raw_samples=512, ) # 使用启发式初始点 from botorch.optim.initializers import gen_one_shot_kg_initial_conditions initial_conditions = gen_one_shot_kg_initial_conditions( acq_function=qEI, bounds=bounds, q=3, num_restarts=10, )4. 与AX框架的深度集成
AX（Adaptive Experimentation）平台提供了更上层的API，简化了贝叶斯优化的应用流程。以下是典型集成模式：
4.1 自定义模型集成
from ax.modelbridge.torch import TorchModelBridge from ax.modelbridge.registry import Models from ax.service.ax_client import AxClient # 创建AX客户端 ax_client = AxClient() # 设置实验参数 ax_client.create_experiment( name="custom_model_experiment", parameters=[ {"name": "x1", "type": "range", "bounds": [0.0, 1.0]}, {"name": "x2", "type": "range", "bounds": [0.0, 1.0]}, ], objective_name="objective", minimize=True, ) # 使用自定义BoTorch模型 model_bridge = Models.BOTORCH_MODULAR( experiment=ax_client.experiment, data=ax_client.get_data(), surrogate=Surrogate(SimpleCustomGP), )4.2 多目标优化实现
AX原生支持多目标优化场景：
from ax.metrics.noisy_function import GenericNoisyFunctionMetric # 定义多目标 metrics = [ GenericNoisyFunctionMetric( name="metric1", f=lambda x: (x["x1"] - 0.5)**2 + (x["x2"] - 0.5)**2, ), GenericNoisyFunctionMetric( name="metric2", f=lambda x: 0.5 - x["x1"] + x["x2"], ) ] # 设置优化策略 ax_client.create_experiment( parameters=[...], objectives={ "metric1": ObjectiveProperties(minimize=True), "metric2": ObjectiveProperties(minimize=False), }, objective_thresholds=[ {"metric": "metric1", "bound": 0.1, "relative": False}, {"metric": "metric2", "bound": 0.5, "relative": False}, ] )5. 高级技巧与性能优化
5.1 处理高维参数空间
对于超过20维的高维问题，传统高斯过程效果会下降。此时可采用：
SAAS先验：
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y)随机嵌入：
from botorch.models.transforms.input import RandomProjection input_transform = RandomProjection(input_dim=50, output_dim=10)
5.2 并行评估加速
BoTorch支持同步评估多个候选点：
from botorch.acquisition import qNoisyExpectedImprovement qNEI = qNoisyExpectedImprovement( model=gp, X_baseline=train_X, sampler=sampler, prune_baseline=True, # 加速计算 ) # 优化获取5个候选点 candidates, _ = optimize_acqf( acq_function=qNEI, bounds=bounds, q=5, # 批量评估5个点 num_restarts=20, raw_samples=1024, )5.3 内存与计算优化
大规模数据集下的优化策略：
使用诱导点近似：
from gpytorch.variational import CholeskyVariationalDistribution from gpytorch.variational import VariationalStrategy from botorch.models import ApproximateGP class SVGPModel(ApproximateGP): def __init__(self, inducing_points): variational_distribution = CholeskyVariationalDistribution( inducing_points.size(0) ) variational_strategy = VariationalStrategy( self, inducing_points, variational_distribution ) super().__init__(variational_strategy) ...启用CUDA加速：
import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_X = train_X.to(device) train_Y = train_Y.to(device) gp = gp.to(device)
在实际项目中，我们发现对输入数据进行标准化处理、合理设置采集函数的超参数（如UCB中的β值）、以及选择适当的蒙特卡洛采样次数，往往能显著提升优化效率。特别是在处理噪声观测时，正确选择高斯过程模型类型（固定噪声vs推断噪声）对最终结果影响巨大。
