Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Rui JiaoJiaqi HanWenbing HuangYu RongYang Liu

article2023AAAI65 citations

Proposes an energy-motivated equivariant pretraining framework that pairs E(3)-invariant force prediction via Riemann-Gaussian denoising with graph-level noise scale estimation to improve 3D molecular representation learning.

Listen

Accurately modeling three-dimensional molecular structures without relying on scarce labeled data is critical for advancing drug discovery, materials design, and molecular simulations. Conventional machine learning pretraining primarily focuses on two-dimensional graph representations, failing to capture essential spatial geometries and physical symmetries, such as invariance to rotations and translations in space.

The article introduces and evaluates a self-supervised pretraining framework termed 3D Equivariant Molecular Graph Pretraining (3D-EMGP). The main objective is to establish whether training an energy-motivated, symmetry-preserving model on unlabeled 3D molecular structures improves downstream molecular property and force predictions compared to existing pretraining approaches.

The authors designed a physics-inspired approach using an equivariant neural network backbone that relates atomic forces to potential energy gradients. The framework combines two self-supervised objectives: a node-level force prediction task framed as position denoising using a rotation- and translation-invariant Riemann-Gaussian distribution, and a graph-level classification task that identifies the magnitude of noise applied to a molecule. The model was pretrained on 100,000 unlabeled molecular conformations from the GEOM-QM9 dataset and subsequently evaluated on two standard benchmarks: MD17 for force and energy simulation and QM9 for molecular property prediction.

The evaluation produced several key findings. First, 3D-EMGP significantly outperformed baseline and existing pretraining methods on force prediction, cutting the average mean absolute error on MD17 to 0.0968—an improvement of more than 22% over the next best approach and a greater than 50% reduction relative to training without pretraining. Second, the framework achieved state-of-the-art results across most quantum-chemical properties in QM9, such as internal energies and orbital gaps. Third, several traditional 2D pretraining techniques exhibited negative transfer on 3D tasks, producing worse results than no pretraining at all. Fourth, ablation studies verified that both the node-level force prediction and graph-level noise detection contributed meaningfully, while the Riemann-Gaussian formulation prevented performance degradation observed under standard Gaussian assumptions. Finally, testing on alternative model backbones demonstrated average error reductions of 6.9% to 36.1%, confirming the versatility of the method.

These results indicate that embedding physical symmetries and 3D geometric tasks into unsupervised representation learning substantially enhances predictive performance and model stability. For research and development organizations, this strategy reduces the costly experimental or computational overhead required to generate labeled quantum-mechanical data, accelerates molecular screening, and provides more physically reliable energy landscapes.

Organizations developing computational chemistry pipelines should integrate symmetry-aware 3D pretraining workflows rather than relying on legacy 2D graph methods when spatial geometry governs target properties. Future efforts should evaluate the framework on larger macro-molecules, such as proteins, and expand pretraining across broader, more diverse conformational datasets. While the experimental evidence strongly supports the method's effectiveness on small organic molecules, stakeholders should exercise caution when applying the model to properties governed purely by global electronic spatial extents, where transfer gains remain limited.

Beijing Academy of Artificial Intelligence Renmin University of China Tencent Tsinghua University bipartite graphs gaussians geometry Energy-Motivated Equivariant Pretraining 3D Molecular Graphs equivariant graph neural networks

arXiv: 2207.08824 jiaor17/3D-EMGP

Paper: E(n) Equivariant Graph Neural Networks, Victor Garcia Satorras et al. (2021). Introduces E(n)-equivariant graph neural networks with coordinate update mechanisms that provide the foundational symmetry-preserving neural architecture used as backbones in 3D molecular modeling.
Paper: Geometric and Physical Quantities improve E(3) Equivariant Message Passing, Johannes Brandstetter et al. (2022). Formulates steerable E(3) equivariant graph message passing incorporating directional geometric and physical vectors such as forces and energies.
Paper: E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials, Simon Batzner et al. (2021). Establishes data-efficient E(3)-equivariant graph neural networks for modeling interatomic potentials and conservative force fields from molecular geometries.
Paper: SchNet: A continuous-filter convolutional neural network for modeling quantum interactions, Kristof Schütt et al. (2017). Pioneers continuous-filter convolutions on 3D coordinates for joint molecular energy and force prediction under physical conservation laws.
Paper: Equivariant Diffusion for Molecule Generation in 3D, Emiel Hoogeboom et al. (2022). Develops equivariant diffusion and coordinate denoising frameworks for 3D molecular structures that motivate geometry-based self-supervised pretraining objectives.
Paper: Strategies for Pre-training Graph Neural Networks, Weihua Hu et al. (2020). Examines multi-level self-supervised pretraining strategies on graph neural networks and analyzes the phenomenon of negative transfer.
Paper: Machine Learning Force Fields, Oliver T. Unke et al. (2020). Provides a comprehensive conceptual overview of physical symmetries, gradient formulations, and force prediction in machine learning force fields.
Paper: ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs, Limei Wang et al. (2022). Presents complete and efficient 3D message-passing schemes capturing interatomic angles and distances on benchmarks like QM9.

