AtomFlow: Generative Prediction of Chemical Reactions in 3D Space

1 文章导读

从目标分子逆推可能的反应物，是逆合成分析的重要任务，也与药物研发和新材料设计密切相关。许多 AI 方法将反应预测建模为 SMILES 字符串之间的翻译，未直接利用三维原子坐标，也不能从模型构造上保证原子守恒。中国科学院计算技术研究所贾伟乐团队提出 AtomFlow，一个直接在三维空间中运行的流匹配生成框架。它学习从产物到反应物的原子坐标变化，并保持各元素的原子类型与数目不变。在不超过7个重原子的小分子基准 Transition1x 上，MatRIS 版本的 Top-10 准确率达到96.2%。研究同时评估了生成结构质量与批量推理效率，为三维逆合成预测提供了新的方法。

关键词：化学反应预测、流模型、原子守恒、三维分子生成、逆合成；

Citation：Enji Li, Siyu Hu, Xiao Tian, Wentao Zhang, Tao Liu, Linwang Wang, Weile Jia*. AtomFlow: Generative Prediction of Chemical Reactions in 3D Space[J]. AI for Science. DOI: 10.1088/3050-287X/ae9e53.

2 本文要点

要点1

把化学反应建模为三维原子的“流动”

AtomFlow 以产物的三维结构为起点，通过图神经网络学习从产物到反应物的坐标映射，并据此构建生成过程中的速度场。原子坐标随生成过程变化，原子类型与数目则保持不变，因此框架从构造上满足原子守恒。这里的输运轨迹描述的是模型的生成过程，不应直接等同于真实反应的动力学路径。模型输出显式三维结构，可进一步进行几何与能量评估。

图 1.AtomFlow 将化学反应表示为速度场中的原子输运过程：从产物结构（t=0）出发，沿学习到的速度场演化至反应物结构（t=1），全程保持原子类型与数目不变；下方为基于图神经网络的核心架构。

要点2

原子守恒与三维结构预测质量

AtomFlow 在生成过程中保持全部原子，满足100%原子守恒这一约束。在 Transition1x 上，采用 MatRIS 骨干的 AtomFlow_P2R_MR 获得79.30%的 Top-1 和96.20%的 Top-10 准确率；Top-10 表示在排名前十的候选中找到了符合评估标准的反应物。论文结构质量评估中，生成结构相对真实结构的 RMSD 中位数为0.0147 Å，均值为0.2007 Å，说明典型误差较小，但部分案例仍存在明显偏差。Transition1x 的反应数量约为常用 USPTO 基准的百分之一；由于数据集、分子规模及评估协议不同，这不能视为相同任务下仅需1%训练数据的证明，也不能据此作受控的模型优劣比较。

图 2.生成质量评估：（a）RMSD 与能量偏差随函数评估次数（nfe）的收敛行为；（b）nfe 为5与100时的误差分布；（c）代表性分子的真实结构（上）与生成结构（下）对比。能量偏差单位为 kcal mol⁻¹ atom⁻¹，RMSD 单位为 Å。

要点3

噪声校准带来候选多样性与鲁棒性

AtomFlow 在训练与推理中引入可控高斯噪声。对同一产物多次带噪采样，再积分同一确定性流场，可以生成多个候选反应物结构；模型对候选去重，并按出现频次排序。实验表明，噪声强度需要校准，适度噪声有助于预测，过强噪声则可能降低准确率。图3展示噪声评估和排序流程，图4展示具体案例：有些真实结构排在第一位，有些需要查看后续候选，也存在未能正确预测的情况。候选出现频次用于排序，并不等同于经过校准的置信概率。

图 3.噪声强度的影响与 Top-k 评估流程：（a）训练与推理噪声的影响；（b）多次采样、去重与频次排序；（c）不同方法的 Top-k 结果。跨数据集曲线仅供参考，最终准确率以论文表2为准。

图 4.代表性候选结构案例。每行给出真实反应物及按频次排名第1、2、5、10的候选；绿色表示与真实结构一致，红色表示不一致。案例同时包含成功与失败情形。

要点4

轻量模型支持高效批量推理

在效率测试中，采用600万参数 GVP 骨干的模型在单张 H100-80G GPU 上完成1,000次预测约需20秒，10,000次预测约需50秒。这些耗时对应论文的小分子批量推理设置；达到96.20% Top-10 的则是2000万参数的 MatRIS 版本，两组结果来自不同模型配置。批量生成三维候选结构有望为高通量筛选和高效逆合成预测提供支持，实际应用仍需结合反应条件、化学可行性和进一步验证。

3 总结与未来展望

AtomFlow 将原子守恒嵌入三维生成模型，在 Transition1x 小分子基准上展示了反应物结构预测能力。原子守恒是必要的基本约束，但并不自动保证预测反应正确、能量合理或实验可行。向更大、更复杂体系推广仍需验证。后续可探索金属体系、溶剂化与电子效应，引入能量感知训练或量子力学精修，并研究过渡态条件下的正向预测。相关代码、模型权重和复现脚本已开源，便于进一步检验和扩展。

4资源链接

代码开源地址：github.com/lienji/AtomFlow

支持单位：