利用Uni-Mol升级版进行高精度量子化学属性预测

今天为大家介绍的是来自Guolin Ke团队的一篇论文。近年来，在深度学习领域的发展已经在加快量子化学（QC）属性预测方面取得了显著进展，其通过消除了昂贵的电子结构计算需求，大大提高了效率。然而，之前的方法主要是基于一维的SMILES序列或二维的分子图进行学习，但这些方法未能实现高精度预测。这是因为QC属性主要依赖于通过电子结构方法优化的三维平衡构象，这与序列型和图形型数据截然不同。

最近的研究表明，使用深度学习加速QC属性预测具有潜力。这涉及到训练一个深度神经网络模型，使用分子输入来预测属性，从而替代繁重的DFT计算。先前的研究主要利用了一维SMILES序列或二维分子图作为分子输入，因为它们容易获得。然而，从一维SMILES和二维分子图预测QC属性可能不够有效，因为大多数QC属性是基于由DFT优化的三维平衡构象来估计的，这与一维/二维分子格式有着内在和显著的差异。为了应对这一挑战，作者在文中提出了一种叫做Uni-Mol+的方法。与直接从一维/二维数据预测QC属性的先前方法不同，Uni-Mol+首先使用如RDKit这样的低成本方法从一维/二维数据生成原始的三维构象。由于原始构象不准确，Uni-Mol+接着使用神经网络迭代更新它，使其趋向于DFT平衡构象，并从学习到的构象中预测QC属性。显而易见，这种学习范式的成功高度依赖于学习到的构象的质量。为了获得准确的预测，作者使用大规模数据集（例如，PCQM4MV2基准）来构建数百万对RDKit生成的原始构象和高质量DFT平衡构象的配对，并从这些监督信息中学习更新过程。通过精心设计的模型架构和训练策略，Uni-Mol+在各种基准测试中显示出了优越的性能。

方法部分

图 1

模型的整体架构如图1所示。在模型方面，作者设计了一个能够同时学习构象和预测QC属性的新型模型。这个模型使用原子特征、边特征和原子的3D坐标，它们都是可以学习的参数。和Uni-Mol模型类似，模型维护了两种表征：原子表征和成对表征，用于模型的不同部分。模型包含多个模块，每个块都包含原子表征和成对表征的更新。位置编码在模型中扮演了重要角色，用于编码3D空间和2D图位置信息。在这个过程中，作者将3D空间位置编码和图位置编码结合起来。相比于Graphormer的多跳图编码方式，作者采用单跳编码，使得这一步的信息处理速度更快。

原子表征是通过原子特征的嵌入来初始化的，并且在每个模块中通过自注意力机制和前馈网络（FFN）来更新。成对表征的更新使用位置编码进行初始化，并通过一系列操作，如外积、三角形更新和前馈网络来更新。此外，研究者在更新3D坐标时采用了迭代方法。作者提出了一种表征可交互网络，原子表征和成对表征在模型网络互相交互，整体架构类似于AlphaFold2。原子表征更新过程会在自注意力机制层部分受到成对表征影响，成对表征也会被原子表征的外积优化。

实验部分

表 1

作者使用两个广泛认可的公共数据集进行了基准测试：PCQM4MV2和OC20。这些数据集有助于评估Uni-Mol+在小型有机分子和催化剂系统中的性能。PCQM4Mv2数据集源自OGB大规模挑战赛，旨在促进和评估用于预测分子量子化学（QC）属性的机器学习模型的发展，特别是被称为”HOMO-LUMO gap “的目标属性。这一属性表示最高占据分子轨道（HOMO）和最低未占据分子轨道（LUMO）能量之间的差距。该数据集包含大约400万分子，通过SMILES符号表示，并为训练和验证集提供HOMO-LUMO gap的标签，但测试集的标签未公开。训练集包含DFT平衡构象，而验证和测试集则不包括。基准测试的目标是在没有DFT平衡构象的情况下，使用SMILES符号在推断过程中预测HOMO-LUMO gap。结果在表1中总结，观察结果如下：(1) Uni-Mol+的所有三个变体都在性能上相较于之前的基准有了显著的提升。(2) 尽管参数更少，但6层的Uni-Mol+在性能上超越了所有之前的基准。(3) 从6层增加到12层，准确性得到了显著的提高，超过所有基准一个很大的幅度。(4) 18层的Uni-Mol+展示了最高性能，超越所有基准。(5) 在排行榜上，单个18层Uni-Mol+模型的表现值得注意，尤其是它在不使用集成或额外技术的情况下超越了以往的最先进方法。相比之下，之前的最先进方法GPS++依赖于一个由112个模型组成的集成，并且在训练时包含了验证集。

表 2

表 3

Open Catalyst 2020（OC20）数据集专门为催化剂的发现和优化的机器学习模型开发而设计。OC20涵盖了三个任务：结构到能量和力（S2EF）、初始结构到松弛结构（IS2RS）和初始结构到松弛能量（IS2RE）。作者聚焦于IS2RE任务。它包含大约46万个训练数据点。虽然训练时提供了DFT平衡构象，但在推断期间不允许使用。此外，与PCQM4MV2数据集不同，OC20 IS2RE任务已经提供了初始构象，无需自行生成初始输入构象。作者展示了在OC20 IS2RE验证集和测试集上各种模型的性能比较，如表2和表3所示。这两个表显示了每个模型的平均绝对误差（MAE）和能量在阈值内的百分比（EwT）。从表格中明显看出，Uni-Mol+在MAE和EwT方面都显著优于所有先前的基准，展示了其卓越的性能。值得注意的是，作者的方法在所有类别中均达到了最低的MAE，包括领域内（ID）、领域外吸附（OOD Ads.）、领域外催化（OOD Cat.）、领域外两者（OOD Both）和平均值（AVG.）。此外，在EwT方面，Uni-Mol+在所有类别中始终取得最高值。

编译 | 曾全晨

审稿 | 王建民

参考资料

Lu, S., Gao, Z., He, D., Zhang, L., & Ke, G. (2023). Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+. arXiv preprint arXiv:2303.16982.