On the Importance of Pretraining Data Alignment for Atomic Property Prediction
Authors: Yasir M. Ghunaim
Hasan Abed Al Kader Hammoud
Bernard Ghanem
Venue: Accepted by TMLR ()

Abstract: This paper challenges the recent paradigm in atomic property prediction that links progress to growing dataset sizes and computational resources. We show that pretraining on a carefully selected task-aligned dataset can match or even surpass large-scale joint pretraining while using only 1/24th of the pretraining budget. We introduce the Chemical Similarity Index (CSI), a simple metric for molecular graphs inspired by the Fréchet Inception Distance in computer vision, which quantifies the alignment between upstream pretraining datasets and downstream tasks. By selecting the most aligned dataset with minimal CSI distance, we show that models pretrained on a smaller, focused dataset consistently achieve better performance on downstream tasks than those pretrained on massive, mixed datasets such as JMP. This holds even when the mixed dataset includes the upstream dataset most aligned with the downstream task. Counterintuitively, we also find that indiscriminately adding more data can degrade model performance when the additional data is poorly aligned with the target task. Our findings highlight that quality often outperforms quantity in pretraining for atomic property prediction.

Published in Transactions on Machine Learning Research (01/2026)
On the Importance of Pretraining Data Alignment for Atomic Property Prediction
Yasir Ghunaim King Abdullah University of Science and Technology (KAUST)
yosir-ghunncinkkuutteduu
Hasan Abed Al Kader Hammoud
King Abdullah University of Science and Technology (KAUST)
Bernard Ghanem
bermani.hhnnem@kanutt.c
King Abdullah University of Science and Technology (KAUST)
Reviewed on OpenReview: https: // opwnreveew.ew.... net/ formm? id= JfD9BsrDTb
Abstract
This paper challenges the recent paradigm in atomic property prediction that links progress to growing dataset sizes and computational resources. We show that pretraining on a care- fully selected task-aligned dataset can match or even surpass large-scale joint pretraining while using only 1/24th of the pretraining budget. We introduce the Chemical Similar- ity Index (CSI), a simple metric for molecular graphs inspired by the Fréchet Inception Distance in computer vision, which quantifies the alignment between upstream pretrain- ing datasets and downstream tasks. By selecting the most aligned dataset with minimal CSI distance, we show that models pretrained on a smaller, focused dataset consistently achieve better performance on rownstreamm tasks than those pretrained on massive, mixed datasets such as JMP. This holds even when the mixed dataset includes the upstream dataset most aligned with the downstream task. Counterintuitively, we also find that indiscrimi- nately adding more data can degrade model performance when the additional data poorly aligned with the target. task. Our findings highlight that quality often outperforms quantity in pretraining for atomic property prediction. The code is publicly available at: github.com com Slhmaai//fficientaattee /fffiiient-atom
1 Introduction
Machine learning is transforming molecular modeling, driving advancements in accurate predictions and simulations of molecular behavior (Chanussot et al. 2021 Tran et al.) 2023 Liao al. 2023). These breakthroughs directly impact the acceleration of progress in crucial fields such as drug discovery (Huang
et al. 2021) and global climate change mitigation (Sriram et al. 2024). The improvements in this field have
been primarily attributed to innovations in model architectures (Liao et al. 2023 Gasteiger et al. 2021
Passaro & Zitnick 2023) and the growing availability large-scale molecular datasets. In recent years, the
sizes of molecular datasets have increased dramatically, from tens of thousands of examples (Christensen
Von Lilienfeld 2020: Chhmiela et al., 2023 Wu al.) 2018)) hundreds of millions (Chamussot al. 2021
Tran et al.] 2023). This rapid growth in scale has also caused a surge in the computational resources required
for pretraining, increasing from few days
single GPU to over thousand GPU-days (Shoghi et al.
2023
Liao et al. 2023). This trend raises the question:
scaling data and resources the only path forward for pretraining in atomic property prediction, or can intelligent data selection achieve similar performance more efficiently?
While data selection strategies for pretraining have been explored in fields like natural language process-
ing (Penedo al. 2024) and computer vision (Hammoud et. al. 2024 Li et al. 2023), this area remains

---

Published in Transactions on Machine Learning Research (01/2026)
rMD17 (Aspirin)
MD22 (Ac- Alc-Ala-NNHR
SPICE (Solvated Amiuo Acids)
QM9 (Un)
-10%
+5%
+6%
22%
MAE. is
6.7
2.6
2.9
5.7
5.4
3.3
3.1
5.2
SoTA (JMP-S)
Ours: 24x less Resources
% indicates relative MAE improvement SoTA
Figure 1: Pretraining on a High-Quality, Task- Aligned Dataset. Pretraining on a carefully selected high-quality dataset achicves comparable or superior mean absolute crror (MAE) across tasks while reducing pretraining budget by a factor of 24 compared JMP-S, which pretrained on all upstream datasets. Lower MAE indicates better performance.
largely underexplored in atomic property prediction, where unique challenges arisc. This gap is especially critical for 3D molecular and material pretraining, which requires far more computation (Liao et al.) 2023 than methods based on 2D representations (Rong et al. 2020). In this paper, we challenge the prevailing assumption that 'bigger better' by exploring whether a smaller, strategically selected dataset can lead comparable or even superior performance while substantially reducing computational demands. We in- troduce pretraining paradigm for 3D molecular and material systems that shifts the focus from data and compute scaling to selecting the most relevant upstream dataset for improved downstream performance.
Our experiments reveal two key insights: (1) Competitive Performance Can Be Achieved with 24x Fewer Resources: Selecting upstrean ream datasets based on downst ream aligmment. matches or exceeds large- scale pretrained models such as JMP (Shoghi et al. 2023), while using only 10M training instances instead of JMP's 240M-instance pretraining budget (a 24x reduction under our budget definition), as shown Figure [ (2) Quality Outperforms Quantity: Expanding the pretraining dataset by incorporating additional data from less aligned sources harm downstream performance rather than improve it.
To explore the potential of dataset selection for pretraining in atomic property prediction, introduce the Chemical Similarity Index (CSI), simple metric inspired by the Fréchet Inception Distance (FID) from computer vision. CST measures the alignment between an upstream dataset and a downstream task, en- abling the selection of chemically relevant pretraining data. By focusing on these highly relevant datasets, we significantly reduce computational costs while maintaining competitive performance and, in many cases, achieving improvements. While large-scale datasets like OC20 (Chanussot et al.] 2021 Tran al.) 2023) and mixed datasets like JMP (Shoghi et al. 2023 popular choices for pretraining in molecular domains (Kol- luru et al. 2022 Shoghi et al. 2023). our findings challenge their universal utility. Surprisingly, pretraining single, carefully selected dataset guided by CSI often outperforms models trained on mixtures, even when those include the most relevant dataset.
The contributions of this paper are threefold: (1) We introduce a novel framework for computationally efficient pretraining of molecular machine learning models, demonstrating that strategic data selection can match or outperform models trained on much larger datasets. (2) We propose the Chemical Similarity Index (CSI), metric for assessing the alignment between upstream and downstream molecular datasets, enabling effective dataset selection. (3) We provide extensive empirical evaluation demonstrating the effectiveness of our approach, offering a practical and efficient alternative to the current trend of ever-incrcasing data and computational costs in molecular machine learning.
2 Related Work
Pretraining for Atomic Property Prediction. Inspired by the success of pretraining in computer vision and natural language processing, pretraining for atomic property prediction has gained significant
attention in recent years. Most approaches in molecular machine learning have focused on self-supervised
pretraining [Rong et. al. 2020
Liu et. al. 2021 Jiao et al. 2023 Chen et al. 2021 Zhou al. 2022a
2

---

Published in Transactions on Machine Learning Research (01/2026)
Ji ct al. 2024 while fewer studies have explored the effectiveness of supervised pretraining Smith et al,
2019 2018 Kolluru et al. 2022). Early self-supervised methods such as GROVER (Rong et al. 2020 target 2D molecular graphs. More recent approaches, including GraphMVP (Liu et al, 2021), Uni-Mol (Zhou ct al.) 2022a) and Uni-Mol2 (Ji et al. 2024), extend to 3D structures and often use denoising objectives that are best suited for equilibrium geometries (Liao et al.) 2024). This limits their applicability to large- scale non-equilibrum datasets common in molccular and materials modeling. In contrast, recent supervised pretraining frameworks such as Joint Multi-domain Pre-training (JMP) (Shoghi al.) (023) jointly pretrain diverse large-scale labeled datasets and operate effectively non-equilibrium data. While powerful, JMP requires substantial compute due to the massive data volume and does not reveal how each pretraining dataset influences downstream performance, Our work addresses this gap by systematically analyzing the link ween pretraining datasel choice and downstream performance. Based on our analysis, we show that selecting a single task-aligned upstream dataset can match or exceed joint pretraining, enabling efficient model dovelopment under constrainod computational budgets.
Computational Budgeting. Recent research highlights the importance of studying model performance under computationally budgeted setups. In continual learning (CL), works by Prabhu et al. (2023) and Chunaim et al. (2023) show that simple basclines often outperform statc-of-the-art methods in compute- constrained settings. TiC-CLIP Garrg et al. 2024) further demonstrates efficient rehearsal- based training for time-contimuous data. For Vision Transformers, Pan et al. (2022) propose a framework to dynamically control model complexity during training, achieving competitive performance under varying budgets. Li et al. (2019) formalize budgeted training, showing that budget-aware learning rate schedules, such as linear decay, are critical for robust performance across tasks like image classification and object detection. Tn multi-domain learning, Berriel et al.] (2019) introduce Budget-Aware Adapters, which reduce computational complexity while maintaining accuracy by selecting relevant feature channels. These findings across domains emphasize the critical need for more efficient approaches that can achieve competitive performance while minimizing computational costs.
Distribution Similarity. A variety of methods exist for quantifying the distance between probability dis- tributions, including Kullback Leibler (KL) divergence (Kullback & Leibler 1951), Jensen-Shannon Shannon (JS) divergence Lin 2002 Maximum Mcan Discrepancy (MMD) (Gretton et al. 2012), , and the Fréchet In- ception Distance (FID) (Heusel et al. 2017). FID has become standard metric in computer vision for comparing real and generated data distributions by evaluating differences in the means and covariances of their feature representations. In the chemical domain, the Fréchet Ceem Distance (FCD) (Preuer et al. 2018) adapts FID for ovaluating generative models that use SaILES-based molecular representations and has been integrated into benchmarking frameworks such as GuacaMol (Brown et al. 2019). While FCD demonstrates the applicability of FID-based metrics in chemistry, it primarily applied to SMILES-basod molecular data within a single chemical domain. Our proposed Chemical Similarity Index (CSI) is also de- rived from FID, but is designed for 3D graph-structured systems and focuses measuring alignment across multiple domains such as molecules and materials. To our knowledge, we are the first to systematically study pretraining datasct alignment in 3D atomic property prediction.
Data Selection. Efficient training through data selection has been explored via two primary approaches: subset selection and dataset distillation. Subset selection aims to identify a representative subset of the training data that matches or even outporforms training on the full datasot. Several methods have been proposed for vision and NLP lasks (Attendu & Corbeil 2023 Killamsetty et al. 2021a b Kaushal et al., 2019 Bairi et al. 2015 Lapedriza et al. 2013). Dataset distillation, introduced by Wang et al (2018). focuses on generating a smaller, synthetic subset of the datasct that preserves performance while reducing training time and storage requirements. Subsequent work has explored techniques such as meta-learning (Zhou et al., 2022b; Nguyen et al. 2021a b), gradient matching (Zhao et al. 2021). 3, and distribution matching (Zhao & Bilen 2023). While most research in distillation has focused on vision tasks, a few studies have extended it to graph data (Jin et al.) 2022b Liu et al.] 2022) Jin et al. 2022a), though primarily targeting knowledge and social graphs rather than molecular graphs.
Two recent vision studies are particularly relevant to our work. First, Hammoud et al. 2024 shows that increasing pretraining data diversity enhances performance only when distribution shifts between upstream and downstream tasks are minimized, Second, Li et al. 2023) introduces a method to dynamically leverage
3

