《Journal of Chemical Information and Modeling》:ChemTorch: A Deep Learning Framework for Benchmarking and Developing Chemical Reaction Property Prediction Models
编辑推荐:
本文介绍开源深度学习框架ChemTorch,旨在解决化学反应预测模型开发中因软件碎片化导致的复用困难、对比不公与可复现性低等问题。该框架通过模块化流水线(数据、表征、模型、训练流程)简化开发流程,内置多种数据拆分器支持分布内/外评估。文章以RDB7数据集上的能垒预测为例,比较了四种典型模型(指纹、序列、图、3D),发现结构感知模型优势明显,并强调了严格基准测试在评估模型外推能力中的重要性。
引言:化学反应建模的挑战与深度学习机遇
理解和预测化学反应如何进行对于合成规划、高通量实验和过程优化等应用至关重要。反应的可行性通常通过实验确定,或使用计算活化能垒的量子力学(QM)方法进行估算。虽然QM方法被认为是金标准,但过渡态(TS)搜索使得其计算成本过高。机器学习(ML)提供了一种替代方案,可以用毫秒级的推理速度替代昂贵的计算。深度学习方法已经超越传统方法,扩展到了更广泛的反应任务范围,例如活化能垒、焓值、产率预测、反应分类、过渡态几何结构预测乃至逆合成分析。
所有这些任务中的一个关键设计选择是反应表征。文献中探索的表征包括带有化学感知分词功能的SMILES字符串、指纹、基于图的方案以及分子构象体的三维(3D)结构。此外,每种表征都可以与一系列架构配对。这种表征和架构的多样性也反映在当前零散的软件生态系统中,这使得跨表征的性能评估变得困难,并影响了代码重用、公平比较和可复现性。为了弥补这一差距,我们推出了ChemTorch,一个用于化学反应深度学习的通用开源软件框架,旨在简化实验并加速模型开发。
设计理念:模块化的核心
在ChemTorch中,反应建模工作流被形式化为一系列模块化组件,每个组件都遵循“契约式设计”原则,实现预定义的接口。这种模块化设计促进了可重用性和可扩展性。主要抽象包括数据流水线、表征、变换、模型、例程和训练器。数据流水线负责摄入原始数据集并进行准备;表征将化学反应转换为机器可读的输入;变换是可选模块,可在建模前修改表征;模型实现神经网络架构;例程处理特定任务的训练逻辑;训练器则协调整个工作流的执行。这种模块化设计减少了工程开销,并支持跨模态和架构的系统性基准测试。该设计理念通过一张概览图清晰地展示出来。
核心功能与特性
ChemTorch的核心目标是使研究流程更简单、快速和可靠。其关键特性包括:
- 1.
分层配置:使用Hydra的分层配置系统控制流水线的所有组件,用户可以轻松地通过覆盖配置项来混合匹配组件,无需修改源代码。
- 2.
命令行界面(CLI):支持通过命令行便捷地启动实验,并支持实时配置覆盖。
- 3.
日志、可视化与运行管理:集成Weights & Biases进行实验跟踪、交互式仪表板和结果协作共享。
- 4.
超参数优化:通过Weights & Biases扫描实现无缝、资源高效的超参数优化。
- 5.
默认可复现性:所有已解析的配置与日志和指标一起存储在Weights & Biases中,任何运行都可以通过ChemTorch的CLI用一个命令精确复现。
案例研究:RDB7数据集上的能垒预测
为了展示ChemTorch的基准测试能力,我们在RDB7数据集上评估了代表指纹、序列、图和3D四种模态的四种架构,进行能垒(活化能)预测。RDB7数据集包含近12,000个有机反应的向前和向后能垒以及过渡态几何结构。研究仅使用向前能垒。
模型与评估策略
我们比较了四种模型:(1) DRFP/MLP:在差分反应指纹上的多层感知机;(2) SMILES/HAN:基于SMILES语法分词的分层注意力网络;(3) CGR/D-MPNN:在缩合反应图(CGR)上训练的定向消息传递神经网络;(4) DimeReaction:一种3D消息传递模型,通过将过渡态(TS)和反应物(R)几何嵌入之间的差异传递到额外的MLP头来预测能垒。
评估采用了多种数据拆分策略:随机拆分用于测试分布内插值能力;此外,还测试了向分布外数据的泛化能力,包括大小(按重原子数排序)、目标(按能垒值排序)、反应核心(按反应过程中断裂或形成的键的最小联合子图分组)和反应物支架(基于反应物的Bemis–Murcko支架)拆分。性能通过平均绝对误差(MAE)和均方根误差(RMSE)评估。
插值性能与结果
在随机拆分上,基于图和3D的模型显著优于基于指纹和序列的方法。DRFP/MLP和SMILES/HAN的平均绝对误差(MAE)分别为14.50和10.36 kcal/mol,而CGR/D-MPNN和DimeReaction将误差分别降至4.10和2.54 kcal/mol。这些结果突显了三点:第一,明确编码化学结构(图或3D几何)的架构始终优于基于分词和指纹的基线;第二,加入3D信息显著提高了预测性能;第三,即使提供了真实的反应物和过渡态几何结构,DimeReaction的误差仍然超过通常被认为是化学精度的1 kcal/mol阈值。四种模型在随机拆分测试集上的预测值与真实值的对比情况通过一幅比对图直观展示。
缩放与数据效率
通过在不同大小的训练数据子集上重新训练模型,研究了性能随数据集大小的缩放行为。学习曲线显示,所有模型在达到最大数据集大小时均未表现出明显的饱和迹象。CGR/D-MPNN和DimeReaction在所有训练规模上都优于SMILES/HAN和DRFP/MLP。值得注意的是,它们使用不到1000个反应(十分之一的训练数据)所达到的误差,比DRFP/MLP和SMILES/HAN使用完整数据集所达到的误差还要低。这表明学习到的表征比基于指纹的描述符更能有效捕获可转移的反应模式。平均绝对误差随训练集大小的变化趋势通过一幅曲线图清晰呈现。
分布外泛化能力
为了评估模型在更现实场景下的泛化能力,我们使用ChemTorch内置的多种数据拆分器,在多个化学相关的分布外(OOD)设置下评估了所有模型。总体而言,当根据化学结构(如反应物支架和反应核心)明确划分反应时,所有模型的性能都显著下降,这反映了泛化到结构或机制新颖的反应的难度增加。DRFP/MLP在反应核心拆分上几乎未表现出相对性能下降,但这反映了其表征局限性而非模型优势。DRFP捕获了反应中哪些局部环境出现或消失,但无法利用特定机制内的转换信息。相比之下,基于反应物支架的拆分会导致DRFP/MLP性能的显著下降。
对于能够利用结构和机制信息的模型,Spiekermann等人引入的反应物支架拆分产生的误差与随机基线相当。这一结果具有误导性,因为测试集中的产物可能与训练集中的反应物共享支架,导致分区之间出现“支架”泄露。因此,对于化学反应建模,建议使用更严格的方案,如反应核心拆分,它能更好地隔离机制新颖性,并提供更现实的分布外性能评估。
所有模型在小分子上的泛化表现相对可靠,但在预测超出训练观察范围的能垒值时都表现出严重的误差增加,这凸显了神经网络模型本质上的插值特性。所有模型在向更大分子的外推中也出现系统性的性能退化。各种模型在不同拆分策略下的分布外性能表现通过一幅雷达图进行综合比较。
结论与展望
我们介绍了ChemTorch,这是一个模块化且可扩展的开源框架,在单一代码库中支持多种化学反应表征。通过将数据流水线、表征、模型和训练例程抽象为可互换的组件,ChemTorch极大地减少了开发新型反应表征和神经网络架构相关的工程开销。作为案例研究,我们使用指纹、序列、图和3D模型将ChemTorch应用于RDB7能垒预测。这次比较展示了如何以最小的努力使用ChemTorch进行跨模态基准测试。我们的结果重现了先前报告中提到的趋势。
具体而言,像D-MPNN这样在CGR上训练的图神经网络(GNN)提供了一个仅需反应物和产物图形子映射的计算高效基线。当反应物和过渡态结构已知时,像DimeReaction这样的3D模型能提供最高的精度。与Vadaddi等人的发现类似,我们观察到模型泛化能力对评估方案的选择高度敏感。总的来说,直接在结构信息(如分子图或3D几何)上学习能带来更好的归纳偏差。然而,不同的架构和训练方法可能会产生不同的性能水平。展望未来,ChemTorch为跨模态和任务系统性地探索反应学习问题奠定了基础。其设计使得整合新的数据集、模型和训练协议变得简单,而其内置的拆分器和超参数优化工具则促进了透明和可复现的研究。