Search at the Cost of Sampling: Nearly-Instant Latent Space Bayesian Optimization¶

会议: NeurIPS 2026
arXiv: 2609.19476
OpenReview: https://openreview.net/forum?id=jqZgN0DS0E
代码: https://github.com/DonneyF/nearly-instant-lsbo
领域: 优化/理论
关键词: 潜在空间贝叶斯优化、高维黑盒优化、线性代理模型、超球面几何、薄壳现象

一句话总结¶

本文提出了一种在生成模型隐空间超球面上构建线性代理模型的高速贝叶斯优化算法，利用球面全局对称性将代理拟合与采集函数优化简化为单次特征分解与一维标量求根，在保持甚至超越非线性高斯过程样本效率的同时，将单轮优化开销降低逾两个数量级，实现了接近纯随机采样的极速搜索。

研究背景与动机¶

在现代小分子设计、多肽合成与蛋白质工程等从头设计（de novo discovery）流水线中，核心范式依赖于两项关键技术：其一是预训练生成模型（如变分自编码器 VAE、流匹配 Flow Matching、潜在扩散模型 LDM），将离散的组合化学或生物序列映射为平滑连续的潜在向量；其二是基于计算的虚拟筛选（virtual screening）或分子动力学模拟管线，仅需数秒即可完成生成样本的性质打分与排序。将优质候选样本的发现转化为生成模型潜在空间上的连续黑盒优化问题，贝叶斯优化（Bayesian Optimization, BO）凭借其利用历史评估数据自适应权衡探索与利用的机制，理应是这类任务的理想算法。

然而，经典贝叶斯优化建立在一个核心假设之上——目标黑盒函数的物理评估成本远高于代理模型的拟合与采集函数优化开销。在虚拟筛选耗时仅需几秒钟的高通量计算场景中，这一假设彻底崩溃。主流潜在空间贝叶斯优化（LSBO）依赖非线性高斯过程（GP），拟合核函数超参数需要进行多次 $\mathcal{O}(n^3)$ 的求逆运算，随后还要通过多起点梯度下降优化非凸的期望提升（EI）或置信区间上限（UCB）采集函数，单轮决策往往耗时数分钟甚至数十分钟，计算开销占整个循环的 99.8% 以上。这种巨大的计算延迟迫使工业界更倾向于采用样本效率极低但能极速并行采样的“散弹式”（shotgun）随机抽样策略。此外，高维各向同性高斯先验具有著名的“高斯典型集”（薄壳现象，thin shell phenomenon），绝大多数概率质量集中在半径为 $R=\sqrt{d}$ 的超球壳上；传统方法在超立方体边界内搜索容易漂移至低似然的“死亡空间”，解码出无意义结构甚至诱发针对虚拟筛选打分器的奖励作弊（如拼凑极长无意义碳链）。

本文的切入角度是：与其在超立方体中挣扎，不如直接将搜索空间显式限制在隐向量自然聚集的超球面上，并搭配线性代理模型。核心 idea：将潜在空间的搜索域严格约束在半径为 $\sqrt{d}$ 的超球壳上，并构建高维线性代理模型，利用无界超球面的完美全局旋转对称性，将原本高昂的高斯过程超参数拟合与采集函数极值求解降解为单次矩阵特征分解加一维标量求根，以接近随机采样的微秒级开销实现 SOTA 样本效率的自适应贝叶斯优化。

方法详解¶

整体框架¶

