SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dingshuo Chen $^{1*}$  Yanqiao Zhu $^{2*}$  Jieyu Zhang $^{3}$  Yuanqi Du $^{4}$  Zhixun Li $^{5}$  Qiang Liu $^{1}$  Shu Wu $^{1\dagger}$  Liang Wang $^{1}$
<sup>1</sup>Center for Research on Intelligent Perception and Computing Institute of Automation, Chinese Academy of Sciences
$^{2}$ Department of Computer Science, University of California, Los Angeles
<sup>3</sup>The Paul G. Allen School of Computer Science and Engineering, University of Washington
$^{4}$ Department of Computer Science, Cornell University
$^{5}$ Department of Systems Engineering and Engineering Management The Chinese University of Hong Kong
Primary contact: dingshuo.chen@cripac.ia.ac.cn
Molecular Representation Learning (MRL) has emerged as a powerful tool for drug and materials discovery in a variety of tasks such as virtual screening and inverse design. While there has been a surge of interest in advancing model-centric techniques, the influence of both data quantity and quality on molecular representations is not yet clearly understood within this field. In this paper, we delve into the neural scaling behaviors of MRL from a data-centric viewpoint, examining four key dimensions: (1) data modalities, (2) dataset splitting, (3) the role of pre-training, and (4) model capacity. Our empirical studies confirm a consistent power-law relationship between data volume and MRL performance across these dimensions. Additionally, through detailed analysis, we identify potential avenues for improving learning efficiency. To challenge these scaling laws, we adapt seven popular data pruning strategies to molecular data and benchmark their performance. Our findings underline the importance of data-centric MRL and highlight possible directions for future research.
The research enthusiasm for Molecular Representation Learning (MRL) is steadily increasing, attributed to its potential in expediting drug and materials discovery processes compared with conventional in vitro and in vivo experiments [1-4]. Within the context of MRL, the central objective is to leverage specific featurizations, or modalities, of molecules in order to learn continuous vector representations. These representations aim to capture comprehensive chemical semantics and exhibit high expressiveness, thereby effectively addressing various downstream tasks [5-17].
A trend in the field is developing neural architectures and training strategies to improve the expressiveness of the learned representations. However, the influence of varying data scales on the performance of MRL under different learning scenarios is yet to be fully understood. To fill this gap, we draw attention to the following questions: What are the neural scaling behaviors of molecular representation learning? Do they align with the previous scaling laws such as the power-law, established in other domains [18-20]? Beyond common research objects in neural scaling laws such as the impact of pre-training and model parameter sizes [21, 18, 22], MRL further presents unique
*These authors made equal contribution to this research.
†Corresponding author: Shu Wu (shu.wu@nlpr.ia.ac.cn).
data-oriented challenges. These include the selection of appropriate modality [23] and issues related to Out-Of-Distribution (OOD) generalization [24].
To provide a comprehensive understanding of these complexities, we investigate the impact of various design dimensions on MRL from a data-centric perspective. Specifically, our exploration spans four core design dimensions: (1) data modalities (molecular featurizations), (2) data splitting, (3) the role of pre-training, and (4) model capacity. We identify five scientific questions and outline our key observations as follows. A succinct summary of contribution of this work is present in Figure 1.
(Section 3.1) How does performance scale with data quantities? We conduct extensive experiments on four large-scale molecular property prediction datasets. These datasets contain a number of classification and regression tasks, both in single-task and multi-task settings, focusing on properties ranging from quantum mechanical properties to macroscopic influence on human body. The experimental results indicate that the model performance generally follows a power-law relationship with data quantities. Compared with the neural scaling laws in Natural Language Processing (NLP) and Computer Vision (CV) domains, there is no apparent performance plateau [21] within the range of datasets we explored, regardless of low-data and high-data regimes, which implies that the performance improvement with increasing dataset size in MRL is highly predictable.
(Section 3.2) How do different molecular modalities influence scaling laws? Since distinct molecular featurizations might carry different semantic meanings and their corresponding neural encoders have different inductive bias, the selection of appropriate modalities (molecular featurizations) in MRL remains an open question. In our investigation, we specifically choose three commonly used modalities (2D graphs [6, 8, 24], SMILES strings [25-28], and Morgan fingerprints [23]) and compare their performance on three classification tasks. We find that different modalities exhibit distinct learning behaviors in MRL; graphs and fingerprints are identified as the most data-efficient modalities, exhibiting the largest power-law exponent. In comparison, SMILES strings demonstrate lower performance improvement with the same data increment.
(Section 3.3) Does pre-training consistently result in positive transfer across all data scales in MRL? While previous studies tend to argue that molecular pre-training can improve performance on downstream tasks, these conclusions are often drawn based on evaluations on the entirety of available datasets [29-33, 30, 34]. Our work further probes into the effects of graph-based pre-training across varying scales of downstream datasets. The results show that pre-training indeed brings improvements in low-data regimes. However, the power-law exponent of the performance curve, reflecting the rate of growth with incremental data sizes, is smaller with pre-training compared to training from scratch. We suppose that pre-training only delivers stable gains when the downstream dataset is relatively small—under 40K samples, for instance. As the downstream dataset scales up, this positive gains diminish and might even revert to a negative transfer in the high-data regime.
(Section 3.4) What influence do dataset splits exert on scaling laws? Out-of-Distribution (OOD) generalization poses a great challenge in MRL [24]. Given that the training distributions, which typically encompass known compounds, often differ from the test distributions containing unknown compounds, this divergence is particularly prevalent in drug discovery. However, the impact of dataset splitting strategies on neural scaling laws remains largely unknown. To bridge this gap, we study three splitting strategies: random, imbalanced, and scaffold splitting. These strategies differ primarily in the degree of overlap between their respective training and test sets. Our empirical findings indicate that while all three splitting schemes adhere to the scaling laws, the exponent associated with random splitting is the smallest, suggesting that as the divergence between training and test distributions increases, the efficiency with which the model utilizes data samples decreases. Such results accentuate the inherent challenges of MRL in practical scenarios, which face with ubiquitous problems of distribution shifts.
(Section 3.5) How does the model capacity affect the scaling laws? The model capacity stands as another crucial factor impacting the performance of MRL models [18, 22] as it significantly influences the requisite amount of training data. To study the effect of model capacity in MRL, we vary the parameter size of Graph Isomorphism Networks (GIN) [5], a widely adopted graph neural network architecture and examine its performance on three classification tasks. We observe that, while model capacity does affect data utilization efficiency, there is not a clear correlation between the dataset size and the optimal model capacity required to obtain saturation performance. Unexpectedly, we find that the optimal model capacity for a task with a smaller training set could be larger than that of a task with a larger training set, which contradicts our intuition that larger models usually need larger
Figure 1: Summary of contributions of this work. We conduct a comprehensive data-centric study to examine the neural scaling laws in molecular representation learning. We explore four dimensions affecting the data utilization efficiency: (1) data modalities, (2) the role of pre-training, (3) dataset splitting, and (4) model capacity. Additionally, we study the subset selection problem by adapting seven data pruning strategies to molecular graphs.
training datasets. It suggests the need for further exploration of the interplay between model capacity and data scale, rather than applying a one-size-fits-all approach.
(Section 3.6) Can a curated subset from the full dataset yield comparable or even superior results? In CV and NLP domains, the utility of data pruning has been extensively explored due to the computational burden imposed by increasingly large models and massive amounts of data [35-37]. In MRL, however, the extend of data redundancy and the potential of pruning strategies to alleviate computational burden remain largely unexplored. To address this gap, we benchmark seven data pruning strategies originally proposed for image data and adapt them to molecular graph models on three classification tasks. The results show that these data pruning methods do not significantly outperform random selection, which highlights the need for developing data pruning strategies specifically tailored to molecular data.
Based on our empirical analysis, we identify several factors that can enhance data utilization efficiency of MRL. These include the utilization of graph and fingerprint modalities as input data, the application of pre-trained models on small-scale downstream datasets, and the design of data pruning strategies specifically tailored to molecular data. To the best of our knowledge, our study is the first to approach MRL from a data-centric perspective. We envision that this work could provide valuable insights that facilitate future explorations in the field.
2 Experimental Settings
In order to answer the six scientific questions through empirical investigation, we conduct a series of experiments on the neural scaling behaviors between model performance and varied data quantities. Specifically, we divide the whole dataset into nine proportional subsets:  $[1\%, 5\%, 10\%, 20\%, 30\%, 40\%, 60\%, 80\%, 100\%]$ , and for each configuration, we randomly select five seeds and report the mean performance. Then, we study how each design dimension of interest in MRL models influences the scaling laws. To demonstrate the trend of neural scaling laws, we employ the least square model to fit the performance curve.
2.1 Datasets and Tasks
Given the potential issues of over-fitting and spurious correlations that may arise with limited samples, we focus on relatively large-scale datasets containing more than 40K molecules in MoleculeNet [38]. Also, these datasets should cover both classification and regression tasks, are diverse in task settings (i.e. single- and multi-task settings), and focus on important biophysics and quantum mechanics properties. As a result, we choose four datasets ranging from molecular-level properties to macroscopic influences on human body for experimental investigation: HIV [39], MUV [40], PCBA [41] and QM9 [42]. Readers can refer to Appendix B for a more detailed discussion on datasets and tasks.
2.2 Implementation Details
In the following, we briefly introduce the implementation details of our experiments. Note that our experiments cover multiple dimensions, and thus all experimental settings will remain consistent except for the design dimension to be studied. Please refer to Appendix A for a detailed introduction of implementation details.
Modalities. Molecular featurizations translate chemical information of molecules into representations that can be understood by machine learning algorithms. Each featurization can thus be regarded as a modality of the molecular data. Here, we consider the following four molecular modalities: (1) 2D topology graphs model atoms and bonds as nodes and edges respectively, (2) 3D geometric graphs include coordinates of atoms into their representation to depict how atoms are positioned relative to each other in 3D space, (3) Morgan fingerprints [23] encode molecules into fixed-length bit vectors which map certain structures of the molecule within certain radius of molecular bonds, and (4) SMILES strings [25] represent chemical structures in a linear notation using ASCII characters, with explicit information about atoms, bonds, rings, connectivity, aromaticity, and stereochemistry.
MRL models. Since our experiments involve four different data modalities, each modality is modeled with its corresponding encoders.
- For 2D graphs, we utilize the Graph Isomorphism Network (GIN) [5] as the encoder. To ensure the generalizability of our research findings, we adopt the commonly recognized experimental settings proposed by Hu et al. [24], with 5 layers, 300 hidden units in each of layer, and  $50\%$  dropout ratio.
- For 3D geometries, we employ the widely-used SchNet model [43] as the encoder. We set the hidden dimension and the number of filters in continuous-filter convolution to 128. The interatomic distances are measured with 50 radial basis functions, and we stack 6 interaction layers in SchNet.
- For fingerprints, we first use RDKit [44] to generate 1024-bit molecular fingerprints with a radius  $R = 2$ , which is roughly equivalent to the ECFP4 scheme [7]. Given the lack of established encoders for fingerprints, we conduct a comparison between two classic encoders, a single-layer MLP and a single-layer Transformer. Please refer to Appendix C.4 for the detailed results. According to the experiments, we choose the Transformer model [45] with 8 attention heads for modeling fingerprints. The bit embedding dimension is set to 64, and the hidden dimension is set to 300.
- For SMILES strings, we employ the same model architecture as the fingerprints to ensure a fair comparison. The only difference is the dictionary size, which will be discussed in Section 3.2.
Training details. We follow the settings proposed by Hu et al. [24]. The dataset is split randomly. For classification tasks, we employ an  $80\% / 10\% / 10\%$  partition for the training, validation, and test sets, respectively. Meanwhile, for regression tasks, the dataset is split into 110K molecules for training, 10K for validation, and another 10K for testing. All models are initialized using the Glorot initialization [46]. The Adam optimizer [47] is employed for training with a batch size of 256. For classification tasks, the learning rate is set at 0.001 and we opt against using a scheduler. For regression tasks, we align with the original experimental settings of SchNet, setting the learning rate to  $5 \times 10^{-4}$  and incorporating a cosine annealing scheduler.
Evaluation metrics. For HIV and MUV datasets, performance is measured using the Area Under the ROC-Curve (ROC-AUC), while report the performance on PCBA in terms of Average Precision (AP) —higher values in both metrics indicate better performance. When assessing quantum property predictions, the Mean Absolute Error (MAE) is used as the performance metric, with lower values indicating better accuracy.
3 Empirical Studies
In this section, we summarize our empirical studies on the aforementioned dimensions of MRL. Firstly, we show the neural scaling laws between model performance and data quantities. Then, we demonstrate the impact of four design dimensions of MRL: data modalities, the role of pretraining, data splits, and model capacity. Lastly, we study data pruning strategies for molecular graphs.
Figure 2: The general neural scaling laws of molecular representation learning. The first row displays the relationship between model performance with respect to varied data sizes in linear coordinates and the second row shows the same on a logarithmic scale.
3.1 General Neural Scaling Laws for Molecular Data
Early studies of both classical learning theory and neural scaling laws [48, 18] show that the test performance  $L(n)$  increases polynomially with the training data size  $n$ :
$$
L (n) = \delta \cdot n ^ {\alpha}, \tag {1}
$$
where  $\delta$  is the coefficient and  $\alpha$  is the exponent of the power law. We start our investigation on whether MRL also adheres to this power law relationship with the most common setting of graph-based MRL. Specifically, on the QM9 dataset, we study GINs on 2D graphs for classification tasks and SchNet on 3D geometric graphs for regression tasks. Figure 2 demonstrates the relationship of model performance with respect to the data size across all datasets, where the first row displays the curves in linear coordinates and the second row shows the same on a logarithmic scale.
Observation 1. It can be observed that figures in the second row generally exhibit a linear relationship, which suggests alignment with the aforementioned power law. Unlike previous findings in the NLP and CV domains [18], there is no apparent performance plateau within the range of datasets we explored. This pattern remains consistent in both low- and high-data regimes, which suggests that the performance of supervised MRL is highly predictable and consistently improves with increasing data quantity. We note that we observe the same phenomenon consistently when applying other modalities, as we will show later in Figure 3. Furthermore, we also validate the neural scaling laws by examining them on a large-scale dataset PCQM4Mv2 [49] from Open Graph Benchmark (OGB) [50], as well as by employing two advanced 3D GNN encoders PaiNN [51] and SphereNet [52] on QM9. Readers of interest may refer to Appendix C.2 for detailed results and analysis.
3.2 Scaling Laws with Different Molecular Modalities
The choice of the most appropriate modalities in MRL remains a topic of continuous discussion. Also, there is a lack of fair comparisons regarding data efficiency across these modalities. In this section, we compare the learning behaviors of the three modalities: 2D graphs, SMILES strings, and fingerprints. To ensure a fair comparison, we maintain roughly equivalent parameter sizes across the modality encoders. Such configurations are selected to maximize the expressiveness of each model. Specifically, we employ a 5-layer GIN model for the 2D graph modality and a 1-layer transformer model for both SMILES strings and fingerprints. It is noting that transformers used for SMILES and fingerprints differ in vocabulary sizes: while the former has a size of 2 for fingerprints due to the binary nature of fingerprint strings, the latter adopts a more expansive vocabulary size of 7,924 for SMILES, following ChemBERTa [26]. We present the results illustrating model performance in relation to different data sizes for the three encoders in Figure 3, with anomalous performance explicitly marked in red.
Observation 2. In general, we observe that different modalities exhibit distinct learning behaviors. Among the three modalities, the graph modality consistently demonstrates superior data utilization efficiency across all three tasks, emerging as the most efficient modality for MRL. Fingerprints, characterized with a lower exponent, deliver a slightly lower performance improvement with equivalent
Figure 3: Neural scaling laws with three encoders: 2D graphs, SMILES strings, and fingerprints. Anomalous performance is explicitly marked in red.
data increments. SMILES strings, on the contrary, are the least data efficient modalities. On the MUV dataset, we even notice counter-intuitive performance degradation that breaks the power law  ${}^{3}$  . However, it is important to recognize that several studies have underscored the efficacy of language models pre-trained on large-scale SMILES string datasets, which exhibit remarkable performance in the downstream tasks [28, 26]. This points towards the superior transferability of the SMILES-based modality and we will leave this performance gap between pre-trained and those trained from scratch for future research.
3.3 Influence of Pre-training on the Scaling Laws
Molecular pre-training studies how to leverage unlabeled molecular data to improve the learned representations [53]. Previous studies have generally suggested that pre-training can induce a positive transfer to downstream tasks [33, 30, 29]. However, we argue that the existing evaluations rely on the entirety of available datasets. We hypothesize that when fine-tuning with datasets of varied sizes, the extent of this positive transfer might fluctuate. Therefore, in this section, we explore whether pre-training consistently leads to positive transfer with different sizes of fine-tuning datasets.
In the experiments, we focus on the state-of-the-art molecular pre-training model GraphMAE [33], which involves masking the atom type of partial atoms in the molecules, re-masking the encoded atom representations from the backbone model, and eventually using a decoder model to reconstruct the original atom features. The GraphMAE model is first pre-trained on PCQM4Mv2 and then fine-tuned on three classification tasks. The results in Figure 4 compares the performance curve with and without pre-training on downstream datasets of varied sizes.
Observation 3. Our empirical results reveal that, while pre-trained models still adhere to a power law as their downstream data size varies, they differ from models trained from scratch by having a higher intercept and a lower exponent in the logarithm-scaled plots. This implies that while pre-trained models benefit from better initializations, they demonstrate reduced data efficiency when fine-tuned on downstream datasets.
Our findings also indicate that the benefits of pre-training are particularly pronounced when fine-tuning with smaller datasets and such advantages start to diminish as the size of downstream datasets increases. This is evidenced by the PCBA dataset, where we observe a crossover in performance at a data scale of 40K. Beyond this threshold, the performance of pre-trained models starts to fall behind its non-pretrained counterpart. This behavior suggests the phenomenon of parameter ossification [54] in pre-trained models, suggesting that pre-training can inadvertently "freeze" the model weights in a
<sup>3</sup>We confirm that the performance degradation is a common phenomenon rather than an outcome of limited model capacity by further altering the number of model layers. Readers can refer to Appendix C for more discussions.
Figure 4: Comparison of scaling laws with and without pre-training on varied downstream data sizes.
way that limits their adaptability to the fine-tuning distribution, especially in large-data scenarios. Consequently, we believe the role and implications of pre-training in MRL warrant careful evaluation.
3.4 Influence of Dataset Splits on the Scaling Laws