分子Transformer：一种用于不确定性校准的化学反应预测模型

学术报告

本研究由来自IBM苏黎世研究实验室的Philippe Schwaller、Teodoro Laino、Théophile Gaudin、Costas Bekas，以及来自英国剑桥大学物理系的Philippe Schwaller、Alpha A. Lee，和来自剑桥大学化学系的Peter Bolgar、Christopher A. Hunter共同完成。该研究以题为“molecular transformer: a model for uncertainty-calibrated chemical reaction prediction”的论文形式，发表于ACS Central Science期刊，于2019年8月30日在线发表（2019, 5, 1572–1583）。

研究背景 本研究的科学领域属于人工智能在化学信息学与合成化学的交叉应用，具体聚焦于化学反应正向预测问题。在药物发现领域，有机合成是获取新分子的关键，但其复杂性构成了主要瓶颈。一个核心挑战是解决正向问题：给定反应物和试剂，预测产物。传统计算方法主要分为基于模板（template-based）和无模板（template-free）两类。基于模板的方法依赖于专家手工编写或从数据集中自动提取的反应规则库，但存在可扩展性差、模板有效性存疑以及依赖预原子映射（atom mapping）等循环依赖问题。而无模板方法，特别是将SMILES（Simplified Molecular-Input Line-Entry System）字符串视为语言序列，将反应预测视为机器翻译任务的序列到序列（seq-2-seq）模型，为克服这些限制提供了新途径。然而，早期的seq-2-seq模型基于循环神经网络（RNN），存在无法有效处理SMILES字符串中长程化学依赖关系的缺陷。此外，现有大多数方法未能可靠地估计自身预测的不确定性，而这对多步合成路径的风险评估至关重要。因此，本研究旨在开发一种更准确、更通用且能进行不确定性校准的化学反应预测模型。

研究流程 本研究的工作流程主要包括数据准备、模型构建、模型训练与评估、以及性能验证与比较四个主要阶段。

第一阶段：数据准备与预处理。 研究使用了四个公开或非公开的化学反应数据集进行训练和测试，以进行公平比较：USPTO_MIT、USPTO_LEF、USPTO_STEREO和Pistachio_2017。这些数据集均来源于专利文本，并使用SMILES字符串表示化学反应。研究特别关注了两种预处理方法：一是“分离式”（separated），即明确用“>”符号分隔反应物（贡献原子到产物）和试剂（如溶剂、催化剂）；二是“混合式”（mixed），即将所有输入分子（反应物和试剂）混合在一起，不加以区分。后者使预测任务更具挑战性，因为模型需要从更多分子中自行判断反应中心，也更符合现实情况（实际应用中区分反应物和试剂并非总是显而易见）。所有SMILES字符串都使用RDKit工具进行规范化（canonicalized），并使用特定的正则表达式进行分词（tokenized）。为了增加数据多样性并提升模型泛化能力，研究采用了SMILES数据增强（data augmentation）技术，通过为训练集中的每个反应生成一个随机SMILES表示的双份副本来扩大训练数据。

第二阶段：模型构建——分子Transformer。 本研究开发的核心模型称为“分子Transformer”（Molecular Transformer），其架构基于谷歌提出的完全基于注意力机制（attention）的Transformer模型（Vaswani et al., 2017），并进行了适应性修改。与之前基于RNN的seq-2-seq模型不同，分子Transformer完全摒弃了循环组件，核心是多头注意力（multihead attention）层。该模型是一个自回归的编码器-解码器架构：编码器通过多头注意力层处理输入的SMILES序列（反应物/试剂），将其编码为隐藏表示；解码器则利用掩码多头注意力层（关注已生成的部分产物序列）和另一个关注编码器输出的多头注意力层，逐步生成产物的SMILES序列。位置编码（positional encoding）被用来注入序列的顺序信息。模型的创新之处在于利用多头注意力机制同时关注输入序列中的多个部分，从而能够捕获反应中可能相隔很远的化学基团之间的长程相关性，这对于理解化学选择性至关重要。研究团队对基础Transformer参数进行了调整，将层数从6层减少到4层，每层大小从512减少到256，从而将可训练参数量从6500万降至1200万。他们发现至少需要4个注意力头才能达到峰值精度，但最终保留了原始的8头配置以获得最佳验证性能。训练中使用了Adam优化器，并采用了动态学习率调整策略。

第三阶段：模型训练、评估与不确定性量化。 模型在单个GPU（NVIDIA P100）上进行训练，最佳单一模型通过训练48小时后，对最后20个检查点（checkpoint）的权重进行平均得到。此外，也探索了集成多个模型以提升性能。评估指标主要关注Top-k准确率（即真实产物出现在模型预测的前k个候选产物中的比例）。一个关键的创新点是模型的不确定性估计能力。 研究者利用模型解码时生成的每个令牌（token）的概率的乘积，作为整个预测序列的置信度分数。通过在不同置信度阈值下，计算模型预测正确与否的分类性能（绘制ROC曲线并计算AUC-ROC面积），他们验证了该置信度分数作为不确定性校准指标的有效性。值得注意的是，研究发现在训练中不使用标签平滑（label smoothing）技术，虽然对最终准确率提升微乎其微，却能显著提升模型区分正确与错误预测的能力，这对于实际应用中的风险评估至关重要。

