Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching

TL;DR

The work tackles the challenge of pretraining molecular representations from 3D geometry by introducing GeoSSL, a coordinate-denoising self-supervised framework. It leverages SE(3)-invariant score matching to convert coordinate denoising into denoising pairwise distances (GeoSSL-DDM), maximizing the mutual information between two perturbed conformers. Empirically, GeoSSL-DDM achieves robust, state-of-the-art performance across 22 downstream tasks spanning quantum mechanics, force prediction, and binding affinity, outperforming a wide range of baselines. This approach demonstrates the value of pure 3D geometry pretraining and provides a general MI-based methodology for structured geometric data, with potential extensions to 2D topology and protein geometries.

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.

Figures (4)

• Figure 1: Illustration on coordinate geometry of molecules. The molecule is in a continuous motion, forming a potential energy surface (PES), where each 3D coordinate (x-axis) corresponds to an energy value (y-axis). The provided molecules, i.e., conformers, are in the local minima (${\bm{g}}_1$). It often comes with noises around the minima (e.g., statistical and systematic errors or vibrations), which can be captured using the perturbed geometry (${\bm{g}}_2$).
• Figure 2: Pipeline for GeoSSL-DDM. The ${\bm{g}}_1$ and ${\bm{g}}_2$ are around the same local minima, yet with coordinate noises perturbation. Originally we want to conduct coordinate denoising between these two views. Then as proposed in GeoSSL-DDM, we transform it to an equivalent problem, i.e., distance denoising. This figure shows the three key steps: extract the distances from the two geometric views, perform distance perturbation, and denoise the perturbed distances. Notice that the covalent bonds in the 3D data are added for illustration only.
• Figure 3: Pipeline for denoising coordinate matching.
• Figure 4: Pipeline for GeoSSL-DDM. The ${\bm{g}}_1$ and ${\bm{g}}_2$ are around the same local minima, yet with coordinate noises perturbation. Originally we want to do coordinate denoising between these two views. Then as proposed in GeoSSL-DDM, we transform it to an equivalent problem, i.e., distance denoising. This figure shows the three key steps: extract the distances from the two geometric views, perform distance perturbation, and denoise the perturbed distances.