Paper: Protein Representation Learning by Geometric Structure Pretraining, Zuobai Zhang et al. (2023). Extends geometric 3D self-supervised representation learning from small molecules to large-scale macromolecular protein structures.
Paper: E3Bind: An End-to-End Equivariant Network for Protein-Ligand Docking, Yangtian Zhang et al. (2023). Applies iterative equivariant 3D coordinate refinement networks to the downstream structural task of flexible protein-ligand docking.
Paper: FlowMM: Generating Materials with Riemannian Flow Matching, Benjamin Kurt Miller et al. (2024). Generalizes Riemannian geometric learning principles to periodic 3D crystal structures and material generation.
Paper: SE(3)-Stochastic Flow Matching for Protein Backbone Generation, Avishek Joey Bose et al. (2024). Builds on SE(3) geometric dynamics to develop stochastic flow-matching frameworks for generating complex 3D macromolecular backbones.
Paper: SE(3) diffusion model with application to protein backbone generation, Jason Yim et al. (2023). Utilizes SE(3) equivariant networks and geometric diffusion on rigid-body transformations to generate macromolecular 3D structures.

Cover for Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

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. Code is available at https://github.com/jiaor17/3D-EMGP.

Topics

bipartite graphs gaussians geometry energy-based models two dimensional directed graphs weighted graph masked language modeling

Keywords

Energy-Motivated Equivariant Pretraining 3D Molecular Graphs equivariant graph neural networks molecular representation learning node-level equivariant force prediction molecular property prediction Graph Contrastive Learning Invariant Noise-scale Prediction MD17 force prediction

Introduction
Related Works
Method Energy-based Molecular Modeling
Node-Level: Equivariant Force Prediction
Graph-Level: Invariant Noise-scale Prediction
Experiments
Experimental Setup
Main Results
Ablation Studies
Visualization
Conclusion
Acknowledgments
References

Knowls

Knowl 1 — Energy-Based Equivariant Molecular Representation Modeling
model/method
To model molecular conformations while preserving 3D physical symmetries, a molecule of $N$ atoms is represented by coordinates $X = [x_1, \dots, x_N] \in \mathbb{R}^{3 \times N}$ and node features $H = [h_1, \dots, h_N] \in \mathbb{R}^{m \times N}$ . An $\mathrm{E}(3)$ -equivariant graph neural network $\varphi_{\text{EGN}}$ produces invariant latent node representations:

$H' = \varphi_{\text{EGN}}(X, H, \mathcal{E})$

where $H' = [h'_1, \dots, h'_N] \in \mathbb{R}^{k \times N}$ and $\mathcal{E}$ denotes molecular graph connectivity. The scalar molecular potential energy $\hat{E}(X) \in \mathbb{R}$ is obtained via node pooling and a multi-layer perceptron projection head $\varphi_{\text{Proj}}: \mathbb{R}^k \to \mathbb{R}$ :

