离散扩散在线适应在分子优化中的设计空间研究

# 离散扩散在线自适应分子优化的设计空间研究

来源：https://arxiv.org/html/2607.02834

Trevor Chen1∗  
Ariel Dai1∗  
Jason Yang1  
Riccardo De Santi3,2  
Daniel Khalil1  
Wenda Chu1  
Nate Gruver4  
Pranav Murugan4  
Alexander F. G. Goldberg4  
Maruan Al-Shedivat4  
Yisong Yue1  

1Caltech  
2ETH Zurich  
3ETH AI Center  
4Genesis Molecular AI  

###### 摘要

分子优化通常从一个预训练的生成模型出发，该模型捕捉了有效分子结构的广泛先验。然而在测试时，目标不是从该先验中采样，而是利用有限的评估预算将生成过程转向任务特定的高奖励分子。我们针对离散扩散模型研究这一自适应问题。每一轮在线循环耦合了多个选择：该循环必须决定评估哪些候选分子、如何将奖励转化为模型更新、重用哪些反馈，以及在预训练先验的基础上偏移多远。这些选择大多被孤立地研究过，尚不清楚它们在完整的在线自适应循环中是否互补、冗余或相互干扰。我们在六个小分子结合亲和力任务和三个蛋白质适应度任务上进行了受控研究。我们发现采集、奖励塑造和模型去偏为获得更高奖励提供了互补途径，特别对小分子而言。重放进一步稳定了学习，而有效性惩罚使小分子探索保持在有效的分子流形上。综合这些发现，我们提出了一种反馈高效的分子优化实用方案：结合采集、奖励塑造、去偏、重放和有效性控制的在线微调。在匹配的评估次数和GPU小时预算下，该方案在性能上超越了离线微调和推理时搜索基线。当高奖励候选分子需要从预训练先验偏移较大时，收益最为显著。

## 1 引言

