From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction
Authors: Nima Shoghi
Adeesh Kolluru
John R. Kitchin
Zachary Ward Ulissi
C. Lawrence Zitnick
Brandon M Wood
Venue: ICLR 2024 poster (2024)
Keywords: atomic property prediction
pre-training
3D atomic pre-training
graph neural networks
multi-task learning
molecules
materials
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We pre-train a large model on multiple chemical datasets in a multi-task learning framework to generate transferable atomic representations that can be fine-tuned for SOTA results across various tasks.

Abstract: 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.

Published as a conference paper at ICLR 2024
FROM MOLECULES TO MATERIALS: PRE-TRAINING LARGE GENERALIZABLE MODELS FOR ATOMIC PROP- ERTY PREDICTION
Nima Shoghi Adeesh Kolluru² John R. Kitchin²
Zachary W. Ulissi!
C. Lawrence Zitnick' Brandon M. Wood
1 Fundamental Al Research (FAIR) at Meta
CCarnegie Mellon University
"Work done while FAIR
Correspondence to: ns ns@niima.. sh, bmwoodemeta.com com
ABSTRACT
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 the challenges of training effective models across multiple chemical domains. To address this, we introduce Joint Multi-domain Pre-training (JMP), supervised pre-training strategy that simul- taneously trains on multiple datasets from different chemical domains, treating each dataset as a unique pre-training task within multi-task framework. Our combined training dataset consists of ~ 120M systems from OC20, OC22, ANI-Ix, 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- pre-training strategies that utilize diverse data to advance property prediction across chemical domains, especially for low-data tasks.
INTRODUCTION
Computing atomic properties accurately and efficiently for vast array of molecules and materials is crucial for range of applications, from drug discovery (Chan al., 2019 Deng al. 2022) to catalyst design (Zitnick et al., 2020). Currently, the quantum chemistry method Density Functional Theory (DFT) is commonly employed for atomic property calculations. Unfortunately, DFT's use limited by its significant computational expense, which can range from hours to days for certain calculations. Machine learning (ML. potentials, which approximate augment DFT, capable of reducing the computational cost by orders of magnitude (Behler) 2016). In recent years, much progress has been made towards this goal (Kolluru al.) (022b), fueled in part by the release of large and diverse DFT-generated datasets for training ML models. While these datasets are incredibly useful, they are also extremely expensive generate, e.g., 4400 million CPU hours for the Open Catalyst 2020 dataset (OC20) (Chanussot et al. 2021). As a consequence, it is impractical to create large dataset for every specific chemistry problem of interest. Similarly, non-ideal to train model from scratch for all use cases, which is common practice currently.
Foundation models (FMs) - large pre-trained models that can fine-tuned for various tasks have achieved remarkable success in domains such as natural language processing (NIP) and computer vision (CV), especially when fine-tuned on low-resource downstream tasks. Several key factors have enabled this effectiveness: (1) the availability of massive datasets, (2) the development of
widely adopted pre- pre-training strategies, and (3) the establishment of diverse benchmarks to rigorously assess the performance of these fine-tuned models. Despite the availability of large DFT-labeled datasets (c.g., OC20) and the existence of a wide and diverse range of downstream tasks (e.g., QM9 (Ruddigkeit al. 2012b), MatBench (Dunn et al. 2020)), the adoption of pre-training in MI.

---