算法围绕预训练生成模型的确定性解码器 $G: \mathcal{Z} \to \mathcal{X}$ 与快速黑盒评估函数 $F: \mathcal{X} \to \mathbb{R}$ 展开。输入为当前观测数据集 $\mathcal{D}_t = \{(\bm{z}_i, y_i)\}_{i=1}^{n_t}$，其中隐向量 $\bm{z}_i \in \mathbb{R}^d$ 均位于超球面 $\mathcal{Z} = \{\bm{z} \in \mathbb{R}^d : \|\bm{z}\|_2 = \sqrt{d}\}$ 上。整体流水线分为四步：首先对历史隐向量构成的散度矩阵进行单次特征分解；随后通过一维标量求根在对数边际似然上剖析求解均值常数与观测噪声方差；接着利用球面对称性导出的几何性质，以闭式投影（TS）或一维标量求根（EI/UCB/PI）直接获取超球面上全局最优候选隐向量 $\bm{z}_{t+1}$；最后将该隐向量送入生成模型解码并调用虚拟筛选打分，扩充数据集进入下一轮迭代。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["历史隐向量与评估集合<br/>D_t = {(z_i, y_i)}, ||z_i|| = sqrt(d)"] --> B["矩阵特征分解预计算<br/>分解 Z^T Z (n >= d) 或 Z Z^T (n < d)"]
    B --> C["一维对数边际似然求根<br/>闭式剖析均值常数 c 求解观测噪声 sigma_n^2"]
    C --> D{"采集策略分支"}
    D -->|汤普森采样 TS| E["后验权重向量球面投影<br/>闭式求解: z* = sqrt(d) beta / ||beta||"]
    D -->|期望提升 EI / UCB| F["拉格朗日乘子一维标量求根<br/>复用特征基向量求解旋转空间根 tau*"]
    E --> G["确定性解码与目标评估<br/>x = G(z*), y = F(x) 扩充数据集"]
    F --> G

关键设计¶

1. 超球面域上的线性代理建模：将薄壳约束转化为几何加速基石

传统 LSBO 将设计空间设为超立方体 $[-b, b]^d$，在高维下会导致严重的维度灾难。一方面，各向同性高斯先验的概率质量严格集中在 $\|\bm{z}\|_2 = \sqrt{d}$ 的超薄球壳上，超立方体绝大部分体积对应解码器训练时未曾见过的死区；另一方面，虽然已有研究尝试将超立方体投影至部分球面，但边界截断破坏了对称性，依然需要耗时的梯度优化。本文直接将可行域定为全超球壳 $\mathcal{Z} = \{\bm{z} \in \mathbb{R}^d : \|\bm{z}\|_2 = \sqrt{d}\}$，并建立高维线性代理模型： $$f(\bm{z}) = \bm{z}^\top \bm{\beta} + c, \quad \bm{\beta} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$$ 该设定在理论上将搜索限定在 $(d-1)$ 维流形上。作者通过理论与实验证明，离散生成模型的贪婪解码器是一个多对一映射（many-to-one），任意高分结构 $\bm{x}$ 的原像集（preimage set）在高维空间中几乎必然与 $\sqrt{d}$ 球面相交。由此，流形限制不仅未损伤有效设计空间的覆盖表达能力，更天然杜绝了解码无意义产物与高范数区域的奖励黑客行为，同时赋予了全局连续对称性。

2. 边际似然快速剖析：单次特征分解实现常数级超参数拟合

