Research map / Graph learning / Graph self-supervised learning
Pushing the Limits of All-Atom Geometric Graph Neural Networks: Pre-Training, Scaling, and Zero-Shot Transfer
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.
Studies scaling behavior of pre-trained all-atom geometric GNNs (ET, ViSNet) via denoising pre-training on PCQM4Mv2 and OrbNet-Denali, and tests them as zero-shot descriptors for VAMP kinetic modeling and protein fold classification.
Problem
Supervised training of Geom-GNN force fields and property predictors suffers from scarce, noisy labels and poor OOD generalization, and it is unclear how Geom-GNNs scale and whether pre-trained all-atom embeddings transfer zero-shot to molecular dynamics and protein tasks.
Approach
Use denoising-pre-trained Equivariant Transformer and ViSNet embeddings as frozen, transferable all-atom descriptors, combined with token mixers and VAMPNet objectives, and characterize their scaling (width, depth, aspect ratio, cutoff) across self-supervised, supervised and unsupervised setups.
Key results
- Pre-trained ViSNet embeddings with MLP-Mixer token mixing improve VAMP-2 over the VAMPNet baseline by 20.4% (Ala2), 50.0% (pentapeptide), and 15.1% (lambda6-85).
- On lambda6-85, MLP-Mixer token mixing gives over 46% higher VAMP-2 than sum pooling.
- Pre-training loss does not follow a clean power law with model size; gains saturate beyond roughly 1e7 parameters, and Geom-GNNs show no predictable scaling on QM9 scaffold split.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use pre-trained all-atom graph embeddings instead of pairwise-distance features | VAMP kinetic modeling · Ala2, Pentapeptide, and λ6-85 | VAMP-2: 20.4%, 50.0%, and 15.1% | Helped |
| Add an MLP token mixer after residue-level pooling (ablation) | PCQM-ViSNet VAMP representations · λ6-85 | VAMP-2: over a 46% | Helped |
| Increase embedding dimension beyond 64 (ablation) | PCQM-ViSNet VAMP representations · Ala2 | VAMP-2 | No clear effect |
| Increase embedding dimension (ablation) | PCQM-ViSNet VAMP representations · Pentapeptide and λ6-85 | VAMP-2 | Helped |
| Pre-train on PCQM before fine-tuning | ViSNet and ET · Scaffold-QM9 | MAE | Helped |
| Pre-train on Denali non-equilibrium structures | ET and ViSNet · xxMD-Temporal | Energy and force MAE | Mixed |
| Use deeper message-passing networks (ablation) | Geom-GNN denoising pre-training · PCQM and Denali | Pre-training loss | Helped |
| Decrease width/depth aspect ratio at similar parameter counts (ablation) | ViSNet denoising pre-training · PCQM and Denali | Pre-training loss | Helped |
| Set the radius cutoff near 5 Å (ablation) | 6-layer ViSNet denoising pre-training · PCQM and Denali | Pre-training loss: 5 Å | Helped |
| Increase model width (ablation) | ViSNet and ET supervised models · Scaffold-QM9 | MAE | Mixed |
| Use PCQM pre-training in data-limited xxMD (ablation) | ViSNet · Random-xxMD stilbene | Energy and force MAE: 10^3 | Helped |
Surprising findings
- Geom-GNNs do not follow power-law scaling during denoising pre-training and show limited predictable scaling on supervised QM9 tasks.
- Pre-training does not universally help: ViSNet sometimes shows negative transfer on xxMD, and pre-trained models can underperform from-scratch training beyond about 1e3 samples.
Limitations
- Only the denoising pre-training objective was studied.
- Protein systems are treated as single graphs, which may cause over-smoothing for larger proteins like lambda6-85.
- Fold classification results are preliminary and not compatible with ProNet data augmentation, so it is not a state-of-the-art claim.
Experimental setup
Datasets: PCQM4Mv2 OrbNet-Denali QM9 xxMD MD17 alanine dipeptide pentapeptide lambda6-85 FOLD
Benchmarks: QM9 xxMD alanine dipeptide VAMP-2 pentapeptide VAMP-2 lambda6-85 VAMP-2 FOLD fold classification
Metrics: VAMP-2 score mean absolute error denoising MSE accuracy
Baselines: VAMPNet with pairwise distance features Mardt et al. VAMPNet TICA ProNet-Amino-Acid from-scratch ET and ViSNet
Builds on: VAMPNet Equivariant Transformer ViSNet ProNet denoising pre-training for molecules
Closest papers by meaning
- Neural P$^3$M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs (NeurIPS 2024 · similarity 0.89)
- Neural Atoms: Propagating Long-range Interaction in Molecular Graphs through Efficient Communication Channel (ICLR 2024 · similarity 0.88)
- Generative Modeling of Full-Atom Protein Conformations using Latent Diffusion on Graph Embeddings (NeurIPS 2025 · similarity 0.88)
- M²N: A Progressive Macro-to-Micro 3D Modeling Scheme for Unveiling Drug-Target Affinity (AAAI 2025 · similarity 0.88)
- On the Scalability of GNNs for Molecular Graphs (NeurIPS 2024 · similarity 0.88)
- Towards Multiscale Graph-based Protein Learning with Geometric Secondary Structural Motifs (NeurIPS 2025 · similarity 0.88)
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.