[论文解读] Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching

Shengchao Liu, Hongyu Guo+1 more|arXiv (Cornell University)|Jun 27, 2022

Machine Learning in Materials ScienceMaterials Science被引用 17

一句话总结

本文提出GeoSSL-DDM，一种新颖的3D分子几何预训练框架，通过SE(3)-不变的去噪距离匹配，利用坐标去噪建模势能面，学习鲁棒的分子表征。通过分数匹配将坐标去噪转化为成对距离去噪，该方法在22项下游几何分子预测任务中达到最先进性能，优于现有的2D与3D预训练基线。

ABSTRACT

Molecular representation pretraining is critical in various applications for drug and material discovery due to the limited number of labeled molecules, and most existing work focuses on pretraining on 2D molecular graphs. However, the power of pretraining on 3D geometric structures has been less explored. This is owing to the difficulty of finding a sufficient proxy task that can empower the pretraining to effectively extract essential features from the geometric structures. Motivated by the dynamic nature of 3D molecules, where the continuous motion of a molecule in the 3D Euclidean space forms a smooth potential energy surface, we propose GeoSSL, a 3D coordinate denoising pretraining framework to model such an energy landscape. Further by leveraging an SE(3)-invariant score matching method, we propose GeoSSL-DDM in which the coordinate denoising proxy task is effectively boiled down to denoising the pairwise atomic distances in a molecule. Our comprehensive experiments confirm the effectiveness and robustness of our proposed method.

研究动机与目标

• 为解决3D分子几何数据缺乏有效预训练方法的问题，而这类数据对准确的分子性质预测至关重要。
• 开发一种自监督代理任务，以捕捉势能面（PES）中分子构象的动态与噪声特性。
• 通过将直接去噪3D坐标这一挑战转化为使用SE(3)-不变分数匹配去噪成对原子距离，克服直接去噪3D坐标的困难。
• 实现对3D空间中旋转和平移扰动不变的鲁棒分子表征学习。
• 通过实证验证所提方法在多样化3D分子预测基准上的优越性。

提出的方法

• 提出GeoSSL，一种通用的几何自监督学习框架，通过最大化干净分子几何与其中噪声扰动之间的互信息。
• 引入GeoSSL-DDM，一种SE(3)-不变的分数匹配方法，将坐标去噪重构为成对原子距离的去噪，以确保旋转与平移不变性。
• 采用受NCSN启发的多层去噪分数匹配目标，实现在多个去噪阶段的渐进式去噪。
• 使用基于能量的建模（EBM）近似真实数据分布，并最大化干净与噪声几何之间的互信息。
• 利用图神经网络（GNNs）编码分子结构并预测去噪距离，损失通过分数匹配进行优化。
• 应用互信息的下界以使优化可行，有效训练模型从噪声输入中重建干净几何。

实验结果

研究问题

• RQ1 基于3D分子几何的自监督预训练框架是否能在下游分子预测任务中超越现有的2D图基方法？
• RQ2 与直接的3D坐标去噪相比，去噪成对原子距离是否是更有效且更易处理的代理任务，适用于3D分子表征学习？
• RQ3 去噪层数量如何影响3D分子预训练中模型的性能与鲁棒性？
• RQ4 SE(3)-不变分数匹配是否能有效捕捉分子的平滑势能面，同时保持对旋转和平移的不变性？
• RQ5 所提方法是否能在多样化的分子数据集和下游任务（如力预测与能量回归）中实现泛化？

主要发现

• GeoSSL-DDM在22项下游几何分子预测任务中达到最先进性能，优于九种预训练基线，包括2D图基方法与其他3D方法。
• 在MD17数据集上，GeoSSL-DDM在力预测任务中达到1.637 ± 0.01的平均绝对误差，为所有对比方法中的最低值。
• 在某些力预测任务中，该方法相比监督基线将平均绝对误差降低高达20%，表明自监督预训练具有强大的泛化能力。
• 实验表明，增加去噪层数量（L > 1）可进一步提升性能，证实多阶段去噪优于单层方法。
• GeoSSL-DDM在所有评估任务中均优于并行工作[76]，尤其在更深的去噪网络中优势更明显，表明多层分数匹配的优势。
• 消融研究证实，去噪成对距离比直接坐标去噪更有效，且SE(3)不变性显著提升了模型的鲁棒性与性能。