以采样为代价的搜索：近乎即时的潜在空间贝叶斯优化

# 以采样为代价的搜索：近乎即时的潜空间贝叶斯优化  
来源：https://arxiv.org/html/2609.19476  
Colin Doumont³ Aleksandra Kalisz⁴ Paul Duckworth⁴ 通讯作者：  
Jacob R. Gardner⁵ Henry Moss⁶ Geoff Pleiss¹,²  
单位：  
¹不列颠哥伦比亚大学 ²Vector Institute ³图宾根AI中心  
⁴InstaDeep ⁵宾夕法尼亚大学 ⁶兰卡斯特大学  

###### 摘要  
生成模型在许多全新发现流程中日益成为核心，这些流程通过大规模生成设计，并经由虚拟筛选过滤出候选方案以进行实验验证。虽然贝叶斯优化（BO）与这类场景天然契合（因其能利用历史评估结果指导后续提议），但当虚拟筛选成本相对较低时，其顺序决策所需的计算开销会成为瓶颈。我们利用线性模型在球形域（高维潜变量集中于此）上的独特组合，使BO在此类场景中变得实用。我们基于近期关于线性代理模型的研究，推导出近乎闭式解的代理建模与采集函数问题，充分利用了球形对称性。结果表明，该方法在分子与图像生成基准测试中，相比最先进基线实现了至少100倍加速，同时达到相当或更优的性能。总体而言，我们的方法使BO成为全新发现流程中实用的即插即用方案，此前因速度过慢而无法考虑的场景现已可行。  

## 1 引言  
近期在结构化高维域（如蛋白质⁷、⁴⁵，肽¹¹、⁴²，小分子⁵²、⁹等）上的全新发现取得成功，依赖两个关键要素。首先，离散或组合域上的预训练生成模型能从连续向量生成结构有效的样本，同时保持功能相似性。其次，虚拟筛选或模拟流程作为昂贵湿实验室实验的廉价代理，通常在数秒内完成样本验证与排序。两者结合实现了新颖设计的快速生成与评分，仅最有前景的方案进入现实验证。  

**图1：目标函数值 vs. 优化壁钟时间**  
在osmbGuacaMol多目标优化问题中，使用维度d=512的COATI-LDM Flow模型，批量大小q分别为5和5000的总评估次数下，我们的方法在达到与所有基线相当的分数的同时，所需壁钟时间减少了数个数量级，且差距随q增大而扩大。当q=1000时，我们的方法几乎与从潜空间“随机采样”一样快，但能生成显著更优的设计。  

给定虚拟筛选管道，识别高分候选方案是对生成模型潜空间进行的连续黑箱优化问题。贝叶斯优化（BO）是自然工具，它使用高斯过程（GP）等概率代理模型将查询导向有前景的区域。近期高维通用BO⁵⁰、²⁰、¹³和潜空间BO（LSBO）⁴³、²⁷、²⁶的进展已在全新发现中取得经验成功。然而，BO方法通常假设评估成本远高于代理模型拟合与采集优化——这一假设在上述虚拟筛选场景中失效。此时，LSBO方法可能对每个候选方案耗费数分钟进行代理模型计算，而实际生成与评分仅需数秒²⁷。  

“随机采样”等高通量策略⁷、⁹尽管样本效率低，却可能在壁钟时间上主导LSBO，尤其当非自适应采样可并行化而BO本质上是顺序进行时。我们引入一种LSBO算法，在保持现有方法样本效率的同时，消除了大部分代理模型开销。通过将非线性GP替换为球形域上的线性代理模型，我们将每次迭代复杂度从样本数量的立方降至线性，并将代理拟合与采集优化简化为闭式操作与一维数值搜索；整个流程每次迭代耗时不足1秒，相比竞争方法的数分钟大幅缩短。  

我们对线性代理模型的使用基于Doumont等人¹³的研究，他们在设计空间的球面投影上进行线性模型BO。他们将超立方体映射到球面的一部分，而我们直接将设计空间限制为潜空间的球形子集。由于映射仅覆盖球面的一部分，其边界破坏了全局对称性（我们的闭式解依赖于此），迫使代理拟合与采集优化都需在完整搜索空间上回退到基于梯度的优化。  

除了计算便利性，球形限制符合潜空间（基础分布）中良好设计的分布。大多数生成模型（变分自编码器、潜扩散模型、流匹配等）共享两个相关特性：首先，其各向同性基础分布引发“薄壳”现象：潜变量密度集中于球面¹⁰，而偏离球面的区域会解码为无意义或无效设计，浪费BO预算³⁰、⁵、²⁹。其次，在离散域中，这些模型是多对一的，且我们经验发现：任意有效设计的原像都能可靠地与球面相交（跨模型族，见第4节）。因此，该约束保留了表达能力，将预算集中于合理设计，并提供了闭式解所需的全局对称性。  

球形几何上的线性模型与新颖闭式解的结合，使LSBO对基于生成模型与虚拟筛选的全新问题变得实用。经验表明，我们的方法在高达16,384维潜空间的分子设计与图像生成任务中，匹配或超越了基线的样本效率。更令人印象深刻的是，其壁钟运行速度快100倍以上，且随着观测预算与潜维度增加、基线变得难以处理，差距进一步扩大。借助并行采集，我们的方法成本几乎等同于“随机采样”，但能返回更高分候选方案（图1）。因此，它在评估成本过低而无法合理使用BO的场景中变得可行。  

