PaperlayerView exact revision on arXiv

Source-linked AI summary

From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction

Nima Shoghi, Adeesh Kolluru, John R. Kitchin, Zachary W. Ulissi, C. Lawrence Zitnick, Brandon M. Wood

arXiv:2310.16802v2cs.LG

TL;DR

Atomic property prediction lacks broadly effective pre-training because chemical domains, structures, and continuous targets vary substantially. JMP addresses this with supervised multi-task pre-training across diverse atomic datasets, then fine-tunes a shared backbone on downstream tasks. It improves performance over training from scratch and matches or sets state-of-the-art on 34 of 40 benchmarks, while supporting larger models in low-data settings.

  • Problem

    Atomic property prediction requires generalization across diverse chemical domains, structures, and continuous targets, but existing models are often trained from scratch.

    Source: §1 · pp. 1–2

  • Method

    JMP simultaneously pre-trains a shared backbone on diverse datasets by treating each dataset as a separate supervised multi-task learning objective.

    Source: §1 · p. 2

  • Results

    JMP sets or matches state-of-the-art on 34 of 40 fine-tuning benchmarks and improves performance across domains, datasets, and tasks.

    Source: §1 · p. 2, §5 · p. 6, §6 · p. 9

  • Takeaways & Limitations

    Pre-training diverse atomic data supports generalizable representations, effective scaling to larger models, and performance gains on low-resource downstream tasks.

    Source: §6 · p. 9

    The study evaluates only the GemNet-OC backbone, discards pre-training heads during fine-tuning, and uses models smaller than the largest NLP and CV models.

    Source: §6 · p. 9

Abstract

from arXiv · show

Foundation models have been transformational in machine learning fields such as natural language processing and computer vision. Similar success in atomic property prediction has been limited due to the challenges of training effective models across multiple chemical domains. To address this, we introduce Joint Multi-domain Pre-training (JMP), a supervised pre-training strategy that simultaneously trains on multiple datasets from different chemical domains, treating each dataset as a unique pre-training task within a multi-task framework. Our combined training dataset consists of $\sim$120M systems from OC20, OC22, ANI-1x, and Transition-1x. We evaluate performance and generalization by fine-tuning over a diverse set of downstream tasks and datasets including: QM9, rMD17, MatBench, QMOF, SPICE, and MD22. JMP demonstrates an average improvement of 59% over training from scratch, and matches or sets state-of-the-art on 34 out of 40 tasks. Our work highlights the potential of pre-training strategies that utilize diverse data to advance property prediction across chemical domains, especially for low-data tasks. Please visit https://nima.sh/jmp for further information.

1 INTRODUCTION

Atomic property prediction is costly and difficult to generalize across diverse chemical domains, motivating JMP, a supervised multi-domain pre-training strategy. JMP trains on over 120 million structures and improves fine-tuning performance across broad benchmarks.

1 INTRODUCTION

  • Atomic property calculations are expensive with DFT, while ML potentials can reduce computational cost by orders of magnitude.Creating large datasets for every chemistry problem is also impractical, and training from scratch for each use case is common.Source: §1 · p. 1
  • Before JMP, models trained from scratch held previous state-of-the-art on 34 of the 40 benchmark tasks evaluated.Scratch-trained models held all previous state-of-the-art results on rMD17, MD22, SPICE, and MatBench, plus 7 of 12 QM9 targets.Source: §1 · p. 2
  • JMP concurrently trains over 120 million equilibrium and non-equilibrium structures by treating each chemical domain as a separate supervised pre-training task.The approach is designed to learn representations that transfer across small molecules, large molecules, and materials.Source: §1 · p. 2
  • JMP sets or matches state-of-the-art on 34 of 40 fine-tuning benchmarks and consistently outperforms training from scratch.Its evaluation spans a diverse benchmark suite across small molecules, large molecules, and materials.Source: §1 · p. 2
  • Pre-training enables a 235M-parameter model to avoid the overfitting associated with training large models from scratch on small datasets.JMP also recovers its expensive pre-training cost through over 12x faster fine-tuning than training from scratch.Source: §1 · pp. 2–3

2 RELATED WORK

Prior atomic modeling has advanced through descriptors, neural networks, graph neural networks, and transfer learning, but broad cross-domain pre-training remains limited. Existing self-supervised approaches also commonly target equilibrium structures despite the prevalence of non-equilibrium data.

