[论文解读] Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Rui Jiao, Jiaqi Han+3 more|arXiv (Cornell University)|Jul 18, 2022

Machine Learning in Materials Science被引用 5

一句话总结

该论文提出3D-EMGP，一种用于3D分子图的自监督预训练框架，利用E(3)-等变消息传递和受物理启发的预训练任务——通过黎曼-高斯去噪实现等价力预测和噪声尺度预测。该方法在MD17和QM9基准上实现了最先进性能，即使在使用3D坐标增强的情况下，也优于2D预训练基线方法。

ABSTRACT

Pretraining molecular representation models without labels is fundamental to various applications. Conventional methods mainly process 2D molecular graphs and focus solely on 2D tasks, making their pretrained models incapable of characterizing 3D geometry and thus defective for downstream 3D tasks. In this work, we tackle 3D molecular pretraining in a complete and novel sense. In particular, we first propose to adopt an equivariant energy-based model as the backbone for pretraining, which enjoys the merits of fulfilling the symmetry of 3D space. Then we develop a node-level pretraining loss for force prediction, where we further exploit the Riemann-Gaussian distribution to ensure the loss to be E(3)-invariant, enabling more robustness. Moreover, a graph-level noise scale prediction task is also leveraged to further promote the eventual performance. We evaluate our model pretrained from a large-scale 3D dataset GEOM-QM9 on two challenging 3D benchmarks: MD17 and QM9. Experimental results demonstrate the efficacy of our method against current state-of-the-art pretraining approaches, and verify the validity of our design for each proposed component.

研究动机与目标

• 为解决现有2D分子预训练方法无法捕捉分子动力学和能量预测等任务所必需的3D几何信息的局限性。
• 开发一种自监督预训练框架，充分利用E(3)-等变模型和3D感知的预训练目标来挖掘3D分子几何信息。
• 在预训练损失中确保对称性不变性（E(3)-不变性），以提升对3D空间变换的鲁棒性和泛化能力。
• 通过在多个3D主干网络和基准上进行消融研究与泛化分析，评估所提组件的有效性。

提出的方法

• 采用基于等变图神经网络（如EGNN）的E(3)-等变消息传递主干网络，预测原子级力，确保在3D旋转、平移和反射下的不变性。
• 提出一种基于力预测的节点级预训练任务，通过黎曼-高斯分布重构为去噪损失，以保持E(3)-不变性。
• 设计一种图级代理任务，用于预测一对干净与噪声3D分子构象中施加的噪声尺度，以提升模型鲁棒性和泛化能力。
• 将两种预训练目标统一为单一框架3D-EMGP，并在包含3D构象的大规模GEOM-QM9数据集上进行训练。
• 采用可微分的能量建模方法：从聚合的节点表示中预测能量，并通过能量梯度计算力，提升全局一致性。
• 通过微调将预训练模型应用于下游3D任务，并添加能量投影头以实现能量景观的可视化与评估。

实验结果

研究问题

• RQ1 3D分子图的自监督预训练框架是否能在3D基准任务上超越现有的2D预训练方法？
• RQ2 通过黎曼-高斯去噪强制实现E(3)-不变性，对预训练损失的影响如何？是否显著提升模型鲁棒性和性能？
• RQ3 结合节点级（力预测）和图级（噪声尺度预测）预训练任务对最终性能的贡献是什么？
• RQ4 基于能量的建模策略——通过预测能量的梯度计算力——是否优于直接力预测？
• RQ5 所提出的预训练框架在不同3D图神经网络主干网络上的泛化能力如何？

主要发现

• 3D-EMGP在MD17力预测基准上实现了0.0876的测试MAE，显著优于次佳基线方法（0.0968），即使在使用3D坐标增强的情况下，也超越了2D预训练方法。
• 消融研究显示，若移除黎曼-高斯去噪方案，MAE上升至0.0912，证明其在保持E(3)-不变性和鲁棒性方面具有关键作用。
• 图级噪声尺度预测任务具有实质性贡献，移除后MAE从0.0876上升至0.0905，表明其在提升泛化能力方面具有重要价值。
• 基于能量的力预测策略（通过能量梯度计算力）的MAE为0.0876，低于直接力预测（0.0914），证实其在捕捉全局几何模式方面的优越性。
• 该方法在不同主干网络上表现出良好的泛化能力：在SchNet和TorchMD-ET上均观察到一致的性能提升，证实其广泛适用性。
• 能量景观的可视化显示，3D-EMGP在构象周围生成平滑、稳定的能量表面，而2D预训练方法则产生崎岖的景观，表明其更符合物理能量分布。