高斯过程拟合的关键瓶颈在于超参数优化时反复对协方差矩阵求逆。对于线性模型，观测噪声为 $\sigma_n^2$ 时的 Gram 矩阵为 $\widetilde{\bm{K}} = \bm{Z}\bm{Z}^\top + \sigma_n^2 \mathbf{I}$。通过对平移散度矩阵 $\widetilde{\bm{C}} = \bm{Z}^\top \bm{Z} + \sigma_n^2 \mathbf{I} \in \mathbb{R}^{d \times d}$ 应用 Woodbury 矩阵逆恒等式与矩阵行列式引理，协方差逆与对数行列式被精确降维转化： $$\widetilde{\bm{K}}^{-1} = \sigma_n^{-2} (\mathbf{I} - \bm{Z} \widetilde{\bm{C}}^{-1} \bm{Z}^\top), \quad \log\det\widetilde{\bm{K}} = \log\det\widetilde{\bm{C}} + (n - d)\log\sigma_n^2$$ 在无先验或共轭先验下，对于任意给定的噪声水平 $\sigma_n^2$，均值常数的最优解存在精确闭式：$c^*(\sigma_n^2) = \frac{\mathbf{1}^\top \widetilde{\bm{K}}^{-1} \bm{y}}{\mathbf{1}^\top \widetilde{\bm{K}}^{-1} \mathbf{1}}$。将 $c^*(\sigma_n^2)$ 代回对数边际似然（LML）后，目标函数仅剩标量未知数 $\sigma_n^2$。通过预先计算散度矩阵的单次特征分解 $\bm{Z}^\top \bm{Z} = \bm{U}\bm{\Lambda}\bm{U}^\top$（当 $n < d$ 时分解 $\bm{Z}\bm{Z}^\top$），任何 $\sigma_n^2$ 下的矩阵求解均坍缩为对角阵 $\bm{\Lambda} + \sigma_n^2 \mathbf{I}$ 的常数时间标量操作。因此，超参数拟合仅需在 $\sigma_n^2 \ge \sigma_{\min}^2$ 上进行极快的一维数值求根，单轮耗时仅为毫秒级。

3. 球面对称采集优化：拉格朗日乘子降维与闭式极值求解

在超球面上，采集函数的最大化问题展现出极其规整的几何结构。对于汤普森采样（Thompson Sampling, TS），在后验权重空间抽样 $\bm{\beta} \sim \mathcal{N}(\bm{\mu}, \bm{\Sigma})$ 后，线性函数在球面上的最优候选解可直接由柯西-施瓦茨不等式得出闭式投影： $$\bm{z}^* = \arg\max_{\|\bm{z}\|_2=\sqrt{d}} \bm{\beta}^\top \bm{z} = \sqrt{d} \frac{\bm{\beta}}{\|\bm{\beta}\|_2}$$ 对于期望提升（Expected Improvement, EI），作者将带球面约束的优化问题建立拉格朗日函数 $\mathcal{L}(\bm{z}, \lambda) = \alpha_{\text{EI}}(\bm{z}) - \frac{\lambda}{2}(\|\bm{z}\|_2^2 - d)$。对 $\bm{z}$ 求驻点并引入变量重参数化 $\tau = \frac{\lambda \sqrt{\text{Var}[f(\bm{z})|\mathcal{D}_t]}}{\sigma_n^2 \phi(\zeta)}$，得到最优解的隐式结构： $$\bm{z}(\tau) = \gamma \cdot \bm{v}(\tau), \quad \bm{v}(\tau) = (\tau \mathbf{I} - \widetilde{\bm{C}}^{-1})^{-1} \mathbb{E}[\bm{\beta} \mid \mathcal{D}_t]$$ 其中 $\gamma$ 仅由标量标准正态分布密度 $\phi(\zeta)$ 与累积分布 $\Phi(\zeta)$ 决定。强行满足球面模长约束 $\|\bm{z}(\tau)\|_2 = \sqrt{d}$ 等价于标量残差函数 $g_{\text{EI}}(\tau) = 0$ 的一维求根。由于已缓存特征向量基 $\bm{U}$，旋转坐标系下的矩阵求逆为纯对角矩阵运算，残差计算开销仅为 $\mathcal{O}(d)$。当样本数 $n_t < d$ 时，作者进一步证明了零空间分量在驻点处的解析结构，通过至多两次一维标量搜索即可完全确定兼顾均值与方差全局最优的候选点。

损失函数 / 训练策略¶