2 RELATED WORK

  • Graph neural networks have surpassed descriptor-based methods across benchmarks spanning small molecules, catalysts, and bulk materials, but one model remains difficult to generalize across all domains.Source: §2 · p. 3
  • Most transfer-learning studies focus on representations transferred within the same chemical domain using limited pre-training data.Source: §2 · p. 3
  • Most existing pre-training work does not explore generalization across multiple chemical domains and often uses self-supervised denoising on equilibrium structures.Source: §2 · p. 3
  • Equilibrium-structure denoising applies to less than 1% of the training data because most atomic-property data is non-equilibrium.Alternative formulations applicable to non-equilibrium data had begun to emerge when this work was conducted.Source: §2 · p. 3

3 DATASETS

The study organizes atomic data into small-molecule, large-molecule, materials, and catalysis domains, using diverse datasets with 3D structures and atomic properties. Pre-training uses four large datasets, while fine-tuning tests transfer within and beyond the pre-training domains.

3 DATASETS

  • The dataset taxonomy includes small molecules, large molecules, materials, and catalysis, with samples containing 3D structures and node-level or graph-level properties.Source: §3 · p. 3
  • Pre-training uses ANI-1x, Transition-1x, OC20, and OC22, selected for their diversity and large size.Source: §3 · p. 3
  • The combined pre-training data contains over 120M energy- and force-labeled examples, with more than 99% from non-equilibrium structures.Different prediction heads and per-dataset energy referencing accommodate differences in DFT theory and software.Source: §3 · p. 3
  • Fine-tuning evaluates small-molecule, large-molecule, and materials datasets, including both in-distribution energy and force labels and out-of-distribution labels such as QM9’s ∆ϵ.Source: §3 · p. 3

4 JOINT MULTI-DOMAIN PRE-TRAINING

JMP frames multi-domain atomic pre-training as supervised multi-task learning with one shared backbone and dataset-specific prediction heads. Its design balances heterogeneous targets, dataset sizes, and system sizes before fine-tuning on downstream tasks.

4 JOINT MULTI-DOMAIN PRE-TRAINING

  • JMP treats every label from each pre-training dataset as a separate task, allowing one backbone to learn simultaneously from multiple chemical domains and labels.Source: §4 · p. 4
  • GemNet-OC computes node and edge representations, with prediction heads producing summed system-level energies and direction-weighted node-level forces.During pre-training, forces use a direct equivariant block for computational efficiency and faster convergence on large datasets.Source: §4 · p. 4

4.1 MULTI-TASK PRE-TRAINING

  • Each dataset receives separate energy and force heads, while target normalization places energies and forces from different datasets on a common scale.Source: §4.1 · p. 4
  • Temperature sampling with T = 2 balances contributions from pre-training datasets ranging from 2 million to 100 million samples.The sampling strategy is intended to prevent high-resource datasets from dominating and low-resource datasets from being underfit.Source: §4.1 · pp. 4–5
  • Structure-wise loss reduction averages force loss per system before averaging across systems, preventing larger systems from dominating the force loss.The method addresses datasets such as Transition-1x and OC22, which average 14 and 80 atoms per system, respectively.Source: §4.1 · p. 5
  • Loss coefficients balance energy and force losses within each dataset without requiring an infeasible grid search over 2 · M hyperparameters.Source: §4.1 · p. 5
  • Fine-tuning discards the pre-training heads, adds randomly initialized downstream heads, and updates the full model on each target task.For evaluation, force labels are computed by differentiating energy with respect to atomic positions.Source: §4.1 · p. 5

5 EXPERIMENTS

JMP is evaluated across diverse in-domain and out-of-domain atomic property benchmarks, with comparisons to scratch training, prior methods, and ablated variants. It improves scaling and accuracy across tasks while reducing downstream fine-tuning compute, though pre-training requires substantial upfront resources.

