ConRetroBert: 基于模板的单步逆合成的EMA稳定双编码器

# ConRetroBert: EMA稳定化双编码器用于基于模板的单步逆合成
来源: https://arxiv.org/html/2605.12736
Mohammad Jahid Ibna Basher 工业工程系 中佛罗里达大学 佛罗里达州奥兰多 32816 mo253203@ucf\.edu &Ali Khodabandeh Yalabadi 工业工程系 中佛罗里达大学 佛罗里达州奥兰多 32816 yalabadi@ucf\.edu Ivan Garibay 工业工程系 中佛罗里达大学 佛罗里达州奥兰多 32816 igaribay@ucf\.edu &Ozlem Ozmen Garibay 工业工程系 中佛罗里达大学 佛罗里达州奥兰多 32816 ozlem@ucf\.edu

###### 摘要

基于模板的单步逆合成通过选择并应用显式反应模板来预测反应物，使每次预测都可追溯到一条化学变换规则。这种可解释性对于合成规划很有用，但基于模板的方法通常被认为不如无模板模型有竞争力，因为模板预测通常被形式化为一个长尾规则库上的全局分类问题。我们认为，这一弱点并非模板本身的固有缺陷，而是学习范式的限制。我们提出**ConRetroBert**，一种双编码器框架，它将基于模板的逆合成重新定义为：先进行产物-模板稠密检索，再进行候选集列表排序。在第一阶段，对比预训练学习产物与反应模板之间的共享嵌入空间。在第二阶段，一个多正例列表排序目标在挖掘出的困难负例候选集上优化模板排序，这比全词汇分类更贴近推理时的决策问题。为了在不破坏困难负例挖掘稳定性的前提下实现模板端的适应性调整，ConRetroBert采用慢速更新的指数移动平均（EMA）模板编码器来构建检索库，同时通过排序损失更新实时模板编码器。在本地USPTO-50k基准上，主要增益来自第二阶段的候选集排序，它将top-1反应准确率从50.5%提升至约61.3%，而EMA稳定化的模板适应性调整则进一步将其提升至约62.4%。本地模型在top-3、top-5和top-10上分别达到81.6%、85.3%和87.8%。我们进一步表明，基于检索的模板预测在稀有模板的长尾部分尤其强大，并且第二阶段的排序相对于仅对比检索，提高了检索模板的适用性。作为一项独立的缩放结果，从无泄漏控制的USPTO-Full检查点微调后，在USPTO-50k上达到了75.4%的top-1准确率。这些结果表明，基于模板的逆合成可以将强大的预测性能与可化学检查的预测结合起来，挑战了高准确率必须放弃显式反应模板的常见假设。代码和数据已上传至GitHub链接 (https://github.com/JahidBasher/ConRetroBert)。

## 1 引言

计算机辅助逆合成旨在为目标产物分子恢复合理的反应物集合，是现代合成规划系统的核心组成部分 (Coley et al., 2017 (https://arxiv.org/html/2605.12736#bib.bib2); Segler and Waller, 2017 (https://arxiv.org/html/2605.12736#bib.bib1))。在实际的多步合成规划中，单步逆合成模型被反复调用来扩展搜索节点，早期断开决策中的错误可能通过整个路线构建过程传播 (Torren-Peraire et al., 2024 (https://arxiv.org/html/2605.12736#bib.bib3))。因此，改进单步逆合成不仅仅是基准准确率的问题。一个有用的模型还应提供可靠、可检查且与下游规划约束兼容的预测。

最近的进展主要由无模板和半模板方法推动，包括序列生成、图编辑、多任务图学习和基于流的生成模型 (Zheng et al., 2019 (https://arxiv.org/html/2605.12736#bib.bib18); Wan et al., 2022 (https://arxiv.org/html/2605.12736#bib.bib15); Sacha et al., 2021 (https://arxiv.org/html/2605.12736#bib.bib11); Zhong et al., 2023 (https://arxiv.org/html/2605.12736#bib.bib12); Chen et al., 2023 (https://arxiv.org/html/2605.12736#bib.bib13); Wang et al., 2023 (https://arxiv.org/html/2605.12736#bib.bib4); Zhao et al., 2025 (https://arxiv.org/html/2605.12736#bib.bib5); Yadav et al., 2025 (https://arxiv.org/html/2605.12736#bib.bib7))。这些方法具有吸引力，因为它们并未显式限制于固定的反应模板库，并且近期几个系统在USPTO-50k上报告了强劲的基准表现。特别是，RetroSynFlow报告了较高的top-k准确率，并将离散流匹配呈现为通向准确且多样化的单步逆合成的灵活途径 (Yadav et al., 2025 (https://arxiv.org/html/2605.12736#bib.bib7))。与此同时，最近关于可解释和结构化逆合成的工作凸显了许多高性能系统的一个持久局限：一个模型可能生成合理的反应物集合，但仅提供有限的机制证据来解释为何选择特定的断开 (Wang et al., 2023 (https://arxiv.org/html/2605.12736#bib.bib4); Zhao et al., 2025 (https://arxiv.org/html/2605.12736#bib.bib5))。

基于模板的逆合成提供了一套互补的优势。基于模板的系统并非直接生成反应物，而是预测一个显式的反应规则，通常编码为SMARTS变换，然后将该规则应用于产物分子 (Coley et al., 2017 (https://arxiv.org/html/2605.12736#bib.bib2); Segler and Waller, 2017 (https://arxiv.org/html/2605.12736#bib.bib1); Dai et al., 2019 (https://arxiv.org/html/2605.12736#bib.bib9))。这使得预测通过构造方式可追溯。每个输出都可以关联到一个具体的变换模式，该模式可被检查、被反应引擎验证、被化学约束过滤，或用作规划器内部的行动。这些特性在科学和规划场景中尤其有价值，用户通常不仅需要理解预测的反应物，还需要理解产生这些反应物的化学操作。

尽管有这些优势，基于模板的方法通常被视为可扩展性差且不如无模板模型有竞争力。我们认为，这一局限并非模板本身的固有缺陷，而是用于模板预测的主流学习范式所致。许多基于模板的系统将问题形式化为一个长尾反应模板库上的全局分类问题 (Dai et al., 2019 (https://arxiv.org/html/2605.12736#bib.bib9); Chen and Jung, 2021 (https://arxiv.org/html/2605.12736#bib.bib6))。这引发了三个耦合问题。首先，频繁模板主导损失函数，使得稀有但有用的变换难以学习。其次，softmax分母将每个未观察到的模板都视为竞争标签，尽管多个未记录的模板可能在化学上适用或可能导致同一产物的其他合理断开。这种标签不完整性并未被我们的方法消除，但候选集排序将问题从全库竞争缩小到在推理时决定决策的化学相似替代方案。第三，输出层随模板库大小增长，使得当模板集合扩展到数十万个变换时，全局分类越来越不具吸引力。

在本工作中，我们从检索学习的角度重新审视基于模板的逆合成。我们提出**ConRetroBert**，一个针对基于模板单步逆合成的双编码器框架，它将反应模板视为可搜索的化学规则，而非大型输出词汇表中的类别标签。产物编码器和模板编码器将产物和反应模板映射到一个共享嵌入空间，因此模板预测变为最近邻检索，随后是排序后的模板应用。ConRetroBert与检索导向的基于模板逆合成方法以及使用困难负例和慢速更新编码器目标的稠密检索训练最为相关；更充分的讨论见附录G (https://arxiv.org/html/2605.12736#A7)。我们并不声称检索本身是新颖之处：先前的如MHNreact和本地模板检索等方法已经使用了产物-模板嵌入或基于邻域的证据。我们的贡献在于将双编码器对比产物-模板预训练、候选集困难负例列表排序以及EMA稳定化的模板端适应性调整整合到一个统一的基于模板的逆合成框架中。

ConRetroBert采用两阶段训练策略。第一阶段使用对称对比预训练学习一个联合的产物-模板表示空间，将每个产物与其观察到的反应模板对齐，同时使用批次中的其他样本作为负例。第二阶段则通过多正例列表排序，在动态挖掘的困难负例候选集上，将这个表示精炼成一个模板提议策略。这第二阶段是该方法的中心步骤。它用候选集排序替代了全词汇竞争，因此模型学习区分观察到的正例模板与那些在测试时实际检索到的化学和几何上相似的替代方案。在我们的实验中，这一变化也是性能提升的主要来源：在本地USPTO-50k设置下，从第一阶段模型转到冻结模板编码器的第二阶段模型，top-1反应准确率提高了10.9个百分点。

当双编码器架构中的两个编码器都被允许调整时，会出现一个关键的技术挑战。产物编码器作为查询模型应该快速更新，而模板编码器则定义了用于构建困难负例检索索引的嵌入空间。如果模板编码器变化太快，索引就会过时，挖掘出的负例可能不再反映模型当前的评分函数。为了解决这个问题，ConRetroBert引入了EMA稳定化的可训练模板编码器。实时模板编码器通过列表排序目标的梯度进行更新，而一个慢速更新的EMA影子编码器则用于在epoch级别重建检索索引。因此，EMA作为一种稳定化机制，使得在候选构建不被破坏的前提下实现模板端的适应性调整。它不应被解释为该模型增益的主要来源；相反，主要的改进来自于第二阶段的候选集排序公式。在本地USPTO-50k设置下，实验表明，第二阶段EMA相对于冻结模板编码器（TE）将平均top-1反应准确率从61.28%提升至62.44%，并且在多次运行中保持稳定。我们还发现，基于检索的模板预测在长尾部分尤其强大，并且第二阶段相对于第一阶段提高了检索模板的适用性。作为一项独立的缩放结果，从无泄漏控制的USPTO-Full检查点微调后，达到75.4%的top-1准确率。这些结果表明，显式反应模板在保留化学解释和合成规划所需的规则级透明度的同时，仍然可以保持竞争力。

我们的贡献如下。首先，我们将基于模板的单步逆合成形式化为稠密产物-模板检索后接候选集列表排序，用更贴近推理时决策的目标替代全局模板分类。其次，我们引入EMA稳定化的可训练模板编码器，在保持用于困难负例挖掘的稳定检索索引的同时，实现可控的模板端适应性调整。第三，我们为基于模板的预测提供了一个仔细的评估流程，包括反应物集排序、模板应用、去重以及从USPTO-Full到USPTO-50k的无泄漏控制迁移。第四，我们通过本地、长尾、可靠性以及匹配预算迁移分析表明，候选集排序显著增强了基于模板的逆合成，同时保留了显式反应规则的可追溯性。

## 2 方法

ConRetroBert将基于模板的单步逆合成重新定义为对显式化学变换规则的检索和排序。给定一个产物分子，模型从学习到的嵌入空间中检索候选反应模板，用列表排序目标对它们进行排序，然后应用得分最高的模板来生成反应物集。本节描述核心模型。训练和推理流程的示意性概览见附录B (https://arxiv.org/html/2605.12736#A2) 的图1 (https://arxiv.org/html/2605.12736#A1.F1)。更多架构、算法、数学、数据筛选和评估细节见附录B (https://arxiv.org/html/2605.12736#A2) 至附录D (https://arxiv.org/html/2605.12736#A4)。

### 2.1 问题形式化

设 \(x\) 表示产物分子，\(\mathcal{T}=\{\tau_1,\ldots,\tau_N\}\) 为包含 \(N\) 个反应模板的模板库。每个模板 \(\tau\) 表示为一个SMARTS变换，可通过反应引擎应用于 \(x\) 以生成零个或多个候选反应物集。对于每个产物 \(x\)，令 \(P(x) \subseteq \mathcal{T}\) 表示训练数据中与该产物关联的观测正例模板集合。该集合可能包含多个模板，因为同一产物可能出现在多个反应中或与多个记录的变换一致。

我们学习一个评分函数

\[
s_\theta(x,\tau) = \langle z_P(x;\theta_P), z_T(\tau;\theta_T) \rangle,
\tag{1}
\]

其中 \(z_P\) 和 \(z_T\) 由两个编码器塔产生的产物嵌入和模板嵌入。在推理时，模板按 \(s_\theta(x,\tau)\) 排序，应用 top 模板到 \(x\)，得到的反应物集被规范化学、去重，并按模板评分排序。

### 2.2 双编码器模板检索

ConRetroBert 使用两个独立的 Transformer 编码器，具有相同的字符级词汇表但参数独立。一个塔编码产物 SMILES 字符串，另一个塔编码模板 SMARTS 字符串。每个塔由词元嵌入、可学习的位置嵌入和 \(L\) 个 Transformer 编码层组成。我们直接将第一个词元的隐藏状态作为序列嵌入：

\[
z_P(x) = f_P(x;\theta_P), \quad z_T(\tau) = f_T(\tau;\theta_T).
\tag{2}
\]

两个嵌入均进行 \(\ell_2\) 归一化，最终得分为温度缩放的余弦相似度：

\[
s_\theta(x,\tau) = \frac{\hat{z}_P(x)^\top \hat{z}_T(\tau)}{\tau_{\mathrm{temp}}}.
\tag{3}
\]

该架构用一个可搜索的模板库替代了全局模板分类器。输出层不再随模板数量增长，训练可以专注于候选级排序而非全库分类。

### 2.3 第一阶段：对比的产物-模板对齐

第一阶段学习共享的产物-模板检索空间。给定一个小批量 \(B\) 个配对样本 \(\{(x_i, \tau_i)\}_{i=1}^B\)，我们计算相似度矩阵 \(S \in \mathbb{R}^{B \times B}\)，其中元素 \(S_{ij} = s_\theta(x_i, \tau_j)\)。我们使用对称对比损失同时训练两个编码器：

\[
\mathcal{L}_{\mathrm{S1}} = \frac{1}{2} \left[ \frac{1}{B} \sum_{i=1}^B -\log \frac{\exp(S_{ii})}{\sum_{j=1}^B \exp(S_{ij})} + \frac{1}{B} \sum_{j}
\]