Research map / Graph learning / Graph self-supervised learning
Equivariant Masked Position Prediction for Efficient Molecular Representation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EMPP, a self-supervised pretraining/auxiliary task that masks an atom's position and predicts it via equivariant spherical-harmonic distributions over neighbor-based directions and radii, avoiding Gaussian-mixture denoising approximations.
Problem
Self-supervised molecular representation learning for GNNs is limited: attribute masking is ill-posed and misses quantum mechanical features, while denoising relies on hard-to-tune Gaussian mixture approximations of the potential energy surface.
Approach
Mask an atom's 3D position (removing its node while keeping its atomic number) and predict its position from neighboring atoms' equivariant embeddings, using a predicted radius distribution and a spherical-harmonic-projected direction distribution trained with KL divergence against soft targets; this makes the task near well-posed and bypasses Gaussian-mixture approximations.
Key results
- Without extra data, EMPP (3-Mask) on Equiformer achieves best results on 11 QM9 tasks, e.g., alpha 0.039 vs 0.046 for Equiformer.
- With PCQM4Mv2 pretraining on TorchMD-Net backbone, EMPP surpasses denoising methods on nine QM9 tasks and is best on seven, e.g., alpha 0.035 and HOMO 13.7 meV.
- On MD17, EMPP (3-Mask) achieves best performance on most molecules, with energy-only variant reaching e.g. 4.8 meV on Aspirin energy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use 3-Mask instead of 1-Mask (ablation) | EMPP with Equiformer · QM9 (no extra pre-training) | α MAE: .041 -> .039 | Helped |
| Use 3-Mask instead of 1-Mask (ablation) | EMPP with Equiformer · QM9 (no extra pre-training) | ε_HOMO MAE: 14 -> 13 | Helped |
| Use 3-Mask instead of 1-Mask (ablation) | EMPP with Equiformer · MD17 test | energy MAE: 5.1 -> 5.0; 2.3 -> 2.1; 2.2 -> 2.0; 3.0 -> 3.0; 4.1 -> 3.7; 4.3 -> 4.1; 3.5 -> 3.3; 4.1 -> 4.0 | Mixed |
| Encode energy only instead of energy and forces (ablation) | EMPP with Equiformer · MD17 test | energy MAE: 5.0 -> 4.8 (Aspirin); 3.7 -> 3.5 (Toluene); 4.0 -> 3.3 (Naphthalene) | Mixed |
| Encode force only instead of energy and forces (ablation) | EMPP with Equiformer · MD17 test | forces MAE: 7.5 -> 7.1 (Benzene); 2.7 -> 2.4 (Ethanol); 5.2 -> 4.9 (Malonaldehyde); 1.9 -> 1.5 (Naphthalene) | Mixed |
| Lower spherical sampling rate from 100² to 20² (ablation) | EMPP · QM9 | ε_HOMO MAE: 14.2 -> 20.4 | Hurt |
| Increase spherical sampling rate from 100² to 150² (ablation) | EMPP · QM9 | α MAE: .041 -> .041 | No clear effect |
| Set EMPP loss weight to 0.1 instead of 1 (ablation) | EMPP auxiliary task · QM9 | α MAE: .041 -> .046 | Hurt |
| Set EMPP loss weight to 5 instead of 1 (ablation) | EMPP auxiliary task · QM9 | ε_HOMO MAE: 14.2 -> 15.8 | Hurt |
| Add denoising with a Gaussian-mixture standard deviation (ablation) | DP-TorchMD-Net · QM9 HOMO and LUMO | generalization performance | Mixed |
| Increase spherical sampling rate (ablation) | EMPP · QM9 | generalization performance | Helped |
Surprising findings
- Denoising with Gaussian-mixture targets yields results similar to or worse than the no-denoising baseline when no external data is used, and is highly sensitive to sigma.
- Encoding only the energy label (energy-only) with single-atom masking sometimes beats three-atom masking that encodes both energy and forces.
Limitations
- Extending to higher-order spherical harmonic representations (l>3) remains open
- Multi-atom masking can make position prediction ill-posed, so sequential single-atom masks are used
- Gains without pretraining are mostly on QM9 and MD17 with moderate margins; pretraining comparisons use a modified TorchMD-Net backbone
Experimental setup
Datasets: QM9 MD17 PCQM4Mv2 GEOM-Drug
Benchmarks: QM9 MD17 GEOM-Drug
Metrics: MAE energy MAE (meV) force MAE (meV/Å)
Baselines: AttrMask DP-TorchMD-Net 3D-EMGP SE(3)-DDM Transformer-M Frad SchNet PaiNN DimeNet++ TorchMD-NET SEGNN Equiformer
Builds on: Equiformer TorchMD-Net AttrMask Symphony e3nn
Closest papers by meaning
- EquiPocket: an E(3)-Equivariant Geometric Graph Neural Network for Ligand Binding Site Prediction (ICML 2024 · similarity 0.87)
- E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products (NeurIPS 2025 · similarity 0.86)
- ELoRA: Low-Rank Adaptation for Equivariant GNNs (ICML 2025 · similarity 0.86)
- Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity (ICML 2025 · similarity 0.86)
- Learning Equivariant Non-Local Electron Density Functionals (ICLR 2025 · similarity 0.86)
- Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.