引言

在药物发现和材料科学领域,寻找新分子的过程常被比作“大海捞针”。然而,这里的“大海”是化学空间,其中包含估计达 \(10^{60}\) 种理论上可能的药物样分子。搜索这个空间是一场离散的、组合爆炸式的噩梦。

为了解决这个问题,机器学习研究人员开发了潜在空间贝叶斯优化 (Latent Space Bayesian Optimization, LSBO) 。 这个想法非常优雅: 我们可以使用变分自编码器 (VAE) 将这些离散的化学结构映射到一个连续、平滑的数值空间 (即“潜在空间”) 中,而不是直接搜索离散结构。然后,我们可以在这个平滑空间中使用标准的优化技术来寻找最佳候选者。

多年来,这一直是标准的操作流程。但一篇题为 “Return of the Latent Space COWBOYS” 的新论文认为,我们要么走错了方向,要么方法不对。作者指出,将生成模型 (VAE) 与优化模型过于紧密地耦合会导致显著的效率低下。

他们提出的解决方案是 COWBOYS (Categorical Optimisation With Belief Of underlYing Structure,基于潜在结构信念的分类优化) ,提供了一种将生成与预测分离的解耦方法。通过让 VAE 专注于结构生成,而让高斯过程专注于预测,他们在严格的预算限制下识别高潜力候选分子方面取得了最先进的结果。

在这篇文章中,我们将解构为什么标准 LSBO 举步维艰,COWBOYS 背后的数学洞察,以及这个新框架如何改变我们要处理结构化优化的方式。

背景: 发现的要素

要理解为什么 COWBOYS 是一项创新,我们首先需要了解结构化空间上标准贝叶斯优化的运作机制。

贝叶斯优化 (BO)

贝叶斯优化是一种寻找昂贵黑盒函数 \(f(x)\) 最大值的策略。在药物发现中,\(f(x)\) 可能是测量分子功效的湿实验室实验。由于我们无法测试每一个分子,BO 会建立一个代理模型 (通常是高斯过程) 来根据之前的测试预测性能。然后,它使用采集函数 (acquisition function) 来决定下一个测试哪个分子,以平衡探索 (尝试新事物) 和利用 (优化我们已知有效的) 。

变分自编码器 (VAE)

VAE 是一种深度生成模型。它由两部分组成:

  1. 编码器 (Encoder) : 将复杂的输入 (如分子图) 压缩成潜在空间中的低维向量 \(z\)。
  2. 解码器 (Decoder) : 从 \(z\) 重构原始输入。

VAE 允许我们将离散分子视为连续向量空间 \(\mathbb{R}^d\) 中的点。

标准方法: LSBO

在潜在空间 BO (LSBO) 中,我们在大量分子数据集上训练 VAE。然后,我们在潜在空间内部执行贝叶斯优化。

Algorithm 1 Latent Space Bayesian Optimisation

如算法 1 所示,优化器在潜在空间中选择一个点 \(z\),将其解码为分子 \(x\),对其进行评估,并更新模型。代理模型 \(\tilde{g}(z)\) 试图学习从潜在坐标到分子属性的映射。

潜在空间 BO 的陷阱

虽然 LSBO 将困难的离散问题转化为可管理的连续问题,但作者指出了阻碍其性能的三个主要病灶。

1. 对齐问题 (The Alignment Problem)

VAE 的训练目的是重构分子,而不是预测它们的属性。潜在空间中两点之间的“距离”对应的是结构相似性,而不一定是功能相似性。潜在空间中的一小步可能会导致我们试图优化的化学属性发生巨大跳变。这使得代理模型的工作变得异常困难,因为目标函数在潜在空间中显得参差不齐且不可预测。

2. 随机性问题 (The Stochasticity Problem)

VAE 是概率性的。解码器不会将点 \(z\) 映射到单个分子 \(x\),而是映射到一个分子分布。

Figure 2.(a)In LSBO,the same latent input (blue dot) will, via the stochastic decoder (grey box),map to different values in structure space (black dots) and so corresponds to multiple objective function values (red dots)— a discrepancy that hinders the learning of accurate surrogate models.(b) In higher-dimensional problems, the area of the latent space supported by the prior of the VAE (blue) concentrates in a thin circular shell.