第四阶段：性能验证与深入分析。 研究进行了广泛的实验以验证分子Transformer的性能。首先，在标准基准测试集USPTO_MIT上，分子Transformer取得了90.4%的Top-1准确率和93.7%的Top-2准确率（使用分离式预处理），显著超越了文献中所有已知的基于模板或无模板（图神经网络或序列模型）的方法。即使在更具挑战性的混合式预处理设置下，其Top-1准确率也达到88.6%，仍然优于其他方法在更简单设置下的性能。在包含立体化学信息的USPTO_STEREO数据集上，分子Transformer也取得了显著优于先前seq-2-seq模型的结果，而先前的图神经网络方法则无法处理立体化学。其次，研究通过按反应模板在训练集中的出现频率（流行度）划分测试集子集，证明分子Transformer不仅在常见反应上表现优异，在罕见反应上也大幅领先于当时的先进图模型，表明模型并非简单记忆，而是学会了泛化。第三，研究展示了分子Transformer在应对化学选择性、区域选择性和立体选择性等复杂化学反应预测挑战上的能力，并给出了具体反应实例（如图3所示），这些例子均不在训练集中。第四，研究将分子Transformer与基于量子化学计算的最新区域选择性预测工具RegioSQM在特定反应类型（杂环芳烃的亲电卤代反应）上进行比较，结果显示分子Transformer（83% Top-1准确率）甚至略优于量子化学方法（81%），且计算成本低得多。第五，在一项与人类化学家的对比研究中，分子Transformer在80个随机选取的反应上的Top-1准确率（87.5%）超过了参与测试的化学家（平均最佳为76.5%）和当时最好的图模型（72.5%）。最后，研究还验证了模型通过训练自发学会了避免“炼金术”错误（例如凭空生成反应物中没有的原子），因为在实施约束波束搜索（禁止预测未出现在反应物中的原子令牌）后，模型准确率没有变化。

主要结果 1. 卓越的预测准确率： 分子Transformer在多个基准测试集上达到了当时最高的Top-1和Top-k准确率，特别是在USPTO_MIT上超过90%，证明了其强大的预测能力。 2. 强大的不确定性校准： 模型产生的置信度分数能够有效地区分正确与错误的预测，其AUC-ROC高达0.89。这意味着模型能够为每个预测提供一个可靠的风险评分，这对于评估多步合成路线的可行性至关重要。 3. 出色的泛化与选择性预测能力： 模型不仅在常见反应上表现好，在罕见反应上优势更明显；能够成功预测涉及复杂化学选择性、区域选择性和立体选择性的反应，展示了其从数据中学习深层化学规律的能力。 4. 超越专家与物理模型： 在与人类化学家的盲测对比中胜出；在特定区域选择性预测任务上，其表现甚至超越了基于第一性原理的量子化学计算模型。 5. 处理的灵活性与鲁棒性： 模型不依赖反应物与试剂的预先区分（混合模式），也无需原子映射信息，且能够处理包含立体化学的SMILES字符串，这使其更具普适性和实用性。 6. 对错误归纳偏置的克服： 通过多头注意力机制，模型克服了基于RNN的旧模型因错误假设“序列中接近的令牌化学上才相关”而产生的系统性预测错误（如图2所示），能够正确关联SMILES字符串中相距遥远但化学上相关的基团。

结论与意义 本研究成功开发并验证了“分子Transformer”——一种基于完全注意力机制的序列到序列模型，用于化学反应的正向预测。该模型在准确性、不确定性估计和通用性方面均设立了新的标杆。其主要科学价值在于：第一，证明了纯粹基于数据驱动的、无需任何手工规则的深度学习模型，能够超越依赖专家知识的传统模板方法以及早期机器学习模型，甚至在某些方面媲美或超越基于物理原理的量子化学计算。第二，模型展示出从海量反应数据中自动学习复杂化学规律（如远程官能团效应、选择性控制）的能力，为人工智能辅助发现新化学知识提供了可能。第三，模型提供的可靠不确定性估计，为将其安全、有效地整合到自动合成规划流程中奠定了基石，允许化学家评估合成路径的风险，实现“快速失败、廉价失败”。该模型自2018年8月起已集成于IBM RXN for Chemistry在线平台，被全球数千名化学家用于进行了超过4万次预测，展示了其实际应用价值。

研究亮点 1. 架构创新： 首次将完全基于注意力机制的Transformer架构成功应用于化学反应预测任务，利用多头注意力有效捕捉化学中的长程依赖关系。 2. 性能突破： 在标准基准上实现了超过90%的Top-1准确率，全面超越了当时所有已知算法。 3. 不确定性量化： 创新性地利用模型生成概率进行不确定性校准，提供了实用的预测置信度指标，AUC-ROC达0.89。 4. 处理能力全面： 模型无需预先区分反应物/试剂、无需原子映射、可处理立体化学信息，适用性更广。 5. 超越性验证： 通过在与人类专家对比、与量子化学方法对比、在罕见反应上表现等多个维度的测试，全面证明了模型的优越性和鲁棒性。

其他有价值的内容 研究还探讨了标签平滑技术对不确定性估计的负面影响，这为未来设计兼具高准确率和高校准能力的模型提供了重要经验。此外，研究中对模型在各类具体化学反应类型（如杂原子烷基化、酰基化、C-C键形成、氧化还原等）上的表现进行了详细分析（表5），为领域内研究者评估模型在不同任务上的潜力提供了细致参考。论文中展示的失败案例（图6）分析也颇具价值，有助于理解模型当前的局限性。