[Paper Review] 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 Science17 citations

TL;DR

This paper proposes GeoSSL-DDM, a novel 3D molecular geometry pretraining framework that leverages SE(3)-invariant denoising distance matching to learn robust molecular representations by modeling the potential energy surface through coordinate denoising. By transforming coordinate denoising into pairwise distance denoising via score matching, the method achieves state-of-the-art performance on 22 downstream geometric molecular prediction tasks, outperforming existing 2D and 3D pretraining baselines.

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.

Motivation & Objective

• To address the lack of effective pretraining methods for 3D molecular geometry data, which are critical for accurate molecular property prediction.
• To develop a self-supervised proxy task that captures the dynamic, noisy nature of molecular conformers in the potential energy surface (PES).
• To overcome the challenge of directly denoising 3D coordinates by transforming the task into denoising pairwise atomic distances using SE(3)-invariant score matching.
• To enable robust molecular representation learning that is invariant to rotational and translational perturbations in 3D space.
• To empirically validate the superiority of the proposed method across diverse 3D molecular prediction benchmarks.

Proposed method

• Proposes GeoSSL, a general geometric self-supervised learning framework that maximizes mutual information between a clean molecular geometry and its noisy perturbation.
• Introduces GeoSSL-DDM, an SE(3)-invariant score matching method that reformulates coordinate denoising as denoising of pairwise atomic distances to ensure rotational and translational invariance.
• Employs a multi-layer denoising score matching objective inspired by NCSN, enabling progressive noise removal across multiple denoising stages.
• Uses energy-based modeling (EBM) to approximate the true data distribution and maximize mutual information between clean and noisy geometries.
• Leverages graph neural networks (GNNs) to encode molecular structures and predict denoised distances, with loss optimized via score matching.
• Applies a lower bound on mutual information to make the optimization tractable, effectively training the model to reconstruct clean geometries from noisy inputs.

Experimental results

Research questions

• RQ1 Can a self-supervised pretraining framework based on 3D molecular geometry outperform existing 2D graph-based methods in downstream molecular prediction tasks?
• RQ2 Is denoising pairwise atomic distances a more effective and tractable proxy task than direct 3D coordinate denoising for 3D molecular representation learning?
• RQ3 How does the number of denoising layers affect model performance and robustness in 3D molecular pretraining?
• RQ4 Can SE(3)-invariant score matching effectively capture the smooth potential energy surface of molecules while remaining invariant to rotations and translations?
• RQ5 Does the proposed method generalize across diverse molecular datasets and downstream tasks such as force prediction and energy regression?

Key findings

• GeoSSL-DDM achieves state-of-the-art performance on 22 downstream geometric molecular prediction tasks, outperforming nine pretraining baselines, including 2D graph-based and other 3D methods.
• On the MD17 dataset, GeoSSL-DDM achieves a mean absolute error of 1.637 ± 0.01 on force prediction, the lowest among all compared methods.
• The method reduces mean absolute error by up to 20% compared to supervised baselines on certain force prediction tasks, demonstrating strong generalization from self-supervised pretraining.
• Experiments show that increasing the number of denoising layers (L > 1) leads to further performance gains, confirming the benefit of multi-stage denoising over single-layer approaches.
• GeoSSL-DDM outperforms the parallel work [76] in all evaluated tasks, especially with deeper denoising networks, indicating the advantage of multi-layer score matching.
• The ablation study confirms that denoising pairwise distances is more effective than direct coordinate denoising, and SE(3) invariance significantly improves robustness and performance.