Elite-Weighted Supervised Fine-tuning用于目标导向分子优化

# 面向目标导向分子优化的精英加权监督微调
来源：https://arxiv.org/html/2609.00189
作者：Shiyun Wa
所属机构：研究部, Biogen公司, 美国马萨诸塞州剑桥02142；Manning信息与计算机科学学院, 马萨诸塞大学阿默斯特分校, 美国马萨诸塞州阿默斯特01003
作者：Anna G. Green
所属机构：Manning信息与计算机科学学院, 马萨诸塞大学阿默斯特分校, 美国马萨诸塞州阿默斯特01003
作者：Simone Sciabola
所属机构：研究部, Biogen公司, 美国马萨诸塞州剑桥02142
作者：Ye Wang††（通讯作者：[email protected]）
所属机构：研究部, Biogen公司, 美国马萨诸塞州剑桥02142

###### 摘要
目标导向优化对于引导分子生成器提出具有所需性质的候选分子至关重要。然而，它通常通过策略梯度强化学习实现，该方法需要计算生成轨迹的对数概率，而其形式取决于模型架构和生成过程。这使得优化器难以在不同架构和条件生成设计中复用。监督微调不需要这些机制，但其更新由固定数据集驱动，因此奖励信号从未进入更新过程。我们引入了精英加权监督微调（EW-SFT），它利用奖励来指导对高分分子的精英选择，并通过该集合上的模型自身预训练损失进行更新。消融实验表明，奖励信息主要通过精英选择传递，而非通过所选集合内的连续加权。由于更新仅消耗已评分的分子和模型自身的损失函数，相同的规则可应用于自回归、掩码扩散和离散流生成器，以及从头生成、基元扩展和连接子设计任务。在两个激酶参考化合物的3D形状对齐预言机调用固定预算下，EW-SFT持续优于相应的原生优化器。它在四个保留参考化合物的2D相似性预言机下进一步改进了目标导向优化，并在样本效率基准测试中达到了与更复杂的轨迹级强化学习方法相当的性能。这些结果证明了EW-SFT在分子生成器、设计约束、参考物和预言机之间具有统一性和有效性。

## 1 引言
生成模型可以直接提出具有所需结构和性质的新颖小分子，而不是筛选预定义的分子库。近期的工作转向通用分子生成器（Noutahi等人，2024；Lee等人，2025；Kaech等人，2026），支持从头生成（Brown等人，2019），即从零构建完整分子，以及片段约束生成（Fialková等人，2021；Tan等人，2022；Guo等人，2023；Rossen等人，2025），即保留固定子结构并只设计剩余部分。然而，优化此类生成器以实现设计目标通常采用策略梯度强化学习（Guo和Schwaller，2024），这需要计算每个采样分子生成轨迹的对数概率。该量取决于架构和生成过程，因此必须为每种架构指定策略梯度目标。对于一些两遍生成过程，发布的实现并未提供覆盖完整生成轨迹的联合策略。同时，保持策略梯度更新稳定需要调参的辅助机制（Engstrom等人，2020；Andrychowicz等人，2020）。相比之下，监督微调通过相同的基于似然的预训练目标更新模型，不需要这些机制。其更新由给定的数据样本驱动，因此奖励信号从未到达模型。为了填补这一空白，我们引入了精英加权监督微调（EW-SFT），它通过训练集的组成让奖励信号进入监督更新。每个生成器的遗传算子在其无条件或片段条件接口下提议分子，一个滚动精英缓冲区保留迄今发现的最高分分子，模型在该选定集合上使用其原生预训练损失执行一次加权微调步骤（图1）。先前的分子爬山和GEGL方法也从选定的高分分子中学习（Neil等人，2018；Ahn等人，2020），但其原始形式针对的是需要精确似然模型重拟合的从头优化。EW-SFT则通过每个生成器的原生损失，将这种基于选择的更新扩展到异构架构和约束生成过程。我们发现，一旦应用了精英选择，对每个选定分子的*强调*可以从连续的优势函数简化为二元指示符。由于EW-SFT直接操作于完成并评分的分子，它既不需要生成轨迹对数概率，也不需要PPO风格的重要性比率、评论家或比率裁剪。它适用于自回归、掩码扩散和离散流生成器，并且无需为两次采样遍历定义联合策略即可优化SAFE-GPT的两遍连接子生成。

图1：EW-SFT的工作流程。三个代表性的分子生成器首先在任务约束下采样以初始化遗传群体。然后每个模型充当遗传算子，重新生成新分子，调用预言机评分并更新群体。预言机分数决定哪些遗传进化分子被选入训练集，而选择的强调方式决定了它们在原生损失更新中的权重。

我们在基于配体的分子设计中评估EW-SFT，其目标是生成与已知参考化合物具有良好3D相似性的分子。这基于一个前提：具有相似3D形状和药效团排列的分子可以结合相同靶点并表现出相似的生物效应。在两个激酶参考物的3D相似性预言机调用固定预算下，EW-SFT在所有九种架构-任务组合中运行，并持续优于相应的原生优化器。它在片段约束生成基准测试（Noutahi等人，2024）的2D相似性预言机下进一步改进了优化，而该基准测试近期生成器仅用于冻结采样。最后，在从头实用分子优化（PMO）基准测试（Gao等人，2022）中，EW-SFT达到了与经过调参的、携带更多强化学习机制的策略梯度方法相当的性能。我们的贡献总结如下：
- •我们提出EW-SFT，一种通用的目标导向更新规则，仅使用已评分的分子和每个模型的原生预训练损失，适用于不同架构和任务。
- •我们表明，精英选择而非集合内的连续加权，是奖励信号传递到更新的有效渠道。
- •我们表明EW-SFT可转移到片段约束设计，其中策略梯度强化学习需要任务特定的表述，或发布的生成过程未提供联合策略，并报告了通用分子生成器在基元和连接子约束下的首个目标导向优化结果。

