Uni-Mol+加速量子化学属性预测，深势科技、北大研究登Nature子刊

编辑 |KX

Uni-Mol 是深势科技于 2022 年发布的一款基于分子三维结构的通用大模型，其性能优越、模型泛化能力强，在小分子性质预测、蛋白靶点预测、量子化学性质预测、MOF 材料吸附性能预测等任务上都超越了现有的解决方案。

今年 3 月，深势科技与清华大学等合作，提出基于 Uni-Mol 的领域专属模型 Uni-MOF，用于预测各类工况下纳米多孔材料对各类气体的吸附性能，预测精度高达 0.98。

近日，深势科技与北京大学合作，推出新一代模型 Uni-Mol+，迭代升级后的模型拥有更大的参数量，更多的预训练数据量，并展现出更强大的通用性。

Uni-Mol+ 是一种利用 3D 构象进行精确量子化学属性预测的深度学习方法。基准测试结果表明，Uni-Mol+ 显著提高了各种数据集中 QC 属性预测的准确性。

相关研究以「Data-driven quantum chemical property prediction leveraging 3D conformations with Uni-Mol+」为题，于 8 月 19 日发布在《Nature Communications》上。

论文链接：https://www.nature.com/articles/s41467-024-51321-w

量子化学 (QC) 性质预测对于计算材料和药物设计至关重要，但依赖于昂贵的电子结构计算，如密度泛函理论 (DFT)。

深度学习方法使用 1D SMILES 或 2D 图作为输入来加速这一过程，但难以实现高精度，因为大多数 QC 性质都依赖于精细的 3D 分子平衡构象。

为了应对这一挑战，深势科技提出了一种名为 Uni-Mol+ 的方法。

Uni-Mol+ 凭借精心设计的模型主干和训练策略，在各项基准测试中表现出优异的性能。

研究的主要贡献可以概括如下：

• 研究人员利用从 RDKit 生成的构象到 DFT 平衡构象的构象优化，开发了一种用于 QC 属性预测的新范例。
• 通过生成伪轨迹（pseudo trajectory）并从中采样策略，基于伯努利分布和均匀分布的混合，创建了一种用于 3D 构象优化的新训练策略。
• Uni-Mol+ 的整个框架具有重要的经验价值，因为它在两个广受认可的基准PCQM4MV2 和 Open Catalyst 2020 (OC20) 上的性能明显优于之前的研究。

Uni-Mol+ 概述

对于任何分子，Uni-Mol+ 首先通过廉价方法（例如来自 RDKit 和 OpenBabel 的基于模板的方法）获得原始 3D 构象。然后，它通过原始构象的迭代更新过程学习目标构象，即由 DFT 优化的平衡构象。在最后一步中，根据学习到的构象预测 QC 属性。

为了有效地学习这个构象更新过程，研究人员提出了一个双轨 Transformer 模型主干和一种新颖的训练方法。

图 1：Uni-Mol+ 的整体架构。（来源：论文）

Uni-Mol+ 的模型主干是一个双轨 Transformer，由一个原子表示轨道和一个对表示轨道组成。

与之前 Uni-Mol 中使用的 Transformer 主干相比，进行了两项重大更新：

（1）通过原子表示的外积（称为 OuterProduct）增强对表示，进行原子到对的通信，并使用三角算子（称为 TriangularUpdate）来增强 3D 几何信息。这两个算子在AlphaFold2 中被证明是有效的。

（2）采用迭代过程不断更新 3D 坐标以达到平衡构象。使用 R 表示构象优化的 rounds数。

为了学习构象更新过程，研究人员提出了一种新颖的训练策略。从 RDKit 生成的原始构象和 DFT 平衡构象之间的轨迹中采样构象，并使用采样的构象作为输入来预测平衡构象。必须注意的是，在许多数据集中，实际轨迹通常是未知的；因此，研究人员使用一种假定两个构象之间存在线性过程的伪轨迹。

此外，还设计了一种采样策略，用于从伪轨迹中获取构象，作为模型在训练期间的输入。该策略混合使用伯努利分布和均匀分布。伯努利分布解决了 (1) 训练和推理之间的分布转变，以及 (2) 增强了从平衡构象到 QC 属性的精确映射的学习。同时，均匀分布生成额外的中间状态作为模型输入，有效地增强了输入构象。

基准测试

研究人员在两个大规模数据集基准PCQM4MV2 和 Open Catalyst 2020 (OC20) 上评估了 Uni-Mol+ 的性能。

首先，将之前提交给 PCQM4MV2 排行榜的模型作为基准。除了默认的 12 层模型外，研究人员还评估了 Uni-Mol+ 的性能，其两个变体分别由 6 层和 18 层组成。这旨在探索当模型参数大小改变时模型性能如何变化。

结果如下：

（1）Uni-Mol+ 在单模型性能验证数据上比之前的SOTA 高出 0.0079，相对提高了 11.4%。