---

Published in Transactions on Machine Learning Research (01/2026)
Dataset Selection for Pretraining
Pretraining under a Limited Budget
Finetuning
(
D i
((9)
Train on V random samples D for & eeochs
u
6(2)
((6)
Dd
u (0))
8(K)
6(1
D(K)
D(1)
Backbone
0
Backbone
02
((K)
argmin (D Da) D ISiSK
8(1)
Dd
Energy
Forces
Task Head
= N
Figure 2: Pipeline Overview. Our paradigm for pretraining and finetuning consists of two new components: (1) Dataset Selection Stage, where a distance metric is employed to identify the dataset that is most similar to downstream lask datasel Da, in this case D. This selected dataset. then used for pretraining the model. (2) Limited Budget Pretraining, where impose training budget by subsampling /V random samples from D and training the model for & epochs. This results in a computational budget of C = & N. The pretrained backbone of is subsequently finetuned on the downstream task dataset DD to obtain the final model parameters 0z
the open web, reducing the distribution gap between upstream and downstream tasks through targeted representation learning. Findings from other domains suggest that aligning upstream datasets may be crucial for effective pretraining.
Comparison to Our Work. To the best of our knowledge, no prior work has specifically explored restream dataset selection for molecular graphs, which present unique challenges due to their structural and chemical complexity. In this work, we take the first step in addressing this gap by focusing on aligning upstream and downstreamn distributions at the dataset level rather than subselecting at sample-wise level or creating a synthetic distilled version of the dataset.
3 Formulation and Setup
In this section, we present our problem setup, notion of computational budget, and the formulation of dataset similarity. We then detail how we adapt the Fréchet Inception Distance (FID) to the molecular domain, yielding the Chemical Similarity Inder (CST). Our setup is illustrated in Figure 2 Throughout this work, use the term 'molecular' broadly to encompass both molecular and materials domains, as well their respective datascts.
3.1 Formal Setting
Upstream and Downstream Datasets. Let. (D.)), - ...., ......(k)) denote a collection of K upsiream (pretraining) datasets, each containing molecular structures paired with relevant atomic propertics (e.g.. energies and forces). In the typical paradigm, upstream datasets are typically aggregated into single pretraining set:
K
D, = D(
(1)
1=1
We further define Da as the dounstream dataset, which focuses on a specific prediction task (e.g., predicting an atomic property).
Multi-task Pretraining. We consider a neural network (::;), where encompasses the shared backbone parameters O, and task-specific head parameters 0. (for energy prediction) and O, (for force prediction). During pretraining, the network is trained simultancously predict energies and forces. Formally, the multi-task pretraining objective over an upstream datasct (( is given by:
g arg min Spretrain (0: ((:))
(2)
4

---

Published in Transactions on Machine Learning Research (01/2026)
where 8 = {8b. 0e, } and
pretrain (0: (.) a Lenergy (0,, 0e; ( + Looree (0p. 0y:
We compute Lenergy using the Mean Absolute Error (MAE) and Ltorces using the mean per-atom Euclidean (L2) distance. Coefficients a and B weight the importance of energy and force tasks, respectively. Following the JMP paper (Shoghi et al. 2023
set B to prioritize accurate force predictions in atomistic modeling. Pretraining can be performed on either the joint upstream dataset Du, similar to JMP (Shoghi
et al. 2023 or on an individual upstream dataset D as in our selective setting.
Fine-Tuning. After the multi-task pretraining phase, the task-specific heads 0. and Of are discarded, and new specifc head 0% is attached the pretrained backbone 0b. The downstream objective then becomes:
0
arg min Liinetune Oni 000 Daa), 0r,0h
where 0f*)* denotes the pretrained backbone parameters from Eq- 2 Intuitively, the downstream training refines the shared backbone parameters 0, and learns the task-specific head 0, capture the target property in Da-
In this paper, we introduce two additional definitions for this selling: (1) compudational budget and (2)
dataset similarity.
Computational Budget. Following Hammoud et al. (2024). we define the compulational budgel C to be the product of the number of epochs and the number of unique samples N in the pretraining dataset:
C
(5)
Hence, the computational budget represents the total number of samples processed over training. naturally splits into two factors: the dataset size (N) and the nunber passes through it (E). The choice of depends on the available computing resources. In our main experiments, we fix C, N, and & ensure fair comparison across different upstream datasets. We also include experiments in which M (and thus C) varies, in order to analyze the impact of dataset size and total compute downstream performance.
Dataset Similarity. A key objective of this work is to estimate how well an upstrcam datasct D, aligns with a downstream dataset Da. We therefore seek distance metric
8(Du, Da)
that quantifies their alignment or "similarity." In principle, a lower value of 8(D.. Da) reflects a higher degree of alignment between the upstream and downstream distributions. Thus, among multiple candidate upstream datasets (D!!" D....(()) (K)) the one that minimizes
argmin 8 (D( Da)
ISiS
should provide the most effective pretraining for Da. In this paper, we empirically test this assumption, examining whether lower 8-values indeed correlate with improved downstream performance. Motivated by this, we use d as a principled metric to guide dataset selection for Eq. ( Instead of aggregating all upstream datasets, we modify the pretraining setup to use only the single dataset D( that best aligns with the downstream task under a fixed computational budget.
3.2 The Chemical Similarity Index (CSI)
Recap of FID. Our proposed Chemical Similarity Inder (CST) draws inspiration from the well-known Fréchet Inception Distance (FID) (Heusel et al. (0017) FID is commonly used in computer vision to compare two sets of images via their feature distributions. Specifically, if one extracts features (e.g., from Inception
5

---

extract features Published in Transactions on Machine Learning Research (01/2026)
ANI-Ix
Transition-Ix
OC20
OC22
300
200
CSI
100
rMD17
MD22
SPICE
QM9
OMat24
Downstream Dataset
Figure 3: Alignment Between Upstream and Downstream Using CSL. We assess how well the extracted representations from each upstream dataset align with downstream tasks using our CST metric, where lower values indicate stronger alignment. ANI-Ix demonstrates the closest feature alignment with most downstream tasks. However, for OMat24, the catalysis datasets OC20 and OC22 exhibit the strongest alignment.
network) for datasets and and denotes their empirical means and covariances by μxx Ex and jy, Σy, then
FID(X,Y) = llxx - -y + Tr + -
(6)
The central idea is to represent each sample in feature space where distances encode semantic similarity and then compare the distributions of these representations for the two datasets.
To adapt FID for graph-structured molccular data, we compute the CSI metric using node emboddings features and apply class-balanced sampling to ensure representative coverage of molecular types in each upstream dataset. For computational feasibility, subsample 10k instances from both the upstream and downstream datasets. To keep the metric fixed and comparable across all evaluations,
using EquiformerV2 (Liao et al. 2023) pretrained on OC20 (Chanussot et al. 2021). We provide results with additional pretrained checkpoints in Appendix feature representation choices in Appendix and blaan-ia procedure in Appendix
CSI Between Upstream and Downstream Results. In Figure 8 we present the CSI values for pairs of upstream and downstream tasks related to energy and force prodictions, with additional details about the datasets and targets provided in Section 4 For the first four downstream datasets in Figure ANI-1x (Smith design goal of maximizing chemical diversity. Transition-Ix (Schreiner al.) 2022), which focuses on transition states, ranks second suggesting that its emphasis high-energy transition states leads to partial overlap with downstream distributions, OMat24 [Barroso-Luque et al. (024)) which contains inorganic materials, is more closely aligned with the catalysis datasots OC20 (Chanussot et al.) 2021 and OC22 (Tran et al. 2023) than with other upstream datasets. While OC20 and OC22 are often lavored for pretraining Shooghi et al. 2023: Kolluru et al. 2022) due to their scale and chemical diversity, our metric suggests they may not align well with most of the considered downstrcam tasks. Next, we examine whether these alignment values correlate with downstream performance.
et al. 2020) consistently achieves the closest alignment, reflecting
4 Experiments
We evaluate the impact of pretraining on different upstream datasets for downstream performance and investigate how well the CSI values in Figure 3 reflect the relevance of these datasets. We begin by defining the datasets, baselines, and evaluation setup.
Upstream Datascts: Following JMP (Shoghi et al. 2023). we perform protraining on upstream datascts of
small molecules, including ANI-Ix (Smith et al.] 2020 and Transition-lx (Schreiner al. 2022). as well
large-scale catalysis datasets, OC20 [Chanussot. al. 2021 and OC22 (Tran et al. 2023 These datasets

---

Published in Transactions on Machine Learning Research (01/2026)
vary in domain focus and graph size, enabling us to examine how these factors impact the generalization of pretraining across downstream tasks. The ground-truth labels, energy and forces, are computed using Density Functional Theory (DFT).
For downstream evaluation, we focus on in-distribution (ID) tasks involving energy or force prediction, following the definition JMP (Shoghi et al. 2023). We discuss out-of-distribution (OOD) tasks in Section 5 To keep the experiments tractable, we evaluate on a single target per downstream dataset, since we run six baselines for each. These targets and their corresponding datasets are: Aspirin
Downstream Datasets:
in rMD17 (Christensen & Von Lilienfeld 2020), Ac-Ala3- NeMe in MD22 (Chmiela et al., 2023 solvated
amino acids in SPICE (Eastman et al. 2023), Uo in QM9 (Wu et al. 2018), and force prediction on the aatuled-330--uubsampledd split of OMat24 (Barros-LLqque et al. 2024).
Baselines: We report the original performance of JMP, where 'JMP-S" and 'JMP-L' correspond to the small and large backbones, respectively. Additionally, we present our reproduced fino-tuning results using the official JMP checkpoints, denoted as 'JMP-S"" and "JMP-L"'.
For our budgeted evaluation, present results in two categories: pretraining on single upstream dataset and pretraining on a joint combination of all upstream datascts. For single-dataset experiments, we randomly sample .N' instances from the original upstream data. For joint pretraining, we construct the training set using two different strategies. (1) Balanced Sampling, where an equal number of samples is drawn from each of the four upstream datasets, totaling N samples; and (2) Temperature- Baasee Sampling, which preserves the dataset proportions used in the full 120M sample set of JMP (Shoghi et al.]
2023).
Evaluation Setup: We pretrain the GemNet-OC-S model (Gasteiger et al. 2022 on each individual
upstream dataset, as well as on joint configurations that combine all upstream datasets, following the baseline setups. For our main experiments, we set fixed computational budget of C = 10M, achieved by training on N = 2M samples for & = 5 epochs. This represents a 24x reduction in pretraining budget compared to that used in JMP (Shoghi et al., 2023). Our budget ensures accessibility and reproducibility, with cach pretraining run completing within 1 to days on an A100 GPU. Additional budget coufigurations are explored in later sections and the appendix. Each pretrained model is then fine-tuned separately on each downstream task.
4.1 Does CSI Correlate with Better Performance?
In Figure presented CSI values quantifying the alignment between cach upstream and downstream dataset. These results raise a critical question:
Can CSI reliably guide the selection of pretraining datasets to achieve optimal performance on specific dounstream tasks?
Table [ summarizes the downstream performance of models pretrained on different datasets in the in- distribution setting. Consistent with the CSI values, the upstream dataset with the lowest CSI for each downstream task also achieves the best fine-tuning performance. Specifically, pretraining on ANI-1x, despite being the smallest dataset, yields the best fine-tuning results on rMD17, MD22, SPICE and QM9, outper- forming all other individual pretraining datasets as well as the joint variants. For instance, on the rMD17, SPICE and QM9 datasets, the model pretrained on ANI-Ix achieves MAEs of 5.2, 5.39 and 3.1, compared to 6.7, 5.71 and 3.3 for JMP-S. This strong performance is achieved with less than 5% of the pretraining budget used by JMP-S. For OMat24, the model pretrained on OC20 achieves the best performance among the limited-budget basclines with an MAE of 87.4, comparable to JMP-S's 84.4 despite its much larger 240M pretraining budget. This outcome is also predicted by CSI, which ranks OC20 as the most similar upstream dataset for this task.
Furthermore, joint pretraining, whether balanced or temperature-based, generally performs worse than in- dividual pretraining on the most CSI-aligned dataset for each lask when both are trained under the same limited- limited-budget setting. Following the JMP formulation, temperature- sampling allocates more samples to OC20 and OC22 than balanced sampling. On OMat24, where OC20 and OC22 are most aligned according to CSI values, temperature-bart sampling outperforms balanced sampling. However, temperature based sampling degrades performance on other tasks where OC20 and OC22 are less aligned with the downstrean
7

---

Published in Transactions on Machine Learning Research (01/2026)
Table 1: In-Distribution Evaluation for Energy and Force Targets. We report test MAE when fine-
tuning downstream targets, detailed in Downstream Datascts (Section D The top section represents models pretrained with the large-scale JMP budget, while the lower two sections show results under a limited budget. JMP-S" denotes reproduced results. We report the average and standard deviation over three random seeds.
Upstream Data
Backbone
rMD17 (meV/A)
MD22 (meV/A)
SPICE (m-V/A)
QM9 (mee)
OMat24 (meV/A)
JMP-L (GemNNet-OC-L)
5.1
1.92
4.75
2.9
240M
Joint (Temperature)
JMP-S (GemNettOOC-))
6.7
2.64
5.71
3.3
JMP-8 S" (GemNet-OC-S)
6.8
3.21
5.60
3.4
84.4
ANI-Ix
5.2 0.2
2.92
0.03
5.39
0.26
3.1
98.1 0.5
Transition-Ix
9.8 0.5
3.70
7.70 0.31
3.6 0.3
100.8
10M
GemNet-OC-S
OC20
13.8 +
5.10
9.67 0.73
5.3 0.4
87.4
OC22
15.3 +
5.34
10.62 0.42
5.7
0.1
Joint (Balanced)
7..8
2.73
6.39
3.5 0.2
90.3 0.11"
10M
Joint (Tempeertture)
GemNet-OC-S
11.0 +: 1.5
4.41 4: 0.59
8.44 + 0.63
4.0 0.4
88.4 + 0.1
dataset. These results suggest that, under a limited budget, mixing upstream datasets with varying CSI values is suboptimal and requires significantly larger pretraining budgets achieve competitive performance.
Takcaway. Our experiments reveal three key insights for in-distribution downstream tasks: (1) Task- aligned upstream datasets outperform larger joint datasets. (2) Joint pretraining can match the benefits of
pretraining on a highly task-aligned dataset, but doing so requires substantially larger pretraining budgets. (3) CSI effectively predicts downstream performance, as lower CSI values consistently correlate with better results.
4.2 What is the Effect of Computational Budget?
Building on our earlier findings, we now investigate how varying the computational budget impacts down- stream performance. Specifically, we ask:
Do our findings about dalaset alignment terms of CSI hold across different budget levels? ?
Figure 4 shows downstream MAE across pretraining budgets of 0.5M, 1M, 2M, and 3M samples (each trained for epochs). We observe that incrcasing the budget for low-CSI datasets (i.c., ANI-1x and Transition-lx 1x) beyond 2M leads to diminishing returns on rMD17 and SPICE. In contrast, increasing the pretraining budget for high-CSI datasets (i.e., OC20 and OC22) often degrades downstream performance more substantially, particularly on rMD17, SPICE, and QM9. These results indicate that allocating more compute to misaligned upstream tasks can harm downstream generalization.
Takeaway. Our findings are consistent across budget levels: the upstream dataset with the lowest CSI yields the best downstream performance.
4.3 What is the Effect of Changing the Backbone Size?
In the previous sections, we used the small variant, GemNet-OC-S, as our backbone. Here, we address the question:
Does the correlation between CSI and downstream performance hold across differrnnt back- bone sizes?
Table 2 reports the downstream performance using the large variant, GemNet-OC-L, as the backbone. We also include our best attempt at reproducing the baseline results using JMP-L pretraining (denoted as "JMP-L.").
Consistent with the results on the small backbone, models pretrained on ANI-Ix achieve the best performance across all downstream tasks, aligning with its low CSI values. Notably, using a small computation budget of C 10M (i.e., 2M samples over 5 epochs), ANI-1x outperforms JMP-L, which was pretrained with C = 240M
8

---

Published in Transactions on Machine Learning Research (01/2026)
Pretraining Dataset
ANI-1x
Transition-ls 1x
OC20
OC22
rMD17
MD22
SPICE
QM9
20
13
15
5
11
MAE
9
4
10
7
3
5-
0.5 1.0
2.0
3.0
0.5
1.0
2.0
3.0
3.0 Pretraining Budgel (Millions)
0.5 1.0
2.0
0.5 1.0
2.0
3.0
Pretraining Budget (Millions)
Pretraining Budget (Millious)
Pretraining Budgel (Millions)
Figure Effect of Computational Budget on Performance. While fixing the number of epochs (E) to 5, we vary the number of training samples across W 0.5M, 1M, 2M, and 3M. Our findings are consistent across budget levels where the upstream dataset with the lowest CSI yields the best downstream performance.
Table 2: Effect of Changing the Backbone Size. We analyze the impact of using larger variant of GemNet-OC and find that, irrespective of backbone size, relevance-based upstream dataset selection consistently outperforms costly large-scale joint pretraining.
C
Upstream Data
Backbone
rMD17
MD22 (meV/A)
SPICE (meV/A)
QM9
(meV/A)
(meV)
JMP-L (Ce-Net-OC-L)
5.1
1.92
4.75
2.9
240M
Joint (Temperature)
JMP-L, (GemNeetOOCLL)
5.3
2.59
4.91
3.0
ANI-1x
4.8
2.54
5.24
2.6
Transition-ls
9.7
3.56
7.42
3.0
10M
GemNet-OC-L
OC20
3.90
9.24
4.6
OC22
12.0
4.14
10.43
4.0
Joint (Balanced)
8.9
3.39
7.57
3.2
10M
GemNet-OC-L
Joint (Temperature)
10.8
4.15
9.83
4.4
on a joint upstream dataset. We obtain state-of-the-arl results with an MAE of 4.8 on Aspirin (rMD17) and 2.6 on Uo (QM9), demonstrating that strong dataset alignment can outweigh large-scale pretraining even with increased model capacity. While larger backbones improve overall performance, the gap between aligned and misaligned upstream datasets persists. High-CSI datasets like OC20 and OC22 still underperform, reaffirming the importance of dataset alignment.
Takeaway. Our findings hold across backbone sizes: scaling up the model does not change the relative utility of upstream datasets. Alignment-basee upstream dataset selection outperforms large-scale dataset mixing, even under high-capacity settings and at significantly lower computational budgets.
4.4 What is the Effect of Changing the Backbone Type?
In the previous experiments, we focused on the GemNet-OC backbone to enable direct and fair comparison with the JMP baseline. We now examine whether our CSI-based alignment analysis extends beyond graph neural networks to a different architectural paradigm. Specifically, we ask:
Does the relationship between CSI and downstream performance persist across different backbone types?

---

Published in Transactions on Machine Learning Research (01/2026)
Table 3: Effect of Changing the Backbone Type. Downstream MAE after fine-tuning EquiformerV. V2 (EqV2) pretrained on different upstream datascts. Despite the architectural shift from graph neural net- works to transformer-basee models, upstream dataset alignment remains strong predictor of downstream performance. Numbers in parentheses indicate parameter count.
C
Pretraining Time
Upstream Data
Backhone
MMD1 (meV /Å)
MD22
SPICE (meV/A)
QM9
(A100 GPU-days)
ImeV/A)
(meV)
1.68
ANI-Ix
9.7
3.23
8.88
6.5
1.60
Transition- 1x
11.0
3.37
10.16
10M
EqV2 (31M)
6.78
OC20
12.5
3.58
10.54
7.0
7.72
OC22
12.6
3.72
11.10
6.7
6.69
ANI-1x
9.0
2.90
6.24
3.5
6.72
Transition-Ix
10.2
3.05
7.80
3.6
10M
EqV2 (87M)
16.52
OC20
11.4
3.25
8.10
4.0
17.16
OC22
11.9
3.28
8.60
4.3
To answer this question, we evaluate EquiformerV2 (Liao et al. 2023 backbone. We pretrain EquiformerV2 on the same set of upstream datasets using the identical pretraining and fine-tuning protocol as in our main experiments. We adopt the same pretraining budget of 2M samples over 5 epochs for both the small (31M parameters) and large (87M parameters) Equiformer V2 variants. Pretraining EquiformeVV2 notably more expensive than GemNet-OC. Even under this limited 10M budget, the large variant requires up 17 A100 GPU-days. We report the results in Table 3
transformer-based equivariant
Across both EquiformerV2 model sizes, performance trends remain consistent with GemNet- OC findings. The datasets with the lowest CST lead to the strongest. downstream performance, with ANT-1x being the top-performer in nearly all cases. While Transition-Ix shows marginal advantage on QM9 for the 31M model, ANI-1x achieves the best results as capacity increases to the 87M variant. This confirms that the benefits of alignment reflect a general property of data relevance that persists across architectural paradigms.
Takeaway. CSI-based alignment reliably predicts downstream performance across both backbone sizes and backbone lypes, highlighting the central role of dataset alignment in determining pretraining performance.
4.5 Is More Diverse Data Always Better?
A common assumption in pretraining that larger and more diverse datasets lead better generalization. This intuition motivates the JMP framework, where large-scale pretraining budget of = 240M led to strong downstream results. However, it remains unclear whether this benefit comes from the size, the diversity, or the alignment of the data with the downstream task. Here, we revisit this assumption through targeted experiment:
Does increasing data diversily by adding less aligned sources improve or harm downstream performance?
To test this, we compare two settings: (1) pretraining N 2M unique samples from ANI-Ix, the most CSI-aligned dataset, and (2) pretraining
mixture of 2M ANI-1x samples and 1M OC22 samples (i.c., N = =3MM), both trained for epochs. As shown in Figure 5 simply adding OC22 results in worse downstteam performance across all four tasks, despite the increase in data volume. This indicates that adding less aligned data may interfere with the knowledge transfer gained from aligned pretraining sources.
Takeaway. Our results challenge the intuitive strategy of adding diversity to pretraining datasets without considering alignment. CSI provides a practical signal for curating upstream data that supports better generalization, especially under constrained budgets.
4.6 Is Alignment Important in Self- Supervised Pretraining?
Self-supervised learning has long been central the scarcity of large-scale labeled data (Rong al. 2020 Liu et. al.] 2021 Jiao et al. 2023 Chen et al.
3D molecular representation learning, initially driven by
10

---

Published in Transactions on Machine Learning Research (01/2026)
ANI-Ix (2M)
ANI-Ix (2M) + OC22 (1M)
7.5
MAR
5.0
2.5
0.0
rMD17
MD222
SPICE
QM9
(meV/A)
(meV/A)
(meV/A)
(meV)
Figure Impact of Adding Less Aligned Pretraining Data. Adding 1M OC22 samples to a 2M- sample ANI-1x baseline worsens downstream performance despite a larger pretraining budget. This highlights the importance of dataset alignment and the value of the CSI metric for effective pretraining.
Table 4: Effect of Changing the Pretraining Objective. We report the test MAE for the EquiformerV? (EqV2) backbone, with 87M parameters, pretrained with the DeNS self-supervised denoising objective. Consistent with our supervised results, pretraining on the most CSI-aligned dataset (i.e., ANI-Ix) achieves superior fine-tuning performance across all downstream molecular tasks.
Upstream Data
Backbone
rMD17
MD22
SPICE (meV/A)
QM9
(meV/A)
(meV/A)
(meV)
ANT-1x Transition-Ix
4.7
2.78
5.95
4.1
7.3
2.95
8.71
4.2
10M
Eq (87M) + DeNS
OC20
8.1
3.32
7.34
5.0
OC22
10.4
3.58
9.39
4.9
Zhou al. 2022a Ji et al, 2024 While recent frameworks such as JMP (Shoghi et al. 2023 emphasize supervised pretraining on energy and forces, modeling structural dependencies directly from molecular geometry remains dominant paradigm. This raises critical question regarding the role of data
2021
selection:
Does dataset alignment matter for self-supervised pretraining?
To answer this, we extend our analysis to the self-supervised setting and adopt the DeNS framework (Liao et al. 2024). DeNS follows the standard denoising paradigm in 3D molecular learning, where models are trained to recover the original 3D atomic positions from noisy inputs. Unlike traditional self-supervised denoising objectives designed for equilibrium geometries (Rong al., [2020] Zhou et al.) 2022a),, DeNS explicitly targets non-cquilibrium molecular and material configurations. This makes it well suited to the non-equilibrium datasets considered in this work. We pretrain the EquiformerV2 (87M) backbone with the DeNS objective on each upstream dataset using a fixed budget of 10M training instances (2M samples over 5 epochs).
The results, summarized in Table 4 demonstrate that CST remains robust predictor of downstream per- formance even in the self-supervieed sctting. Across all tasks, the upstream datasct with the lowest CSI (i.e., ANI-1x) yields the strongest transfer performance. Remarkably, the model pretrained ANI-Ix using self-supervised denoising achieves MAE of 4.7 meV/À on Aspirin (rMD17), outperforming the large- scale supervised JMP-L baseline (5.1 meV/A) while using 24x less pretraining budget. This confirms that dataset alignment remains relevant for 3D representation learning under both supervised and self-supervised objectives.
Takeaway. Datasct alignment is objoctive-gnnosic. Even when the pretraining objective is changed model structural dependencies in sele-super manner, CSI-aligned datasets consistently achieve superior performance.
11

---

Published in Transactions on Machine Learning Research (01/2026)
ANI-Ix
Transition-Ix
OC20
OC22
200
150
CSI
100
50
0
QMOF
Bench
Downstream Dataset
Figure 6: CSI Between Upstream and OOD Downstream Tasks. CSI values predict that ANI-Ix is the best pretraining choice for QMOF, while OC20 and OC22 are best for MatBench.
Table 5: OOD Task Performance Across Upstream Sources. We compare the CSI-predicted best upstream sources with actual downstream performance on OOD tasks (QMOF, MatBench, and QM9's
A.). While CSI aligns well with QM9's OOD label, it mispredicts the best source for MatBench. Joint pretraining gencrally improves performance, highlighting the benefits of diverse upstrcam sources for OOD generalization.
C
Upstream Data
Backbone
QM9 [Ac]
QMOF (cV)
MatBench [fold0 / mean] (em
(meV)
JMP-S (GemNNet-OCS))
23.1
0.18
26.60 / 22.77
240M
Joint (Temperature)
JMP-S (GemNettOOCS))
24.0
0.19
24.77 / 21.48
ANI-1x
24.5
0.22
30.09 / 29.60
Transition-lx
25.3
0.22
52.22
38.56
10M
GemNet-OC-S
OC20
30.8
0.22
37.52
30.88
OC22
35.6
0.22
32.78 / 27.55
Joint (Balanced)
27.3
0.21
26.11
/ 24.87
10M
GemNet-OC-S
Joint (Temperature)
27.9
0.21
26.63
25.61
5 Beyond In-Distribution
Recall that our pretraining process is conducted on upstream tasks involving molecules and catalysts, with energy and force targets. We classify in-distribution (ID) any downstream task that uses the same labels (energy and/or forces), and as out-of-distribattion (OOD) those that involve different target properties, such as band gap (e.g,, QMOF) or phonon properties (e.g., Mathench). While our main results focused on ID evaluation, here we explore our metric's applicability to OOD tasks. Specifically, we examine three cases:
the Band Gap property from QMOF (Rosen et al. 2021 Phonons from MatBench Dunnn et al. 2020 and A. from QM9.
In Figure dataset-level CSI values as in Figure 3 We observe that QMOF exhibits pattern similar to other ID domains shown in Figure 3 However, MatBench displays distinct pattern, showing strong correlation with OC20 and OC22, followed by ANI-Ix and Transition-Ix. Next, we analyze the correlation between CSI and downstream performance under OOD evaluation.
present the CSI values for OOD domains; for QM9, the OOD label (A.) uses the same
Table shows that A, QM9 aligns with the CSI pattern, similar to ID evaluation, suggesting that CST remains effective when the downstream domain (i.c., molocules, as QM9) is similar to the upstream tasks, even if the label type is OOD. In QMOF, the different upstream sources achieve similar performance, which
12

---

Published in Transactions on Machine Learning Research (01/2026)
lags behind the full pretraining by JMP. For MatBench (evaluated over 5 folds), OC22 achieves the best mean performance while OC20 lags behind, despite our metric predicting both to be equally suitable. Additionally, for both QMOF and MatBench, joint pretraining variants generalize better than individual sources. This suggests that when both the downstream domain and the label type differ from all upstream sources, mixing diverse upstream domains provides the best performance.
While CSI reliably guides dataset selection for in-distribution tasks, its effectiveness in OOD scenarios reveals important limitations of datasct-level alignment. This may stem from the limited diversity of the backbone used for feature extraction, which was pretrained only energy and force targets. Future work could explore using backbones pretrained on broader sets of chemical properties or incorporating more diverse upstream domains to better caplure variation across OOD tasks. Another promising direction is to leverage foundation models trained on multi-mode or multi-objective tasks, which may offer more transferable representations for similarity assessment across varied downstream domains.
6 Conclusion
This paper challenges the prevailing trend of scaling data and computational resources in atomic property prediction by demonstrating that strategic data selection based on dataset alignment can achieve comparable or superior performance with significantly fewer resources. We introduce the Chemical Similarity Index (CSI), simple metric that quantifies the alignment between upstream pretraining datasets and downstreem tasks, enabling the selection of high-quality, task-aligned pretraining data. Our experiments reveal that smaller, focused dalasets often outperform larger, mixed ones, and that indiscriminately adding data can degrade performance when relevance is low. These findings highlight that aligmment, rather than scale alone, is the koy to effective pretraining, and they point toward a more principled, efficient, and sustainable direction for future research in atomic property prediction.
Acknowledgment
The authors thank Mohamed Elhoseiny, Firas Laakom, and Abdelrahman Eldesokey for their valuable feedback. Yasir Ghunaim was supported by Saudi Aramco. This work is supported by the King Abdullah University Science and Technology (KAUST) Centor of Excellence for Generative AI under award number 5940 and SDAIA-KAUST Center of Excellence Data Science, Artificial Intelligence (SDAIA-KAUST AI). The computational resources are provided by IBEX, which managed by the Supercomputing Core Laboratory at KAUST.
References
Jean-Michel Attendu and Jean-Philippe Corbeil, Nlu on data diets: Dynamic data subset selection for alp classification tasks. 2023.
Ramakrishna Bairi, Rishabh Tyer, Ganesh Ramakrishnan, and Jeff Bilmes. Summarization of multi-document topic hicrarchies using submodular mixtures. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp.
2015.
Luis Barroso-Luque, Muhammed Shuaibi, Xiang Fu, Brandon M Wood, Misko Dzamba, Meng Gao, Ammar Rizvi, C Lawrence Zätnick, and Zachary W Ulissi. Open materials 2024 (oma124) inorganic materials dataset and models. arXie preprint arXiv:24 10.771,,,0.11711 2024.
Rodrigo Berriel, Stephane Lathuillere, Moin Nabi, Tassilo Klein, Thiago Oliveira-Santos, Nicu Sebe, and Elisa Ricci. Budget-aware adapters for multi-domain learning. In ICCV, pp. 382-391, 2019.
Nathan Brown, Marco Fiscato, Marwin HS Segler, and Alain C Vaucher. Guacamol: benchmarking models for de novo molecular design. Journal of chemical information and modeling, 59(3):1096-11 -108, 2019.
13

---

Published in Transactions on Machine Learning Research (01/2026)
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11(10):6059 6072, 2021.
Dong Chen, Kaifu Gao, Duc Duy Nguyen, Xin Chen, Yi Jiang, Guo- Wei, and Feng Pan. Algebraic graph-assisted bidirectional transformers for molecular property prediction. Nature commumicatiinss, (1):3521, 2021.
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, (()):eadf0873, 2023.
Anders S Christensen and O Anatole Von Lilienfeld. On the role of gradients for machine learning of molecular energies and forces. Machine Learning: Science and Technology, 144):045018, 2020.
Alexander Dunn, Qi Wang, Alex Ganose, Daniel Dopp, and Anubhav Jain. Benchmarking materials property prodiction methods: the matbench test set and automatminer reference algorithm. npj Computational Materials, 6(1):138, 2020.
Peler Eastman, Pavan Kumar Behara, David L Dotson, Raimondas Galvelis, John E Herr, Josh T Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
Saurabh Garg, Hadi Pour Ansari, Mehrdad Farajtabar, Sachin Mehta, Raviteja Vemulapalli, Oncel Tuzel, Vaishaal Shankar, and Fartash Faghri. Tic-clip: Continual training of clip models. In ICLR, 2024. URL
htttp://arxiv.. wwwp://////////////ee 2310.162266.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural
networks for molccules. NeurIPS, 1667990-802,, 2021.
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C Lawrence Zitnick, and Abhishek Das. Gemnet-oc: developing graph neural networks for large and diverse molecular simulation datasots. arXiu preprint ar arXiv:2204.02788 2022.
Yasir Ghunaim, Adel Bibi, Kumail Alhamoud, Motasem Alfarra, Hasan Abed Al Kader Hammoud, Ameya Prabhu, Philip HS Torr, and Bernard Ghanem. Real-time evaluation in online continual learning: A new hope. In CVPR, pp. 11888 11897, 2023.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The journal of machine learning research, 13(1):723-77) 773,7 2012.
Hasan Abed Al Kader Hammoud, Tuhin Das, Fabio Pizzati, Philip Torr, Adel Bibi, and Bernard Ghanem. On pretraining data diversity for self-supervised learning. 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. NewrlPS, 30, 2017.
Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W. Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. In NeurIPS Datasets and Benchmarks Track, 2021. URL https:///penrreview
Xiaohong Ji, Zhen Wang, Zhifeng Gao, Hang Zheng, Linfeng Zhang, Guolin Ke, and Weinan E. Exploring htw.pppaereeeeeeeeeeeee/uouaaaa. net/forummid molecular pretraining model at scale. In NewrIPS, 2024. URI. https://openrevierr
64V40K2fDv
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivated equivariant pretraining for 3d molecular graphs. In AAAT, volume 37, pp. 8096-80 0104,, 2023.
14

---

Published in Transactions on Machine Learning Research (01/2026)
Wei Jin, Xianfeng Tang, Haoming Jiang, Zheng Li, Danqing Zhang, Jiliang Tang, and Bing Yin. Condensing graphs via one-step gradient matching. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 720 730, 2022a.
Wei Jin, Lingxia Zhao, Shichang Zhang, Yozen Liu, Jiliang Tang, and Neil Shah. Graph condensation for graph neural networks. In ICLR, 2022b.
Vishal Kaushal, Rishabh lyer, Suraj Kothawade, Rohan Mahadev, Khoshrav Doctor, and Ganesh Ramakrish- nan. Learning from less data: A unified data subset selection and active learning framework for computer vision. In WACV, pp. 1289 1299. IEEE, 2019.
Krishnateja Killamsetty, Sivasubramanian Durga, Ganesh Ramakrishnan, Abir De, and Rishabh Iyer. Grad- match: Gradient matching based data subset selection for efficient deep model training. In ICML, pp. 5464-50 5474. PMLR, 2021a.
Krishnateja Killamsetty, Durga Sivasubramanian, Ganesh Ramakrishnan, and Rishabh Iyer. Glister: Gener- alization based data subset selection for efficient and robust learning. In AAAT, volume 35, pp. 8110-8118, 2021b.
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, 156(18), 2022.
Solomon Kullback and Richard Leibler. On information and sufficiency. The annals of mathematical slalistics, 22(1):79 86, 1951.
Agala Lapedriza, Hamed Pirsiavash, Zoya Bylinskii, and Antonio Torralba. Are all training examples equally valuable? Xix preprint or arXh::3311.6510, 2013.
Alexander C Li, Ellis Brown, Alexei A Efros, and Deepak Pathak. Internet explorer: Targeted representation learning on the open web. In JCML, PMLR, 2023.
Mengtian Li, Ersin Yumer, and Deva Ramanan. Budgeted training: Rethinking deep neural network training under resource constraints. arXie preprint Xiv: 1905. 205.04770 04753, 2019.
Yi-Lum Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In ICLR, 2023.
Yi-Lam Liao, Tess Smidt, Muhammed Shuaibi, and Abhishek Das. Generalizing denoising to non-equilibrium structures improves equivariant force fields. Transaction Machine Learning Research, 2024.
Jianhua Lin. Divergence measures based on the shannon entropy. IEEE Transactions on Information theory, 37(1):145 151, 2002.
Mengyang Liu, Shanchuan Li, Xinshi Chen, and Le Song. Graph condensation via receptive field distribution matching. arXiv preprint arXiv:2206. 13697, 2022.
Shengehao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. arXiv preprint 1rX:v:2110.07728, 2021.
Timothy Nguyen, Zhourong Chen, and Jachoon Lec. Dataset meta-learning from kernel ridgo-regression. In ICLR, 2021a.
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jaehoon Lee. Dataset distillation with infinitely wide convolutional networks. NewrIPS, 1451886-119, 2021b.
Xuran Pan, Xuan Jin, Yuan He, Shiji Song, Gao Huang, et al. Budgeted training for vision transformer. In ICLR, 2022.
Saro Passaro and Lawrence Zitnick. Reducing so (3) convolutions to so (2) for efficient equivariant gnns. In ICML, pp. 27420 27438. PMLR, 2023.
15

---

Published in Transactions on Machine Learning Research (01/2026)
Guilherme Penedo, Hynck Kydliček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In NeurlPS Datasets and Benchmarks Track, 2024. URL https///ppareview.nettt wew.//runneoooooodd-- n6SCkn2QaC
Ameya Prabhu, Hasan Abed Al Kader Hammoud, Puneet K Dokania, Philip HS Torr, Ser-Nam Lim, Bernard Ghanem, and Adel Bibi. Computationally budgeted continual learning: What does matter? In CVPR, pp. 3698-3707 2023.
Kristina Preuer, Philipp Renz, Thomas Unterthiner, Sepp Hochreiter, and Gunter Klambauer. Fréchet chemmet distance: metric for generalive models for molecules in drug discovery. Journal of chemical information and modeling, 58(9):1736 119-11733-1771,, 2018.
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self- supervised graph transformer on large-scale molecular data. Advances in neural informalion processing systems, 13:12559-12577 2020.
Andrew S Rosen, Shaelyn M lyer, Debmalya Ray, Zhenpeng Yao, Alan Aspuru-Guzik Laura Gagliardi, Justin M Notestein, and Randall Q Smurr. Machine learning the quantum-chemical propertics of metal- organic frameworks for accelerated materials discovery. Matter, 4(5):1577-155) 1597, 2021.
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European signal processing conference, pp. 606-610. IEEE, 2007.
Mathias Schreiner, Arghya Bhowmik, Tejs Vegge, Jonas Busk, and Ole Winther. Transition lx-a datasot for building generalizable reactive machine learning potentials. Scientific Dala, 9(1):779, 2022.
Claude E Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3): 979-423, 1948.
Nima Shoghi, Adeesh Kolluru, John R Kitchin, Zachary Ward Ulissi, C Lawrence Zitnick, and Brandon M Wood. From molecules to materials: Pre-training large generalizable models for alomic properly prediction. In ICLR, 2023.
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.
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 ceneral-purpose neural network potential through transfer learning. Nature communications, (1):2903, 2019.
Justin S Smith, Roman Zubatyuk, Benjamin Nebgen, Nicholas Lubbers, Kipton Barros, Adrian E Roitberg, Olexandr Isayev, and Sergei Tretiak. The ani-lecx and ani-lx data sets, coupled-cluster and density functional theory properties for molecules. Scientific data, 7(1):134, 2020.
Anuroop Sriram, Sihoon Choi, Xiaohan Yu, Logan M Brabson, Abhishek Das, Zachary Ulissi, Matt Uyt- tendaele, Andrew J Medford, and David Sholl. The open dac 2023 dataset and challenges for sorbent discovery in direct air capture, 2024.
Richard Tran, Janice Lan, Muhammed Shnaibi, Brandon M Wood, Siddharth Goyal, Abhishek Das, Javier Heras-Doming Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) dataset and challenges for oxide cloctrocatalysts. ACS Catalysis, (3(5))3066-30)) 3084, 2023.
Tongzhou Wang, Jun Yan Zhu, Antonio Torralba, and Alexei A Efros. Dataset distillation. arXiv preprint arXiv:1811. 1811:1811.00952 10959, 2018.
16

---

Published in Transactions on Machine Learning Research (01/2026)
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Ancesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9(2):513 530, 2018.
Bo Zhao and Hakan Bilen. Dataset condensation with distribution matching. In WACV, pp. 6514 6523, 2023.
Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen. Dataset condensation with gradient matching. In ICLR, 2021.
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. 2022a.
Yongchao Zhou, Ehsan Nezhadarya, and Jimmy Ba. Dataset distillation using neural feature regression. NewrIPS, 35:9813-B8 90279,8777 2022b.
17

---

Published in Transactions on Machine Learning Research (01/2026)
A More Epochs or More Data?
To extend the findings presented in the main paper, we explore the trade-off between increasing the number of training epochs and expanding the dataset size under fixed computational budget. Specifically, we aim to answer the following question:
Given fixed computational budget, it more effective train on smaller dataset for more epochs or train on a larger dataset for fewer epochs?
Sctup. To investigate this question, we compare two scenarios under the computational budget of 10M samples: (1) training 2M samples for 5 epochs, and (2) training 1M samples for 10 epochs. We evaluate the performance of models pretrained on ANI-Ix, Transition-ls OC20, and OC22, and fine-tune them on the downstream datasets: rMD17, MD22, SPICE, and QM9. For comparison, we also include the results JMP-L and JMP-S, which use 120M samples for 2 epochs.
Results. Table G presents the downstream performance for the two scenarios. Across all datasets, ANI-1x consistently achieves the best performance, regardless of whether the model is trained 2M samples for 5 epochs or 1M samples for 10 epochs. For example, rMD17, ANI-1x achieves test error of 5.4 in both scenarios, outperforming JMP-S (6.7) and JMP-T (5.1). Similarly, SPICE, ANI-Ix achieves a test error of 5.08 (2M samples, 5 epochs) and 5.04 (1M samples, 10 epochs), compared to 5.71 for JMP-S and 4.75 for JMP-L.
Interestingly, increasing the number of epochs from 10 while reducing the dataset size from 2M to 1M does significantly degrade performance for ANI-Ix. This suggests that for highly aligned datasets like ANI-1x, training fewer samples for more epochs be effective training more samples for fewer epochs. In contrast, for less aligned datasets such OC20 and OC22, increasing the mumber of epochs only partially compensates for the reduced dataset size, as some tasks show similar performance while others experience noticeable degradation.
Takeaway. Our findings indicate that the trade-off between more epochs and more data depends on the alignment of the pretraining dataset with the downstream task. For highly aligned datasets like ANI-1x, training on fewer samples for more epochs can yield comparable performance. In contrast, for less aligned datasets, increasing the dataset size tends to be more beneficial. These results further show the importance of dataset quality and alignment, as quantified by CSI, in determining effective pretraining strategy.
Table 6: Trade-off between increasing the number of samples and the number of epochs. We report the MAE for various downstream tasks while varying the pretraining sample count and epoch count simultaneously. C. N, and E denote the computational budget, number of samples, and mumber epochs, respectively.
N E
Upstream Data
Backbone
rMD17
MD22
SPICE
QM9
(meV/A A)
(meV/A)
(meV/A)
(meV)
ANI-1x
5.4
2.90
5.13
2.9
Transition-lx
10.1
3.73
7.55
3.2
10M
2M
5
GemNel-OC-S
OC20
14.6
4.53
8.74
4.8
OC22
16.0
5.20
10.73
5.7
ANIIx
5.4
2.88
5.04
2.9
Transitionix
10.6
3.79
7.50
3.1
10M
1M
10
GemNet-OCs
OC20
14.8
4.67
10.16
4.9
OC22
17.3
5.24
11.06
5.4
18

---

Published in Transactions on Machine Learning Research (01/2026)
B Impact of Changing the Feature Extractor for CSI
In the main paper, we used an EquiformerV2 model pretrained on the OC20 datasct as our primary feature extractor for computing the CSI values. In this section, we explore alternative open-source pretrained EquiformerV2 checkpoints, including: ODAC23 (Siraam et al.) 2024). MPtrj (Barroso-Luque et al. 2024) and MPtrj (DeNS) (Barroso-Luque et al. 2024 DeNS (Liao et al. 2024 refers to a denoising auxiliary task that extends traditional denoising to non-equiibrrim structures to improve learned representations.
We report the results in Figure Our comparison is based relative rankings since the absolute values are influenced by the choice of feature extractor. Notably, for out of 5 datasets (rMD17, MD22, SPICE, and OMat24), all feature extractors agree on the lowest-CSI upstream dataset: ANI-1x for rMD17, MD22, and SPICE, and OC20 for OMat24. For QM9, OC20 and MPtrj (DeNS) rank ANI-Ix as the most aligned, whereas ODAC23 and MPtrj place ANI-1x second to Transitionlx, which is also a strong upstream candidato for QM9.
ANI-1x
Transition-Ix
OC20
OC22
Feature Extractor: OC20
Feature Extractor: ODAC23
300
10000
8000
200
6000
0
4000
100
0
rMD17
MD22
SPICE
QM9
OMat24
rMD17
MD22
SPICE
QM9
OMat24
Downstream Dalaset
Downstream Datasel
Feature Extractor: MPtrj
Feature Extractor: MPtrj (DeNS)
1500
CSI
1000
500
500
0
rMD17
MD22
SPICE
QM9
OMat24
xMD17
MD22
SPICE
QM9
OMat24
Downstream Dataset
Downstream Dataset
Figure 7: Dataset Alignment Using Feature Extractors Pretrained Various Datasets. We compare CSI relative values computed using feature extractors pretrained on four different datasets: OC20, ODAC23, MPtrj, and MPtrj (DeNS). The lowest CSI upstream dataset is consistent across feature extractors for out downstream tasks, namely rMD17, MD22, SPICE, and OMat24.
19

---

Published in Transactions on Machine Learning Research (01/2026)
C Feature Representation Choices for CSI
A key challenge when adapting the Fréchet Inception Distance (FID) to 3D molecular graphs is handling their variable number of aloms. Our leature extractor, EquilormerV2, produces variable-sized node-level features in Rxxd, where n is the number of atoms and d is the feature dimension. A straightforward approach is to aggregate these features using mean pooling, producing a graph-level representation in RNxd where is the number graphs. This setup is directly analogous to the image-based FID. The central question, however, is whether this aggregation loses valuable information.
For our CSI metrie to accurately measure aligmment across various domains, it critical that the feature expressivity is maximized. To quantify this expressivity, use the Effective Rank (Roy & Vetterli 2007) of the covariance matrix of these features. The Effective Rank is similar to the traditional matrix rank but is more robust to noise.
Formally, as in Roy Vetterli (2007) let C E Rddd be the covariance matrix with eigenvalues {}g}ー=1 and Pj = >// 2g Am. The Shannon (1948) entropy H -E Py log byg and the Effective Rank is:
erank(C) = exp(H).
Higher crank indicates that information is distributed across more independent directions, signifying richer and less redundant feature representation.
To compare node-level and graph-level expressivity in a paired manner, we employ a bootstrap protocol. For a chosen number of graphs k, we proceed as follows:
1. Sample k graphs from the dataset.
2. Construct two matrices in kxd from the sampled graphs: (i) a graph-level matrix, where each row is the pooled representation of one graph, and (ii) node-level matrix, where each row corresponds to a randomly selected node representation from each graph.
3. Standardize each feature (z-score per dimension) in both matrices.
4. Compute the covariance C' and its effective rank crank(C) for both variants.
We repeat this process 10 times and report the mean effective ranks. We perform this comparison for & {5,000, 10,000, 15,000} and show the results in Figure S We observe that across all tested sample sizes, the nodc-level features consistently have a significantly larger effective rank. This finding indicates that the aggregation process for graph-level features results in a loss of intrinsic dimensions that likely significant for dataset alignment studies. Therefore, to preserve maximum feature expressivity, features in our CSI design.
node-level
ANI-Ix
Transition-Is
OC20
OC22
Effective Rank
650
650
800
600
500
650
550
500
500
400
5k
10k
15k
5k
10k
15k
6k
10k
15k
5k
10k
15k
Subset size (k)
Subset size (k)
Subset size (k)
Subset size (k)
Node-level
Graph-level
Figure 8: Node-Level Graph-Level Effective Rank. Effective rank of features from (ANI-1x, Transition-Ix, OC20, OC22) datasets using both node-level (green) and mean-pooled graph-level (orange) representations. The consistently higher effective rank of node-level features demonstrates that mean pooling removes informative directions, which may be inadequate for analyzing alignment between datasets.
20

---

Published in Transactions on Machine Learning Research (01/2026)
To complement our effective rank analysis, compare node-level and graph-level CSI variants in Figure OC20 and OC22 are both catalysis-focused datasets and share similar chemical domains. The node-level features reflect this similarity by yielding comparable scores between the two datasets. In contrast, the graph- level variant suggests greater divergence between the two, which may be due to information loss introduced by mean pooling.
ANI-Ix
Transition-ls lx
OC20
OC22
Node-level (no pooling)
300
200
CSI
100
=MD17
MD22
SPICE
QM9
OMst21
Downstream Dataset
Graph-level (mean pooling)
150
100
CSI
=MD17
MD22
SPICE
QM9
OMa124
Downstream Dataset
Figure 9: CSI values for Node-Level vs Graph-Level Feature Representation
21

---

Published in Transactions on Machine Learning Research (01/2026)
D Class Balancing for Feature Extraction
The upstream datascts analyzed in our main paper are large in size. For example, OC20 contains over 133 million training samples, while Transition-Ix contains around 9.6 million. To make the computation of our CST metric tractable, we sample 10K instances from each upstream dataset for feature extraction. To ensure these subsets are representative of the full distribution, we perform class-balanced sampling. For ANI-Ix and Transition-Ix, a class defined by the molecular formula of the molecule in the trajectory. For OC20 and OC22, a class is defined by the bulk structure. Each class treated bin, and sample equal number of instances from each bin.
Figure 10 compares random sampling and class-balanced sampling in terms of class coverage. Random sampling leads to significant underrepresentation of less frequent classes, while class-balance balanced sampling ensures uniform coverage across all classes.
ANIIX
Transitionix
OC20
OC22
Total Unique 2151
Total Unique: 156
Total Uniques 8128
4003
Total Uniquee ST16
55000
150
8000
5600
Unique Classes
9000
2000
00000
75
10000
1500
50
1000
300
X000
Bolanced
Figure 10: Impact of Sampling Strategics on Subset Construction for Feature Extraction. We
show the differences in class coverage between random and claas-balanced sampling when sampling 10K instances for each upstream task.
22

---

Published in Transactions on Machine Learning Research (01/2026)
E Implementation Details
For both pretraining and fine-tuning experiments, we primarily follow the JMP hyperparameters. However, due to resource constraints requiring smaller batch sizes compared to JMP, we adjusted the learning rate to ensure training stability, as detailed below.
For pretraining, we use a batch size of 20 and a learning rate (LR) of lc-4 for the small backbone (Ge-NettOCC-S). For the large backbone (GemNet-OC-L)) the batch size is reduced 12 to fit GPU memory. Additionally, when training with the OC22 dataset the large backbone, LR of 1e-4 caused gradient instability, thus used LR of 1e-5 for that particular run. Unless otherwise specified, each experiment is run for five epochs on the specified number of samples for each section of the paper. The best checkpoint is selected based on the performance in the validation set. To handle the large size of the upstream validation sets, validation is performed on a smaller subset of 2,000 samples.
For finetuning, we use the batch size specified in the JMP codebase and a default learning rate (LR) of 8e-5, except for cases where adjustments were needed stabilize training. Specifically, use 5e-5 for QMOF, 8c-4 for Mathench when pretrained on Transitionlx.
23


Reviews:
[{"review_id": "b09UNePYYg", "reviewer": "Reviewer_1xGw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n1. this paper conduct in-depth study of the effects of alignment between pretraining and finetuning datasets on downstream task performance. \n\nWeaknesses\n1. FID-like metric has been proposed to quantify dataset similarity for molecules [1] and used in molecular benchmarks [2]. \n2. The conclusion of this paper, similar data for pretraining and finetuning yields better results, seems to be well known and trivial. What is the best approach to quantify molecular dataset similarity is valuable, but this paper did not compare with other similarity metrics. \n\n[1] Preuer, Kristina, et al. \"Fréchet ChemNet distance: a metric for generative models for molecules in drug discovery.\" Journal of chemical information and modeling 58.9 (2018): 1736-1741. \n[2] Brown, Nathan, et al. \"GuacaMol: benchmarking models for de novo molecular design.\" Journal of chemical information and modeling 59.3 (2019): 1096-1108.\n[3]", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigate the importance of data selection between pretraining and downstream tasks. The authors proposed a data alignment metric, called Chemical Similarity Index, to quantify dataset level similarity. According to this study, selecting small but aligned data for pretraining can achieve similar or better performance in downstream tasks than using large but less related data in in-distribution tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. In the details of CSI implementation, the authors mention the node embeddings are used as features and class-balanced sampling is used to ensure representative coverage of molecular type.  1) Is the feature of a molecule represented as $\\mathbf{X} \\in \\mathbb{R}^{n \\times d} , n=\\text{num of nodes}, d=\\text{feature dimension}$? If that is the case, when generating the feature matrix of the entire dataset and calculating the mean and covariance for FID calculation, larger graphs and smaller graphs will have unbalanced contribution. 2) How is class-balanced sampling achieved? \n\n2. In Section 4 downstream datasets part, the author says \"Given the large number of pairs for ID evaluation, we focus on the first molecule for force tasks and the energy target for the multi-property dataset QM9\". What does this sentence mean? \n\n3. In Section 4 baselines part, JMP-S* and JMP-L* are reproduced finetuning results using the official JMP checkpoints. Are they finetuned using the JMP training data?\n\n4. Can you compare CSI with other dataset similarity metric such as Fréchet ChemNet Distance?", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xxgnnS8oyF", "reviewer": "Reviewer_rtE1", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is very clear and well-organized, presenting specific scientific questions and supporting them with proper experimental evidence. The background information and definitions provide a foundation for understanding the paper.\n2. The study addresses an important issue: how to conduct effective and inexpensive pretraining for molecular models to achieve optimal results.\n3. The experiments are well-designed and investigate how different variables, such as computational budget, model size, and task relevance, impact performance on downstream tasks.\n\nWeaknesses:\n1. The code is not provided, which limits the applicability of the presented method. In particular, the implementation of the CSI would help practitioners find the most relevant pretraining datasets for their projects. As the paper may have broader implications, the code would increase the reproducibility of this study and enable the confirmation of these results on other datasets.\n2. The features extracted for calculating the CSI are based on a model trained on OC20, one of the evaluated pre-training datasets. Using feature extractors trained on different datasets would be necessary to make sure the results are not biased towards this dataset.\n3. Based on Figure 3, all downstream datasets share a similar CSI ranking with respect to pre-training datasets. For all datasets, ANI-1x is the most similar dataset, and the two OC datasets are the least similar. It is unclear whether the good performance of the models pretrained on ANI-1x stems from high similarity or other qualities of this particular dataset.\n4. No confidence intervals are provided for the results in Table 1, which makes it difficult to estimate the significance of these results.\n5. Why are the results of pre-training on the combined datasets, with a small budget, not provided in Table 2? In this case, the results of the models pre-trained on the joint dataset with a 240M budget outperform those of models pre-trained on a single dataset (first row, MD22 and SPICE).\n6. The results in Table 3 again suggest that ANI-1x is a generally good pre-training dataset, yielding good performance even when the downstream dataset is not similar to it (MatBench). This contradicts the main claim of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work aims to determine whether, in atomic property prediction, we can achieve equal or better results by pretraining on much smaller datasets that are more closely aligned with the target chemical space. The results indicate that pretraining on high-quality and task-aligned datasets can produce similar outcomes when training on only a fraction of the data used currently to train large atomic models. The authors introduce a chemical similarity index (CSI) inspired by the FID score from computer vision. The GemNet-OC-S model is pretrained on four molecular datasets, where its task is to predict energy and per-atom forces, and then fine-tuned on four downstream datasets. The CSI between the upstream and downstream datasets correlates with the model's performance. Additionally, the authors explore how changing the backbone model size and training on three less related tasks affect the results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the weaknesses above. The most important changes that would improve the impact of the paper are:\n1. Publishing the code to enhance the reproducibility of this research.\n2. Finding more examples of downstream datasets where ANI-1x is not the most similar dataset. Currently, some upstream datasets consistently perform well, while others perform poorly.\n3. Repeating the experiments on multiple folds and providing confidence intervals would strengthen the results.\n4. Considering also using other models to compute CSI since the current model is pretrained on one of the upstream tasks.", "broader_impact_concerns": "The paper does not mention any broader impact concerns. Since it is a study on pretraining machine learning models to enhance performance on atomic property prediction tasks, a detailed discussion of ethical issues is not required.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WBvEy4kRQ7", "reviewer": "Reviewer_yRFh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThe paper shows that small but task‑aligned data beats scaled pre-training dataset, using just 2 M ANI‑1x molecules to rival a 120 M‑sample baseline while reducing compute. It introduces the simple, model‑agnostic metric Chemical Similarity Index (CSI) that predicts transfer quality, and validates its claims across budgets, model sizes and both in‑ and out‑of‑distribution tasks. The paper reveals how misaligned extra data can hurt.\n\n## Weaknesses\n\nIn my opinion, **the pre-training objective is wrong**. I know that there are previous works using the same pre-training scheme, but it does not necessarily validate it's theoretical correctness. From my perspective, the pre-training objective should focus on modeling the structural dependencies within the molecules, rather than performing any property prediction. The authors could check out the pre-training methods adopted by Grover or Uni-Mol [1,2]. This is a fundamental flaw of this work and it makes the soundness of all subsequent discussion questionable.\n\nAdditionally, the backbone model architecture is somewhat outdated. Most of the SOTA works are using Transformers, which have been proven to be more effective in capturing the data features. Experimenting on such backbone models might deliver different results.\n\nOne other concern is related to the usage of the pre-trained models. In both NLP and materials science domains, people tend to directly apply fine-tuning on open-source pre-trained models, rather than manually go through the pre-training process themselves. With this pipeline, the demand on computation resources is not as significant as what have been proposed in the paper.\n\n\n[1] Zhou, Gengmo, et al. \"Uni-mol: A universal 3d molecular representation learning framework.\" (2023).  \n[2] Rong, Yu, et al. \"Self-supervised graph transformer on large-scale molecular data.\" Advances in neural information processing systems 33 (2020): 12559-12571.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "For atomic‐scale property prediction, the work claims that carefully choosing a pre‑training dataset that is chemically aligned with the downstream task can outperform blindly scaling to hundreds of millions of structures. Using 2M ANI‑1x molecules, the authors match or beat SOTA MAE on 4 standard benchmarks with lower pre-training computation resources with certain pre-training schema and backbone model structures. Ablation studies show that lower distance from training and test datasets correlates with better in‑distribution accuracy at any compute budget.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Using a proper pre-training objective (e.g., those used in the referred papers).\n- Using Transformers as the backbone model.\n- Adding experiments and analysis on open-source pre-trained molecular representation models.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)