离散扩散模型已成为分子生成的有力工具，适用于小分子 [19](https://arxiv.org/html/2607.02834#bib.bib17), [42](https://arxiv.org/html/2607.02834#bib.bib3) 和蛋白质 [57](https://arxiv.org/html/2607.02834#bib.bib4)，因为它们能直接模拟真实的化学结构同时产生多样化的候选分子。然而在分子设计中，无条件生成很少有用。预训练生成器 \( p_{\theta_0} \) 编码了一个广泛的分子先验，其概率质量近似于预训练分布，而下游设计任务则定义了一个不同目标：采样能够针对任务特定评估器 \( r(x) \) 获得高奖励的分子。该评估器通常评估成本高昂，可能是一个学习的结合亲和力模型 [16](https://arxiv.org/html/2607.02834#bib.bib34)、一个结构预测管道、一个对接管道 [30](https://arxiv.org/html/2607.02834#bib.bib6) 或一个湿实验。在固定评估预算 \( B \) 下，分子优化的典型目标是生成一个最佳优化给定奖励函数的分子 \( x^* \)：

\[
x^* \in \arg\max_{x \in \mathcal{X}} \; r(x)
\tag{1}
\]

其中 \( \mathcal{X} \) 表示所考虑的有效的化学空间。该公式对应于对未知奖励函数 \( r \) 的 top-1 优化，也称为最优臂识别问题（例如 [18](https://arxiv.org/html/2607.02834#bib.bib65)）。为了解决此类生成优化问题，必须将生成器导向化学空间中与任务相关的区域，同时保持先验有效性 [13](https://arxiv.org/html/2607.02834#bib.bib64), [23](https://arxiv.org/html/2607.02834#bib.bib44), [55](https://arxiv.org/html/2607.02834#bib.bib47)。实现这一点的有效方法是通过在线模型自适应 [60](https://arxiv.org/html/2607.02834#bib.bib22), [47](https://arxiv.org/html/2607.02834#bib.bib33), [54](https://arxiv.org/html/2607.02834#bib.bib43), [32](https://arxiv.org/html/2607.02834#bib.bib11), [48](https://arxiv.org/html/2607.02834#bib.bib39), [52](https://arxiv.org/html/2607.02834#bib.bib18), [60](https://arxiv.org/html/2607.02834#bib.bib22), [46](https://arxiv.org/html/2607.02834#bib.bib16), [47](https://arxiv.org/html/2607.02834#bib.bib33), [22](https://arxiv.org/html/2607.02834#bib.bib46), [29](https://arxiv.org/html/2607.02834#bib.bib48), [4](https://arxiv.org/html/2607.02834#bib.bib15)，如下所述。

**在线模型自适应优化：循环结构。** 给定一个预训练模型，在线模型自适应的每一轮 \( t \) 交替进行两个高层步骤（图1）。第一步，从 \( p_{\theta_t} \) 中采样一批候选分子，可能通过推理时引导，然后选择用于评估的候选分子。第二步，将观察到的奖励转化为训练信号，将模型更新为 \( p_{\theta_{t+1}} \)，然后进入下一轮。

**一个紧密耦合的设计问题。** 在线自适应耦合了通常单独研究的选择：哪些候选分子接受评估、奖励如何转化为模型更新、以及如何推动生成器超越预训练先验中的偏差。汤普森采样可以通过不确定性感知的探索改善评估分配 [54](https://arxiv.org/html/2607.02834#bib.bib43), [18](https://arxiv.org/html/2607.02834#bib.bib65)；top-1/K 或 CVaR 风格的目标可以将更新集中在奖励尾部（例如 [40](https://arxiv.org/html/2607.02834#bib.bib30), [58](https://arxiv.org/html/2607.02834#bib.bib66)）；而负得分正则化可以使生成器远离过度占据的模式 [7](https://arxiv.org/html/2607.02834#bib.bib67), [40](https://arxiv.org/html/2607.02834#bib.bib30), [39](https://arxiv.org/html/2607.02834#bib.bib45)。目前尚不清楚这些机制在完整的在线自适应循环中是否互补、冗余或相互干扰。为揭示这一问题，我们对基于离散扩散的分子与蛋白质优化进行在线训练循环的综合研究，以评估代表性设计选择如何相互作用：采集、反馈重用、奖励塑造和先验去偏的选择是提供独特价值、相互增强，还是组合后变得冗余。

**图1：在线测试时反馈循环用于分子优化。** 每一轮从当前离散扩散生成器 \( p_{\theta_t} \) 中采样一个大小为 \( M \) 的候选池 \( C_t \)，选择大小为 \( K \) 的批次 \( S_t \) 进行评估，将得到的配对 \( (x, r(x)) \) 添加到已评估集合中，将分数转换为训练奖励 \( \tilde{r}(x) \)，对更新进行正则化，并在下一轮前将 \( p_{\theta_t} \) 训练为 \( p_{\theta_{t+1}} \)。

我们在涵盖小分子结合亲和力与 3 个蛋白质适应度优化的 9 个任务上的实验表明，这些机制提供了互补的益处，尤其当高奖励候选分子需要从预训练生成先验偏移较远时。对于小分子，采集、CVaR 奖励塑造和密度熵正则化提供了最清晰的互补增益：采集改善了评估分配，CVaR 将更新集中在奖励尾部，密度熵正则化推动探索远离过度占据的先验模式。重放和有效性控制更具稳定性，将有用反馈保留在更新分布中，并防止探索性更新离开有效的分子流形。在蛋白质适应度任务上，消融效应较弱，这与家族特异性先验已经覆盖高适应度区域一致。总体而言，我们的贡献包括：

*   我们提供了在固定评估和运行时预算下，离散扩散在线自适应用于分子优化的受控研究，以及基于离线微调和推理时搜索的诊断基线。
*   我们揭示了添加到在线优化循环中的简单*即插即用*组件之间的相互作用提供了互补益处：基于奖励模型的汤普森采样改善了评估分配，重放稳定学习，CVaR 奖励塑造将更新集中在高奖励候选分子上，密度熵正则化鼓励探索新模式，无效输出惩罚保持分子有效性。
*   我们将这些组件组合成一个统一的在线主动微调方案，在匹配评估次数预算下，该方案在小分子和蛋白质优化中均优于所有基线。我们还阐明，对于距离原始生成先验分布更远的分子设计解决方案，性能提升更大。

## 2 在线自适应的“即插即用”设计组件

我们将图1中的循环框架化为一个*控制台*：一个固定的在线调度，在反馈收集、转换并用于微调的点上具有可交换的选择。每一轮生成候选分子，选择子集进行评估，将评估分数转换为训练奖励，对更新进行正则化，并微调生成器。该控制台与微调器无关：任何使用每个样本的对数奖励和离线策略缓冲区的目标函数 \( \mathcal{L}_{\text{ft}} \) 都可以使用；我们使用 DDPP-LB [32](https://arxiv.org/html/2607.02834#bib.bib11) 和 VIDD [48](https://arxiv.org/html/2607.02834#bib.bib39) 进行实例化。

我们根据设计选择在反馈循环中的作用来组织它们。*引导策略* 影响哪些候选分子接受评估（第2.1节）；*自适应策略* 决定观察到的奖励如何塑造模型更新（第2.2节）；*目标失配校正* 保持优化与累积反馈和有效分子一致，因为探索会改变采样分布 \( p_{\theta_t} \)（第2.3节）。对于每一类，我们从先前工作中实例化一个代表性选择（表1）。第4节中的逐个旋钮消融测试这些选择在组合时是否仍保留独特价值，或在完整在线循环中变得冗余。

**表1：本文评估的可组合控制台设计选择。** 每个旋钮叠加在在线 DDPP-LB 之上，并在第4节中独立消融，然后组合成第4.1节的最终算法。还使用了仅搜索和离线微调基线（第2节）。

### 2.1 引导设计选择

**图2：汤普森采样。** 尽管 \( c_A \) 具有更高的平均预测，但 \( c_B \) 可能被选中，因为其更大的不确定性使得它有非零概率抽取到高于贪婪阈值。

**汤普森采样采集。** 由于评估预算有限，循环必须从生成的候选分子中选择一个小子集进行评估。最简单的规则是评估在代理模型下预测奖励最高的候选分子。这个贪婪规则纯粹是剥削性的：它将评估调用花在代理已经预测高奖励的地方。我们还考虑使用汤普森采样进行不确定性感知的采集 [51](https://arxiv.org/html/2607.02834#bib.bib49), [36](https://arxiv.org/html/2607.02834#bib.bib50)。给定一个集成代理，对于候选分子 \( x_i \in C_t \) 具有均值 \( \mu_i \) 和分歧度 \( \sigma_i \)，汤普森采样抽取一个采样奖励估计：

\[
\hat{r}_i = \mu_i + \varepsilon_i \sigma_i, \quad \varepsilon_i \sim \mathcal{N}(0,1),
\]

并按 \( \hat{r}_i \) 选择 top-K 候选分子，定义选中的批次 \( S_t \)。图2描述了汤普森采样如何平衡探索/利用权衡。具有高不确定性的候选分子可以超越具有更高预测均值但较低不确定性的候选分子，即该方案激励探索看似质量较低但不确定性高的候选分子。

### 2.2 自适应设计选择

**CVaR 风格奖励塑造。** 第二个杠杆作用于奖励本身，在评估之后、进入梯度之前。我们将得到的塑形训练奖励记为 \( \tilde{r}(x) \)。由于公式 (1) 中的设计目标由高尾而非平均生成分子控制，我们通过 CVaR 风格 [34](https://arxiv.org/html/2607.02834#bib.bib52), [35](https://arxiv.org/html/2607.02834#bib.bib53) 目标对模型进行自适应（图3(a)），这在先前的扩散模型自适应工作中已有引入 [40](https://arxiv.org/html/2607.02834#bib.bib30), [58](https://arxiv.org/html/2607.02834#bib.bib66)。该目标允许牺牲平均样本奖励以提升 top 样本的质量：

\[
\mathrm{CVaR}_{1-\tau}[r(x)] = \mathbb{E}\!\left[r(x) \mid r(x) \geq Q_{\tau}(r)\right].
\]

**(a) CVaR 塑形**

**(b) 密度熵正则化**

**图3：奖励塑造设计选择。** (a) CVaR 高尾塑形：奖励分布中高于 \( \tau \) 分位数（虚线）的右尾被着色，红色实线标记了 \( \mathrm{CVaR}_{1-\tau} = \mathbb{E}[r \mid r \geq Q_{\tau}] \)，即条件于在该尾部的期望奖励。优化将梯度集中在这个尾部期望上，而不是硬 top-k 选择。(b) 密度熵正则化（模型去偏）：从对数奖励中减去当前生成器的扩散内部似然 \( \log p_{\theta_t}(x) \)，将梯度质量从密集占据的模式重新分配到较少访问的区域，暴露了未塑形奖励会使其探索不足的次要模式。

**密度熵正则化。** CVaR 奖励塑造将更新集中在高尾，但它不区分生成器已经频繁采样的高奖励区域和仍然探索不足的高奖励区域。为了独立于 CVaR 来区分这一点，从业者可以应用一个密度熵正则化器 [7](https://arxiv.org/html/2607.02834#bib.bib67), [40](https://arxiv.org/html/2607.02834#bib.bib30), [39](https://arxiv.org/html/2607.02834#bib.bib45)，该正则化器使用当前模型密度作为生成器下区域占据程度的度量，并塑形原始对数奖励为：

\[
\log \tilde{r}_{\mathrm{debias}}(x) = \log r(x) - \gamma \log p_{\theta_t}(x),
\tag{2}
\]

如图3(b)所示，其中 \( \log p_{\theta_t}(x) \) 是边际分布。在连续扩散分析 [7](https://arxiv.org/html/2607.02834#bib.bib67), [40](https://arxiv.org/html/2607.02834#bib.bib30) 中，相应的负得分梯度 \( \nabla_x \log p_{\theta_t}(x) \) 可由可用的得分网络直接近似，但在掩蔽离散扩散中，它通过边际化估计：

\[
p_{\theta_t}(x_0) = \mathbb{E}_{t, x_t \sim q(x_t \mid x_0)}\!\bigl[\, p_{\theta_t}(x_0 \mid x_t) \,\bigr] \cdot \text{(噪声调度权重)},
\]

这是一个指数和...