5 EXPERIMENTS

  • 59% average relative improvement over GN-OC-L shows that JMP-L elevates performance across all domains, datasets, and tasks.JMP-L also achieves a 21% average gain over JMP-S, while scratch-trained GN-OC-L performs 8% worse than GN-OC-S on average.Source: §5 · p. 6, Figure 2 · p. 6
  • JMP-L achieves state-of-the-art performance on nearly all QM9 targets, with R2 as the sole exception.The QM9 evaluation uses dedicated prediction heads, and prior work similarly reports less pronounced pre-training benefits for R2.Source: §5 · p. 6, Table 2 · p. 7
  • JMP-L achieves state-of-the-art performance across QMOF and all evaluated MatBench tasks despite out-of-domain materials structures and target labels.Most pre-training labels are energies and forces, whereas these downstream tasks include different materials targets.Source: §5 · p. 5, Table 3 · p. 7, Table 4 · p. 7
  • JMP achieves state-of-the-art results across all rMD17 molecules, all MD22 molecules, and all SPICE tasks, including unseen large-molecule domains.For rMD17, JMP-L reaches state of the art on 5 molecules and on 6 of 10 targets in the 50-sample subset; direct force heads are used for two memory-intensive MD22 structures.Source: §5 · pp. 7–8, Table 5 · p. 8

5.1 ABLATION STUDIES

  • Ablations identify structure-wise loss reduction and regularization as the most important improvements, with edge dropout reaching 13.2% improvement over the base model.Temperature sampling with T = 2 provides a moderate gain, while higher temperatures have diminishing returns; EMA does not improve performance, and OC20-only pre-training performs substantially worse than the diverse multi-task base.Source: §5.1 · pp. 8–9, Table 6 · p. 8

5.2 COMPUTATIONAL COST ANALYSIS

  • Fine-tuning JMP-L requires about 275 GPU hours versus 3,300 GPU hours for scratch training, a 12x reduction in downstream compute.JMP-L pre-training itself requires about 34,400 GPU hours on 128 V100 GPUs, making the compute trade-off substantial upfront but efficient downstream.Source: §5.2 · p. 9

6 CONCLUSION

JMP’s diverse-data pre-training improves transfer across chemical domains and low-data settings, while the study identifies computational and architectural limits that motivate future work.

6 CONCLUSION

  • The study’s limitations include evaluating only the GemNet-OC backbone, discarding pre-training heads before fine-tuning, and using models smaller than leading NLP and CV systems.The authors identify different backbones, retaining prediction heads, and larger models as directions for future research.Source: §6 · p. 9

A ADDITIONAL PRETRAINING COMPARISONS

B.1 EXPANDED RMD17 RESULTS: TRAINING ON 50 EXAMPLES

  • JMP-L achieved state-of-the-art performance on 6 of 10 rMD17 molecules when trained with only 50 examples per molecule.The pre-trained model also outperformed GN-OC-L in this reduced-data setting.Source: §B.1 · pp. 16, 18, Table 10 · p. 18

B.2 PDBBIND V2013 BINDING AFFINITY RESULTS

  • JMP achieved state-of-the-art performance on the PDBBind v2013 binding-affinity dataset containing 154 protein-ligand complexes.The evaluation used a random 80%/10%/10% split and RMSE reporting.Source: §B.2 · p. 18, Table 11 · p. 18

B.3 ADDITIONAL ABLATION STUDIES

The authors evaluate dataset balancing, threshold losses, task-gradient surgery, and fine-tuning learning-rate components. Temperature-2.0 scaling and the complete learning-rate schedule are retained, while threshold loss and PCGrad do not improve performance.

B.3 ADDITIONAL ABLATION STUDIES

  • 2.61% mean relative improvement from fully balancing datasets is similar to 2.29% with temperature-2.0 scaling, but requires more training steps.The authors therefore use temperature-2.0 dataset scaling for full-dataset models.Source: §B.3 · p. 18, Table 13 · p. 19
  • Threshold loss hurts performance in the final ablation results despite initially appearing promising.It applies dataset-specific physically motivated margins to predictions beyond estimated DFT-label noise thresholds.Source: §B.3 · p. 19, Table 12 · p. 19
  • PCGrad provides no performance improvement over structure-wise loss reduction, consistent with evidence that regularization can reduce the need for complex multi-task optimization.PCGrad projects conflicting task gradients onto normal planes.Source: §B.3 · p. 19

B.4 FINE-TUNING LR SCHEDULING

  • 31% lower force MAE results from the full learning-rate schedule versus cosine decay alone on Aspirin in rMD17.Adding each scheduling component improves performance, and the complete strategy combines warmup, cosine decay, LLRD, and ReduceLROnPlateau.Source: §B.4 · p. 20, Table 14 · p. 20