如图 2(a) 所示,潜在空间中的单个点 (蓝点) 可以解码成多种不同的分子结构 (黑点) ,每种结构具有不同的属性值 (红点) 。在潜在空间中运行的标准高斯过程 (GP) 很难处理这种情况。它将这种变化解释为噪声,导致预测不佳和过拟合。

3. 几何问题 (“甜甜圈”效应)

标准 LSBO 将搜索限制在一个方框内,通常是 \([-\delta, \delta]^d\)。直觉上,有效的分子位于潜在空间的中心附近 (通常是标准正态分布) 。

然而,高维几何是反直觉的。根据高斯环定理 (Gaussian Annulus Theorem) , 高维高斯分布的几乎所有概率质量都位于距离原点 \(\sqrt{d}\) 的一个薄壳 (环) 中,而不是在中心。

图 2(b) 直观地展示了这一点。如果你将搜索限制在一个方框内,你搜索的体积主要由缺乏有效先验概率的区域组成,同时可能切断了有效分子实际存在的高概率“壳层”。这种不匹配造成了“死区”,在这些区域解码器会产生无效或无意义的分子。

核心方法: COWBOYS

作者提出了 COWBOYS 来解决这些问题,从根本上改变了 VAE 和代理模型的交互方式。

解耦哲学

在 LSBO 中,代理模型存在于潜在空间中。在 COWBOYS 中,代理模型存在于结构空间 (即实际的分子空间) 中。

Figure 1. Unlike LSBO where GPs are fit in a VAE’s latent space, COWBOYS’s GP is fit in structure space,decoupled from its VAE. acquisition routines can be employed. Candidate points selected by the optimiser are then decoded back into the original structured domain to yield new query points.

如图 1 所示:

  • LSBO (左) : GP 从 \(z\) 预测 \(y\)。
  • COWBOYS (右) : GP 直接从 \(x\) 预测 \(y\)。VAE 仅用作生成器 (先验) ,而不用作优化的坐标系。

第一步: 结构化代理

通过将代理模型移回结构空间,COWBOYS 可以利用领域知识。作者没有使用潜在向量上的通用欧几里得核,而是使用了 Tanimoto 核 , 这是一种专门为分子指纹设计的相似性度量。

Equation for Tanimoto Kernel

该核 (\(K_T\)) 基于分子中特定子结构 (\(\pmb{m}\)) 的存在与否来计算相似性。这使得 GP 能够比从任意潜在坐标中更有效地学习化学属性。

第二步: 搜索策略

如果我们不在方框内优化坐标,我们如何找到新分子?COWBOYS 将其表述为一个采样问题。

我们希望生成的分子既在化学上有效 (在 VAE 先验 \(p_\theta(x)\) 下具有高概率) ,又具有高预测性能 (在 GP 代理下具有高概率) 。

在算法上,COWBOYS 不是最大化采集函数,而是从条件分布中采样:

Equation for conditional sampling

在这里,我们要求模型生成一个分子 \(x\),使得在给定目前观察到的数据 (\(D_n^{\pmb{\chi}}\)) 的情况下,其预测值 \(f_x\) 大于当前观察到的最佳值 \(f^*\)。

第三步: 通过 MCMC 进行实际实现

直接从这个条件分布中采样是很困难的。然而,作者使用了一个聪明的技巧。他们将搜索映射回潜在空间,仅仅是为了采样,但他们使用结构化 GP 来评估条件。

他们用确定性映射 \(h_\theta(z)\) (取最可能的解码) 来近似随机解码器。采样问题随后变成了寻找一个满足以下条件的潜在代码 \(z\):

Equation for latent sampling with deterministic decoding

利用贝叶斯规则,该后验分布正比于先验乘以改进的可能性:

Equation 5: Bayes decomposition

这个方程非常优雅。

  • \(p(z)\) 是 VAE 的标准高斯先验 (我们知道如何对其进行采样) 。
  • \(p(g_{\theta, z} > f^* | D)\) 类似于改进概率 (PI) 采集函数,GP 可以轻松计算出来。

为了高效地从该分布中采样,作者使用了预处理 Crank-Nicolson (PCN) MCMC。与标准随机游走不同,PCN 专为高维高斯先验设计。它自然地停留在“环” (前面讨论的高概率壳层) 上,避开了困扰方框约束 LSBO 的死区。

算法总结

完整的 COWBOYS 流程 (算法 2) 与 LSBO 显著不同:

Algorithm 2 COWBOYS

  1. 初始设计: 纯粹从 VAE 先验中采样分子。
  2. 循环:
  • 使用所有分子-值对 \((x, y)\) 更新结构化 GP。
  • 使用 MCMC 采样一个新的潜在变量 \(z\),该变量有可能改进最佳结果 \(f^*\)。
  • 将 \(z\) 解码为 \(x\),评估 \(f(x)\),然后重复。

实验与结果

作者将 COWBOYS 与广泛的基线进行了比较,包括使用复杂启发式方法在优化过程中微调 VAE 的最先进 LSBO 方法。

低数据效率

药物发现中最关键的要求之一是样本效率。湿实验室实验非常昂贵;你无法评估数百万个候选者。

图 3 显示了六个不同分子优化任务的性能。x 轴代表预言机调用 (实验次数) ,y 轴代表奖励 (分子质量) 。

Figure 3.Average performance ( \\(\\pm\\) standard error) of COWBOYS over 10 runs on problems considered by (Chu et al.,2024).COWBOYS achieves a substantial improvement in sample efficiency over all existing LSBO methods…

关键结论: COWBOYS (橙色星星) 比竞争方法更快地获得了高奖励。在像 Amlodipine MPO (图左下) 和 Osimertinib MPO (图左上) 这样的任务中,它在不到 500 次评估内就达到了接近最佳的性能,而其他方法则落在后面。值得注意的是,它优于微调 VAE 的方法 (如 LOL-BO 和 InvBO) ,且无需重新训练神经网络的计算开销或不稳定风险。

高维性能

作者还在高维离散序列优化任务上测试了 COWBOYS (表 1) 。

Table 1 - Results

在这个综合基准测试中,COWBOYS 始终获得最高的平均分数 (粗体) 。它击败了进化算法 (GA,HillClimbing) 和其他 BO 方法。这证明该方法不仅适用于小分子图,也具有普适性。

解耦的好处

为了严格证明解耦策略是改进的源头,作者将 COWBOYS 与同样微调 VAE 的 LSBO 方法进行了比较。

Figure 4.Average performance ( \\(\\pm\\) standard error) over 2O repetitions with an log-scaled \\(\\mathbf { X } ^ { }\\) -axis,demonstrating that,among LSBO methods that cannot fine-tune their latent space, COWBOYS provides significant improvemnt in efficiency.

图 4 使用对数刻度的 x 轴来显示更长范围内的性能。即使给标准的 LSBO 方法 (如 TURBO-L) 多几个数量级的评估次数,它们也很难赶上 COWBOYS 的效率。这凸显了在结构空间 (适用化学规则的地方) 拟合 GP 从根本上优于在非为回归设计的潜在空间中拟合 GP。

消融实验

最后,一项消融实验 (表 2) 检查了该方法的鲁棒性。

Table 2.Average performance \\(( \\pm \\mathrm { s . d . } )\\) over 5 repetitions of COWBOYS…

最右边的一栏特别能说明问题: “COWBOYS with Latent GP” 。 在这个实验中,他们保留了 COWBOYS 的采样策略,但强制 GP 使用潜在空间 (回退到 LSBO 代理) 。几乎所有任务的性能都显著下降。这证实了结构化核 (Tanimoto) 是成功的关键组成部分。

结论与启示

“Return of the Latent Space COWBOYS” 为简化生成模型与优化的结合方式提出了令人信服的论据。作者表明,与其强迫代理模型去学习 VAE 潜在空间中复杂且往往不平滑的几何结构,我们应该让每个模型做它最擅长的事情:

  1. VAE 定义搜索空间并确保我们生成有效的结构。
  2. 高斯过程 直接在结构空间中对目标函数进行建模,使用捕获领域知识的核函数。

通过基于原则的贝叶斯更新规则和高效的 MCMC 采样将这两者连接起来,COWBOYS 消除了对任意搜索边界和昂贵的 VAE 微调的需求。

这项工作的影响超越了分子设计。它为结构化数据的优化提出了一个通用蓝图: 将生成器与预测器解耦。 无论是设计蛋白质、计算机代码还是机械零件,该框架都允许从业者将领域专业知识 (通过核函数) 直接注入优化循环,将潜在空间优化的“黑盒”变成一个更透明、更高效的过程。