Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

TL;DR

This paper introduces a novel 3D molecular pretraining method using an equivariant energy-based model that captures 3D geometry and improves performance on downstream tasks.

Contribution

It proposes a new equivariant energy-based pretraining framework with force and noise scale prediction tasks for 3D molecular graphs.

Findings

Outperforms existing pretraining methods on MD17 and QM9 benchmarks.

Effectively captures 3D geometric information for molecular representations.

Demonstrates robustness and invariance in 3D molecular tasks.

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.