C EMBEDDING AND ACTIVATION VISUALIZATIONS

The visualizations examine whether JMP-L learns representations organized across datasets and chemical elements. Aggregated node and edge embeddings are plotted by source dataset, while atom-embedding similarities show periodic-table structure.

C EMBEDDING AND ACTIVATION VISUALIZATIONS

  • JMP-L aggregates node- and edge-level embeddings across randomly selected structures from pre-training and fine-tuning datasets for t-SNE visualization.Each point represents one structure, colored by its source dataset.Source: C EMBEDDING AND ACTIVATION… · p. 20, Figure 5 · p. 21
  • Nearby periodic-table elements have more similar learned atom embeddings for H, C, N, O, Co, and Cu.The observed similarity pattern aligns with known elemental properties and indicates chemically meaningful embedding structure.Source: Figure 6 · p. 21, C EMBEDDING AND ACTIVATION… · p. 22

C.1 EXACT PERCENTAGES FOR FIGURE 2

  • Figure 2 reports per-dataset averages spanning QM9, MD17, MD22, SPICE, Matbench, and QMOF across three reported result groups.The supplied percentages include group (a) values from -25.6% to 2.2%, group (b) values from 9.4% to 31.6%, and group (c) values from 32.0% to 77.9%.Source: §C.1 · p. 22

D.1 OPTIMAL TASK WEIGHTING AND REGULARIZATION

The paper uses unitary scalarization with regularization for multi-task pre-training and documents implementation choices for force computation and structure-wise loss reduction. These choices address task weighting and unequal system sizes across datasets.

D.1 OPTIMAL TASK WEIGHTING AND REGULARIZATION

  • Unitary scalarization sums task losses, while weight decay, edge dropout, and exponential moving averages provide regularization.The cited settings are weight decay 0.1, edge-dropout probability 0.1, and exponential-moving-average decay 0.99.Source: §D.1 · p. 22

D.2 FINE-TUNING FORCE COMPUTATION

  • For force fine-tuning, the evaluations compute forces as energy gradients with respect to atomic positions rather than using the direct equivariant block.The authors report that both approaches work well but choose the gradient-based approach as standard in the literature.Source: §D.2 · p. 22

E STRUCTURE-WISE LOSS REDUCTION CODE

F.1 PRE-TRAINING

Pre-training uses AdamW with linear warmup followed by cosine decay, while fine-tuning uses a separate optimizer setup and a three-stage learning-rate schedule. Fine-tuning also applies layerwise decay and explicit stopping criteria.

F.1 PRE-TRAINING

  • Pre-training uses AdamW at learning rate 0.0003 with 2,000-step warmup from 0.2·LR and cosine decay to 0.1·LR over two epochs.The optimizer uses beta1 = 0.9, beta2 = 0.95, and weight decay 0.1.Source: §F.1 · p. 24

F.2 FINE-TUNING

  • Fine-tuning uses AdamW at learning rate 0.00008 with five-epoch warmup, 32-epoch cosine decay, and ReduceLROnPlateau thereafter.Layerwise learning-rate decay is used during warmup and cosine decay, and the final learning rate after cosine decay is shared across layers.Source: §F.2 · p. 24
  • Fine-tuning stops after 50 patience epochs, 500 epochs or seven days, or a learning rate below 10^-8.For rMD17, the patience is 1,000 and the maximum is 100,000 epochs because its training set is small.Source: §F.2 · p. 24

G MODEL TRAINING TIMES AND CO2 IMPACT

JMP’s pre-training combines diverse catalytic and molecular datasets, while the resulting checkpoints reduce downstream training costs despite substantial upfront computation.

G MODEL TRAINING TIMES AND CO2 IMPACT

  • Approximately half the fine-tuning resources were required by pre-trained models compared with models trained from scratch, while maintaining superior performance.The pre-training investment can be reused across multiple applications.Source: G MODEL TRAINING TIMES AND CO2 IMPACT · p. 24
  • 46,400 computation hours on Tesla V100-SXM2-32GB hardware produced an estimated 6,013.44 kgCO2eq, all directly offset.The infrastructure’s carbon efficiency was 0.432 kgCO2eq/kWh.Source: G MODEL TRAINING TIMES AND CO2 IMPACT · pp. 24–25

