Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

TL;DR

The paper tackles the challenge of 3D molecular pretraining by introducing 3D-EMGP, an $E(3)$-equivariant energy-based framework that leverages both per-atom forces and global energy signals. It couples a node-level Equivariant Force Prediction (EFP) objective, formulated as a denoising task under a doubly $E(3)$-invariant Riemann-Gaussian distribution, with a graph-level Invariant Noise-scale Prediction (INP) objective to capture global geometric structure. The model is pretrained on GEOM-QM9 and evaluated on MD17 and QM9, where it consistently outperforms state-of-the-art 2D and 3D baselines, particularly in force prediction. The results demonstrate that incorporating energy and force signals with $E(3)$-equivariant backbones yields robust, transferable 3D molecular representations suitable for MD simulations and quantum property prediction, with ablations validating each component’s importance.

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.

Figures (6)

• Figure 1: An overview of our 3D-EMGP. It consists of two tasks: node-level equivaraint force prediction and graph-level invariant noise scale prediction. ${\bm{X}},\tilde{{\bm{X}}}$ are the original and perturbed coordinates. $H$ is the input node feature and $H',\tilde{H}'$ are the output features of the original and perturbed graph. $\text{Rie}_{\sigma}(\tilde{{\bm{X}}}\mid{\bm{X}})$ is the proposed Riemann-Gaussian distribution in Eq. (\ref{['eq:rg']}).
• Figure 2: Illustration of different distributions. For typical Gaussian, a data point (in dashed circle) is a specific conformation ${\bm{X}}$, while for Riemann Gaussian, it is a set of conformations with the same geometry $[{\bm{X}}]\coloneqq\{g\cdot{\bm{X}}\mid g\in \text{E}(3)\}$.
• Figure 3: MAE on MD17 with different backbones.
• Figure 4: Energy landscape of different pretrained models.
• Figure 5: Additional energy landscape visualizations on MD17.

Theorems & Definitions (9)

• Definition 1
• Proposition 1
• Proposition 2
• Proposition 1
• Lemma 1
• proof
• proof
• Proposition 3
• proof