SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hanchen Wang $^{1\star}$  Jean Kaddour $^{2\star}$  Shengchao Liu $^{3,4\star}$  Jian Tang $^{3,5,6}$  Joan Lasenby $^{1}$  Qi Liu $^{7}$ $^{1}$ Cambridge  $^{2}$ UCL  $^{3}$ MILA  $^{4}$ UdeM  $^{5}$ HEC  $^{6}$ CIFAR  $^{7}$ HKU  $^{\star}$  Equal Contribution
Graph Self-Supervised Learning (GSSL) provides a robust pathway for acquiring embeddings without expert labelling, a capability that carries profound implications for molecular graphs due to the staggering number of potential molecules and the high cost of obtaining labels. However, GSSL methods are designed not for optimisation within a specific domain but rather for transferability across a variety of downstream tasks. This broad applicability complicates their evaluation. Addressing this challenge, we present "Molecular Graph Representation Evaluation" (MOLGRAPH), generating detailed profiles of molecular graph embeddings with interpretable and diversified attributes. MOLGRAPH offers a suite of probing tasks grouped into three categories: (i) generic graph, (ii) molecular substructure, and (iii) embedding space properties. By leveraging MOLGRAPH to benchmark existing GSSL methods against both current downstream datasets and our suite of tasks, we uncover significant inconsistencies between inferences drawn solely from existing datasets and those derived from more nuanced probing. These findings suggest that current evaluation methodologies fail to capture the entirety of the landscape.
Learning neural embeddings of molecular graphs has become of paramount importance in computer-aided drug discovery [1, 2]. For instance, a molecular property prediction (MPP) model can expedite and economise the design process by reducing the need for synthesising and measuring molecules. Thereby, such models can be immensely useful in the hit-to-lead and early lead optimisation phase of a drug discovery project [3]. However, obtaining labels of molecule properties is expensive and time-consuming, especially since the size of potential pharmacologically active molecules is estimated to be in the order of  $10^{60}$  [4, 5].
Graph Self-Supervised Learning (GSSL) paves
the way for learning molecular graph embeddings without human annotations that are transferable to various downstream datasets. Unfortunately, the evaluation of such general-purpose embeddings is fundamentally complex. Different proxy objectives will place different demands on them, and no single downstream dataset can be definitive. Moreover, many of the previously proposed GSSL works are disconnected in terms of the tasks they target and the datasets they use for evaluation, making direct comparison difficult.
Figure 1: Overview of MOLGRAPH. Given molecular graphs, we train GNNs to predict SSL proxy objectives. We then extract embeddings of (possibly unseen) graphs using pre-trained models, which form the inputs for probe models, trained and evaluated on the designed metrics.
Contributions. Our goal is to unbiasedly evaluate molecular graph embeddings obtained by GSSL methods on existing downstream tasks and a new suite of probe tasks (Fig. 1). We summarised some key findings based on a total of 90,918 probe models and 1,875 pre-trained GNNs.
1) On MPP tasks, we observe every GSSL method can introduce substantial performance gains. Yet, there is a significant difference in the rank depending on whether we fine-tune the pre-trained network on the downstream dataset or not. Also, the pre-training configurations for obtaining the optimal embeddings or initialisation for fine-tuning are different; see - Finding 2.
2) Several discrepancies between MPP tasks and MOLGRAPH evaluate demonstrate how the latter complements GSSL evaluation with novel insights fostering future work:
- While embeddings from a randomly initialised GNN perform poorly on MPP tasks, they sometimes outperform on topological properties (which can be useful in some molecular tasks), indicating that GSSL methods are not universally better, see - Finding 4.
- While contrastive GSSL methods perform better than most other methods on MPP tasks, it might attribute to their superiority in identifying the crucial substructures, see - Finding 8.
- In contrast to previous work [6], we find that feature distribution uniformity is not always a strong indicator for MPP performance. For instance, maximising the mutual information between multi-scale graph representations (INFOGRAPH) results in the most uniform distributions, yet, it ranks  $7^{\text{th}}$  among the 9 GSSLs on the downstream MPP tasks, see - Finding 9.
Graph SSL (GSSL) can be divided into contrastive and generative methods [7-9]. Contrastive GSSL [10-12] construct multiple views of the same graph via augmentations and then learn embeddings by aligning positive samples against negative ones. Generative GSSL [11, 13-15] yields embeddings by reconstructing input graphs. Zhu et al. [16] conduct an empirical analysis of contrastive GSSL methods and their components. In contrast, we investigate both generative and contrastive GSSL methods and propose a novel suite of tasks to probe the learned embeddings' attributes.
Probe models and benchmarks on graphs. Probe models, trained exclusively on embedding vectors from pre-trained models, serve as an effective tool for evaluating the quality of learned embeddings [17]. Their effectiveness has been demonstrated across various domains such as language [18-23], vision [24-28], relational tables [29], and science [30-32]. While there exist benchmarks for graph learning [33-37], applying probe models to GSSL remains an unexplored frontier.
Graph. A graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  consists of a set of nodes  $\nu$  and edges  $\mathcal{E}$ . In molecular graphs, nodes are atoms, and edges are bonds. We use  $\pmb{x}_u$  and  $\pmb{x}_{uv}$  to denote the feature of node  $u$  and the bond feature between nodes  $[u,v]$ , respectively. For notation simplicity, we use an adjacency matrix  $\mathbf{A}\in \mathbb{R}^{|\mathcal{V}|\times |\mathcal{V}|}$  to represent the graph, where  $\mathbf{A}[u,v]\neq 0$  if the nodes  $(u,v)$  are connected.
GNN. Graph neural networks (GNNs) give rise to learning molecular graph embeddings [14, 38-41]. A prototypical GNN relies on messaging passing [39], which updates atom-level embeddings based on their neighbourhoods. Given an input atom  $\pmb{h}_u^0 = \pmb{x}_u$ , we compute its embedding by:
$$
\boldsymbol {m} _ {u} ^ {t + 1} = \sum_ {v: \mathbf {A} [ u, v ] \neq 0} M _ {t} \left(\boldsymbol {h} _ {u} ^ {t}, \boldsymbol {h} _ {v} ^ {t}, \boldsymbol {x} _ {u v}\right) \quad \boldsymbol {h} _ {u} ^ {t + 1} = U _ {t} \left(\boldsymbol {h} _ {u} ^ {t}, \boldsymbol {m} _ {u} ^ {t + 1}\right) \tag {1}
$$
where  $M_{t}$  and  $U_{t}$  are the "message" functions and "vertex update" functions, respectively. Repeating message passing for  $T$  steps, the embedding of each atom contains their  $T$ -hop neighbourhood information. A readout function  $R$  is then used to pool node-level embeddings for graph-level representations:  $\hat{y} = R(\{\pmb{h}_u^\top \mid u \in \mathcal{V}\})$ . Following previous GSSL methods on molecular graphs, we adopt the Graph Isomorphism Network (GIN) [42] as the backbone model and incorporate edge features during message passing following [11].
Pre-Training. We inspect nine GSSL methods (1,875 configurations in total): EDGEPred [13], INFOGRAPH [10], GPT-GNN [15], ATTRMASK [11], CONTEXTPred [11], GROVER [43], GRAPHCL [12], JOAO [44], and GRAPHMVP [45]. We use all qualified molecules (around 0.33 million, i.e., leave out the molecules that appeared in downstream datasets) from the GEOM dataset [46] to pre-train the GIN backbone. As many of these pre-training methods are not primarily designed for molecular graphs, we perform the grid search over the hyperparameter space and save
Table 1: Evaluating GSSL methods on molecular property prediction tasks. For each downstream dataset, we report the mean and standard deviation of the ROC-AUC scores over three random scaffold splits. The best and second best scores are marked **bold** and **bold**, respectively. The performance scores are based on the fixed pre-trained embeddings with linear probe models, we also report the average ROC-AUC scores with fine-tuned pre-trained GNN on MPP tasks ("Avg (FT)"). For each pre-training method, we report the highest scores in the table and their corresponding hyperparameter configurations in Tables 7 and 9 in Appendix B.
BBBP	Tox21	ToxCast	Sider	ClinTox	MUV	HIV	Bace	Avg	Avg (FT)
# Molecules	2,039	7,831	8,575	1,427	1,478	93,087	41,127	1,513	-	-
# Tasks	1	12	617	27	2	17	1	1	-	-
RANDOM	50.7 ±2.5	64.9 ±0.5	53.2 ±0.3	53.2 ±1.1	63.1 ±2.3	62.1 ±1.3	66.1 ±0.7	63.4 ±1.8	59.60	66.16
EDGE PRED	54.2 ±1.0	66.2 ±0.2	54.4 ±0.1	56.1 ±0.1	65.4 ±5.0	59.5 ±0.9	73.6 ±0.4	71.4 ±1.2	62.59	68.16
ATTRMASK	62.7 ±2.7	65.7 ±0.8	56.1 ±0.2	58.3 ±1.5	61.9 ±6.4	60.9 ±1.8	65.5 ±1.4	64.8 ±2.6	61.99	69.20
GPT-GNN	62.0 ±0.9	64.9 ±0.7	55.4 ±0.2	55.3 ±0.8	55.0 ±5.1	61.2 ±1.5	71.2 ±1.5	61.0 ±1.2	60.74	67.58
INFOGRAPH	65.9 ±0.6	65.8 ±0.7	54.6 ±0.1	57.2 ±1.0	61.4 ±4.8	63.9 ±1.9	71.4 ±0.6	67.4 ±4.9	63.44	68.92
CONT.PRED	55.5 ±2.0	67.9 ±0.7	54.0 ±0.3	57.1 ±0.5	67.4 ±4.3	60.5 ±0.9	66.2 ±1.5	54.4 ±3.2	60.36	69.40
GROVER	67.0 ±0.3	63.9 ±0.3	53.6 ±0.4	59.9 ±1.7	65.0 ±6.4	62.7 ±1.4	67.8 ±1.0	69.0 ±4.7	63.62	69.97
GRAPHCL	64.7 ±1.7	69.1 ±0.5	56.2 ±0.2	59.5 ±0.9	60.8 ±3.0	60.6 ±1.8	72.5 ±1.4	77.0 ±1.7	65.04	70.33
JOAO	66.1 ±0.8	68.1 ±0.2	55.1 ±0.4	58.3 ±0.3	65.3 ±6.1	62.4 ±1.2	73.8 ±1.2	71.1 ±0.8	65.05	69.75
GRAPHMVP	69.2 ±1.8	63.8 ±0.3	55.5 ±0.3	58.6 ±0.4	58.7 ±1.9	63.8 ±1.3	68.6 ±1.0	73.3 ±4.7	63.92	70.06
the optimal settings. For these nine GSSL methods, we have pre-trained 1,875 GNNs with different configurations, as elaborated in Appendix B. We extract embeddings using the pre-trained weights, select the optimal hyperparameter sets based on their downstream MPP performance and use these optimal embeddings for further probing tasks.
Probe. We use probe models [18] to study whether self-supervised learned embeddings encode helpful structural information about graphs. Concretely, we extract embeddings from a pre-trained GNN and train a linear model to predict the probe tasks with node and graph embeddings as inputs. As the first work that designs probe methods on graph embeddings, we follow previous works on computer vision and natural language processing. We mainly compute and compare the quality of pre-trained embeddings using linear probe models. We have also experimented MLPs with one hidden layer as the probe models, as this architecture is utilised in some previous works. We observe similar findings with both probe architectures and reported the results of MLP probes in Appendix B.4. We use scaffold split to partition data into  $80\% / 10\% / 10\%$  for the training/validation/testing set. The training procedure runs for 100 epochs with a fixed learning rate of 0.001. We report the test results based on the best validation scores. To account for statistical significance, we average all experimental results over three independent runs. We find that different data splits are the primary cause for performance variations ( $\sim 2\%$ ), instead of initialising probe models with different random seeds ( $< 0.01\%$ ).
4 Benchmarking GSSL on MPP
We first conduct a rigorous empirical investigation of the GSSL methods' effectiveness in predicting the biochemical properties of molecules. Following previous work [11, 12], we consider eight molecular datasets consisting of 678 binary property prediction tasks [47, 48]. Unless explicitly stated otherwise, we extract the node/graph embeddings from the last GNN layer. We devise two settings: (i) fixed embeddings, where we train the probe models with fixed embeddings extracted from pre-trained GNNs; (ii) fine-tuned embeddings ("FT"), where we update weights of both the pre-trained GNNs and the probe models. Setting (i) follows the procedures in previous probing literature, while (ii) is widely utilised as the "pre-training, then fine-tuning" paradigm. We use Adam optimiser with no weight decay, set the batch size as 256, and apply identical pre-processing procedures for all experiments.
Findings. Table 1 notes the results, and we summarise the following findings, some of which contrast with those drawn from the concurrent study [49].
1) All GSSL methods perform better than RANDOM. By carefully exploring the pre-training hyperparameters, all GSSLs substantially improve the MPP tasks for both fixed and fine-tuned embeddings. Contrastive-based GSSL methods (i.e., GRAPHCL, JOAO and GRAPHMVP) achieve the overall best performance. As [49] declares that molecular graph pretraining is ineffective; however, we find that their conclusions are based on a few selected finetuning datasets and fixed
pre-training hyperparameters. We further observe that such improvements will reduce when the number of molecules in downstream datasets increases. Specifically, for MUV [50], a dataset designed for validating virtual screening (used in drug discovery to find how likely molecules that bind to a drug target), the average performance gain brought by pre-training is  $-0.3\%$ ; while for BBBP, it is  $12.3\%$ . The number of molecules in BBBP is only  $2\%$  of the MUV's.
2) Rankings differ between probing and fine-tuning. The rank correlation between the fixed and fine-tuned embeddings is 0.77 (p-value=9e-4), indicating that we cannot utilise the rank of fixed embeddings as a definite indicator for fine-tuning performance, though they are positively correlated. Part of this observation has been spotted in a study on masked visual transformers [25]. In the context of molecular property prediction, embeddings pre-trained with JOAO achieve the best score with fixed scenarios but perform the fourth after end-to-end fine-tuning. The reason is unclear and should be investigated by future work.
3) The optimal sets of pre-train hyperparameters for fixed and fine-tuned embeddings vary. We observe that the optimal pre-training hyperparameters on fixed and fine-tuned embeddings differ. Only two out of nine GSSLs (INFOGRAPH and EDGEPRED) share the same set of optimal parameters, as detailed in Tables 7 and 9 in Appendix B. This suggests that probing the fixed embeddings might not truly reflect pre-trained models' performance on downstream MPP tasks, as it ignores the consequent improvements induced by fine-tuning. In Fig. 2, we visualised the hyperparameter space of the ATTRMASK pre-trainer, the local minima in the hyperparameter space distribute differently. As shown in Fig. 2 also Table 9, the best pre-training configuration for probing is "mask rate=0.85 and learning rate=1e-4", while in terms of fine-tuning scores, the optimal setting is "mask rate=0.50 and learning rate=5e-4". Also, it can be inferred that the optimal pre-training hyperparameters for different pre-training datasets vary; therefore, using reported hyperparameters without carefulness and concluding "graph pretraining is ineffective in molecular domain" is not convincing [49].
5 Molecular graph representation evaluation
The goal of GSSL for molecular graphs is to obtain embeddings that capture generic information about the molecule and its properties. However, there is no free lunch [51]: different training objectives optimise for different properties, and evaluating the extracted embeddings on only a handful of downstream datasets does not provide the whole picture (as we confirm empirically in Sec. 6). Also, from Sec. 4, we know the probing and fine-tuning performance are positively correlated, yet their optimal pre-training configurations are largely diverging. In the Appendix, we also provide results on the worst pre-training configurations, some of which
Figure 2: Hyperparameter space of ATTR-MASK (mask rate  $\times$  learning rate), coloured by MPP test scores. Above: fixed; below: fine-tuned. Note the colour bars are different.
cause negative transfer due to initialising the encoders into local bad minima. Investigations are required to understand what kind of property makes the pre-trained encoders differ.
To this end, we propose MOLGRAPH, which encompasses a variety of carefully-selected probe tasks, categorised into three classes: (i) generic graph properties, (ii) molecular substructure properties and (iii) embedding space properties. In the upcoming subsections, we explain the tasks in more detail and why they are essential for molecular graph embeddings.
5.1 Generic graph properties
Topological property statistics are often used as features in machine learning pipelines on graphs that do not rely on neural networks [52]. Based on their scale, they can be divided into {node-, pair-, and graph-} level statistics. For molecular graphs, topological metrics have been widely used as molecular descriptors in cheminformatics for decades [53-56], metrics at different scales will facilitate different tasks.
Node-level statistics accompany each node with a local topological measure, which could be used as features in node classification [52]. Concretely, node-level information such as degree [57] can reflect the reaction centres [58]; thus, it can aid in discovering chemical reactions [59].
- Node degree  $(d_u)$  counts the number of edges incident to node  $u$ :  $d_u = \sum_{v \in V} \mathbf{A}[u, v]$ .
- Centrality  $(e_u)$  represents a node's importance. The eigenvector centrality is determined by a relation proportional to the average centrality of its neighbours:  $e_u = \left( \sum_{v \in V} \mathbf{A}[u, v] e_v \right) / \lambda, \forall u \in \mathcal{V}$ .
- Clustering coefficient  $(c_{u})$  measures how tightly clustered a node's neighbourhood is:  $c_{u} = (|(v_{1}, v_{2}) \in \mathcal{E}: v_{1}, v_{2} \in \mathcal{N}(u)|) / d_{u}^{2}$ , i.e., the fraction of closed triangles in neighbourhood [60].
Graph-level statistics summarise global topology information and are helpful for graph classification tasks. For molecules, graph-level statistics can be used, e.g., to classify a molecule's solubility [61]. We briefly describe their intuitions; formal definitions can be found, e.g., in [52].
- Diameter is the maximum distance between the pair of vertices (i.e., longest path in a molecule).
- Cycle basis is a set of simple cycles that forms a basis of the graph cycle space. It is a minimal set that allows every even-degree subgraph to be expressed as a symmetric difference of basis cycles.
- Connectivity is the minimum number of elements (nodes or edges) that need to be removed to separate the remaining nodes into two or more isolated subgraphs.
- Assortativity measures the similarity of connections in the graph with respect to the node degree. It can be seen as the Pearson correlation coefficient of degrees between pairs of linked nodes.
Pair-level statistics quantify the relationships between nodes (atoms), which is vital in molecular modelling. For example, molecular docking techniques aim to predict the best matching binding mode of a ligand to a macro-molecular partner [62]. For predicting such binding compatibility, connectivity and distance awareness (how close a pair of atoms can be) are important. In our implementation, we randomly select a fixed number (i.e., 10) of atom pairs from each molecular graph. These pairs are then categorised based on their originating molecule, ensuring all pairs from a single molecule are designated to a singular split: either train, validation, or test.
- Link prediction tests whether two nodes are connected or not, given their embeddings and inner products. Based on the principle of homophily, it is expected that embeddings of connected nodes are more similar compared to disconnected pairs:
$$
\mathbf {S} _ {\text {L i n k}} [ u, v, \mathbf {x} _ {u} ^ {T} \mathbf {x} _ {v} ] = \mathbb {1} _ {\mathcal {N} (u)} (v) \tag {2}
$$
- Jaccard coefficient seeks to quantify the overlap between neighbourhoods while minimising the biases induced by node degrees [63]:
$$
\mathbf {S} _ {\text {J a c c a r d}} [ u, v ] = | \mathcal {N} (u) \cap \mathcal {N} (v) | / | \mathcal {N} (u) \cup \mathcal {N} (v) | \tag {3}
$$
- Katz index is a global overlap statistic defined by the number of paths between a pair of nodes:
$$
\mathbf {S} _ {\text {K a t z}} [ u, v ] = \sum_ {i = 1} ^ {\infty} \beta^ {i} \mathbf {A} ^ {i} [ u, v ] \tag {4}
$$
where  $\beta \in \mathbb{R}^{+}$  determines the weight between short and long paths.  $\beta < 1$  reduces the weight of long paths, in implementations we set  $\beta = 1$  to give all paths equal importance.
5.2 Molecular substructure properties
Table 2: ROC-AUC scores of classifiers predicting molecular properties.
LINEAR REGRESSION	RANDOM FOREST	XGBOOST	RANDOM (FIX/FT)	JOAO (FIX)	GRAPHCL (FT)
59.91	61.95	62.31	59.60 / 66.16	65.05	70.33
Molecular substructures often serve as reliable indicators of biochemical properties [64-67]. For instance, molecules with benzene rings typically share consistent physical properties, such as solubility, as well as chemical characteristics like aromaticity [68].
Substructures. We investigate 24 substructures from three groups: rings (Benzene, Beta lactams, ..., Thiophene); functional groups (Amides, Amidine, ..., Urea); and redox active sites (Allylic). We provide chemical knowledge on how they relate with molecular properties in Appendix D.
How predictive are substructures? To demonstrate that substructures are quite predictive of molecular properties, we utilise counts of substructures within a molecular graph as the input for
classic ML methods (linear regression, random forest, and XGBoost) to predict the molecular properties on eight downstream datasets. Table 2 shows the results. For ease of comparison, we add the performance of RANDOM, GRAPHCL (FIX), and JOAO (FT) from Table 1. Notably, even basic models trained exclusively on substructures yield performance akin to the randomly trained GNN baseline. This underscores the profound correlation between substructures and MPP task efficacy.
5.3 Embedding space properties
Beyond MPP metrics, we assess domain-agnostic properties of the embedding space produced by pre-trained graph encoders. These properties correlate positively with downstream generalisation [6]. Hence, they can serve as proxies for embedding degeneration, especially in label-scarce scenarios. We explore three such properties in MOLGRAPH:
- Alignment quantifies how similar produced embeddings are for similar samples [6]. Ideally, two samples with the same (or very similar) semantics should be mapped to nearby features, thus mostly invariant to unneeded noise factors. To examine alignment, we construct positive and negative molecule pairs. Positive pairs in a dataset are those that share identical molecular properties, whereas negative pairs are those that differ in their properties. A better alignment represents a better nearest neighbourhood formulation, which has been especially useful for tasks such as compound potency prediction [70].
- Uniformity measures how uniformly the embeddings are distributed on the unit hypersphere [6]. A more uniformly distributed embedding space is expected to be with better generalisation under some mild assumptions.
- Dimensional collapse refers to the problem of embedding vectors spanning a lower-dimensional subspace instead of the entire available embedding space [69, 71]. Following Jing et al. [69], one simple way to test the occurrence of dimension collapse is to inspect the number of non-zero singular values of a matrix stacking the embedding vectors  $\mathbf{Z} = \{\mathbf{z}_i\}_{i=1}^N$ .
6.1 Generic graph properties
For node- and graph-level topological properties, we benchmark GSSL methods on all the molecular graphs from each dataset; for pair-level metrics, we bootstrap  $10k$  node pairs from each dataset for evaluation. The reported
test scores are averaged over three runs. We plot the distribution of these metrics in Appendix E.
Figure 3: Embedding space property.
(a) Alignment [6]
(b) Uniformity [6]
(c) Dimensional collapse [69]
Findings. Table 3 shows the results, and we summarise these findings.