H.1 PRE-TRAINING DATASETS

  • The pre-training mixture comprises OC20, OC22, ANI-1x, and Transition-1x, combining catalytic and small-molecule structures with energy and force labels.The datasets were selected for their substantial sizes and structural diversity.Source: §H.1 · p. 25
  • Dataset-wise linear referencing and normalization standardize energy and force labels despite differing DFT functionals and engines.Energy labels are divided by dataset standard deviations, while force labels are divided by training-set force-component RMS values.Source: §H.1 · p. 25
  • OC20 contributes 130 million relaxation-trajectory examples, with additional AIMD and rattled data bringing its cumulative size to 189 million entries.The large training configuration uses 100 million entries and a 2 million-entry development set.Source: §H.1 · p. 25
  • OC22 supplies 8 million oxide-focused training examples generated with the PBE + U functional, complementing OC20’s broader catalyst coverage.The full OC22 training set is used for the large run, with 200,000 examples for development.Source: §H.1 · p. 25
  • ANI-1x provides 5 million small-molecule DFT calculations, while Transition-1x contributes nearly 10 million reaction-pathway configurations with denser potential-energy-surface sampling.ANI-1x contains C, H, N, and O molecules; Transition-1x samples configurations around Nudged Elastic Band reaction pathways.Source: §H.1 · p. 26

H.2 FINETUNING DATASETS

Fine-tuning evaluates JMP across materials, small-molecule, and large-molecule domains, with scaled-down development sets supporting ablation studies.

H.2 FINETUNING DATASETS

  • Fine-tuning spans six datasets across three domains: QMOF and Matbench for materials, MD17 and QM9 for small molecules, and SPICE and MD22 for large molecules.This selection tests transfer across varied structures and atomic-property tasks.Source: §H.2 · p. 26
  • QMOF predicts band gaps for approximately 15,000 experimentally synthesized metal-organic frameworks, using mean pooling because band gap is an intensive property.The training split contains 10,000 systems.Source: §H.2 · p. 26
  • Matbench contributes eight structure-based regression tasks spanning optical, thermal, electronic, thermodynamic, tensile, and elastic properties.JMP-L and JMP-S are averaged across five folds, while other comparisons use fold 0; max pooling is used for phonons.Source: §H.2 · p. 26
  • MD17 uses 950 training and 50 validation samples, with force predictions modeled as energy gradients and an additional 50-sample split for few-shot evaluation.The revised dataset contains eight small organic molecules with DFT-computed energies and forces.Source: §H.2 · p. 26
  • QM9 covers organic molecules with up to nine heavy atoms and supplies properties including atomization energy, HOMO-LUMO gap, dipole moment, and polarizability.The experiments use atomwise references, standard-Gaussian normalization, and sum pooling.Source: §H.2 · pp. 26–27
  • SPICE fine-tuning focuses on dipeptides and solvated amino acids, while MD22 tests molecules of 42–370 atoms across eight large-molecule classes unseen during pre-training.SPICE contains over one million conformers, and MD22 provides energies and forces from finite-temperature molecular dynamics.Source: §H.2 · p. 27

H.3 DEVELOPMENT SPLITS FOR ABLATION STUDIES

  • Scaled-down development sets use a 2-million-example OC20 subset, proportionally selected companion datasets, and one target or molecule per fine-tuning dataset for ablations.The reduced configurations address the high computational cost of full pre-training and fine-tuning.Source: §H.3 · p. 27, Table 17 · p. 27

I DATASET OVERLAP

The study examines overlap between pre-training and QM9 and documents the model and optimization configurations used for small and large variants.

I DATASET OVERLAP

  • ANI-1x and QM9 share some or similar molecules, but differ in DFT level, labels, and potentially 3D structures, leaving the overlap’s effect on fine-tuning uncertain.The strictest overlap definition matches systems by composition.Source: I DATASET OVERLAP · p. 28
  • QM9 results are hard to distinguish across overlapping-only, overlap-removed, and full test sets, indicating the fine-tuned model is not simply memorizing similar pre-training examples.The comparison is reported in Table 18.Source: I DATASET OVERLAP · p. 28, Table 18 · p. 28

J MODEL AND OPTIMIZATION HYPERPARAMETERS

Loading 2310.16802v2…