## 2 背景  
##### 潜空间贝叶斯优化旨在最大化黑箱函数  
F:𝒳→ℝ，其中设计空间𝒳通常复杂且离散，方法是通过确定性预训练生成模型  
G:ℤ→𝒳  
的潜空间ℤ⊆ℝᵈ  
替代对𝒳的直接搜索。自然随机生成器（如扩散模型或自回归VAE）可通过使用确定性采样器（如³⁹）或贪婪解码变为确定性。详见附录B。  

我们在复合函数  
F∘G  
上放置概率代理模型f，观测  
y = f(z) + ε，其中 ε∼𝒩(0,σₙ²)  
。数据  
𝒟ₜ = {(zᵢ,yᵢ)}ᵢ₌₁ⁿₜ  
被迭代收集，其中f拟合已有观测，并用于提议下一个待观测点  
zₜ₊₁  
。每个LSBO轮次包含三个步骤：  

1. **(i) 代理模型拟合**  
常用代理模型为  
f∼𝒩(μ,k)，其中均值函数μ与核函数k编码先验平滑性信念。GP超参数（核长度尺度、观测噪声）通过最大化对数边际似然（基于梯度优化）拟合。流行核函数如RBF核  
k_RBF(z,z')=exp(-½‖Dℓ⁻¹z - Dℓ⁻¹z'‖²)  
（对角长度尺度矩阵Dℓ）需要每次梯度步进行  
O(nₜ³)  
线性求解。在中等规模问题（nₜ≈2000, d≈512）中，即使使用GPU加速，该过程因数千次线性求解需20分钟以上。许多LSBO流程进一步联合优化生成模型参数与GP超参数²⁷、²⁶，使壁钟时间增加一倍以上²⁷、¹⁸。  

2. **(ii) 采集函数优化**  
下一个查询点  
zₜ₊₁  
最大化采集函数  
α(z; 𝒟ₜ)  
，该函数权衡探索与利用。常见选择如期望改进（EI）²⁸、²¹、汤普森采样（TS）⁴¹、上置信界（UCB）缺乏闭式优化器。EI与UCB使用重启梯度下降；尽管初始成本为  
O(nₜ³)  
，但其非凸特性常需数百次迭代，在中等规模问题上每轮成本达2分钟。TS则在离散候选集上优化，成本可能相当³³、¹⁵。  

3. **(iii) 评估**  
最后一步是解码  
xₜ₊₁ = G(zₜ₊₁)  
并评估  
yₜ₊₁ = F(xₜ₊₁)  
。当F为计算筛选或模拟⁴⁸、⁵²（而非湿实验室实验）时，即使考虑生成模型本身成本，这也是成本最低的阶段：在我们的分子设计任务中，每轮最多耗时3秒，约占LSBO循环的0.2%。  

##### 薄壳现象  
BO需要紧致域以使采集函数优化适定。大多数BO算法将设计空间限制在超立方体  
ℤ = [-b,b]ᵈ  
（某b>0）。虽然非潜问题的边界由问题指定，但LSBO中b的选择不明确。大多数生成模型（变分自编码器、归一化流、扩散模型）将各向同性基础分布（如𝒩(0,I)）变换为目标数据分布。薄壳现象⁴⁴、¹⁰、²⁵或“高斯典型性”³¹使潜变量在d→∞时集中于  
‖z‖₂² = d  
的球面，但分布具有非紧支撑。b过小会排除潜在高分潜变量，b过大会使搜索体积随d指数膨胀；[-b,b]ᵈ的大部分区域包含低似然潜变量，处于解码器支持较差的“死区”¹²、⁴、⁴⁷、⁴⁶、²。结果输出通常无效或低质量（见图2与图8），但可能从虚拟筛选器中产生虚高的目标值（见第4节）。  

LSBO加剧了该问题，因为采集优化倾向于集中在ℤ的边界³⁷，需要仔细调整b²⁷、解码器有效性感知惩罚¹⁷、范数或球面感知潜先验³⁸、²⁴作为部分补救措施。  

**图2：偏离薄壳的潜变量解码为不真实图像**  
给定提示词“一只狗在草地上用牵引绳行走的高质量照片”，我们在Stable Diffusion 1.5（d=16,384）上解码选定半径R与单位向量z的潜变量Rz。d√d的图像视觉效果佳，而低范数图像缺乏细节，高范数图像过曝且过度锐化，表明球形潜域主要排除不太可能的输入。  

## 3 近乎即时的LSBO流程  
现在我们提出针对生成模型潜空间的快速BO流程。我们的基础模块是Doumont等人¹³的近期工作，他们提出对输入z到单位球面的一部分进行投影后使用线性代理模型进行BO：  
f(z) = P_θ(z)ᵀβ + c, β∼𝒩(0,I), ℤ=[-b,b]ᵈ, P_θ:ℝᵈ→𝒮ᵈ  
其中b>0指定设计空间，P_θ是从ℝᵈ到𝒮ᵈ子集的双射映射，c与θ为可学习超参数。该模型等价于应用于P_θ(z)的线性核GP。尽管使用线性模型，该设置在经验上匹配了最先进高维BO算法的性能¹³。作者证明球面投影P_θ对其成功至关重要，因为它使f在‖z‖上非单调，从而避免边界搜索行为。尽管该方案具有线性模型的  
O(nₜd²)  
渐近复杂度，模型拟合与采集优化仍