$\hat{E}(X) = \varphi_{\text{Proj}}\left(\sum_{i=1}^N h'_i\right)$

Because atomic forces point in the direction of steepest potential energy decrease, the predicted 3D force matrix $\hat{F}(X) \in \mathbb{R}^{3 \times N}$ acting on all atoms is derived as the negative spatial gradient of the predicted energy:

$\hat{F}(X) = -\nabla_X \hat{E}(X)$

Under Euclidean group transformations $g \in \mathrm{E}(3)$ (comprising 3D orthogonal transformations $O \in \mathbb{R}^{3 \times 3}$ and translations $t \in \mathbb{R}^3$ ), the predicted energy $\hat{E}(X)$ is $\mathrm{E}(3)$ -invariant ( $\hat{E}(g \cdot X) = \hat{E}(X)$ ), while the derived force field $\hat{F}(X)$ is $\mathrm{O}(3)$ -equivariant and translation-invariant ( $\hat{F}(OX + t) = O\hat{F}(X)$ ).
Knowl 2 — Doubly E(3)-Invariant Riemann-Gaussian Denoising Formulation
model/method
Under a Boltzmann energy distribution $p(X) = \frac{1}{Z} \exp\left(-\frac{E(X)}{kT}\right)$ , the score function corresponds directly to molecular force: $\nabla_X \log p(X) \propto -\nabla_X E(X) = F$ . Score matching enables learning this force without ground-truth labels by minimizing denoising error against perturbed coordinates $\tilde{X} \sim p(\tilde{X}|X)$ :

$\mathcal{L}_{\text{EFP-DN}} = \mathbb{E}_{G \sim \mathcal{G}, \tilde{X} \sim p(\tilde{X}|X)} \left[ \|\hat{F}(\tilde{X}) - \nabla_{\tilde{X}} \log p(\tilde{X}|X)\|_F^2 \right]$

Standard Gaussian noise violates the requirement that geometric perturbation probabilities must be independent of spatial orientation and position ( $p(g_1 \cdot \tilde{X} \mid g_2 \cdot X) = p(\tilde{X} \mid X)$ for all $g_1, g_2 \in \mathrm{E}(3)$ ). To satisfy double $\mathrm{E}(3)$ -invariance, the perturbation is modeled via a Riemann-Gaussian distribution with noise scale $\sigma$ :

$p_\sigma(\tilde{X}|X) = \text{Rie}_\sigma(\tilde{X}|X) = \frac{1}{Z(\sigma)} \exp\left(-\frac{d^2(\tilde{X}, X)}{4\sigma^2}\right)$

where the distance metric is defined over zero-mean centered coordinates $Y = X - \mu(X)$ and $\tilde{Y} = \tilde{X} - \mu(\tilde{X})$ (with $\mu(X) \in \mathbb{R}^3$ being the mean atom position):

$d(X_1, X_2) = \|Y_1^\top Y_1 - Y_2^\top Y_2\|_F$

The analytical score function of this Riemann-Gaussian distribution with respect to $\tilde{X}$ is:

$\nabla_{\tilde{X}} \log p_\sigma(\tilde{X}|X) = -\frac{1}{\sigma^2}\left[(\tilde{Y}\tilde{Y}^\top)\tilde{Y} - (Y\tilde{Y}^\top)Y\right]$

which has computational complexity $\mathcal{O}(N)$ for $N$ atoms.
Knowl 3 — Double E(3)-Invariance of the Riemann-Gaussian Conformation Distribution
theoretical result
Let $X, \tilde{X} \in \mathbb{R}^{3 \times N}$ denote molecular atomic coordinate matrices. Let $Y = X - \mu(X)$ and $\tilde{Y} = \tilde{X} - \mu(\tilde{X})$ denote coordinates translated to zero mean, where $\mu(X) = \frac{1}{N} \sum_{i=1}^N x_i$ . Define the conformation distance metric as:

$d(\tilde{X}, X) = \|\tilde{Y}^\top \tilde{Y} - Y^\top Y\|_F$

and the corresponding Riemann-Gaussian conditional probability density with scale parameter $\sigma$ as:

$p_\sigma(\tilde{X}|X) = \frac{1}{Z(\sigma)} \exp\left(-\frac{d^2(\tilde{X}, X)}{4\sigma^2}\right)$

For any rigid Euclidean transformations $g_1, g_2 \in \mathrm{E}(3)$ where $g_k \cdot X = O_k X + t_k$ ( $O_k \in \mathrm{O}(3)$ , $t_k \in \mathbb{R}^3$ ), the distribution satisfies double $\mathrm{E}(3)$ -invariance:

$p_\sigma(g_1 \cdot \tilde{X} \mid g_2 \cdot X) = p_\sigma(\tilde{X} \mid X)$

Additionally, the distance $d(\tilde{X}, X)$ and density $p_\sigma(\tilde{X}|X)$ are permutation invariant with respect to the ordering of columns (atoms) in $\tilde{X}$ and $X$ .
Knowl 4 — Multi-Scale Equivariant Force Prediction (EFP) Loss
equation
The node-level Equivariant Force Prediction (EFP) training loss across multiple noise scales $\{\sigma_l\}_{l=1}^L$ is defined as:

$\mathcal{L}_{\text{EFP-Final}} = \mathbb{E}_{G \sim \mathcal{G}, l \sim \mathcal{U}(1,L), \tilde{X} \sim p_{\sigma_l}(\tilde{X}|X)} \left[ \sigma_l^2 \left\| \frac{1}{\sigma_l} \hat{F}(\tilde{X}) - \frac{1}{\alpha} \nabla_{\tilde{X}} \log p_{\sigma_l}(\tilde{X}|X) \right\|_F^2 \right]$

where $\mathcal{G}$ is the pretraining dataset of molecular conformers, $\mathcal{U}(1, L)$ is a discrete uniform distribution over $L$ noise levels, $\hat{F}(\tilde{X}) = -\nabla_{\tilde{X}} \hat{E}(\tilde{X})$ is the predicted force vector matrix, $\nabla_{\tilde{X}} \log p_{\sigma_l}(\tilde{X}|X) = -\frac{1}{\sigma_l^2}[(\tilde{Y}\tilde{Y}^\top)\tilde{Y} - (Y\tilde{Y}^\top)Y]$ is the Riemann-Gaussian score, and $\alpha$ is a doubly $\mathrm{E}(3)$ -invariant normalization coefficient given by:

$\alpha = \frac{1}{2}\left(\|\tilde{Y}\tilde{Y}^\top\|_F + \|Y\tilde{Y}^\top\|_F\right)$
Knowl 5 — Graph-Level Invariant Noise-Scale Prediction (INP) Loss
model/method
To complement node-level force prediction with global structural awareness, a graph-level classification task is defined to identify the magnitude of perturbation applied to an input molecule. Given original coordinates $X$ and perturbed coordinates $\tilde{X} \sim p_{\sigma_l}(\tilde{X}|X)$ at noise level index $l \in \{1, \dots, L\}$ , the equivariant graph neural network backbone $\varphi_{\text{EGN}}$ generates invariant graph-level representations $u = \sum_{i=1}^N h'_i$ and $\tilde{u} = \sum_{i=1}^N \tilde{h}'_i$ .

A classification head $\varphi_{\text{Scale}}$ operates on the concatenated graph embeddings to predict logits over all $L$ noise levels:

$p = \varphi_{\text{Scale}}(u \mathbin{\Vert} \tilde{u}) \in \mathbb{R}^L$

The invariant noise-scale prediction loss is the multi-class cross-entropy between the predicted logits and the one-hot target $\mathbb{I}[l]$ :

$\mathcal{L}_{\text{INP}} = \mathbb{E}_{G \sim \mathcal{G}, l \sim \mathcal{U}(1,L), \tilde{X} \sim p_{\sigma_l}(\tilde{X}|X)} \left[ \mathcal{L}_{\text{CE}}(\mathbb{I}[l], p) \right]$
Knowl 6 — 3D Equivariant Molecular Graph Pretraining (3D-EMGP) Framework
model/method
The 3D-EMGP framework unifies local physical force estimation and global noise perception into a joint pretraining objective:

$\mathcal{L} = \lambda_1 \mathcal{L}_{\text{EFP-Final}} + \lambda_2 \mathcal{L}_{\text{INP}}$

where $\lambda_1$ and $\lambda_2$ are balancing hyperparameter coefficients, $\mathcal{L}_{\text{EFP-Final}}$ is the node-level multi-scale equivariant force prediction loss, and $\mathcal{L}_{\text{INP}}$ is the graph-level invariant noise scale prediction loss.

Pretraining proceeds on unlabeled 3D molecular conformations sampled from large datasets (e.g., GEOM-QM9). Conformations $\tilde{X}$ are sampled from the Riemann-Gaussian distribution using Langevin dynamics. Once pretrained, the backbone parameters $\varphi_{\text{EGN}}$ are transferred to downstream 3D molecular property prediction or force field simulation benchmarks (such as QM9 and MD17) and fine-tuned with task-specific projection heads.

Knowl 7 — MD17 Benchmark Force Prediction Performance

data/table

The table below reports the Mean Absolute Error (MAE) for atomic force prediction across 8 organic molecules in the MD17 dataset. All evaluated methods utilize the identical EGNN 3D backbone architecture.

Method	Aspirin	Benzene	Ethanol	Malon.	Naph.	Salicylic	Toluene	Uracil	Average
Base	0.3885	0.1861	0.0599	0.1464	0.3310	0.2683	0.1563	0.1323	0.2086
AttrMask	0.3643	0.2277	0.0567	0.1456	0.1773	0.3890	0.1093	0.1560	0.2032
EdgePred	0.4707	0.2036	0.0743	0.1268	0.2310	0.3400	0.1854	0.1933	0.2281
GPT-GNN	0.4278	0.2492	0.0703	0.1484	0.2080	0.3609	0.1541	0.2219	0.2301
InfoGraph	0.6578	0.2743	0.1257	0.2647	0.2860	0.5793	0.3821	0.4238	0.3742
GCC	0.3996	0.2346	0.0662	0.1484	0.2798	0.4263	0.3378	0.2369	0.2662
GraphCL	0.2333	0.1845	0.0503	0.0852	0.0966	0.1587	0.0725	0.1167	0.1247
JOAO	0.3646	0.2331	0.0642	0.1029	0.2017	0.3020	0.1322	0.1683	0.1961
JOAOv2	0.3447	0.2198	0.0568	0.0981	0.1889	0.2753	0.1001	0.1850	0.1836
GraphMVP	0.3198	0.2800	0.0629	0.0788	0.2350	0.2641	0.0903	0.1339	0.1831
3D Infomax	0.4592	0.1914	0.0705	0.1263	0.2642	0.3401	0.2032	0.1836	0.2298
GEM	0.3994	0.2105	0.0871	0.1161	0.1489	0.2344	0.1193	0.1827	0.1873
PosPred	0.3050	0.2023	0.0519	0.0937	0.0971	0.2481	0.0945	0.1270	0.1525
3D-EMGP	0.1560	0.1648	0.0389	0.0737	0.0829	0.1187	0.0619	0.0773	0.0968

3D-EMGP achieves the lowest force MAE across every individual molecule and an overall average MAE of 0.0968, representing a substantial improvement over the second-best baseline GraphCL (0.1247) and the un-pretrained Base model (0.2086).

Knowl 8 — QM9 Chemical Property Prediction Benchmark Performance

data/table

The table below presents the Mean Absolute Error (MAE) across 12 target properties on the QM9 dataset, evaluated with an EGNN backbone across all pretraining strategies.

Method	$\alpha$	$\Delta\epsilon$	$\epsilon_{\text{HOMO}}$	$\epsilon_{\text{LUMO}}$	$\mu$	$C_v$	$G$	$H$	$R^2$	$U$	$U_0$	ZPVE
Base	0.070	49.9	28.0	24.3	0.031	0.031	10.1	10.9	0.067	9.7	9.3	1.51
AttrMask	0.072	50.0	31.3	37.8	0.020	0.062	11.2	11.4	0.423	10.8	10.7	1.90
EdgePred	0.086	58.2	37.4	31.9	0.039	0.038	14.5	14.8	0.112	14.2	14.7	1.81
GPT-GNN	0.103	54.1	35.7	28.8	0.039	0.032	12.2	14.8	0.158	24.8	12.0	1.75
InfoGraph	0.099	72.2	48.1	38.1	0.041	0.030	16.5	14.5	0.114	14.9	16.4	1.69
GCC	0.085	57.7	37.7	32.3	0.041	0.034	12.8	14.5	0.104	13.2	13.1	1.66
GraphCL	0.066	45.5	26.8	22.9	0.027	0.028	10.2	9.6	0.095	9.7	9.6	1.42
JOAO	0.068	46.0	28.2	22.8	0.028	0.030	10.5	10.0	0.076	9.9	10.1	1.48
JOAOv2	0.066	45.0	27.8	22.2	0.027	0.028	9.9	9.2	0.087	9.8	9.5	1.43
GraphMVP	0.070	46.9	28.5	26.3	0.031	0.033	11.2	10.4	0.082	10.3	10.2	1.63
3D Infomax	0.075	48.8	29.8	25.7	0.034	0.033	13.0	12.4	0.122	12.5	12.7	1.67
GEM	0.081	52.1	33.8	27.7	0.034	0.035	13.2	13.3	0.089	12.6	13.4	1.73
PosPred	0.067	40.6	25.1	20.9	0.024	0.035	10.9	10.2	0.115	10.3	10.2	1.46
3D-EMGP	0.057	37.1	21.3	18.2	0.020	0.026	9.3	8.7	0.092	8.6	8.6	1.38

3D-EMGP achieves the best performance across 10 of the 12 property targets (such as $\alpha = 0.057$ , $\Delta\epsilon = 37.1$ , $\epsilon_{\text{HOMO}} = 21.3$ , $\epsilon_{\text{LUMO}} = 18.2$ , $\mu = 0.020$ , $C_v = 0.026$ , $G = 9.3$ , $H = 8.7$ , $U = 8.6$ , $U_0 = 8.6$ , and $\text{ZPVE} = 1.38$ ), exhibiting consistent advantages over 2D and 3D pretraining baselines.

Knowl 9 — Ablation Study of 3D-EMGP Pretraining Components

data/table

The ablation study below evaluates the impact of each design component in 3D-EMGP on MD17 average energy and force prediction MAE:

Configuration	EFP	INP	RG	Energy Head	Energy MAE	Force MAE
Base					0.1191	0.2086
Full 3D-EMGP	✓	✓	✓	✓	0.0876	0.0968
INP only		✓	✓	✓	0.0974	0.1350
EFP only	✓		✓	✓	0.0905	0.1193
Gaussian noise	✓	✓		✓	0.0912	0.1060
Distance denoising	✓	✓			0.0931	0.1292
Direct force	✓	✓	✓		0.0914	0.1267

The results show that:

Removing either the node-level force prediction (EFP only vs. INP only) degrades performance relative to joint training (Full 3D-EMGP).
Replacing the Riemann-Gaussian (RG) noise with standard Euclidean Gaussian noise increases force MAE from 0.0968 to 0.1060 due to broken double $\mathrm{E}(3)$ -invariance.
Predicting forces directly via equivariant vector output (Direct force) rather than via the negative gradient of a pooled energy head (Energy Head) worsens force MAE from 0.0968 to 0.1267.

Knowl 10 — 3D-EMGP Generalization Across Geometric GNN Backbones
empirical result
When applied to different 3D graph neural network architectures, 3D-EMGP pretraining consistently improves downstream force prediction on MD17:
1. For SchNet, 3D-EMGP reduces the average MD17 force prediction MAE by 36.1% compared to training SchNet from scratch.
2. For TorchMD-ET (Equivariant Transformer), 3D-EMGP reduces the average MD17 force prediction MAE by 6.9% compared to training TorchMD-ET from scratch.
These gains indicate that the benefits of 3D-EMGP pretraining are architecture-agnostic and transfer across diverse equivariant and invariant 3D GNN designs.

Coverage note — None was omitted; all key contributions including the energy-based equivariant model, Riemann-Gaussian formulation, EFP and INP pretraining losses, benchmark results (MD17 and QM9), ablations, and backbone generalizations are covered.

References

1.Anderson, B.; Hy, T. S.; and Kondor, R. 2019. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32.
2.Axelrod, S.; and Gomez-Bombarelli, R. 2022. GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1): 1–14.
3.Boltzmann, L. 1868. Studien uber das Gleichgewicht der lebenden Kraft. Wissenschafiliche Abhandlungen, 1: 49–96.
4.Chmiela, S.; Tkatchenko, A.; Sauceda, H. E.; Poltavsky, I.; Schutt, K. T.; and M¨ uller, K.-R. 2017. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5): e1603015.
5.Fang, X.; Liu, L.; Lei, J.; He, D.; Zhang, S.; Zhou, J.; Wang, F.; Wu, H.; and Wang, H. 2022. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 1–8.
6.Finzi, M.; Stanton, S.; Izmailov, P.; and Wilson, A. G. 2020. Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, 3165–3176. PMLR.
7.Fuchs, F.; Worrall, D. E.; Fischer, V.; and Welling, M. 2020. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. In Larochelle, H.; Ranzato, M.; Hadsell, R.; Balcan, M.; and Lin, H., eds., Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
8.Gilmer, J.; Schoenholz, S. S.; Riley, P. F.; Vinyals, O.; and Dahl, G. E. 2017. Neural message passing for quantum chemistry. In International conference on machine learning, 1263–1272. PMLR.
9.Hamilton, W.; Ying, Z.; and Leskovec, J. 2017. Inductive representation learning on large graphs. Advances in neural information processing systems, 30.
10.Han, J.; Rong, Y.; Xu, T.; and Huang, W. 2022. Geometrically Equivariant Graph Neural Networks: A Survey. arXiv preprint arXiv:2202.07230.
11.Hu, W.; Liu, B.; Gomes, J.; Zitnik, M.; Liang, P.; Pande, V. S.; and Leskovec, J. 2020a. Strategies for Pre-training Graph Neural Networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
12.Hu, Z.; Dong, Y.; Wang, K.; Chang, K.; and Sun, Y. 2020b. GPT-GNN: Generative Pre-Training of Graph Neural Networks. In Gupta, R.; Liu, Y.; Tang, J.; and Prakash, B. A., eds., KDD '20: The 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Virtual Event, CA, USA, August 23-27, 2020, 1857–1867. ACM.
13.Hutchinson, M. J.; Le Lan, C.; Zaidi, S.; Dupont, E.; Teh, Y. W.; and Kim, H. 2021. Lietransformer: Equivariant self-attention for lie groups. In International Conference on Machine Learning, 4533–4543. PMLR.
14.Kearnes, S.; McCloskey, K.; Berndl, M.; Pande, V.; and Riley, P. 2016. Molecular graph convolutions: moving beyond fingerprints. Journal of computer-aided molecular design, 30(8): 595–608.
15.Kipf, T. N.; and Welling, M. 2016. Variational Graph Auto-Encoders. NIPS Workshop on Bayesian Deep Learning.
16.Liu, S.; Wang, H.; Liu, W.; Lasenby, J.; Guo, H.; and Tang, J. 2021. Pre-training Molecular Graph Representation with 3D Geometry. In International Conference on Learning Representations.
17.Liu, Y.; Wang, L.; Liu, M.; Lin, Y.; Zhang, X.; Oztekin, B.; and Ji, S. 2022. Spherical Message Passing for 3D Molecular Graphs. In International Conference on Learning Representations.
18.Luo, S.; and Hu, W. 2020. Differentiable manifold reconstruction for point cloud denoising. In Proceedings of the 28th ACM international conference on multimedia, 1330–1338.
19.Luo, S.; Shi, C.; Xu, M.; and Tang, J. 2021. Predicting Molecular Conformation via Dynamic Graph Score Matching. Advances in Neural Information Processing Systems, 34.
20.Qiu, J.; Chen, Q.; Dong, Y.; Zhang, J.; Yang, H.; Ding, M.; Wang, K.; and Tang, J. 2020. Gcc: Graph contrastive coding for graph neural network pre-training. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 1150–1160.
21.Ramakrishnan, R.; Dral, P. O.; Rupp, M.; and Von Lilienfeld, O. A. 2014. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1): 1–7.
22.Rong, Y.; Bian, Y.; Xu, T.; Xie, W.; Wei, Y.; Huang, W.; and Huang, J. 2020. Self-Supervised Graph Transformer on Large-Scale Molecular Data. In Larochelle, H.; Ranzato, M.; Hadsell, R.; Balcan, M.; and Lin, H., eds., Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
23.Rosenstein, M. T.; Marx, Z.; Kaelbling, L. P.; and Dietterich, T. G. 2005. To transfer or not to transfer. In In NIPS'05 Workshop, Inductive Transfer: 10 Years Later.
24.Said, S.; Bombrun, L.; Berthoumieu, Y.; and Manton, J. H. 2017. Riemannian Gaussian distributions on the space of symmetric positive definite matrices. IEEE Transactions on Information Theory, 63(4): 2153–2170.
25.Satorras, V. G.; Hoogeboom, E.; and Welling, M. 2021. E (n) equivariant graph neural networks. In International Conference on Machine Learning, 9323–9332. PMLR.
26.Schlick, T. 2010. Molecular modeling and simulation: an interdisciplinary guide, volume 2. Springer.
27.Schütt, K.; Kindermans, P.-J.; Sauceda Felix, H. E.; Chmiela, S.; Tkatchenko, A.; and Müller, K.-R. 2017. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30.
28.Shi, C.; Luo, S.; Xu, M.; and Tang, J. 2021. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, 9558–9568. PMLR.
29.Song, Y.; and Ermon, S. 2020. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33: 12438–12448.
30.Stärk, H.; Beaini, D.; Corso, G.; Tossou, P.; Dallago, C.; Günnemann, S.; and Liò, P. 2022. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, 20479–20502. PMLR.
31.Sun, F.; Hoffmann, J.; Verma, V.; and Tang, J. 2020. InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
32.Sun, M.; Xing, J.; Wang, H.; Chen, B.; and Zhou, J. 2021. MoCL: Data-Driven Molecular Fingerprint via Knowledge-Aware Contrastive Learning from Molecular Graph. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, KDD '21, 3585–3594. New York, NY, USA: Association for Computing Machinery. ISBN 9781450383325.
33.Thölke, P.; and De Fabritiis, G. 2021. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations.
34.Thomas, N.; Smidt, T.; Kearnes, S.; Yang, L.; Li, L.; Kohlhoff, K.; and Riley, P. 2018. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219.
35.Velickovic, P.; Fedus, W.; Hamilton, W. L.; Lio, P.; Bengio, Y.; and Hjelm, R. D. 2019. Deep Graph Infomax. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
36.Vincent, P. 2011. A connection between score matching and denoising autoencoders. Neural computation, 23(7): 1661–1674.
37.Wallach, I.; Dzamba, M.; and Heifets, A. 2015. AtomNet: a deep convolutional neural network for bioactivity prediction in structure-based drug discovery. arXiv preprint arXiv:1510.02855.
38.Xu, K.; Hu, W.; Leskovec, J.; and Jegelka, S. 2019. How Powerful are Graph Neural Networks? In International Conference on Learning Representations.
39.You, Y.; Chen, T.; Shen, Y.; and Wang, Z. 2021a. Graph Contrastive Learning Automated. In Meila, M.; and Zhang, T., eds., Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, 12121–12132. PMLR.
40.You, Y.; Chen, T.; Shen, Y.; and Wang, Z. 2021b. Graph contrastive learning automated. In International Conference on Machine Learning, 12121–12132. PMLR.
41.You, Y.; Chen, T.; Sui, Y.; Chen, T.; Wang, Z.; and Shen, Y. 2020. Graph Contrastive Learning with Augmentations. In Larochelle, H.; Ranzato, M.; Hadsell, R.; Balcan, M.; and Lin, H., eds., Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
42.Zheng, L.; Fan, J.; and Mu, Y. 2019. Onionnet: a multiple-layer intermolecular-contact-based convolutional neural network for protein–ligand binding affinity prediction. ACS omega, 4(14): 15956–15965.
43.Zhou, G.; Gao, Z.; Ding, Q.; Zheng, H.; Xu, H.; Wei, Z.; Zhang, L.; and Ke, G. 2023. Uni-Mol: A Universal 3D Molecular Representation Learning Framework. In The Eleventh International Conference on Learning Representations.
44.Zhu, J.; Xia, Y.; Wu, L.; Xie, S.; Qin, T.; Zhou, W.; Li, H.; and Liu, T.-Y. 2022. Unified 2d and 3d pre-training of molecular representations. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2626–2636.

