[论文解读] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+

Shuqi Lu, Zhifeng Gao+3 more|arXiv (Cornell University)|Mar 16, 2023

Machine Learning in Materials Science被引用 6

一句话总结

Uni-Mol+ 提出了一种深度学习框架，通过首先从1D/2D输入生成低成本的3D分子构象，然后使用双轨Transformer模型迭代优化为DFT优化的平衡构象，从而预测量子化学（QC）性质。该方法在PCQM4MV2和OC20基准测试中达到最先进性能，在OC20测试集上实现0.3745 eV的MAE和11.29%的EwT，显著优于先前方法。

ABSTRACT

Recent developments in deep learning have made remarkable progress in speeding up the prediction of quantum chemical (QC) properties by removing the need for expensive electronic structure calculations like density functional theory. However, previous methods learned from 1D SMILES sequences or 2D molecular graphs failed to achieve high accuracy as QC properties primarily depend on the 3D equilibrium conformations optimized by electronic structure methods, far different from the sequence-type and graph-type data. In this paper, we propose a novel approach called Uni-Mol+ to tackle this challenge. Uni-Mol+ first generates a raw 3D molecule conformation from inexpensive methods such as RDKit. Then, the raw conformation is iteratively updated to its target DFT equilibrium conformation using neural networks, and the learned conformation will be used to predict the QC properties. To effectively learn this update process towards the equilibrium conformation, we introduce a two-track Transformer model backbone and train it with the QC property prediction task. We also design a novel approach to guide the model's training process. Our extensive benchmarking results demonstrate that the proposed Uni-Mol+ significantly improves the accuracy of QC property prediction in various datasets. We have made the code and model publicly available at \url{https://github.com/dptech-corp/Uni-Mol}.

研究动机与目标

• 解决由于1D/2D分子表示与DFT计算所需的3D平衡构象之间不匹配，导致的量子化学（QC）性质预测差距。
• 克服先前深度学习模型直接从SMILES或2D图预测QC性质的局限性，这些方法无法捕捉对精度至关重要的构象依赖关系。
• 开发一种方法，通过神经网络学习将原始3D构象（如RDKit生成）逐步优化为DFT平衡构象，从而实现高精度性质预测。
• 通过在训练过程中引入伪轨迹采样策略，减轻分布偏移问题，提升在分布外设置下的泛化能力和鲁棒性。
• 建立一个可扩展的端到端框架，用于基于深度学习的分子性质高通量筛选，减少对昂贵DFT计算的依赖。

提出的方法

• 使用如RDKit等低成本工具，从1D SMILES或2D分子图生成初始3D分子构象。
• 采用双轨Transformer主干网络：一轨用于原子表示，另一轨用于成对表示，两者均通过3D空间和图位置编码增强。
• 通过学习的神经网络迭代更新3D坐标，将原始构象逐步映射至DFT优化的平衡构象。
• 采用新颖的训练策略，沿从RDKit生成构象到DFT优化构象的伪线性轨迹采样中间构象，以监督优化过程。
• 通过外积、三角更新和前馈网络的多头注意力机制，反复更新成对表示以实现精细优化。
• 使用L1损失进行QC性质预测，结合去噪目标进行构象优化，端到端训练模型，利用大规模配对原始构象与DFT构象数据集。

实验结果

研究问题

• RQ1 深度学习模型是否能通过学习将3D分子构象优化为DFT优化结构，在量子化学性质预测中实现最先进精度？
• RQ2 将构象优化过程建模为一系列迭代更新，相比直接使用1D/2D输入建模，能否显著提升预测精度？
• RQ3 在训练过程中使用伪轨迹和采样中间构象，能在多大程度上增强泛化能力并减少分布偏移？
• RQ4 与标准GNN或Transformer相比，采用增强成对表示和空间编码的双轨Transformer架构是否在3D分子几何任务上表现更优？
• RQ5 所提出方法是否能有效泛化至分布外分子数据集，例如OC20基准中的数据？

主要发现

• Uni-Mol+ 在OC20 IS2RE测试集上实现0.3745 eV的平均绝对误差（MAE），显著优于此前SOTA方法（DRFormer）的0.3865 eV MAE。
• 在PCQM4MV2基准测试中，Uni-Mol+ 在所有类别中均达到最低MAE，包括0.3795 eV（域内，ID）、0.4526 eV（域外吸附，OOD Ads.）和0.4021 eV（域外两者，OOD Both）。
• Uni-Mol+ 在OC20测试集上达到11.29%的EwT（能量加权Top-1准确率），较此前SOTA（DRFormer为9.18%）提升超过2个百分点。
• 消融研究证实，双轨Transformer主干网络和伪轨迹采样策略对性能至关重要，任一组件的移除均导致精度显著下降。
• 该模型在分布外数据上表现出强鲁棒性，在OC20和PCQM4MV2基准的所有OOD类别中均取得最高EwT得分。
• 该方法不仅在域内设置下达到最先进结果，且在具有挑战性的分布外设置中也表现优异，表明其具备强大的泛化能力。