模型选择通过优化对数边际似然实现。对于去中心化标签数据 $\bm{r} = \bm{y} - c^*(\sigma_n^2)\mathbf{1}$，一维求根所优化的目标函数形式为： $$L(c^*(\sigma_n^2), \sigma_n^2) = -\frac{1}{2\sigma_n^2}\left(\bm{r}^\top \bm{r} - \bm{r}^\top \bm{Z} \bm{U} (\bm{\Lambda} + \sigma_n^2 \mathbf{I})^{-1} \bm{U}^\top \bm{Z}^\top \bm{r}\right) - \frac{1}{2}\sum_{i=1}^d \log(\lambda_i + \sigma_n^2) - \frac{n-d}{2}\log\sigma_n^2 - \text{const}$$ 在全流程优化中，每次迭代仅需执行一次维度为 $\min\{n_t, d\}$ 的矩阵特征分解，总时间复杂度严格约束在 $\mathcal{O}(n_t d \min\{n_t, d\} + \min\{n_t, d\}^3)$。对于确定性生成模型，使用 DDIM 等确定性常微分方程（ODE）求解器或贪婪解码（如 SELFIES 的 token 贪婪映射），确保隐空间向量到实体结构的映射单射且完全可复现。

实验关键数据¶

主实验¶

实验覆盖分子发现（GuacaMol 4 项多属性目标 MPO：osmb, pdop, rano, zale；生成模型采用 $d=512$ 的 COATI-LDM 与 $d=256$ 的 SELFIES-VAE）、高维图像生成对齐（Stable Diffusion v1.5, $d=16{,}384$；优化 CLIP 分数与 ImageReward）以及蛋白质构象优化（Boltz-2 结构基础模型, $d=3549$；优化与钙调蛋白 1CLL 的 TM-score）。所有虚拟筛选评测单次均在 5 秒以内完成。下表汇总了在 osmb 分子设计基准与蛋白质 TM-score 任务中的优化表现及耗时对比：

评估任务	算法配置	采集批大小 $q$	评价次数	最优目标值 (越高越好)	优化端总耗时 (s)	相对 Vanilla BO 加速比
GuacaMol (osmb, d=512)	Random Sampling (Shotgun)	1000	5000	0.68 ± 0.02	120 s	720× (基线无自适应)
GuacaMol (osmb, d=512)	Vanilla BO (BoTorch RBF)	1	5000	0.81 ± 0.02	86,400 s (>1天)	1.0× (基准)
GuacaMol (osmb, d=512)	TuRBO-1	1	5000	0.80 ± 0.02	18,200 s	4.7×
GuacaMol (osmb, d=512)	Linear (Warped)	1	5000	0.79 ± 0.03	5,400 s	16.0×
GuacaMol (osmb, d=512)	本文方法 (EI, 顺序式)	1	5000	0.83 ± 0.01	680 s	127×
GuacaMol (osmb, d=512)	本文方法 (TS, 批量式)	1000	5000	0.82 ± 0.01	180 s	480×
蛋白质 Calmodulin (d=3549)	Unguided Sampling	1	1000	0.42 ± 0.03	5,200 s	-
蛋白质 Calmodulin (d=3549)	本文方法 (EI)	1	1000	0.68 ± 0.02	5,700 s	代理开销仅占 8.8%

消融与分析实验¶

为了严谨验证“将隐空间限制在 $\sqrt{d}$ 超球面上是否会损伤生成多样性”以及“球面几何是否能够防御奖励黑客”，作者在 GuacaMol 数据集上抽取 20,000 个分子进行超球面原像重构（求解公式 $\min_{\|\bm{z}\|_2=\sqrt{d}} \mathcal{L}(G(\bm{z}), \bm{x})$），并在易受攻击的 LogP 极大化任务上测试了非受限立方体域与超球面域的生成有效性。

生成模型	隐空间维度 $d$	评估分子样本量 $n$	超球面重构成功率 (%)	LogP 优化下的分子平均碳链长度	LogP 优化化学有效性
SELFIES-VAE	256	1,273,077	99.987% (1,272,917/1.27M)	-	-
COATI-LDM Flow	512	20,000	99.990% (19,998/20,000)	14.2 (真实合理结构)	99.8% (有效药效团)
COATI-LDM Diffusion	512	1,000	100.00% (1,000/1,000)	13.8 (真实合理结构)	100.0% (无异常聚集)
Vanilla BO (未约束超立方体)	512	- (优化轨迹生成)	-	68.5 (病态超长聚乙烯链)	12.3% (严重奖励作弊)
Linear Warped (部分球面投影)	512	- (优化轨迹生成)	-	54.1 (病态长碳链)	26.5% (严重奖励作弊)

