Research map / Label-efficient and robust learning / Recognition and long-tail
From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction
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 JMP, a supervised multi-task pre-training of a GemNet-OC backbone on ~120M DFT-labeled structures from four datasets, then fine-tunes it on QM9, rMD17, MatBench, QMOF, SPICE and MD22.
Problem
Foundation-model-style pre-training has seen little use for atomic property prediction; models trained from scratch hold most SOTA results, and existing pre-training (mostly denoising on equilibrium structures) does not transfer across chemical domains or handle non-equilibrium data.
Approach
Treat each dataset and label as a separate task in a multi-task supervised pre-training of one GemNet-OC backbone, with temperature sampling across datasets, structure-wise force loss averaging, per-dataset loss coefficients, and strong regularization (weight decay, edge dropout), then fine-tune on out-of-domain tasks.
Key results
- JMP-L achieves an average 59% relative improvement over training GN-OC-L from scratch across 40 fine-tuning tasks.
- JMP sets or matches state-of-the-art on 34 of 40 tasks, including all MatBench regression tasks and QMOF.
- Fine-tuning JMP-L matches scratch-trained GN-OC-L performance using about 1/12 of the GPU hours (275 vs 3,300).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| pre-train on diverse chemical-domain datasets with supervised multi-task learning instead of training from scratch | GemNet-OC Large · QM9 ε_HOMO | MAE (meV): 22.7 -> 8.8 | Helped |
| use structure-wise force loss averaging rather than atom-wise averaging (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 0% -> 7.7% | Helped |
| use temperature sampling with T=2.0 (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 0% -> 2.2% | Helped |
| use fully balanced dataset sampling (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 2.61% | Helped |
| increase weight decay from 0.01 to 0.1 (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 11.4% | Helped |
| apply dropout p=0.1 on atom update layers (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 11.4% | Helped |
| drop edges with p=0.1 during pre-training (ablation) | JMP multi-task pre-training · Ablation fine-tuning tasks | mean relative improvement: 11.4% -> 13.2% | Helped |
| fine-tune on EMA weights (ablation) | JMP · Ablation fine-tuning tasks | mean relative improvement: 13.2% -> 12.4% | Hurt |
| pre-train on OC20 only instead of diverse multi-task data (ablation) | JMP · Ablation fine-tuning tasks | mean relative improvement: 0% -> -9.9% | Hurt |
| add threshold loss (ablation) | JMP multi-task pre-training with SWL · Ablation fine-tuning tasks | mean relative improvement: 7.68% -> 5.38% | Hurt |
| add PCGrad (ablation) | JMP multi-task pre-training with SWL · Ablation fine-tuning tasks | mean relative improvement: 7.68% -> 6.48% | No clear effect |
| add layerwise learning rate decay and ReduceLROnPlateau to warmup plus cosine scheduling (ablation) | JMP-L fine-tuning · rMD17 Aspirin validation | force MAE (meV/Å): 7.8 -> 5.1 | Helped |
Surprising findings
- Training OC20-only on a 120M subset performed substantially worse than the base multi-task model (-9.9% mean RI), indicating diversity matters.
- Threshold loss and PCGrad did not improve performance over the chosen recipe in ablations.
Limitations
- Only the GemNet-OC backbone was explored due to pre-training cost.
- Discarding pre-training prediction heads before fine-tuning may be suboptimal for datasets with labels similar to pre-training.
- Model size (~235M) is small compared with largest NLP/CV models.
Experimental setup
Datasets: OC20 OC22 ANI-1x Transition-1x QM9 rMD17 MatBench QMOF SPICE MD22
Benchmarks: QM9 rMD17 MatBench QMOF SPICE MD22
Metrics: MAE force MAE relative improvement over scratch RMSE
Baselines: GN-OC trained from scratch TorchMD-Net Equiformer MACE Allegro MODNet coGN sGDML NequIP ET+NN (pretrained denoising)
Builds on: GemNet-OC temperature sampling PCGrad layerwise learning rate decay
Closest papers by meaning
- MORE: Molecule Pretraining with Multi-Level Pretext Task (AAAI 2025 · similarity 0.87)
- ChemPile: A 250 GB Diverse and Curated Dataset for Chemical Foundation Models (NeurIPS 2025 · similarity 0.87)
- Benchmarking Compositional Generalisation for Machine Learning Interatomic Potentials (NeurIPS 2026 · similarity 0.87)
- Pin-Tuning: Parameter-Efficient In-Context Tuning for Few-Shot Molecular Property Prediction (NeurIPS 2024 · similarity 0.86)
- MIPT: Multilevel Informed Prompt Tuning for Robust Molecular Property Prediction (ICML 2025 · similarity 0.85)
- Physical Consistency Bridges Heterogeneous Data in Molecular Multi-Task Learning (NeurIPS 2024 · similarity 0.85)
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.