Published as a conference paper at ICLR 2024
Catalysis
Small Molecules
Small Molecules
Molecules
Materials
OC20
OC22
AN11X
Trans1x
QMB
SPICE
Matbenchh
Leamed Representations
MD17
MD22
QMOF
Data Normalization Temperature Sampling
Fine-tune
Backbone Model
Pre-train
Knowledge Transfer
Backbone
Model
Multi-Task Learning
E SPICE
From
QM9
E
E
OC20
OC22
AN
Trans1x
MD17
MD22
QMOF
Figure 1: An overview of the Joint Multi-domain Pre-training (JMP) method. Left: JMP's pre- training setup, where a single model is simultaneously trained on set of diverse pre-training datasets using multi-task learning. Center: JMP's fine-tuning process, where the pre-trained JMP backbone equipped with new prediction heads and trained on downstream tasks. Right: t-SNE visualizations of JMP's node-level (h) embeddings for randomly selected structures from all datasets.
for atomic property prediction has been noticeably less prevalent. This under-utilizztion becomes evident when noting that most of the lfte-of-thhearr (SOTA) results downstream tasks come from models trained from scratch. More specifically, prior to our work, all previous SOTA results on the from models trained form scratch. For QM9,
rMD17, MD22,, SPICE, and MatBench datasets models trained from scratch hold SOTA status for 7 of the 12 targets. In total, out of the 40 total tasks explored in this work's evaluation benchmark, models trained from scratch hold the previous SOTA on 34 tasks.
At its core, the challenge of pre-training for atomic property prediction lies in the complexity and diversity of the underlying chemical space. Target applications vary from drug design to catalysis, the data ranges from small molecules with only atoms to periodic crystals with hundreds, and even the properties of interest for each application vary from various energies to forces to phonon peaks. Furthermore, the nature of atomic properties imposes a unique set of challenges. Unlike NLP or CV, where the data is often discrete and finite, atomic properties are continuous and can span several orders of magnitude. This requires models to robust to outliers and capable of predicting highly variable outputs. Further, existing pre- pre-training strategies (e.g.. Zaidi al. (2022 : ZZou et al. (2023b) are designed with equilibrium systems in mind and are not directly applicable to non-equilibrim systems, which are common in DFT datasets (e.g., 99.7% of OC20's training data comprises of non-equilibrium structures). These challenges motivate the need for flexible and generalizable pre-training strategy that can be adapted different applications and datasets.
In this work, we introduce Joint Multi-domain Pre-training (JMP), supervised pre-training strategy tailored the challenges and opportunities of machine learning for atomic modeling. JMP concur- rently trains over 20 million diverse equilibrium and non-equilibrium atomic structures by framing each chemical domain as separate pre- pre-training task multi-task framework. This large-scale pre-trae pre-training enables learning generalizable representations of atomic interactions. The contributions of our work are summarized as follows: First, we introduce the JMP method, shown Figure and demonstrate its powerful generalization ability by evaluating its fine-tuning performance across diverse benchmark suite spanning small molecules, large molecules, and materials. Our results show that JMP consistently outperforms training from scratch and sets matches the state-of-the-ae on 34 out of the 40 fine-tuning benchmarks. Second, we show that JMP enables efficient scaling to larger models that would normally overfit if trained from scratch small datasets. Pre-training acts as a strong regularizer, allowing us to train a 235M parameter model that sets new state-of of-the-art performance multiple low-data benchmarks. Finally, conduct detailed analysis of JMP's
2

---

Published as a conference paper at ICLR 2024
computational requirements. While expensive upfront, show JMP's pre-training cost is recovered by enabling over 12x faster fine-tuning compared to training from scratch. By pre-training large models on diverse chemical data, we believe JMP represents an important step towards the goal of universal MI potential, and that the continued growth of available data and compute power will only improve JMP's ability to learn transferable atomic representations.
RELATED WORK
There has been significant progress in developing ML models for atomic property prediction. Initial approaches focused on descriptor-based methods, where these descriptors were hand-fitted physically meaningful analytical functions (González, 201T Sundius 2002, Dinur and Hagler 1991). These functions were incorporated into gaussian process models (Chmiela al. 2017) 7) neural networks (Behler and Parrinello 2007). Recent advances in graph neural networks (GNNs) have shown to be a promising approach for these tasks, surpassing descriptor- based methods (Gasteiger et al. [2000] (Schitt et al., 2017] Batzzer et al.) 2021 Batatia al.) 2022) on multiple benchmarks across the atomic domains of small molecules, catalysts, and bulk materials. While much progress has been made, it remains difficult for single model perform well across all chemical domains.
Machine learning potentials:
Pretraining and transfer learning on 3D atomic systems: The concept of transfer learning, where representations are learned on one dataset and transferred to another, has been successfully applied to a number of atomic modeling tasks (Kolluru et al.) 2022a] Cai et al. 2020; Tsubaki and Mizoguchi 2021 Smith et al. 2018). However, most of the focus in this area has been transferring representations within the same chemical domain with a limited amount of pre-training data (Smith)
al. 2019 Yamada et al.) 2019 Pesciullesi et al.) 2020). There are beginning to be more dedicated works on pre-training (Zhu et al. 2022 Liu et al. 2021 Jiao et al. 2022 Zhou et al. 2023b), but most do not explore generalization across multiple chemical domains. Many of these works focus self-supervised pre- trainini on molecular graphs and/or 3D atomic structures. Recent supervupee methods have focused on denoising methods (Song et al., 2020) applied to equilibrium structures
-ie. the per atom forces are close to zero Zaidi et al., 2022, Feng et al. [023b) Liu et al. (022). The original formulation of denoising equilibrium structures applicable less than 1% of training data because most of the atomic properties data is non-equilibrium. This is an active area of
research and since the beginning of our present work, alternative formulations that could apply to non-equilibrium data have started to emerge (Feng et al., 2023a Zheng et al. (003))
3 DATASETS
We separate the atomic space into four domains for the purposes of this manuscript including, small molecules (1-20 atoms), large molecules (more than 20 atoms), materials, and catalysis (contains material surfaces with molecules). Each dataset sample contains a 3D atomic structure (positions and atomic numbers) and set of atomic properties. The atomic properties can be either node-level (c.g., forces) or graph-level (c.g., energy). The datasets are summarized in Table [ with additional information, including details on train, validation, and test splits, in Appendix H
To study the ability of pre-trained models to generalize across domains and tasks, we only pre-train on small molecule and catalysis datasets, and fine-tune on small molecule, large molecule, and materials datasets, Our pre-training datasets include the ANI-Ix (Smith et al. 2020 and Transition- 1x (Schreiner al. 2022) small molecule datasets and the OC20 (Chanussot al. 202T] and OC22 (Tran et al. 2022) catalysis datasets. These datasets were chosen due to their diversity and large size.
The combined pre-training dataset contains over 120M training examples with energy and force labels, with the majority of the data (> 99%) coming from non-equilibrium structures. Due to the difference in underlying DFT theory and software used across the datasets, we utilize different prediction heads for each dataset. We also use per-dataset linear referencing scheme for the energies. For fine-tuning, we use smaller datasets from three domains to evaluate how pre-trained models perform in similar (small molecule) and unseen domains (large molecule and materials). These datasets may contain in n--istribution (ID) (i.c., energies and forces) or out-of-distribution (OOD) labels (c.g., QM9's A,).
3

---

Published as a conference paper at ICLR 2024
Pataset
Domain
Label
Elements
Arg size
Train Set
Deseription
Pretraining Paatasets
OC20
Catalyst
E,F
(3-(-225)
100M
Canalyst relaxations
Catalyst
(17-228)
Oxide catalyst relaxations
ANI-Ix
Small Molocule
E.F
H.C.N.O
15544-55)
2M
MD simulations
Transtion- Ix
Small Molocule
E.
H. N,C,N,
(4434))
10M
Reactions database
Finetuning
Muthench
Materials (OOI)
ID/
-600-1
Materaal properties
QMOF
Materials (OOB)
0DD
77
-109 (17, 500)
10k
MGF properties
MD17
Small Mols. (II)
ID
H,C.N,O
-13((9-11)
1k
MD simulation
QM9
Small Mols, (ID)
ID/00D
H.C.N,O
~130k
QM properties
SPICE
Large Mols, (ODD)
ID
H.C.N.O.S
(26-6)
1300,
MD sinulations
Lange Mols. (OOD)
N.C.N.O
(42:370)
.-600
simulations
Table 1: Summary of datasets and their properties, including the domain, target labels, atomic elements present, their sizes and brief description.
4 JOINT MULTI-DOMAIN PRE-TRAINING
Joint Multi-domain Pre-training (JMP), shown in Figure T is based on the intuition that pre-training on diverse set of chemical domains should lead to better representation learning and thus better generalization through fine-tuning. The pre- pre-training task is framed as a multi-task supervised learning problem, where each label of each pre-training dataset is treated as a separate task. This allows us to pre-train a single backbone model on multiple chemical domains and labels simultaneously.
We use the following notation throughout this section. Let D (D1.... DM be the set of M datasets that we pre-train on. Each dataset, D,. is a set systems (e.g., molecules or crystals), where each system is tuple of atomic numbers (Z), atomic positions (RR), and target (i.e., ground-truth) energy (É) and forces (F)). For a given mini-bated of B systems, W, is the index of the dataset that system b E B belongs to, and N₂ is the number of atoms in system
Notation:
Model Architecture: Our goal in this work is to design model-agnostic strategies for supervised pre-trainii traininnngg..... For our backbone model architecture, we chose GemNet-OC (Gasteiger et al., (202)) for its effectiveness across wide spectrum of chemical domains well as large scales (Sriram eet al. 2022). GemNet-OC is a message- message-passing neural network that computes a node representation h; for each atom i and edge representation m,y for pairs of nearby atoms, and j- Using these representations, prediction heads compute desired target properties. System-level scalar predictions, such as energy, are computed by summing the node representations, E = Σ MLP((.). Node- level vector predictions, such as forces, are computed by summing the edge direction unit vectors, weighted by the edge representations, F; Σj=1 (MLP(m.y) Piz). During pre-training, we compute forces using a direct equivariant block, similar to Kliepera et al. (2021)'s 's model setup for the OC20 dataset. This is for two reasons: (1) direct force prediction is much more computationally efficient than gradient-based force prediction, as the latter needs to perform a secondary backward pass compute the gradient of the energy with respect the atomic positions and (2) previous works (Gasteiger et al. 2022) have shown that for larger datasets, direct force prediction shows much faster convergence while producing similar converged accuracies to gradient-based force prediction.
4.1 MULTI-TASK TRAITIGINNNG
In the multi-task setting, each dataset has its own energy and force prediction heads, as shown in [ This allows us to train a single model on multiple datasets simultaneously. In the following sections, we describe each of these imbalances and our proposed solutions detail.
When pre-training multiple datasets, we first need to normalize the targets to make sure they are on a common scale across datasets. Since our pre-training task is energy and force prediction, for each dataset we first linearly reference the total energies and then normalize them to unit Gaussian distributions. We normalize the forces by dividing them by component-wise RMS foree. This puts the energy and forces for each dataset on common scale.
Data Normalization:
Dataset Size Imbalance: Our pre- tre-training datasets vary greatly in size, from 2 million to 100 million training samples, for total of 120M samples. To maintain a proper balance between the total contribution of large, high-resource and small, low-resource pre-training datasets and to prevent overfitting to high-resource datasets and underfitting on low-resource datasets, we use temperature
4

---

Published as a conference paper at ICLR 2024
sampling (Devlin et al.) 2018) during batch construction. Specifically, we sample each dataset i with probability p. x (Σ, 12, 1/T", where [i] the number of samples in dataset and T is the temperature hyperparameter. Inspired by |Shaham et al. (2023), which shows that T = 2 optimizes model performance on high and low-resource languages for large models, we use T = 2.
System Size Imbalance: The number of atoms per system varies greatly across our pre-training
datasets. For example, Transition- Ix has 14 atoms per system on average, while OC22 has 80 atoms per system average. The naive loss reduction method shown in the non-teal terms of Equation ( which is the default behavior of most machine learning libraries, computes an atom-level force loss and then averages the force loss across all atoms in the batch. This leads to datasets with more atoms per system dominating the force loss. To address this issue, we propose a structure wise loss reduction strategy which first computes the average force loss for each system and then computes the average force loss across all systems. This ensures that the relative importance of the force loss is roughly equal across datasets, regardless of the number of nodes per system. Equation ( the updates to the naive formulation of the loss function are shown in teal and removed terms are red. This simple change leads to a significant improvement in model performance, as shown in Section 5.1
B B
No
C
= b=0
1
E
Es
Es
+
B
BE
No 6
No
Fo.s - Fboll
(I)
i=0
Energy Loss (Cg)
Force Loss (CF)
Loss Imbalance Within Single Dataset: In the single-dataset setting, Äg and Xp are typically tuned by grid search, but this approach is not feasible in the multi-dataset setting, as there M hyperparameters to tune, and changing one hyperparameter affects the optimal values of the others. Therefore, we need a simple heuristic to determine the loss coefficients for each dataset that provides reasonable balance between the energy and force losses. Inspired by Tran et al. (2022) size invariant force loss, which computes a dynamic Xf based on the number of atoms in each system of the input batch, we fix X 1 and X D., where D, is the average number of atoms per system in the ith dataset, D;. This provides a reasonable balance between energy and force loss within each dataset.
Once we have a fully pre-trained model, we can fine-tune it on downstream tasks. Our fine-tuning procedure is very similar to other fine-tuning procedures in the machine learning literature (Devlin et al. 2018 Zaidi et al., 2022): We discard the pre-training prediction heads, add new randomly initialized prediction heads for the downstream task, and fine-tune the entire model on the downstream task. This procedure is illustrated in Figure [ (middle). For fine-tuning tasks with force labels, we have the option of using the directly computed forces (i.e., using the direct equivariant block) or computing the forces by taking the gradient of the energy with respect to the atomic positions. Our initial experiments showed that JMP works well with both methods. In our evaluations, however, we chose to compute forces conservatively by taking the gradient of the energy with respect to the atomic positions, as this is the standard approach in the literature.
Fine- Tine-Tuning:
5 EXPERIMENTS
We benchmark our pre-trained models on a diverse set of atomic ML tasks. In previous related works, evaluations are commonly restricted downstream datasets that align closely with the pre-training dataset's 's domain (Zaidi et al., 2022 Zhou et al. 2023a). We posit that true success in pre-training models for atomic machine learning tasks requires adeptness at out-of-domail extrapolation. To test this hypothesis, our benchmark uniquely spans across diverse domains including small molecules (QM9 and rMD17), large molecules (MD22 and SPICE), and materials (MatBench and QMOF).
We compare our fine-tuned models (JMP) to randomly initialized models trained from scratch (GN- OC) to demonstrate the effectiveness of JMP. We also compare to previous state-of-the-arr models where available. For each task, we present results for both a small (~30M parameters, labeled with the -S suffix) and large (~~230M parameters, labeled with the -L suffix) pre-trained model to probe the impact of the model size. These two variants utilize the GN-OC Base and Large backbone architectures from Gasteiger et al. (2022), respectively. Finally, we conduct ablation studies to understand the impact of various components of JMP. More information on the datasets used for
5

---

Published as a conference paper at ICLR 2024
QM9
MD17
MD22
SPICE
Matbench
QMOF
(a) Scrach Large Scratch Small
(b) Fine-tuned Large Fine-tuned Small
Fine-tuned Large Scratch Large
Relative Improvement (%)
75
50
0
-25
-75
Fine-tuning Tasks
Figure 2: Relative performance improvement across across all tasks of all fine-tuning datasets, in percentages, of (a) Scratch Large (GN-OC-L.) over Scratch Small (GN-OC-S), (b) Fine-tuned Large (JMP-L. over Fine-tuned Small (JMP-S), and (c) Fine-tuned Large (JMP-L) over Scratch Large (GN-OC-L). GN-OC shows poor scaling to large models, a clear sign of overfitting, WhereasIMP reverses this, exhibiting much improved scaling dynamics. JMP also consistently outperforms GN-OC across all domains, datasets, and targets. The shaded rectangles indicate the average relative performance across all tasks for each dataset. The exact percentages can be found in Appendix C.I
Details on the pre-training and fine-tuning setup, such as the optimizers, learning rate schedules, and early stopping information, can be found
pre-training and fine-tuning can found in Section 3
in Appendix F Exact hyperparameters can be found in Appendix J
Common Observations: We begin by highlighting some com- mon observations across all experiments. First, when training from scratch, GN-OC-L performs 8% worse average than GN-OC-S, as shown in Figure (a). This is a clear indication of overfitting and has been consistently observed in low-data regimes (Gasteiger et al.] 2022). Second, this problem of over- fitting nearly eliminated by JMP. illustrated in Figure 2 (b). On average, JMP-L exhibits an impressive 21% relative perfor- mance gain over JMP-S. This indicates that the JMP training procedure is able to effectively leverage the additional capacity of the large model, even low-data regimes. Third, we observe that pre-training with JMP elevates performance across all do- mains, datasets, and tasks (Figure (c)), with an average relative improvement of 59% for JMP-L over GN-OC-L.
GNS-TAT-NN
JMP-S
ET-NN
JMP-L
ET-OREO
ELUIO
EHOND
20
40
60
Results on Small Molecules- - QM9 and rMD17: For each target of QM9 (Wu et al. 2018), we fine-tune dedicated model using a simple prediction head with sum pooling for all targets. For R². we use the same prediction head formulation as Thölke and De Fabritiis 2022). Our results can be found in Table 2 coompaee
Relative Improvement (%)
Figure 3: Relative improvement, over training from scratch, of dif- ferent pre-training methods on QM9's eLUMO and éHоMO-
against previous state-of-the-arr works (Liao and Smidi 2022
Batatia et al. 2022, Musaclian et al.) 2023 Feng et al.) 2023a;
Zaidi
et al. 2022). h With the sole exception of R², our JMP-L model achieves state-of-the -he--rt
results on all QM9 targets. For the R2 target, a similar phenomenon has been observed in previous
pre-training works (Zaidi et al. 2022) where the benefits of using pre-trained models are not as pronounced.
In addition their impressive performance, our JMP-S and JMP-L models demonstrate large im- provement relative to their scratch-trained counterparts. Figure 3 compares this relative improvement measured on the éLUMO and 6номо targets to other state-of-the-art pre-training and transfer learning methods for QM9. As shown. JMP outperforms all previous methods by a significant margin. This is a strong signal that our pre- pre-training approach is effective at learning generalizable representations for small molecules. We also report additional pretraining comparisons all our finetuning benchmarks with a pre-trained model from (Zaidi et al. (0022) and demonstrate significant improvements on all tasks in Appendix A
Data overlap: Due to the limited complexity of small molecules, there is some data overlap between our pre-training datasets (ANI-Ix and Transition- 1x) and QM9. To check the impact of this overlap

---

Published as a conference paper at ICLR 2024
on results, we evaluate the fine-tuning performance of our JMP-L QM9 dataset that excludes the overlapping molecules. Using molecular compositions to identify overlaps, we observe that the exclusion of overlapping molecules has negligible impact on results (see Appendix D.).
Turget (Uniss)
TorchMD-
Equi- fermer
Allegro
Pretrained KI-OREO
Pretruined GNS-TAT-NN
GN-OC-
GN-OC-
JMP.
(27)
0.011
0:111
0.015
0.016
0.020
0.233
0,008
0.059
0.946
0,388
0,05
"aceao
16.8
Fusao (meV)
14.7
17.3
(mu))))
30.0
22.0
38.5
23.1
R2 (og)
0,033
0.251
0.210
0 440
0.210
0.171
0,200
0,163
ZPVE (inel)
1.8
1.3
1.2
1.0
1.2
1.2
1.0
0.9
Uo (mue)))
6.2
4.1
4.7
5.8
7.2
9.4
3.1
2,9
(ime V
6.4
5.8
6.9
9.7
3.3
2.8
V)
6.2
C. (CalMMuK)
0,026
0.024
0,018
Table 2: MAE test split results on all targets of the QM9 dataset. SOTA results are bolded.
For rMD17, compute forces by taking the negative gradient of the energy with respect to the atomic positions. shows our force prediction results on the rMD17 dataset. Similarly to QM9, observe that JMP consistently outperforms GN-OC across all rMD17 targets. Our JMP-L model achieves state of the art performance in 5 molecules and is very competitive on the rest. Appendix B.] also shows that JMP achieves SOTA on 6/10 targets on the few-shot 50-sample subset of rMD17.
Molccules
MACK
Allegro
GN-OCS
GN-OC-L
IMP-
JMP-L
Aspirin
1.3
24.7
Benzene
0.3
0.2
1,0
1.0
0.3
Ethanol
2.1
2.1
13.0
13.3
2.8
2,0
Malomaidehyck
4,1
4.1
21,1
25.7
4,0
Naphthalene
Salicylic acid
3.1
2.9
16.7
15.1
4.6
3.4
Tolurne
1.5
1.8
6.8
7.2
2.3
1.5
Umail
2.1
1,8
12.0
12.9
2,5
Paracetamol
18.4
Anobenzeee
3.0
2.6
11.1
11.4
33
Table 3: Force MAE results in me V/À on the test split of the rMD17 dataset. SOTA bolded. Results Materials MatBench and QMOF: In the materials domain, fine-tune the MatBench (Dunn et al. 2020) and QMFF datasets (Rosen et al. 2021). For MatBench, we evaluated all regression tasks that utilize 3D structure as input and compared them with competitive models on the leaderboard (De Breuck et al. 2021 Ruff et al. 2023). For QMOF, we predict the band gap target on 10k split, similarly to Kang et (2022) ; Cao al. (023). We use mean pooling for all experiments, except MatBench's 's phonons, which is the measure of frequency of the highest frequency optical phonon mode peak and thus uses max pooling. Our results can be found Table 4 We observe that JMP-L achieves SOTA performance across QMOF and on all MatBench tasks. These two datasets contain diverse out-of-domain chemical structures (materials) and off-of-domain target labels (i.e., not energies and forces), relative to the pre-training datasets. JMP's impressive performance is yet another positive signal indicating that JMP is learning generalizable representations.
Materials (Units)
MODNet mcam)
coGGN
GN-OC-S (Sold()
GN-OC-I. (folde)
JMP-8 moean)
(Bold0 mcan)
uscan)
JDFT2D (yne /atom))
33.20
22.25
26.19
25.34
20.72/
29.94
Phonons (cans
34.77/
32.12/29.77
93.45
26.6/ 23.77
Dielectrie (umitlesa)
0.169 /0.271
0.178/ 0..309
01.225
0.211
0.135/0.255 0.252
0.119/ 6.249
Log GVRH (loglo(
/OLT7
0.068/ (0,009
0.082
01082
0.966 0.062
0.457 0.059
KVRH (log10(GPA.)
0.005/0.05
0.052/0.055
0.051
0.033
0,044 0.006/0.046
0,845 0.005
Perovskites (e//unitoel)
0.093/0.99L
0.027/0.027
0.045
Gap (el)
0.215/0.220
0.153/0.156
0.228
0.235
0.119/0121
0.089 0.091
MP R Form (mee V//tom)))
40.2 44.8
174411
31.4
$3.1
13,6/19. 13.3
10.3/ 10.1
PT OGCNN
PT MOFTrasformmr
QMOF
0.28
0.27
0.25
0.34
0.18
0.16
Table 4: MAE test split results on different targets the materials domain. SOTA is bolded.
To further investigate the impact of pre-training on unseen domains, we evaluate two large molecule datasets, MD22 (Chmiela et al.) (2023) and SPICE (Fastman al. 2023) and compare our results the previous state-of-the-art (Kovacs et al. 2023)
Results on Large Molecules MD22 and SPICE:
7

---

Published as a conference paper at ICLR 2024
Molecule
SGDMLL
MACE
Allegro
GN-OC-S
GN-OC-L
JMP-S
JMP-L
Ac-Ala3-NHMe
34,55
3,80
4.63
5.07
6.27
2.64
1.92
DHA
32.41
2,80
3.17
2.87
3.95
2.01
1.37
Stachyose
29.24
2.22
AT-AT
29.97
4.30
4.13
5.38
5.96
3.02
1.98
AT-AT-CG-CG
30.48
5,00
5.55
5.80
5.62
3.28
2.11
Buckyball Catcher
29.57
Double Walled Nanotubes
22.68
12.00
11.20
9.61
8.36
6.17
Solvated Amino Acids
6.20
8.02
Dipeptides
2.46
2.99
1.32
1,02
Table 5: Force MAE results in V/À test splits of large molecule datasets. SOTA is bolded.
For SPICE, we only use the large molecule sub-tasks: solvated amino acids and dipeptides. Similar to rMD17, we compute forces by taking the negative gradient of the energy with respect to the atomic positions. However, for MD22's Buckyball Catcher and Double- Walled Nanotubes, we were unable to fit these large structures in memory with using gradientbbased force predictions; therefore, we used direct force prediction heads instead. Our results can be found Table Once again, our model demonstrates Sate-offthhear results across all molecules of MD22 and all tasks of SPICE.
5.1 ABLATION STUDIES
Our ablations demonstrate the impact of various changes JMP the downstream fine-tuning performance. We performed pre-training experiments including dataset sampling strategies, loss formulation, and regularization strategies, and observed their impact on fine-tuning. Given the computational cost of training models on the full pre-training dataset, ablation experiments were conducted on scaled-down version of the full pre-training dataset containing a randomly selected ~2.5 examples. All pre- training models are trained for 10 epochs. Similarly, fine-tuning for these experiments was run on only one task from each of the fine-tuning datasets (MD17: Aspirin, MD22: Stachyose, QM9: Ae, MatBench: MP E Form, QMOF: Band Gap, and SPICE: Solvated Amino Acids). Additional ablations, including using fully balanced (T oo) sampling, threshold regression loss for energies and forces, and automatic task weighting strategies such as PCGrad (Yu al. (0200) are explored in Appendix Table shows the mean improvement, relative to the base, across all the fine-tuning tasks described above. A summarized insight of each ablation study follows:
Base (B): Base refers to the naive implementation of a multi-task pre-training model without temperature pling, srructurewwise loss reduction, or additional regular- ization. This model serves as the baseline for comparison. Temperature Sampling: 'Temperature-basee sampling with T' provides moderate improvement, while higher values (c.g., T' oc) show diminishing returns. This is consistent with Shaham et al. (2023). which shows that for large-enough models, T = 2 provides ideal per- formance across both low and high resource datasets. Structure- Wise Loss Reduction (SWL): The application of the structure-win wise loss reduction strategy proved to substantial improvement on the model's performance, with T2 + SWL offering a 7.7% improvement over Weight Decay (WD): Elevating the weight decay regular- ization parameter to 0. (from the default 0.01) brings the collective improvement
Ablations
(R)())
Base (Temperature 1.0) [B]
B Temperature 2.0 [T2]
2.2%
B Temperature nc [Too]]
2.6%
T2 SW Loss Averaging (SWL)
7.7%
SWL Weight Decay [WD]
11.4%
SWL Dropout [DO]
11.4%
Edge Dropout [ED]
13.2%
WD
Weights [EMA]
12.4%
EMA OC20 Only (OC20)
-9.9%
Table 6: Ablation results demonstrating the mean relative improvements of each method relative the base method (B), averaged over ablation subsplits.
1..4 over B.
Dropout (DO): Using dropout with = 0.1 on atom update layers yielded similar uplift to WD. Edge Dropout (ED): For this ablation, we drop p 0.1 of the edges at every step. We then scale the embeddings of remaining edges by factor of 1_ This yielded a small improvement over WD, I-p.
increasing the collective improvement to 13.2% over B.
Exponential Moving Average (EMA): Fine-tuning on EMA weights did not improve performance. OC20 Only (OC20): To understand the impact of multi-task pre-training, we trained model on the OC20 dataset only. We selected a 120M subset of OC20 to match the number of examples in the full JMP pre-trainne mre-trainin dataset. Note that this means that the dataset used in the OC20 ablation contains 48x more data points than the rest of our ablations. Despite this, OC20 performed substantially worse than B, indicating that diverse multi-task pre-training is important for generalization.
8

---

Published as a conference paper at ICLR 2024
Based on our ablation results the two most important changes from B were SWL and regularization methods like WD, DO, and E.D. These results are consistent with Kurin et al. (022)) which demonstrates the effectiveness of regularization in multi-task learning. Our final model integrates temperature-basee sampling (T2), structure-wise loss reduction strategy (SWL), an amplified weight
decay regularization parameter of 0.1 (WD), edge dropout with p = 0.1 (ED), and EMA despite not
showing a performance boost as it has been standard for training GemNet (Gasteiger et al. 2022)
5.2 COMPUTATIONAL COST ANALYSIS
Time to Milestone (GPU Hours)
150
JMP
GN-OC
Pre-training JMP-L required significant computa- tional resources, which is typical for foundation model approaches. We pre-trained JMP-L on 128x V100 32GB GPUs for 2 epochs, which took around 34,400 GPU hours in total (sce Appendix G for ex- act training times and CO2 impact). While this is
100
50
a substantial upfront investment, it enables efficient fine-tuning across a diverse set of downstream tasks.
QM9
MD17
MD22
SPICE
Matbench
OMOF
We evaluated JMP-L fine-tuning performance versus training models from scratch (i.e., GN-OC-L). Train- ing GN-OC-L. on the downstream tasks until con- vergence based on our stopping criteria took around 3,300 GPU hours in total across all tasks. In con- trast, fine-tuning JMP-L on the same tasks took only around 275 GPU hours total to match the perfor- mance of the models trained from scratch. This 12x reduction in compute demonstrates the significant benefits of pre-training. Figure 4 shows this difference compute requirements, averaged for each fine-tuning dataset.
Figure 4: The number of GPU hours, aver- aged for each dataset, required to train GN- OC-L to convergence and to fine-tune JMP-I. to match GN-OC-L's performance. Overall, fine-tuning JMP-L was able to match GN-OC- L's performance in & the time. T2
CONCLUSION
Our findings demonstrate the promise of pre-training strategies that leverage diverse data to improve atomic property prediction. Our source code and pre-trained models are publicly available on GitHut" We hope this will spark further research and innovation in this area to accelerate progress towards a foundation model for chemistry. It is essential to acknowledge the potential for bias in our models and datasets, which can impact their prediction capability and quality. Additionally, given the potential for misuse of atomic machine learning models, we emphasize that these models must be used responsibly,
There are several limitations of this study that provide fertile ground for future research. Due to the large computation expense of pre-training new models, our study only experimented with the GemNet-OC backbone model. Further exploration of different backbone models is warranted. Additionally, the current methodology of discarding pre-training prediction heads before fine-tuning may not be optimal, particularly for datasets with similar labels to the pre-training sets. Lastly, the model size used in this study, though substantial, dwarfed by the largest models in NL.P and CV. We anticipate that employing larger models, trained with the help of recent data and model parallelism techniques (Sriram et al. 2022), could enhance performance.
In summary, this work presents Joint Multi-domain Pre-training (JMP), novel atomic pre-training strategy that leverages diverse atomic datasets to learn rich representations through multi-task regression. By effectively formulating and regularizing the multi-task learning problem, we achieve remarkable performance on various in-domain and out-of-domain fine-tuning tasks. Scaling up to large model with over 235 million parameters leads to improved performance on all downstream tasks, even low-resource ones like rMD17. This suggests that pre-training not only enhances accuracy but also facilitates effective scaling, allowing models benefit from increased capacity without overfitting, Additionally, we establish a comprehensive set of fine-tuning benchmarks across various chemical domains and tasks. Building off the results present here to create larger, more general, and more accurate ML potentials will remain significant challenge for the field moving forward.

---

Published as a conference paper at ICLR 2024
Acknowledgementt:: We would like to thank Dr. Abhishek Das, who provided valuable insights and feedback on the research direction and manuscript. We would also like to thank Anuroop Sriram and Dr. Johannes Gasteiger for their valuable insight. Finally, we would like to thank the entire FAIR Chemistry team for their support and feedhack.
10

---

Published as a conference paper at ICLR 2024
REFERENCES
llyes Batatia, Dávid Péter Kovács, Gregor NC Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. arXiv preprint urXiv:2206.07697 2022.
Simon Batzmer, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korbluuh, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. Se (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXiv:2101.00 164, 2021.
Jörg Behler. Perspective: Machine learning potentials for atomistic simulations. The Journal of chemical physics, 145(17):170901, 2016.
Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-eneryy surfaces. Physical review letters, 98( 14)441444401, 2007.
Chenjing Cai, Shiwei Wang. Youjun Xu, Weilin Zhang, Tang, Qi Ouyang, Luhua Lai, and Jianfeng Pei. Transfer learning for drug discovery. Journal of Medicinal Chemistry, 63(16):8683-866 2020.
Zhonglin Cao, Rishikesh Magar, Yuyang Wang, and Amir Barati Farimani. Moformer: self-supervised transformer model for oetal-organic framework property prediction. Journal of the American Chemical Society, 145(5):2958- 29688829967224 2023.
HC Stephen Chan, Hanbin Shan, Thamani Dahoun, Horst Vogel, and Shuguang Yuan. Advancing drug discovery via artificial intelligence. Trends in pharmacological sciences, 40(6):592-604, 2019.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras- DomaisDo Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 11(10)6059-6072. 2021.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e 1603015 2017.
Stefan Chmiela, Valentin Uussile--aaliddr Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexan- dre Tkatchenko, and Klaus-Robers Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 92):aaddf087,, 2023.
Pierre-Paul De Breuck, Geoffroy Hautier, and Gian-Marco Rignanese. Materials property prediction for limited datasets enabled by feature selection and joint learning with modnet. upj Computational Materials, 7(1):83, 2021.
Jianyuan Deng, Zhibo Yang, Iwao Ojima, Dimitris Samaras, and Fusheng Wang. Artificial intelligence in drug discovery: applications and techniques. Briefings in Bioinformaties, 23(1), 2022.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wxXiv:1810.0480. 2018.
Uri Dinur and Amold T Hagler. New approaches to empirical force fields. Reviews in computational chemistry, 2:99-164 164, 1991.
Alexander Dunn, Qi Wang. Alex Ganose, Daniel Dopp, and Anubhav Jain. Benchmarking materials property prediction methods: the matbench test set and automatminer reference algorithm. upj Computational Materials, 6(1):138, 2020.
Peter Eastman, Pavan Kumar Behara, David L Dotson, Raimondas Galvelis, John E Herr, Josh Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
Rui Feng, Qi Zhu, Huan Tran, Binghong Chen, Aubrey Toland, Rampi Ramprasad, and Chao Zhang May the force be with you: Unified force-centric pre- voe-training for molecular conformations. arXiv preprint arXiv:2308 14759, 2023a.
11

---

Published as a conference paper at ICLR 2024
Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Wei- Ying Ma. Fractional denoising for 3d molecular pre-trate pre-training, In International Conference on Machine Learning, pages PMLR. 2023b.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations, 2020. URL https://
openreview net/ F
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C Lawrence Zitnick, and Abhishek Das. Gemnet-oc: developing graph neural networks for large and diverse molecular simulation datasets. Transactions on Machine Learning Research, 2022.
MA González. Force fields and molecular dynamics simulations. École thématique de la Société Française de la Neutronique, 22:16-2000, 2011.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. arXiv preprint arXiv: 00111000116666.... 2018
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: large-scale challenge for machine learning graphs. arXiv preprint arXiv:2103.0P4 2021.
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. 3d equivariant molecular graph pretraining. arXiv preprint (rXXi:2207.088244 2022.
Yeonghun Kang. Hyunsoo Park, Berend Smit, and Jihan Kim. Moftransformer: A multi-modal pre-training transformer for universal transfer learning in metal-organic frameworks. 2022.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7482-7491, 2018.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 2014.
Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. arXiv preprint arXiv:2106.080000 2021.
Adeesh Kolluru, Nima Shoghi, Muhammed Shuaibi, Siddharth Goyal, Abhishek Das, C Lawrence Zitnick, and Zachary Ulissi. Transfer learning using attentions across atomic systems with graph neural networks (taag). The Journal of Chemical Physics, (55(18):180) 84702, 2022a.
Adeesh Kolluru, Muhammed Shuaibi, Aini Palizhati, Nima Shoghi, Abhishek Das, Brandon Wood, C Lawrence Zitmick, John R Kitchin, and Zachary W Ulissi. Open challenges in developing general- izable large scale machine learning models for catalyst discovery. arXiv preprint arXiv:2206.02005, 2022b.
David Peter Kovacs, Ilyes Batatia, Eszter Sara Arany, and Gabor Csanyi. Evaluation of the mace force field architecture: from medicinal chemistry to materials science. arXiv preprint arXiv:2305.14244 2023.
Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and Pawan K Mudigonda. In defense of the unitary scalarization for deep multi-task learning. Advances in Neural Information Processing Systems, 35:12169-12188 2022.
Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and Thomas Dandres. Quantifying the carbon emissions of machine learning. arXiv preprint 1rXiv:1910.09700 2019.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990 2022.
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre- training molecular graph representation with 3d geometry. arXiv preprint arXiv: 2021.
12

---

Published as a conference paper at ICLR 2024
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with se (3)-invari denoising distance matching. arXiv preprint arXiv:2206.13602. 2022.
Zhihai Liu, Yan Li, Li Han, Jie Li, Jie Liu, Zhixiong Zhao, Wei Nie, Yuchen Liu, and Renxiao Wang. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31(3): 405-412, 2015.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint AXXv::6608.03983, 2016.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communicutions, 14(1):579, 2023.
Giorgio Pesciullesi, Philippe Schwaller, Teodoro Laino, and Jean-Louis Reymond. Transfer learning enables the molecular transformer predict regio-and stereoselective reactions on carbohydrates. Nature communications, 11 11844:8874, 2020.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. arXiv preprint arXiv:1907.10000 2019.
Andrew S Rosen, Shaelyn M Iyer, Debmalya Ray, Zhenpeng Yao, Alan Aspuru- Aspuru-Guzik, Guzik, Laura Gagliardi, Justin M Notestein, and Randall Q Snurr. Machine learning the quantum-chemical properties of metal-organic frameworks for accelerated materials discovery. Matter, 4(5):1578- 1597, 2021.
Lars Ruddigkeit, Ruud Deursen, Lorenz C. Blum, and Lean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of Chemical Information and Modeling, (2(11):2864-22
2012a. doi: 10. 10.1021/ci300415d URL
PMID: 23088335.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz Blum, and Jean-Louis Reymond. Enumeration 166
billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 11:2264-28875 2012b.
Robin Ruff, Patrick Reiser, Jan Stühmer, and Pascal Friederich. Connectivity optimized nested graph networks for crystal structures. arXiv preprint 2xXiv:2302.14102. 2023.
Mathias Schreiner, Arghya Bhowmik, Tejs Vegge, Jonas Busk, and Ole Winther. Transition lx-a dataset for building generalizable reactive machine learning potentials. Scientific Data, 9(1):779, 2022.
Kristof Schütr, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems, volume 30.
Curran Associates, Inc., 2017. URL httpss://prcceeddggg nwwrips....mmm ccwpappppccpppppr.r.pco fi file// wwesspapper/0000 (6 wtw.:aabbbb6c22299baaabbbb
pddf
Uri Shaham, Maha Elbayad, Vedanuj Goswami, Omer Levy, and Shruti Bhosale. Causes and cures for interference multilingual translation. In Proceedings of the 6lst Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15849-15863,, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. 10.186533/1/2023.ccco 883. URL htt.stts:t/////tttc http://aaaaaacoooo..oooooooooooa wwwplllyyyyoo./////// ott.pp/pes...llaaaaooooooooooooo
883
Justin S Smith, Benjamin T Nebgen, Roman Zubatyuk, Nicholas Lubbers, Christian Devereux, Kipton Barros, Sergei Tretiak, Olexandr Isayev, and Adrian Roitberg. Outsmarting quantum chemistry through transfer learning. 2018.
13

---

Published as a conference paper at ICLR 2024
Justin S Smith, Benjamin T Nebgen, Roman Zubatyuk, Nicholas Lubbers, Christian Devereux, Kipton Barros, Sergei Tretiak, Olexandr Isayev, and Adrian E Roitberg. Approaching coupled cluster accuracy with a ceneral-purposs neural network potential through transfer learning. Nature communications, 10(1):1-8, 2019.
Justin S Smith, Roman Zubatyuk, Benjamin Nebgen, Nicholas Lubbers, Kipton Barros, Adrian E Roitberg, Olexandr Isayev, and Sergei Tretiak. The ani-Iccx and ani-Ix data sets, coupled-cluster and density functional theory properties for molecules. Scientific data, 7(1):1-10, 2020.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint (XKv:200000055000 2020.
Anuroop Sriram, Abhishek Das, Brandon M Wood, and C. Lawrence Zitnick. Towards training billion parameter graph neural networks for atomic simulations. In International Conference on Learning Representations, 2022. |RL hppp//eepnreeeeeeeeeeeeeeee /ppettttes//oeeeeeviieeee net. wtt.r//wwwweeeeeueeeeeeeeeeeeeee
a T Sundius. Scaling of ab initio force fields by molvib. Vibrational Spectroscopy, 29(1-)):89-95 2002.
Philipp Thölke and Gianni De Fabritüi. Torchmd-net: equivariant transformers for neural network based molecular potentials. arXiv preprint arXiv:2202.02541, 2022.
Richard Tran, Janice Lan, Muhammed Shuaibi, Siddharth Goyal, Brandon M Wood, Abhishek Das, Javier HerisnDo Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysis. arXiv preprint arXiv:2206.08911 2022.
Masashi Tsubaki and Teruyasu Mizoguchi. Quantum deep descriptor: Physically informed transfer learning from small molecules to polymers. Journal of Chemical Theory and Computation, 17(12): 7814-7821, 2021.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: benchmark for molecular machine learning. Chemical science, 9()):513-530, 2018
Hironao Yamada, Chang Liu, Stephen Wu, Yukinori Koyama, Shenghong Ju, Junichiro Shiomi, Junko Morikawa, and Ryo Yoshida. Predicting materials properties with little data using shotgun transfer learning. ACS central 5(10n 1117717777777 2019.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824-5836, 2020.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Manchez-Gonzalee Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.0011 2022.
Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng. Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, et al. Towards predicting equilibrium distributions for molecular systems with deep learning. arXiv preprint arXiv:2306.0554 2023.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023a. URL https 7/ 7/openreviee n www.://ntt...ttttettt Sttrummm.dd=K//MMwwwss
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: universal 3d molecular representation learning framework. 2023b.
Jinhua Zhu, Yingce Xia, Lijun Wu, Shufang Xie, Tao Qin, Wengang Zhou, Houqiang Li, and Tie-Ye Yan Liu. Unified 2d and 3d pre-training of molecular representations. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 2626-2636, 2022.
Lawrence Zitnick, Lowik Chanussot, Abhishek Das, Siddharth Goyal, Javier Heras- Peeras-Domingo.. Caleb Ho, Weihua Hu, Thibaut Lavril, Aini Palizhati, Morgane Riviere, et al. An introduction to electrocatalyst design using machine learning for renewable energy storage. arXiv preprint wrXi::2010.09435, 2020.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
Table of Contents:
Additional Pretraining Comparisons
B
Additional Results and Experiments
C
Embedding and Activation Visualizations
D
Additional Method Information
Structure-Winc Loss Reduction Code
Training Setup
Model Training Times and CO2 Impact
Additional Dataset Details
Dataset Overlap
Model and Optimization Hyperparameters
A ADDITIONAL PRETRAINING COMPARISONS
We compare our finetuning results on all benchmarks with the publicly available checkpoint from Zaidi al. (2022) that performs pretraining via denoising. The only checkpoint available with TorchMDNet et backbone and not their best-periorming model, therefore we just add results for pretraining ET+NN where NN stands for noisy nodes as described in the original work. In this method, they pretrain on the PCQM4Mv2 (Hu et al. (221)) dataset which is a 3D small molecule dataset. We perform significantly better on all datasets (including the small molecule domain) showing the impact of our supervised joint pretraining approach. A model trained just small molecule dataset denoising setup focused on equilibrium systems performs worse than the model trained supervised setup non-equilibrium combining datasets from multiple chemical domains. Our Figure in the main paper also demonstrated the relative improvement coming from pretraining relative to models performed from scratch.
Molecules
Pretrained
JMP-S
JMP-L
ET+NN
Aspirin
15.1
6.7
5.1
Benzene Ethanol
1.0
0.7
0.3
8.9
2.8
2.0
Malonaldehyde
13.0
5.3
4.0
Naphthalene
4.5
2.2
1.4
Salicylic acid
9,8
4.6
3.4
Toluene
4.8
2.3
1.5
Uracil
6.4
4.0
2.5
Paracetamol
12.6
5.3
4.0
Azobenzene
7.6
4.5
3.3
Table 7: Force MAE errors in v/À on the test split of the rMD17 dataset.
B ADDITIONAL RESULTS AND EXPERIMENTS
B. EXPANDED RMD17 RESULTS: TRAINING ON 50 EXAMPLES
Following the idea presented by Batatia et al. 2022 we reduced the rMD17 per molecule training set size from 1000 50 examples and evaluated the performance of the JMP-L model. The test set remained unchanged. Table shows these results. The large pre-trained model provided ofate-of-the. art performance on 6 out of 10 molecules and gives competitive results the rest. Most notably,
15

---

Published as a conference paper at ICLR 2024
Target (Units)
Pretrained ET+NN
JMP-S
JMP-I
(D) a (ag)
0.015
0.010
0.008
0.069
0.037
0.032
EHOMO (meV)
23.1
11.1
8.8
LLUMO (meV)
19.1
10.8
8.6
AE (meV)
39.8
23.1
19.1
R² (az)
0.556
0.200
0.163
ZPVE (meV)
1.1
1.0
0.9
Uo (meV)
6.0
3.3
2.9
U (meV)
6,0
3.3
2.8
H (meV)
6.1
3.3
2.8
G (meV)
6.9
4.5
4.3
Cv (cal/mol K)
0.021
0.018
0.017
Table 8: Average absolute error results all targets of the QM9 dataset.
Molecule
Pretrained ET+NN
JMP-S
JMP-L
Ac-Ala3-NHMe
8.53
2.64
1.92
DHA
7.23
2.01
1.37
Stachyose
10.12
2.69
1.73
AT-AT
10.35
3.02
1.98
AT-AT-CG-CG
8.81
3.28
2.11
Buckyball Catcher
3.08
2.26
Double Walled Nanotubes
8.36
6.17
Solvated Amino Acids Dipeptides
7.11
1.60
1.33
6.56
1.32
1.02
Table 9: Force MAE results in me V/Â on the test splits for large molecule datasets (MD22 and SPICE). Missing numbers are due to runs failing or hitting NaNs.
16

---

Published as a conference paper at ICLR 2024
that the fine-tuned model performs significantly better than the GN-OC-L model. These results offer preliminary evidence that reasonable few-shot performance may be achievable with a large pre-trained atomic model.
Molecules
NequIP
MACE
GN-OC-L
JMP-L
Aspirin
52.0
43.9
119.7
36.8
Benzene
2.9
2.7
8.8
2.8
Ethanol
40.2
32.6
95.3
22.2
Malonaldehyde
52.5
43.3
159.2
42.9
Naphthalene
10.0
9.2
39.5
Salicylic acid
35.0
28.4
108.5
26.4
Toluene
15.1
12.1
53.3
12.4
Uracil
40.1
25.9
135.5
25.8
Paracetamol
39.7
31.5
66.4
27.3
Azobenzene
20.0
17.7
139.0
17.8
Table 10: Average absolute force prediction errors in meV/À the test split of the revised MD17 dataset comparing the performance of fine-tuning our pre-trained model with GN-OC trained from scratch, as well as other state-of-the-art results. For all runs, the train was limited to only 50 examples.
B.2 PDBBIND V2013 BINDING AFFINITY RESULTS
We also evaluate the performance of our pre-trained model on the core of the PDBBind v2013 dataset (Liu et al. 2015), provided in the MoleculeNet benchmark (Wu et al. 2018). which contains 1 154 protein-ligand complexes. The dataset contains 3D structures protein- ligand complexes, as well as the binding affinity of the ligand to the protein. Following MoleculeNet, we use random 80%/10%//0%% train/validationntestt split, and we report the RMSE metric. Table Table[[showw shows the results of our pre-trained model compared to the state-of-the-art results. We see that our pre-trained model achieves state-of of--oh--thttte the-art performance on this dataset, further demonstrating that our pre-trained model can be used for wide range of molecular tasks, including those that involve protein-ligand complexes.
Method
log Ka K;
RMSE
JMP-L
1.36
DeepBindGCN _G_
1.49
SE-OnionNet
1.69
DeepBindRG
1.81
GraphBAR (dataset 4, best)
1.63
BAPA
1.45
Table 11: Binding affinity RMSE results on the PDBBind v2013 dataset.
B.3 ADDITIONAL. ABLATION STUDIES
Fully balanced datasets: Finding the optimal dataset scaling is important for pre-training with datasets of variable size. While we investigated dataset scaling with temperature of 2.0 in the main ablation results, we additionally tested the impact of fully balancing the datasets (high temperature). The results for this ablation can found in Table [3] labeled as B + Fully balanced. Compared to the base pre-trained model where there mean relative improvement (RI) of 2.61% on the ablation fine-nuning tasks. This is very similar to with temperature 2.0 (mean RI over base of 2.29%). but requires more training steps to reach an epoch of the largest dataset. As a result, we used temperature 2.0 dataset scaling for all our models trained the full pre-training dataset.
dataset scaling, the fully balanced run provided
the improvement
17

---

Published as a conference paper at ICLR 2024
Threshold loss: DFT is approximate method and result the labels (e.g. energy and forces) in DFT datasets have some error associated with them. In our multi-dataset case, many datasets are computed using different DFT engines and levels of theory, leading distinct noise distributions for each dataset. To address this, we experiment with a threshold loss which measures the distance between the model's prediction and the ground-truth (i.e., DFF-computed) lahel, but only penalizes predictions that are outside of a given threshold. To implement this threshold loss, we modify the loss original functions, Ci by incorporating the predefined physically motivated margins for each The equation for the threshold loss is shown in Equation (
dataset as shown in Table 12
C. v) - (0.3) if (Ú. 8) 2 margin otherwise
(2)
where C is the original loss function for dataset i, and M; is the metric used to compute the margin for dataset For all our datasets, we use the mean absolute error (MAE) the metric for computing the margin. Table 13 shows ablation studies for model which used this threshold loss mechanism, labeled B + T2 + SWI. + Threshold Loss. While our initial experiments indicated that the threshold loss was a promising modification, our final ablation results show that the threshold loss hurts performance.
Dataset
Energy (eV)
Forces (eV/Å)
ANI-1x
0.043
0,01
Transition-
0.043
0.01
OC20
0.03 0.1
OC22
0.1
0.03
Table 12: Summary energy and force threshold values for different datasets.
Automatic Task Weighting with PCGrad: PCGrad Yu et al. 2020) is "gradient surgery" method that attempts to address the optimization challenges in multi-tase task learning by projecting each task's gradient onto the normal plane of the gradient of any other task that has conflicting gradient. We evaluate the usage of PCGrad alongside our structure-wise loss reduction strategy. Our experiments showed that PCGrad (PCG) does not offer performance improvements over SWL. These results are consistent with Kurin et al. (2022), which shows that adequate regularization can mitigate the need for complex multi-task optimization methods.
Ablations
Mean RI (%)
Base (B) B + Fully balanced
0%
2.61 %
+ Temperature 2.0 (T2)
2.29%
T2 SW Loss Averaging (SWL)
7.68%
SWL + Threshold Loss
5.38 %
SWL + PCGrad (PCG)
6.48%
SWL + Weight Decay (WD)
11.37%
SWI. + Dropout (DO)
11.41%
SWL + Edge Dropout (ED)
13.18%
SWL + WD EMA Weights (EMA)
12.38%
EMA + OC20 Only (OC20)
-9.88 %
Table 13: Additional experiments with all ablation results
B.4 FINE-TUNING LR SCHEDULING
Learning rate scheduling has a large impact the downstream fine-tuning performance. We utilize with a robust learning rate scheduling strategy that we employ in all our fine-tume fine-tuning experiments. Namely, combine the ideas behind linear warmup. cosine decay (Loshchilov and Hutter 2016).
18

---

Published as a conference paper at ICLR 2024
Layerwise Learning Rate Decay (LLRD) (Howard and Ruder 2018), and ReduceLROnPlateent Let a be our learning rate, Nw be the number of warmup epochs, Ju be the warmup initial learning rate coefficient, Nº. be the number of cosine decay epochs, f. be the cosine decay final learning rate coefficient, Np be the number of patience epochs for SecduceLRRnPlann learning rate coefficient, and D; be the L.RD decay coefficient for layer Then, our learning rate scheduling strategy for layer is described below:
fp be the patience
1. Linear Warmup:
During the initial phase training, the learning rate a begins at fw a D2 and gradually escalates to reach a D; over N warmup epochs. This strategy aids in preventing substantial gradients at the beginning of training, thereby ensuring a stable optimization process.
2. Cosine Decay:
After the warmup phase, we transition to cosine decay strategy. The learning rate starts from D; and decays fe a over Ne epochs. This phase facilitates the gradual reduction of the learning rate, enabling efficient model convergence. The final learning rate after the cosine decay phase the same for all layers, meaning LLRD only influences the warmup and cosine decay phases.
3. ReduceL ROnPlateau:
The final phase is governed by the Reducel.ROnPlateee ROOnPlteeau strategy- It commences with the learning rate set at fe a. and if the validation loss does not decrease for N, epochs, the learning rate is decreased by multiplying it with fp- This dynamic adjustment of the learning rate based on validation loss performance assists fine-tuning the model parameters towards the end of the training process.
This scheduling strategy provides a robust method to adjust the learning rate across different phases of the model training, thereby balancing need for both rapid learning and careful optimization as the model converges to the best solution. Table 14 shows the fine-tuning performance of the JMP-L model when fine-tuned on the Aspirin molecule of the rMD17 datasct using different components of our LR scheduling strategy. We observe that the performance of the model improves we add each component of the LR scheduling strategy. We can see that all components of our scheduling strategy contribute the performance of the model. Overall, we see a massive improvement of 31 % in the force MAE when we use our full LR scheduling strategy when compared to simply using cosine decay strategy. This demonstrates the importance of a robust learning rate scheduling strategy for fine-tuning the model on downstream tasks. For our final runs, the specific values of the LR scheduling hyperparameters can be found in Appendix[]
LR Schedule
Forces MAE
Warmup Cos
7.8
Warmup + Cos LLRD
6.4
Warmup Cos LLRD RLP
5.1
Table 14: Impact of learning rate schedules Aspirin Force MAE (meV/À) when evaluated on the validation set. Cos: Cosine Decay, LLRD: Layerwise Learning Rate Decay, RLP: ReduceL.ROn- Plateau
C EMBEDDING AND ACTIVATION VISUALIZATIONS
In this section, illustrate the potency of our pre-trained model's embeddings through visualizations. The t-SNE plot in Figure rrepresents the node-level (h) and edge-level (m) GN-OC embeddings of the JMP-L model, computed, averaged across the entire system shown Equation (3), where N is the number of nodes in the system, E is the number of edges in the system, h E (N, Dr)) is the GN- OC node embedding tensor for the system, and E (E, D.,)) the GN-OC edge embedding tensor for the system. We randomly select structures across all pre- pre-training and fine-tuning development datasets (as described in Appendix HH3)) and compute these aggregated system-level embeddings. Fach data point corresponds to unique structure, and its color indicates the dataset from which the structure originated.
N
E
h and m
m
(3)
N
E
1=0
e=0
19

---

Published as a conference paper at ICLR 2024
OC20
Translx
QM9 SPICE
Matbench
MD17
OC22
ANI1x
QMOF
MD22
(a) Node Embeddings
(b) Edge Embeddings
Figure 5: t-SNE visualizations of the node-level (h) and edge-level (m) JMP-L embeddings for randomly selected structures from pre-training and fine-tuning development datasets. Each point represents a structure, and the color indicates the dataset from which the structure was sampled.
H
Co
Cu
-04
-08
Figure Visualization of the cosine similarity between the learned atom embeddings H, C, N, O. Co, and Cu with all other elements in the dataset. Note that nearby atoms in the periodic table have more similar embeddings, which consistent with known element properties.
20

---

Published as a conference paper at ICLR 2024
Furthermore, Figure 6 provides intriguing visualization of the cosine similarities between the learned atom embeddings for a selection elements (H, C, N, O, Co, and Cu) with all other elements in the dataset. The plot reveals that atoms located adjacently the periodic table exhibit more similar embeddings, a result that aligns well with their known elemental properties. This underlines the model's s ability accurately capture and encode fundamental chemical properties within its learned embeddings.
C.1 EXACT PERCENTAGES FOR FIGURE 2
The per-dataset averages values (with translucent shading) for Figure 2 are: (a) QM9 -10.7%, MD17 -5.4%, MD22 15.4%, SPICE -25.6%, Matbench 0.4%, QMOF 2.2%; (b) QM9 14.3%, MD17 31.6%, MD22 31.1%, SPICE 19.5%, Matbench 9.4%, QMOF 9.5%; and (c) QM9 49.5%; MD17 77.9%;
MD22 61.0%; SPICE 74.6%; Matbench 45.0%; QMOF 32.0%.
D ADDITIONAL METHOD INFORMATION
D.1 OPTIMAL TASK WEIGHTING AND REGULARIZATION
The problem of optimal task weighting is active of research in multi-task learning. Existing works range from dynamically weighting tasks based on uncertainty estimates computed from their loss (Kendall et al. 2018) performing "gradient surgery" on per-task gradients to alleviate conflict- ing gradients. However, recent work by Kurin et al. (2022) has shown that unitary sealarization, or simply summing the losses across all tasks, matches outperforms most existing methods, provided that adequate regularization is used. Inspired by these findings, we use unitary scalarization for pre-training, and use the following regularization techniques: We a weight decay of 0.1, edge dropout (Rong et al., 2019) with p = 0.1, and exponential moving average with a decay of 0.99 on the model weights.
D.2 FINE- TUNING FORCE COMPUTATION
For fine-tuning tasks with force labels, have the option of using the directly computed forces (i.e., using the direct equivariant block) or computing the forces by taking the gradient of the energy with respect to the atomic positions. Our initial experiments showed that JMP works well with both methods. In our evaluations, however, we chose to compute forces conservatively by taking the gradient the energy with respect the atomic positions, as this is the standard approach in the literature.
E STRUCTURE- SIISE LOSs REDUCTION CODE
In this section, provide the code for the structure-wise loss reduction (SWL) loss function, well as the original GemNet-OC loss function for comparison. The code is written Py and is shown in Listings and 2 Table below maps the notation used the sample code to the notation used in Equation in the main text.
def gemnet_oc_loss(
8: Tensor, (batch_size,
F: Tensor, (batch sizesnum atoms), 3)
E. target:: Tensor, # (batch_ad sincc__siiee)
F tttarget: Tensor, # (bbatch_sizernnn size=nun atoms)) 3) batch_idx: Tensor,
(ibatch_sizoxuum_att
lambda_E: Tensor, scalar
lambda Tensor, scalar
"The batcc_attcccccea idx tensor is product of PyTorch Geometric's sparse graph batching, which constructs batched graph as one large graph with disconnected components. bat patchh_on tensor maps each node in the batched graph to its corresponding component in the original batch. In the notation of the paper, we do not use this explicitly to simplify the notation. Instead, we use subscripts on node-level quantities to denote the component that the node helongs to. For example, Fo, denotes the force vector of the i-th atom in the b-th component of the batched graph, N, denotes the number of atoms in the b-th component of the batched graph.
21

---

Published as a conference paper at ICLR 2024
Sample Code Notation
Paper Notation
F
B_target
E
F F_target
F
Lambda_ E
Ag
lambda_F
Ay
batch patch_idx
Not presen
dataset_idx
W
Table 15: Mapping of notation used the sample code to the notation used Equation (1 in the main text.
10
Energy loss computation
11
(E E_target) 2 (batch_size, )
12
energy_ = lambda E forch.mean/
systemwise _eerrrrgseeeee_yeeeeeee_ loss
scalar
15
16
Force loss computation atomvisc_F_loss = torch.norm
17
F F_target,
p-2, dim--1
20
21
} + baatchssizewune atoms,)
22
force_lcse Lambda F torch mean (atomwise Loss) scalar
return energy loss force loss 24
Listing 1: Original GemNet-OC loss function
def
jmp loss
E: Tensor, (batch_size) F: Tensor,
beattsh_sizeenu_attoo 3)
K_target: Tensor, (batch_size,)
F_target: Tensor, ((aatch_sizennu__ttom 3) batch_iixt Tensor,
dataset idx: Tensor, f (batch_size) lambda_ E: Tensor, # (mum_datasets)
lambda_F: Tensor, (num_datasets,)
11
Energy computation
12
Energy computation is the same as the original GemNet-OC: loss function, with the only di fiffennnce being that have task-specific energy loss coefficients.
13
16
systemwise_enettee eeweennergyy__ energy_ loss = (
(E E_target) 2 baatch_size,
17
18
systemwiise_enerqy__
lambda_ [dabddaa_[[attaet__ido
20
scalar
21
22
Force loss computation
atomvise___loss = torch.norn
F F_target,
25
p=2,
26
dim=-1
27
) t
The only difference is that
perform strusture-level averaging,
22

---

Published as a conference paper at ICLR 2024
instead of atom-level averaging. structurowiss_forree.loss scatter_mean
31
32
atomwise_F_lose,
33
batch_idx
(batch_size,)
Similarly to the energy less, we have
36
task-specific force loss coefficients.
37
force_lose = torch.mean.
38
Lambda dataset_idaa idx|
40
) scalar
41
42
return energy_loss force_less
Listing Our proposed loss function
F TRAINING SETUP
F.1 TRAITRIINING
Optimizer We use the AdamW (Kingma and Ba, 2014) optimizer with a learning rate of 0.0003, 31 = 0.9, 32 0.95, and a weight decay value of 0.1. Learning Rate Scheduling. During pre-training, we use linear- warmup with cosine decay learning rate schedule. The linear warmup starts with 0.2 LR and warms up to LR over 2000 steps. The cosine decay reduces the learning rate to 0.1 LR over 2 epochs.
F.2 PINE-TUNING
Optimizer We use the AdamW (Kingma and Ba 2014) optimizer with a learning rate of 0.00008,
31 0.9, 32 0.95, and weight decay value of 0.1.
Loss Function. We use the MAE loss function for scalar targets (e.g., energy, band-gap) and the L2 distance loss function vector targets (e.g.. forces).
Learning Rate Scheduling. Across all fine-tuning tasks, we utilize the following learning rate schedule: (1) Warmup over epochs, (2) cosine decay over 32 epochs, and (3) reduce on plateau for the remainder of training. We also utilize Layer-wise Learning Rate Decay (LLRD) (Howard and Ruder 2018 during phases (1) and (2). LLRD employs higher initial learning rates for final, more task-specific layers. The final learning rate after phase (2), however, is the same for all layers. See Appendix(B.I more information on the fine-tumin fine-tuning LR scheduling.
Early Stopping. All our fine-tuning runs use the following stopping criteria: (1) Early stopping with a patience of 50 epochs, (2) maximum of 500 epochs or 7 days of training, or (3) the learning rate dropping below 10-5, For the rMD17 dataset, due the small size of the train set, of 1000 and maximum of 100,000 epochs, similar to Musaclian et al. (2023).
patience
G MODEL TRAINING TIMES AND CO2 IMPACT
Table 16 shows the total training time (in GPU-hours) of pre-training our model on the pre-training datasets, fine-tuning the model on all fine-tuning datasets, and training baseline scratch models for each fine-tuning dataset. The substantial computational resources required for pre-training the models cannot be understated. These numbers, however, are significantly offset when considering the subsequent fine-tuning phase. Remarkably, the fine-tuned models, starting from the pre-trained checkpoints, demonstrated superior performance while utilizing approximately half of the compu- tational resources required by the models trained from scratch. This efficacious when applying our technique substantiates its viability. Morcover, the initial investment of resources in pre-training the models is well compensated these models, once trained, can be reused across multiple applications, thereby amplifying their utility and cost-effectiveness.
of resources
All experiments conducted using private infrastructure, which has carbon efficiency of 0.432 kgCOzeq/khh. A cumulative of 46400 hours computation was performed on hardware of type Tesla V100-SXMP-32GBB (TDP of 300W). Total emissions estimated to be 6013.44 kgCOzeq of
23

---

Published as a conference paper at ICLR 2024
which 100 percents were directly offset. Estimations were conducted using the Machine Learning Impact calculator presented in Lacoste et al. (2019).
Model
Time (GPU-hours)
CO2 Emissions
JMP-L Pre- Pre-Training
34400
4458.24
JMP-S Pre-Training
5700
738.72
JMP-L Fine-Tuning
3000
388.8
Scratch GN-OC-L Training
3300
427.68
Table 16: Training times of different models this work and their corresponding CO2 emission estimates
H ADDITIONAL DATASET DETAILS
H.1 PRE-TRAINING DATASETS
For the purpose of pre-training, we selected two datasets from the catalysis domain, namely OC20 and OC22, as well as two small molecules datasets, ANI-Ix and Transion- 1x. These datasets were chosen due to their substantial training sizes and the diversity of structures they offer. Note that all of the pre-training datasets we used contain energy and forces labels.
Linear Referencing and Normalization: The underlying DFT functional and DFT engine (e.g. VASP Orca) used differs between our pre-training datasets, resulting in variations in energy magnitudes across them. To address these differences and establish a consistent energy reference, we first compute an element specific energy reference for each of these datasets independently. This is achieved by calculating a linear reference across the training split of each dataset to find the per-element contribution (Musaelian et al.) (023). Subsequently, we normalize the energy and force labels dataset-wise. Specifically, we divide the energy labels by their respective standard deviations, ensuring a standardized scale for comparison. The force labels are normalized by dividing them by the poot-mea--see -quure of the force components in the corresponding training set.
The Open Catalyst 2020 (OC20) dataset (Chanussot et al., 2021 large and diverse catalyst dataset comprising training set of 130 million examples, including 55 elements, 82 adsorbates, and catalysts consisting of unaries, binaries, and ternaries. The OC20 dataset consists of DFT relaxation trajectories, where the atom positions are iteratively updated based on the forces to minimize the energy. There are a total of 640k relaxations with an average trajectory length of 200, which makes the total training data ~ 130M examples. In order to get better sampling across the potential energy surface, the dataset also contains ab initio molecular dynamics (AIMD) data and rattled data where atom positions amounts to 189 million entries. All the density functional theory (DFT) calculations conducted in this study utilized VASP with RPBE functional. Due to the considerable size of the pre-training dataset, we needed to establish practical training parameters. Consequently, we designated a training size of 100 million entries for our large training set, along with 2 million entries for our development set. These sizes were chosen to strike a balance between computational feasibility and the inclusion of substantial amount of data for training purposes.
Open Catalyst 2020:
randomly perturbed. The cumulative size of all these data
Open Catalyst 2022: The Open Catalyst 2022 (OC22) dataset, as presented in the work by Tran et al. (Tran al.) 2022), shares similarities with OC20 in terms of optimization trajectories involving adsorbates and catalyst surfaces. However, OC22 is more specialized and specifically focuses on oxide materials. While may not possess the same level of diversity as OC20, it still provides valuable information within this specific context. The dataset comprises a total of 8M training data, which we employ in its entirety for our large run. Furthermore, we utilize 200k data from this dataset for our development run, ensuring similar ratio of dataset sizes as our large run. It is important to note that for this study, a different DFT functional was employed compared to OC20. Specifically, the PBE + U functional account for the specific characteristics and properties oxide catalysts, providing a more accurate representation of their behavior.
utilized in this work. The decision to use this functional was made to
24

---

Published as a conference paper at ICLR 2024
The ANI-Ix dataset (Smith et al., 2020) is small molecule conformation dataset containing C. H. N, and O atoms, created using the Gaussian software with the wb97 / (d) electronic structure method. This is diverse organic dataset generated through active learning through pre-training MI. potential. This dataset has a total of 5M DFT calculations. We screened out organic molecules that have less than 4 atoms from the training data our backhone GemNet model can't calculate quadruplets for these. We split the entire data into train, val, and test such that val and test have molecules that are not present in the train. We subsample 80k split for our development set out of the 44M training data.
ANI-1x:
Transition-Ix: The Transition 1x dataset (Schreiner et al., 2022> contains close to 10M DFT calculations of forces and energies of molecular configurations at the w B97x/6 - 31G(d) level of theory (similar to ANI-Ix). The configurations in this dataset are on and around reaction pathways generated by running Nudged Elastic Band (NEB) on 10k organic reactions. Therefore, this dataset contains more dense sampling of PES for every system as compared to the ANI-Ix dataset. The train, val, and test splits for this dataset are pre-defined and we use the same splits. For our development set, we subsample 200k split out of the total 9M train data.
H.2 FINETUNING DATASETS
To demonstrate the ability of our pre-traines models to generalize over diverse set of fine-tuning tasks, we selected two datasets from three different atomic domains. We include QMOF and Matbench from the materials domain, MD17 and QM9 from the small molecules domain, and SPICE (dipeptides and solvated amino acids subsets) and MD22 from the large molecules domain.
The QMOF dataset (Rosen al. 2021) is database of approximately 15,000 experimentally synthesized metal organic frameworks (MOFs). We use a training dataset of 10,000 systems and split the remaining into validation and test sets. The band gap, which determines the electrical conductivity of material, is an important property for identifying materials for electrocatalysio and energy applications, so we use as the label for our models. It worth noting that the non-referenced energy predictions for the other datasets (OC20, OC22, and ANI-Ix) are extensive properties (i.c., the energy values depend on the size of the system), while the band gap an intensive property (it does not depend on the size of the system). Therefore, we take mean pooling of embeddings across nodes to calculate this scalar property.
QMOF:
MatBench: Mathench (Dunn et al. 2020) is benchmark for predicting the properties of inorganic bulk materials. There are a total of 13 tasks in this benchmark that have samples that range in size from 312-132k samples. Tasks include predicting optical, thermal, electronic, thermodynamic, tensile and clastic properties given a material's 's composition and/or crystal structure. For our work, since we give structure input our model and demonstrate finetuning on regression tasks, we restrict to tasks in Matbench. Each task has 5 folds and predefined test splits. We report an average across 5 folds for our JMP-L and JMP-S models but due to the compute cost, stick to only reporting on fold 0 for all other comparisons. For prediction of all material properties across these tasks, we pooling to work better for phonons as the
use mean pooling except for phonons. We observe vibrational frequencies aren't 't intensive or extensive properties.
MD17: The MD17 (Chmiela al. 2017) dataset, a collection of eight small organic molecules for which energies and forces are computed using ab-initio Molecular Dynamics (MD) simulations with Density Functional Theory (DFT). The revised MD17 data is a recomputed version of the original MD17 with improved numerical accuracy. For this dataset, we use 950 samples for train, 50 samples for validation, and the remainder of the data as test set. This the same number used by other models benchmarked on this dataset. All of our force predictions are modeled as the gradient of energies to achieve improved performances. Additionally, we also demonstrate results on a training split of to demonstrate the few-shot learning capabilities of our pre-trained model (shown in Appendix B))
The QM9 dataset (Ramakrishnan et al., 2014) is a widely used benchmark dataset in the field of quantum chemistry and machine learning. It comprises a collection of quantum mechanical calculations for organic molecules containing up to nine heavy atoms from the GDB-17 database (Ruddigkeit al., 2012b). The dataset provides essential molecular properties, including atomization energy, HOMO-LUMO gap, dipole moment, polarizability, and more. All molecules are modeled using the P/6 31 G1d p) DFT functional. We take an atomwise reference for all properties
QM9:
25

---

Published as a conference paper at ICLR 2024
and then normalize those standard Gaussian for predictions. We empirically find sum pooling to work better for all the property predictions.
SPICE: SPICE (Eastman et al. 2023) is large and diverse dataset with the goal of training potentials relevant to simulating drug-like small molecules with proteins. It contains over a million conformers. restricted our finetuning results to dipeptides and solvated amino acids subset. The dipeptides subset covers a full range of covalent interactions found in naturally occurring proteins and the solvated amino acids subset includes critical non-covalent interaction of protein-water and water-water.
We were interested in finetuning on the domain of larger molecules,
MDD2: MD22 (Chmiela al. 2023) is benchmark dataset for large molecules and includes molecules with sizes from 42 to 370 atoms. This dataset includes MD simulations of 8 molecules which include proteins, lipids, carbohy drates, nucleic acids, and supramolecules. All of these domains are not seen during pre-training. The trajectories for all the systems were sampled at temperatures between 400 and 500 K at a resolution of fs, with corresponding potential energy and forces calculated at PBE + M (B611, 62) level of theory.
H.3 DEVELOPMENT SPLITS FOR ABLATION STUDIES
Due to the high computational cost of training on our full pre-training and fine-tuning sets we made scaled-down development versions for our ablation studies. The pre-training subset includes 2M split of OC20 as Gasteriger correlate with the larger training split. Further, we chose splits of the other pre-training datasets that keep roughly the same ratio as present in the full pre-training set to enable the study of dataset imbalances. Additionally, for finetuning datasets, we pick single target molecule from each
(Gasteiger et al.) 2022) demonstrate that results on this split
dataset. The development datasets are summarized in Table I7
Datasets
Task
Dev. train split
Pre-training datasets
OC20
E, F
2M
OC22
E, F
200k
ANI-Ix
E, F
80k
Transition-lx 1x
E, F
200k
Fine-tuning datasets
Matbench
MP Form
10k
QMOF
Band gap
10k
MD17
Aspirin Forces
1k
QM9
Ac
1 10k
SPICE
Solvated Amino Acids
1
MD22
Stachyose
8k
Table 17: Scaled down development sets for ablation studies.
26

---

Published as a conference paper at ICLR 2024
1 DATASET OVERLAP
Target (Units)
JMP-L Overlapping
JMP-L Non-overlapping
JMP-L Combined
P (D) a (ag)
0.008
0.006
0,008
0.032
0.030
0.032
HOMO (meV)
9.1
7.7
8.8
ELUMO (mee)
8.5
8.9
8.6
A (meV)
19.4
18.1
19.1
Rd (o))
0.162
0.164
0.163
ZPVE (meV)
0.901
1.043
0.9
Uo (meV)
3.0
2.6
2.9
(meV)
2.9
2.5
2.8
H (meV)
2.9
2.5
2.8
G (meV)
4.4
4.2
4.3
C₂ (cal/mol K)
0.017
0.018
0.017
Table 18: QM9 results where the test set is partitioned into overlapping and non-overlapping fractions based on composition.
As mentioned the main text there some overlap between the pre-training and fine-tuning small molecule datasets. This will likely always be the case for small molecules because under certain number of heavy atoms nearly all possible molecules can be enumerated (Ruddigkeit et al. (2012a) and many datasets draw from this distribution. In particular, there is overlap between the molecules in ANI-Ix and QM9, While there are some of the same or very similar molecules present, the level of DFT used generate the data is different (wb97x vs B3LYP), the labels are not identical (although the thermodynamic properties are closely related to the total electronic energy), and the 3D structures may not be identical, given these differences it is unclear how the overlap will impact fine-tuning performance. To examine this further, we evaluated pre-trained model that has been fine-tuned on QM9 (JMP-L) with multiple versions of test set, one with only the overlapping systems, one with all overlapping systems removed, and finally the combined or full test set. We consider the strictest case, where overlap determined by composition i.e. the same atoms are present. The results are hard to differentiate across all test splits, shown in Table [8] indicating that model is learning not simply memorizing similar examples in pre- trainini.
J MODEL AND OPTIMIZATION HYPERPARAMETERS
Table 19 shows the model hyperparmeters for the small and large variants of the GemNet-OC backbone. Table 20 shows the training and optimization hyperparameters for the pre-training, fine- tuning, and scratch baseline training runs. Comma-separated hyperparameter values indicate that multiple values were evaluated in our experiments. In this case, an exhaustive grid search is conducted across all possible hyperparameters, and the hyperparameters that produce the best results (i.e., the lowest validation MAE scores) are selected. Notably, we use ReduceLROnPlatee 3 and factor of 0.8 for all Scratch GN-OC runs and all GN-OC Small fine-tuning runs. For the FT the learning rate scheduling strategy described in Appendix [.4 [.4] For rMD17 fine-tuning runs, we found that using longer cosine decay duration of 128 epochs, with lower cosine final LR factor of 1e-2, produces better results.
with a patience of
GN-OC Large runs,
27

---

Published as a conference paper at ICLR 2024
GemNet-OC Hyperparameters
Small
Large
No. spherical basis
7
7
No. radial basis
128
128
No. blocks
4
6
Atom embedding size
256
256
Edge embedding size
512
1024
Triplet edge embedding input size Triplet edge embedding oulput size Quadruplet edge embedding input size Quadruplet edge embedding output size Atom interaction embedding input size Atom interaction embedding output size Radial basis embedding size
64
64
64
128
32
64
32
32
64
64
64
64
16
32
Circular basis embedding size
16
16
Spherical basis embedding size
32
64
No. residual blocks before skip connection
2
No. residual blocks after skip connection
2
2
No. residual blocks after concatenation
1
4
No. residual blocks in atom embedding blocks
3
No. atom embedding output layers
3
3
Cutoff
12.0
12.0
Quadruplet cutoff
12.0
12.0
Atom edge interaction cutoff
12.0
12.0
Atom interaction cutoff
12.0
12.0
Max interaction neighbors
30
30
Max quadruplet interaction neighbors
8
Max atom edge interaction neighbors
20
20
Max atom interaction neighbors
1000
1000
Radial basis function Circular basis function Spherical basis function Quadruplet interaction Atom edge interaction
Gaussian
Gaussian
Spherical harmonics Spherical harmonics
Legendre Outer
Legendre Outer
True
True
True
True
Edge atom interaction Atom interaction
True
True
True
True
Table 19: Model hyperparameters for the small and large variants the GemNet-OC backbone model
Table 20: Optimization hyperparameters across different training runs.
Pre- Teeiriiing JMP Small/Large
Batch Size
1024/768
Optimizer
Adam
Weight Decay
EMA
0.99
Initial LR
2.00c-4
LR Scheduler
Warmup Cos
Warmup Duration
2000 steps
Warmup Starting Factor
0.2
Cos Duration
2 epochs
Con Final LR Factor
1.0e-1
Max Training Epochs
2
28

---

Published as a conference paper at ICLR 2024
Fine- JMP Small
Optimizer
A
Weight Decay
0.1
EMA
0.99
Initial LR
8.00e-5
LR Scheduler
Warmup + Cos + LLRD + RLP
Warmup Duration
5 epochs
Warmup Starting Factor
1.00e-1
Cos Duration
32 epochs
Cos Annealing
false
Con Final LR Factor
1.0e-1
LLRD Embedding Block Initial LR Factor
0.30
LLRD Block Initial LR Factor
0,35
LLRD Block 2 Initial LR Factor
0,40
LLRD Block 3 Initial LR Factor
0.55
LLRD Block 4 Initial LR Factor
0.625
RLP Patience
5
RLP Factor
0.1
Fine- ToninT JMP Large
Optimizer
AdamW
Weight Decay
0.1
EMA
0.99
Initial LR
8.00e-5
LR Scheduler
Warmup Cos LLRD RLP
Warmup Duration
5 epochs
Warmup Starting Factor
1.00e-1
Cos Duration
32 epochs
Cos Annealing
false
Con Final LR Factor
1.0e-1
LLRD Embedding Block Initial LR Factor
0.30
LLRD Block 1 Initial LR Factor
0.55
LLRD Block 2 Initial LR Factor
0.40
LLRD Block 3 Initial LR Factor
0.30
LLRD Block 4 Initial LR Factor
0.40
LLRD Block 5 Initial LR Factor
0.55
LLRD Block 6 Initial L.R Factor
0.625
RLP Patience
5
RLP Factor
0.1
Fine- Tinn-Tuning JMP Large MD17
Optimizer
AdamW 0.1
Weight Decay
EMA
0.99
Initial LR
8.00e-5
LR Scheduler
Warmup + Cos + LLRD RLP
Warmup Duration
5 epochs
Warmup Starting Factor
1.00e-1
Cos Duration
32 epochs
Cos Annealing
false
Con Final LR Factor
1.0e-2
LLRD Embedding Block Initial LR Factor
0.30
LLRD Block Initial LR Factor
0.55
LLRD Block 2 Initial LR Factor
0.40
LLRD Block Initial LR Factor
0,30
LLRD Block 4 Initial LR Factor
0,40
LLRD Block 5 Initial LR Factor
0.55
LLRD Block 6 Initial LR Factor
0.625
RLP Patience
5
29

---

Published as a conference paper at ICLR 2024
RLP Factor
0.1
Scratch GN-OC Small and Large
Optimizer
AdamW
Weight Decay
0.01
EMA
0.99
Initial LR
le-4, 2e-4, 5e-5 RLP
LR Scheduler
Patience
Factor
0.8
30


Reviews:
[{"review_id": "zjfh4fFCFL", "reviewer": "Reviewer_XFKt", "summary": "This paper explores the possibility of pre-training a foundation-style model over multiple chemical domains to generate transferable atomic representations for downstream fine-tuning tasks. The Joint Multi-domain Pre-training (JMP) strategy utilizes data from multiple chemical domains and achieves state-of-the-art results across many targets of various datasets. The paper establishes a comprehensive set of fine-tuning benchmarks across various chemical domains and tasks.", "questions": "- Can you provide more details about the scaling strategy of the model architecture? The message-passing paradigm suffers from over-smoothing a lot and it is notorious of hard to make the network deep. I would like to understand more about this and how this method overcome the issues.\n\n- Since for larger foundation models, we would like to include far more parameter. But the molecular graphs often just include a very small number of vocabularies. A better strategy could be adding full attention mechanism to make the parameter space larger. And the model would be a transformer-based architecture or so-called \"graph transformer\". Did the authors try this in their experiments? How did the two frameworks perform?\n\n- This paper claims that they can deal with OOD challenge and cross-domain adaption better and they can benefit drug discovery, etc. So I would like to see more results on how this method perform on other therapeutics data [4], e.g., the molecule property predictions (toxicity, solubility, lipophilicity, etc.) on MoleculeNet [5].\n\n[4] Huang, K., Fu, T., Gao, W. et al. Artificial intelligence foundation for therapeutic science. Nat Chem Biol 18, 1033–1036 (2022). https://doi.org/10.1038/s41589-022-01131-2\n\n[5] Wu, Zhenqin, et al. \"MoleculeNet: a benchmark for molecular machine learning.\" Chemical science 9.2 (2018): 513-530.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Comprehensive set of fine-tuning benchmarks: The paper establishes a comprehensive set of fine-tuning benchmarks across various chemical domains and tasks, which enables researchers to evaluate the performance of their models against a standardized set of benchmarks.\n\n- State-of-the-art results: The Joint Multi-domain Pre-training (JMP) strategy achieves state-of-the-art results across many targets of various datasets, which demonstrates the effectiveness of the proposed approach.\n\n- Large and diverse molecular datasets: The paper highlights the importance of large and diverse molecular datasets in enabling the development of accurate and efficient models for atomic property prediction.\n\n- Different model sizes: The paper provides multiple model sizes with pretrained checkpoints that can benefit real-world deployment at different resource levels and potentially accelerate research progress in related fields.", "weaknesses": "- The paper claims that the proposed pre-training method is model-agnostic. However, the only evaluated architecture backbone is GemNet-OC. It would be better to have a variant pre-trained using other types of model backbones to conduct further comparison and analysis.\n\n- The novelty is a bit limited. I admit that this paper contributes on providing the empirical evidence that pre-training cross-domain molecule data can benefit multiple downstream tasks. However, the techniques used in this paper are either introduced by other literature or very simple and straightforward. No novel methods/theories are introduced. I would recommend this paper submit to more domain-related or comprehensive journals. \n\n- This paper claims that many previous efforts on pretraining molecules focus on a specific domain which ignores the information provided by other domains and generalizability. It would be better to provide more empirical evidence that compares the proposed model with other pre-training methods [1].\n\n- More supervised SOTA should be compared. E.g., for the materials domain, there are two more recent papers [2, 3] that can be reported against the proposed method.\n\n- I understand the concerns of releasing the code and models before acceptance. However, in terms of reproducibility, it would be better to provide an anonymized repo including some demo models for testing purpose.\n\n[1] Xia, Jun, et al. \"A Systematic Survey of Chemical Pre-trained Models.\" IJCAI, 2023.\n\n[2] Yan, Keqiang, et al. \"Periodic graph transformers for crystal material property prediction.\" Advances in Neural Information Processing Systems 35 (2022): 15066-15080.\n\n[3] Lin, Yuchao, et al. \"Efficient Approximations of Complete Interatomic Potentials for Crystal Property Prediction.\" arXiv preprint arXiv:2306.10045 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "87ZnUEQ6eL", "reviewer": "Reviewer_KadD", "summary": "This paper pretrains a large generalizable model for atomic property prediction, which outperforms SOTA methods on many downstream tasks.", "questions": "1. Does the use of JMP contribute to the prediction of molecular properties? If so, what specific types of molecular properties show improvements?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The model achieves exceptional performance.\n2. The experiments conducted in this study are comprehensive and thorough, covering various aspects such as hyper-parameter settings, ablation studies, and downstream tasks.\n3. The authors carefully study the balance between different datasets and pre-training tasks, ensuring a comprehensive analysis of their impact.", "weaknesses": "1. The disscussion of the correlation between pre-training tasks and downstream tasks is missing. I am wondering what kinds of downstream tasks can be promoting by pre-training? \n2. The main contribution of this work is implemental and not suprising.\n3. This paper confuses the prediction of atomic properties with the prediction of molecular properties.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cRynoQ0OEe", "reviewer": "Reviewer_R7em", "summary": "Authors explore the application of machine learning in predicting atomic properties across a wide array of applications, from healthcare to climate change. The authors introduce Joint Multi-domain Pre-training (JMP), a supervised pre-training strategy that leverages a vast dataset comprising approximately 120 million examples from multiple chemical domains. The primary goal of JMP is to generate transferable atomic representations that can be fine-tuned for diverse downstream tasks, addressing the challenge of generalizing across the extensive and complex space of molecular interactions.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Innovative Approach: The paper introduces Joint Multi-domain Pre-training (JMP), a novel supervised pre-training strategy that leverages a massive dataset from multiple chemical domains. This approach is innovative in its attempt to generate transferable atomic representations for a wide array of downstream tasks, addressing the challenge of generalizing across diverse molecular interactions.\n\nCreative Combination of Ideas: The authors draw inspiration from successful practices in Natural Language Processing (NLP) and Computer Vision (CV), creatively applying the concept of large-scale pre-training to the domain of atomic property prediction. This cross-disciplinary innovation enhances the originality of the work.\n\nBroad Applicability: The paper’s contributions have broad applicability across various domains, ranging from drug discovery to material science. The ability of JMP to generalize across diverse chemical domains signifies its potential to drive advancements in multiple fields.\nAddressing a Critical Challenge: The paper tackles the critical challenge of generating transferable atomic representations in the vast and complex space of molecular interactions. By addressing this challenge, the paper makes a significant contribution to the field of machine learning for atomic modeling.\n\nComputational Efficiency: The computational efficiency achieved through JMP, with over 12x faster fine-tuning compared to training from scratch, is a notable strength. This efficiency is crucial for practical applications, making the paper’s contributions highly significant.", "weaknesses": "Need for Broader Ablation Studies:\nIssue: While the paper includes ablation studies to analyze the impact of different JMP components, these studies could be broadened to provide a more comprehensive understanding of the model’s behavior and the contributions of individual components.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IQbAyXa8J1", "reviewer": "Reviewer_avot", "summary": "The paper introduces a supervised pre-training strategy named JMP, which pretrains GemNet-OC on various small molecule data from multiple chemical domains. The pre-trained network acts as a foundational model, further finetuned for downstream atomic property prediction tasks. The authors’ major contribution is demonstrating that by pretraining exclusively on small molecules, the network can be finetuned on large molecule datasets, achieving state-of-the-art results. The primary technical advancements concentrate on refining each component of the standard deep learning pipeline, including data preprocessing and hyperparameter tuning. Experiments across 40 fine-tuning benchmarks are conducted to showcase the effectiveness of the method.", "questions": "1. In Sect. 4.1, concerning Data Normalization, the authors have chosen to normalize the property values per dataset. A lingering question is how the output of the NN is transformed. Is the transformation still dependent on each specific dataset? If that is the case, it seems impractical for real world applications where a novel molecule is given, and it would be indeterminable as to which \"dataset\" it inherently belongs to and how to transform its output.\n\n2. Regarding Dataset Size Imbalance, I was wondering if the authors considered utilizing loss reweighting as opposed to data reweighting. By loss reweighting, I am referring to the approach of uniformly sampling the data but adjusting the coefficients of each sample to p_d (ensuring normalization across each batch).\n\n3. I devoted a significant amount of time attempting to digest whether each term in Eq.1 is a novel contribution or a previously introduced one. It would be beneficial if the authors could provide clearer definitions of each symbol used, elaborate more distinctly on the novel improvements introduced in this paper, and add a period to the end of the equation.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The paper is articulate and easily comprehensible. The authors provide an essential level of detail in describing their pipeline, facilitating a clear understanding of the processes. Although source code is not provided, the comprehensive details included in the paper should enable straightforward reimplementation.\n\n+ The experiments presented in the paper are exhaustive and meticulous. A diverse array of molecular property datasets, encompassing both large and small molecules, has been utilized for the experiments. The authors have conducted extensive ablation studies, offering valuable insights into the influence of various hyperparameters used in the pipeline.", "weaknesses": "- My primary concern lies in the paper’s technical contribution. The concept of building a foundational model for molecular property prediction tasks isn’t novel. A significant challenge is bridging the substantial domain gaps across various chemistry domains. The authors seem to emphasize that the proposed pipeline can effectively bridge the molecule size gap, allowing it to work efficiently on larger molecules even when only pretrained on smaller ones. However, molecule size is just one of several apparent factors—and likely among the simpler ones—causing the domain gap. More complex factors, such as intrinsic differences in the distribution of graph structures and issues related to data availability, are not addressed in the paper. Consequently, it is challenging to be convinced that the proposed method significantly contributes to foundational models or represents \"an important step for universal ML potential,\" as claimed in the introduction.\n\n- The improvement brought by the proposed method appears to be mainly attributed to hyperparameter tuning. The network architecture itself isn’t novel, and the loss function closely resembles commonly used ones, albeit with slight modifications to some coefficients. My overarching impression is that the authors engage extensively in manual hyperparameter tuning, which doesn’t offer substantial insights to propel further research advancements. While I acknowledge the empirical enhancements demonstrated through comprehensive experiments, it is still challenging to bestow a favorable overall evaluation on the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)