关键发现¶

计算瓶颈彻底逆转：在高维分子与蛋白质设计任务中，本文方法的单轮代理拟合加采集优化时间从传统高斯过程的数百秒被压缩至 0.05–0.5 秒以内。在 Boltz-2（$d=3549$）实验中，优化器本身的计算延迟仅占整个循环的不到 9%，黑盒评估与生成模型本身首次成为耗时的主导环节。
高维极端尺度的可行性飞跃：在 Stable Diffusion v1.5（$d=16{,}384$）超高维图像对齐实验中，常规的高斯过程与 TuRBO 因内存与逆矩阵计算复杂度彻底无法运行。本文方法是唯一能够流畅运行的自适应 BO 算法，生成图像的 ImageReward 与文本对齐显著超过随机采样。
天然免疫奖励黑客现象：在惩罚 LogP 优化实验中，非受限优化器倾向于沿着隐向量模长发散的方向搜索，利用虚拟筛选计算器对长线性碳氢链的偏好获得虚高分数；而超球面约束强制候选向量停留在生成模型的典型集分布内，在无任何启发式正则惩罚的情况下产出高化学合理性分子。

亮点与洞察¶

将几何约束转化为代数闭式解的精妙构思：传统观点常将等式非凸约束视为优化的额外负担；本文敏锐地洞察到各向同性高斯隐空间的流形本质，主动添加超球面约束，反向激活了正交旋转对称性，将高维非凸采集优化化解为极度平滑的单调/单峰一维求根，将理论直觉与代数技巧结合得浑然天成。
重新定义高通量虚拟筛选场景下的实用 Pareto 前沿：打破了学术界“贝叶斯优化必然沉重缓慢、海量候选只能靠暴力随机采样”的思维定势，证明了高维线性代理模型在局部与方向性搜索上的强大归纳偏置，以仅仅 1.5 倍于纯采样的墙钟时间换来了数十倍的样本效率提升。
通用于连续隐空间的即插即用范式：只要生成模型具备近似各向同性的基分布且解码器单射或具有足够容量的原像流形，该算法即可跨模态直接应用于分子、多肽、单克隆抗体结构及多模态扩散生成，迁移门槛极低。

局限与展望¶

局部平滑假设的适用边界：线性代理模型依赖于隐空间方向向量与目标函数响应之间具备基本的一阶平滑性。若生成模型的隐空间存在剧烈的不连续断裂、多模态跳变或极端“大海捞针”（needle-in-a-haystack）式的孤立奇异极值点，线性近似可能会过早陷入局部次优。
高度异构离群样本生成的探索受限：超球面薄壳约束假设高价值候选样本必然在基分布的高似然典型集之内。若下游某些特定任务极度偏好完全脱离训练分布的特异结构（out-of-distribution），强制留在 $\sqrt{d}$ 球面上可能会对激进的外推探索构成物理限制。
未来方向：探索在超球面上引入低阶球谐函数或多项式核的闭式扩展，以在保持计算速度的同时提升对局部非线性曲率的建模能力；并将该方案与无分类器引导（CFG）的动态退火进一步深度整合。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙利用各向同性隐空间的超球面对称性，首次推导出 EI/TS/UCB 的近乎闭式一维求根解，构思极其惊艳。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖分子、蛋白质及 16K 超高维图像生成，且详尽评估了原像存在性、典型性似然与奖励黑客实验，证据链极具说服力。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严丝合缝，几何直觉交代极为清晰，实验与图表表达精准。
价值: ⭐⭐⭐⭐⭐ 彻底解决了虚拟筛选场景下 LSBO 因优化开销过大无法落地的关键工业痛点，极具实用普及价值。