## 2 预备知识
#### 分子生成与优化。分子生成模型长期以来是针对单个任务构建的。基于简化分子线性输入规范（SMILES）的生成器（Weininger，1988）通常为从头设计（Olivecrona等人，2017；Loeffler等人，2024）而训练，并且为骨架修饰（Fialková等人，2021）和连接子设计（Guo等人，2023；Tan等人，2022）引入了独立架构。基于序列的附着片段嵌入（SAFE）（Noutahi等人，2024）将分子表示为可附着片段的序列，使得一个预训练生成器可以支持多个设计任务。基于此表示构建的通用生成器在架构上有所不同：SAFE-GPT（Noutahi等人，2024）是自回归的，GenMol（Lee等人，2025）是掩码扩散模型，InVirtuoGen（Kaech等人，2026）是离散流模型。它们的原生预训练目标因此不同：SAFE-GPT使用精确的自回归负对数似然，而GenMol和InVirtuoGen使用似然下界目标。这些模型通过冻结采样（Noutahi等人，2024）或在预言机调用预算下的目标导向优化（Brown等人，2019；Gao等人，2022）进行评估。现有优化器与特定的生成器架构和生成过程配对，不能在不同之间直接互换使用。
#### 策略梯度强化学习。近期的分子优化器通常使用轨迹级策略梯度（Loeffler等人，2024；Guo和Schwaller，2024；Kaech等人，2026；Zhu等人，2026），它用分配给生成轨迹的对数概率的梯度乘以其奖励或优势函数（Williams，1992）。PPO风格的变体还使用当前策略和行为策略之间的重要性比率来重用采样轨迹（Schulman等人，2017）。对于自回归生成器，生成轨迹在令牌决策上分解。对于迭代生成器（Kaech等人，2026；Zhu等人，2026），轨迹及其概率取决于架构。生成过程可能增加额外的依赖关系：SAFE-GPT连接子设计在两次独立遍历中采样两个锚点约束的序列并将其连接，因此策略梯度强化学习需要覆盖完整两遍生成轨迹的联合策略。即使可以形成策略梯度目标，保持其更新稳定也严重依赖于辅助设计选择。代码级别的选择，如价值函数裁剪和奖励缩放，可能比训练算法选择对最终奖励的影响更大，并且其最佳设置取决于环境（Engstrom等人，2020；Andrychowicz等人，2020）。在分子优化中，对基于REINFORCE方法的系统研究发现，常见的先验似然正则化不一定能改善优化（Thomas等人，2025）。先验锚点、KL惩罚、多样性过滤器和回放可以改善特定设置，但引入了许多超参数。它们的最优点随目标变化，调参消耗有限的预言机预算（Gao等人，2022），并且最终的更新仍然可能崩溃或需要大量调参（Thomas等人，2022；Guo和Schwaller，2024）。
#### 监督微调。监督微调则是在完成的分子上，使用生成器自身的原生预训练损失进行更新。因此，它不需要轨迹对数概率或策略梯度表述（Segler等人，2017；Gupta等人，2018）。但对于给定数据集，更新没有预言机反馈。一个高分分子和一个低分分子的贡献是相同的，数据的来源是更新之外的决策（Neil等人，2018）。
#### 奖励引导的监督微调。这些观察启发了一种奖励引导的原生损失更新。在候选分子被提议和评分之后，共享的微调更新应该仅取决于完成的分子、它们的预言机分数以及每个生成器的原生预训练损失，而不是产生它们的生成轨迹。加权的原生损失更新提供了这样的接口。这样就剩下三个设计决策——训练分子来自哪里，选择其中的哪些，以及每个选定分子的权重有多强。我们的消融实验分别测试了这三个作用，并表明遗传搜索提供了超出无引导采样的提议，精英选择携带了有用的奖励信号，而连续缓冲区内加权是不必要的。

## 3 方法
EW-SFT实例化了第2节中确定的加权原生损失更新。它通过每个生成器的原生遗传算子GG和预训练损失lθ\\ell\_{\theta}适应各个生成器，同时应用相同的奖励引导选择和加权表述。算法1给出了带精英选择的GenMol实例化，其他两个骨干网络仅在GG和lθ\\ell\_{\theta}上有所不同（附录B）。
### 3.1 问题设置
给定一个在结构约束C下预训练的生成模型pθp\_{\theta}，以及对f进行N次预言机调用的预算，我们的目标是在此评估预算内生成在f下得分最高的化合物。遵循Noutahi等人（2024）的分类，我们使用三种生成设计任务：从头设计（C=∅C=\varnothing），基元扩展（C是一个必须保留的片段），和连接子设计（C是两个必须都保留并连接的片段）。我们用MC\\mathcal\{M\}\_{C}表示符合C的分子可行集。冻结采样，即从冻结的pθp\_{\theta}中绘制N个样本并在θ\\theta不变的情况下对其进行评分，是当前生成器在约束下原生支持的方法。
### 3.2 精英加权监督微调
我们将EW-SFT应用于三个具有不同生成机制和基于似然的原生预训练损失lθ\\ell\_{\theta}的预训练分子生成器：SAFE-GPT（Noutahi等人，2024）（自回归，下一令牌交叉熵），GenMol（Lee等人，2025）（掩码扩散，负证据下界（NELBO）损失（Sa