（2）Uni-Mol+ 的所有三种变体都比之前的基线表现出显著的性能提升。（3）尽管 6 层的 Uni-Mol+ 的模型参数少得多，但它的表现优于所有之前的基线。

（4）将层数从 6 层增加到 12 层可显著提高准确度，并以相当大的优势超越所有基线。

（5）18 层的 Uni-Mol+ 表现出最高的性能，以显著的优势超越所有基线。这些发现强调了 Uni-Mol+ 的有效性。

（6）单个 18 层 Uni-Mol+ 模型在排行榜（测试开发集）上的表现值得关注，特别是因为它超越了之前最先进的方法，而无需使用集成或其他技术。相比之下，之前最先进的 GPS++ 依赖于 112 个模型集成，并包括验证集进行训练。

Open Catalyst 2020 (OC20) 数据集专门用于促进催化剂发现和优化的机器学习模型的开发。在该研究中，重点关注始结构到松弛能量（IS2RE）任务。

研究人员对 OC20 IS2RE 验证和测试集上的各种模型进行了性能比较，如表 2 所示。从表中可以看出，Uni-Mol+ 在平均绝对误差 (MAE) 和阈值内能量 (EwT) 方面都明显优于所有之前的基线。这证明了 Uni-Mol+ 的卓越性能。研究结果强调了 Uni-Mol+ 在捕捉材料系统中复杂相互作用方面的有效性，以及它在各种计算材料科学任务中广泛应用的潜力。

消融研究

研究人员将对 Uni-Mol+ 进行了全面的消融研究。对 PCQM4Mv2 数据集进行了消融研究，采用默认的 12 层 Uni-Mol+ 配置。研究结果总结在表 3 中，其中 No.1 是默认设置，No.2–7 重点检查模型主干，No. 8–No. 17 重点检查训练策略。

研究结果如下：

（1）比较 No. 8、No. 9 和 No. 10，发现仅从一种构象中采样效果不佳。

（2）通过比较 No. 8、No. 9 和 No. 11，可以推断出从 RDKit 和目标构象的混合中采样会产生令人满意的结果（有效 MAE 为 0.0697）。但是，如果仅从目标和中间构象（No. 12）采样，结果并不令人满意（有效 MAE 为 0.0753）。这一结果表明，从 w1.0中采样是必要的，因为它减少了训练和推理之间的分布偏移。

（3）从三种构象类型中采样的默认策略（No. 1）表现出最佳性能。

（4）改变混合分布的权重（No. 13–17）不会导致比默认策略更好的性能。此外，随着 w0.0的减少，性能会变差。这表明默认加权方案适合这项任务。

（5）比较 No.18 和 No.1 的结果后，很明显，Noisy Nodes（No.18，有效 MAE 为 0.0760）的性能明显低于 Uni-Mol+（No.1，有效 MAE 为 0.0696）。这种巨大的性能差距（0.0760 vs. 0.0696），凸显了所提出的训练策略比以前采用的策略更高效。

（6）对比 No.19 和 No.18，发现在使用噪声节点策略时，之前研究中采用的模型结构比使用 Uni-Mol+ 的主干结构产生的结果更差。这一发现进一步证明了 Uni-Mol+ 的主干结构优于之前提出的模型架构。

总之，消融研究证明了 Uni-Mol+ 中采用的默认采样策略的有效性，强调了利用不同构象混合物来实现卓越性能的重要性。

构象学习的可视化分析

除了 QC 性质预测外，Uni-Mol+ 还可以预测平衡构象。虽然该研究主要集中在 QC 属性预测上，并且证明了 Uni-Mol+ 的有效性，但可视化的结果可以帮助更好地理解 Uni-Mol+ 的工作原理。因此，研究人员还为 PCQM4MV2 数据集中 Uni-Mol+ 的构象学习提供了两个额外的分析。

第一个分析评估预测的构象。如图 2 所示，Uni-Mol+ 可以有效地预测平衡构象。此外，随着更新迭代次数的增加，RMSD 变小，进一步证明了所提出的迭代坐标更新的有效性。

图 2：Uni-Mol+ 预测构象的可视化。（来源：论文）

第二个分析旨在证明 Uni-Mol+ 可以预测较低能量的构象，接近平衡构象。如图 3 所示，Uni-Mol+ 可以预测能量较低的构象。此外，初始构象和预测构象之间的能量差分布，与初始构象和平衡构象之间的能量差分布密切一致。这种相似性证明了 Uni-Mol+ 在准确预测平衡构象方面的有效性。

图 3：δ 能量分布。（来源：论文）

上述结果为所提出的 Uni-Mol+ 的有效性提供了额外的证据，因为它确实可以预测较低能量的构象，并迭代接近目标 DFT 构象。

总之，该研究提出了一种新颖的方法，能够通过辅助任务——构象优化，准确预测量子化学性质。这种方法有望提高高通量筛选的效率，并促进创新材料和分子设计。