Citation

MLA

Jiao, R., et al. “Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs”. arXiv, 2022, http://arxiv.org/abs/2207.08824v4.

APA

Jiao, R., Han, J., Huang, W., Rong, Y., & Liu, Y. (2022). Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs. arXiv. http://arxiv.org/abs/2207.08824v4

Chicago

Jiao, R., J. Han, W. Huang, Y. Rong, and Y. Liu. 2022. “Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs”. arXiv. http://arxiv.org/abs/2207.08824v4.

Harvard

Jiao, R. et al. (2022) “Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2207.08824v4.

Vancouver

1. Jiao R, Han J, Huang W, Rong Y, Liu Y (2022) Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs. arXiv

BibTeX

@article{jiao2022energy,
  title = {Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs},
  author = {Jiao, Rui and Han, Jiaqi and Huang, Wenbing and Rong, Yu and Liu, Yang},
  year = {2022},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2207.08824v4},
  eprint = {2207.08824}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by-sa/4.0/

Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Energy-Based Equivariant Molecular Representation Modeling

Knowl 2 — Doubly E(3)-Invariant Riemann-Gaussian Denoising Formulation

Knowl 3 — Double E(3)-Invariance of the Riemann-Gaussian Conformation Distribution

Knowl 4 — Multi-Scale Equivariant Force Prediction (EFP) Loss

Knowl 5 — Graph-Level Invariant Noise-Scale Prediction (INP) Loss

Knowl 6 — 3D Equivariant Molecular Graph Pretraining (3D-EMGP) Framework

Knowl 7 — MD17 Benchmark Force Prediction Performance

Knowl 8 — QM9 Chemical Property Prediction Benchmark Performance

Knowl 9 — Ablation Study of 3D-EMGP Pretraining Components

Knowl 10 — 3D-EMGP Generalization Across Geometric GNN Backbones

References

Citation

BibTeX

Source Code

Access the Paper