SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kha-Dinh Luong, Ambuj Singh
Department of Computer Science  
University of California, Santa Barbara  
Santa Barbara, CA 93106  
{vluong, ambuj}@cs.ucsb.edu
Property prediction on molecular graphs is an important application of Graph Neural Networks (GNNs). Recently, unlabeled molecular data has become abundant, which facilitates the rapid development of self-supervised learning for GNNs in the chemical domain. In this work, we propose pretraining GNNs at the fragment level, a promising middle ground to overcome the limitations of node-level and graph-level pretraining. Borrowing techniques from recent work on principal subgraph mining, we obtain a compact vocabulary of prevalent fragments from a large pretraining dataset. From the extracted vocabulary, we introduce several fragment-based contrastive and predictive pretraining tasks. The contrastive learning task jointly pretrains two different GNNs: one on molecular graphs and the other on fragment graphs, which represents higher-order connectivity within molecules. By enforcing consistency between the fragment embedding and the aggregated embedding of the corresponding atoms from the molecular graphs, we ensure that the embeddings capture structural information at multiple resolutions. The structural information of fragment graphs is further exploited to extract auxiliary labels for graph-level predictive pretraining. We employ both the pretrained molecular-based and fragment-based GNNs for downstream prediction, thus utilizing the fragment information during finetuning. Our graph fragment-based pretraining (GraphFP) advances the performances on 5 out of 8 common molecular benchmarks and improves the performances on long-range biological benchmarks by at least  $11.5\%$ . Code is available at: https://github.com/lvkd84/GraphFP.
The rise of Graph Neural Networks (GNNs) has captured the interest of the computational chemistry community [11]. Representing molecules and chemical structures as graphs is beneficial because the topology and connectivity are preserved and can be directly analyzed during learning [13], leading to state-of-the-art performance. However, there is a problem in that while highly expressive GNNs are data-hungry, the majority of molecular datasets available are considerably small, posing a critical challenge to the generalization of GNNs. Following the steps that were taken in other domains such as text and images, a potential solution is to pretrain GNNs on large-scale unlabeled data via self-supervised pretraining [3, 5]. How to do so effectively is still an open question since graphs are more complex than images or text, and the chemical specificity of molecular graphs needs to be considered when designing graph-based pretraining tasks.
A large number of pretraining methods have been proposed for molecular graphs [38, 40]. For example, [15] uses node embeddings to predict attributes of masked nodes, [17, 27] predicts graph structural properties, and [16] generatively reconstructs graphs via node and edge predictions. Another
popular direction is contrastive learning [26, 32, 35, 45, 46], in which multiple views of the same graph are mapped closer to each other in the embedding space. These views are commonly generated via augmentations to the original graphs [35, 45, 46]. However, without careful adjustments, there is a great risk of violating the chemical validity or changing the chemical properties of the original molecular graph after applying the augmentations. Other methods such as [23, 31] avoid this problem by contrasting 2D-graphs against 3D-graphs. Unfortunately, privileged information such as 3D coordinates is often expensive to obtain. Node versus graph contrastive learning has also been investigated [32], however, this approach may encourage over-smoothing among node embeddings.
In general, the majority of works pretrain embeddings at either node-level or graph-level. Node-level pretraining may be limited to capturing local patterns, neglecting the higher-order structural arrangements while graph-level methods may overlook the finer details. As such, motif-based or fragment-based pretraining is a new direction that potentially overcomes these problems [27, 48, 49]. Still, existing fragment-based methods use either suboptimal fragmentation or fragmentation embeddings. GROVER [27] predicts fragments from node and graph embeddings, however, their fragments are k-hop subgraphs that cannot account for chemically meaningful subgraphs with varying sizes and structures. MICRO-Graph [48] contrastively learns subgraph embeddings versus graph embeddings; however, these embeddings may not effectively capture global patterns. MGSSL [49] utilizes graph topology via depth-first search or breadth-first search to guide the fragment-based generation of molecules; however, their multi-step fragmentation may overly decompose molecules, thereby losing the ability to represent higher-order structural patterns.
In this paper, we propose GraphFP, a novel fragment-level contrastive pretraining framework that captures both granular patterns and higher-order connectivity. For each molecule, we obtain two representations, a molecular graph and a fragment graph, each of which is processed by a separate GNN. The molecular GNN learns node embeddings that capture local patterns while the fragment GNN learns fragment embeddings that encode global connectivity. The GNNs are jointly trained via a contrastive task that enforces consistency between the fragment embedding and the aggregated embedding of the corresponding atoms from the molecular graphs. Unlike previous work, we pretrain the fragment-based aggregation of nodes instead of individual nodes. With the fragment embeddings encoding the global patterns, contrasting them with the aggregation of node embeddings reduces the risk of over-smoothing and allows flexibility in learning the appropriate node latent space. Since aggregation of nodes in this context can also be considered a kind of fragment representation, our framework essentially contrasts views of fragments, each of which captures structural information at a different scale. Moreover, as the molecular graphs and the fragment graphs are chemically faithful, our framework requires no privileged information or augmentation.
Exploiting the prepared fragment graphs, we further introduce two predictive pretraining tasks. Given a molecular graph as input, the molecular GNN is trained to predict structure-related labels extracted from the corresponding fragment graph. These tasks enhance the structural understanding of the molecular GNN and can be conducted together with contrastive pretraining.
To generate a vocabulary of molecular fragments, we utilize Principal Subgraph Mining [20] to extract optimized prevalent fragments that span the pretraining dataset. Compared to fragmentations used in previous works, this approach can produce a concise and diverse vocabulary of fragments, each of which is sufficiently frequent without sacrificing fragment size.
We evaluate GraphFP on benchmark chemical datasets and long-range biological datasets. Interestingly, both the molecular and the fragment GNNs can be used in downstream tasks, providing enriched signals for prediction. Empirical results and analysis show the effectiveness of our proposed methods. In particular, our pretraining strategies obtain the best results on 5 out of 8 common chemical benchmarks and improve performances by  $11.5\%$  and  $14\%$  on long-range biological datasets.
2.1 Representation Learning on Molecules
Representation learning on molecules has made use of fixed hand-crafted representations such as descriptors and fingerprints [30, 43], string-based representations such as SMILES and InChI [22, 28], and molecular images [47]. Currently, state-of-the-art methods rely on representing molecules as graphs and couple them with modern graph-based learning algorithms like GNNs [11]. Depending
on the type of compounds (small molecules, proteins, crystals), different graph representations can be constructed [11, 39, 44]. Compared to graph-based representation, molecular descriptors and fingerprints cannot encode structural information effectively while string-based representations require an additional layer of syntax which may significantly complicate the learning problem.
Given a graph  $G = (V, E)$  with node attributes  $x_v$  for  $v \in V$  and edge attributes  $e_{uv}$  for  $(u, v) \in E$ , GNNs learn graph embedding  $h_G$  and node embedding  $h_v$ . At each iteration, a node updates its embedding by gathering information from its neighborhood, including both the neighboring nodes and the associated edges. This process often involves an aggregating function  $M_k$  and an updating function  $U_k$  [11]. After the  $k$ -th iteration (layer), the embedding of node  $v$  is,  $h_v^{(k)} = U_k(h_v^{(k-1)}, M_k(\{(h_v^{(k-1)}, h_u^{(k-1)}, e_{uv}) | u \in N(v)\}))$ , where  $N(v)$  is the set of neighbors of  $v$  and  $h_v^{(0)} = x_v$ . The aggregating function  $M_k$  pools the information from  $v$ 's neighbors into an aggregated message. Next, the updating function  $U_k$  updates the embedding of  $v$  based on its previous embedding  $h_v^{(k-1)}$  and the aggregated message. An additional readout function  $R$  combines the final node embeddings into a graph embedding  $h_G = R(\{h_v^{(K)} | v \in V\})$ , where  $K$  is the number of iterations (layers). Since there is usually no ordering among the nodes in a graph, the readout function  $R$  is often chosen to be order-invariant [41]. Following this framework, a variety of GNNs has been proposed [12, 34, 41], some specifically for molecular graphs [9, 29].
2.3 Pretraining on Graph Neural Networks
To alleviate the generalization problem of graph-based learning in the chemical domain, graph pretraining has been actively explored [15, 23, 26, 31, 32, 35, 45, 46] in order to take advantage of large databases of unlabeled molecules [10].
In terms of the pretraining paradigm, existing methods can be categorized into being predictive [15, 27, 37], contrastive [23, 26, 31, 35, 45, 46], or generative [16, 49]. Predictive methods often require moderate to large labeled datasets for pretraining. In the case of an unlabeled large dataset, chemically or topologically generic labels can be generated. Despite the simple setup, predictive methods are prone to negative transfer [15]. On the other hand, contrastive methods aim to learn a robust embedding space by using diverse molecular views [23, 31, 35, 45, 46]. Generative methods intend to learn the distribution of the components constituting molecular graphs [16, 49].
At the pretraining level, methods can be sorted into node-level [16, 27], graph-level [15, 23, 26, 31, 35, 45, 46], and more recently motif-level [27, 48, 49]. Node-level methods only learn chemical semantic patterns at the lowest granularity, limiting their ability to capture higher-order molecular arrangements, while graph-level methods may miss the granular details. Motif-based or fragment-based pretraining has emerged as a possible solution to these problems [27, 48, 49]; however, existing methods use suboptimal fragmentation or suboptimal fragment embeddings. In this work, we learn fragment embeddings that effectively capture both local and global topology. We exploit the fragment information at every step of the framework, from pretraining to finetuning.
In this section, we introduce the technical details of our pretraining framework GraphFP. We begin with discussing molecular fragmentation and fragment graph construction based on an extracted vocabulary. Then, we describe our fragment-based pretraining strategies, including one contrastive task and two predictive tasks. Finally, we explain our approach for combining pretraining strategies and the pretrained molecular-based and fragment-based GNNs for downstream predictions.
3.1 Molecule Fragmentation
Graph fragmentation plays a fundamental role in the quality of the learning models because it dictates the global connectivity patterns. Existing methods rely on variations of rule-based procedures such as BRICS [4] or RECAP [21]. Though chemistry-inspired, the extracted vocabulary is often large, to the order of the size of the pretraining dataset, and contains unique or low-frequency fragments,
Figure 1: Fragment-based contrastive pretraining framework.  $\mathrm{GNN}_M$  processes molecular graphs while  $\mathrm{GNN}_F$  processes fragment graphs. The fragment-based pooling function FRAGPOOL aggregates node embeddings into a combined embedding that forms a positive contrastive pair with the corresponding fragment embedding. Notice that the two -OH groups, blue and light pink, are considered distinct. Therefore, the aggregated node embedding corresponding to the blue fragment and the embedding of the light pink fragment form a negative pair.
posing a significant challenge for pattern recognition. To overcome these drawbacks, some works further break down fragments [49], even to the level of rings and bonds [18]. However, reducing the sizes of fragments forfeits their ability to capture higher-order arrangements.
3.1.1 Principal Subgraph Extraction
In this work, we borrow the Principal Subgraph Mining algorithm from [20] to alleviate the above-mentioned problems. Given a graph  $G = (V,E)$ , a subgraph of  $G$  is defined as  $S = (\tilde{V},\tilde{E})$ , where  $\tilde{V} \subseteq V$  and  $\tilde{E} \subseteq E$ . In a vocabulary of subgraphs,  $S$  is a principal subgraph if  $\forall S'$  intersecting with  $S$  in any molecule, either  $S' \subseteq S$  or  $c(S') \leq c(S)$  where  $c(\cdot)$  counts the occurrences of a fragment among molecules. Intuitively, principal subgraphs are fragments with both larger sizes and more frequent occurrences. The algorithm heuristically constructs a vocabulary of such principal subgraphs via a few steps:
Initialize: Initialize the vocabulary with unique atoms.
Merge: For each molecular graph, for all pairs of overlapping fragments in the graph, merge the fragment in the pair and update the occurrences of the resulting combined fragment.
Update: Update the vocabulary with the merged fragment with the highest occurrence. Repeat the last two steps until reaching the predefined vocabulary size.
For a more detailed description, we refer readers to [20]. We investigate the effect of vocabulary sizes on the performance of pretrained models in Section 4.4 and Appendix C. Similarly, we investigate the effect of fragmentation strategies in Section 4.2.
3.1.2 Fragment-graph Construction
Given an extracted vocabulary of principal fragments, for each molecular graph, we construct a corresponding fragment graph. Let  $F = \{S^{(0)}, S^{(1)}, \dots, S^{(m)}\}$  be the fragmentation of a molecular graph  $G_{M} = (V_{M}, E_{M})$ , where  $S^{(i)} = (\tilde{V}^{(i)}, \tilde{E}^{(i)})$  is a fragment subgraph,  $\tilde{V}^{(i)} \cap \tilde{V}^{(j)} = \emptyset$  and  $\cup_{i=1}^{m} \tilde{V}^{(i)} = V_{M}$ . We denote the fragment graph as  $G_{F} = (V_{F}, E_{F})$ , where  $|V_{F}| = |F|$  and each node  $v_{F}^{(i)} \in V_{F}$  corresponds to a fragment  $S^{(i)}$ . An edge exists between two fragment nodes of  $G_{F}$  if there exists at least a bond interconnecting atoms from the fragments. Formally,
$E_{F} = \{(i,j)|\exists u,v,u\in \tilde{V}^{(i)},v\in \tilde{V}^{(j)},(u,v)\in E\}$ . For simplicity, in this work, we withhold edge features in the fragment graph. As a result, a fragment graph purely represents the higher-order connectivity between the large components within a molecule. Fragment node features are embeddings from an estimizable lookup table. Vocabulary and fragment graphs statistics are provided in Appendix C.
3.2 Fragment-based Contrautive Pretraining
Figure 1 illustrates our contrastive framework. We define two separate encoders,  $\mathrm{GNN}_M$  and  $\mathrm{GNN}_F$ .  $\mathrm{GNN}_M$  processes molecular graphs and produces node embeddings while  $\mathrm{GNN}_F$  processes fragment graphs and produces fragment embeddings. Since GNNs can capture structural information, the node embeddings encode local connectivity of neighborhoods surrounding atoms. Similarly, the fragment embeddings encode global connectivity and positions of fragments within the global context.
We apply contrastive pretraining at the fragment level. Contrastive pretraining learns robust latent spaces by mapping similar views closer to each other in the embedding space while separating dissimilar views. Learning instances are processed in pairs, in which similar views form the positive pairs while dissimilar views form the negative pairs. In this case, a positive pair consists of a fragment from a fragment graph and the collection of atoms constituting this fragment from the corresponding molecular graph. On the other hand, a fragment and any collection of atom nodes constituting a different fragment instance form a negative pair. Notice that different occurrences of the same type of fragment in the same or different molecules are considered distinct instances because when structural contexts are taken into account, the embeddings of these instances are dissimilar. The -OH groups in Figure 1 illustrate one such case.
To obtain the collective embedding of atom nodes corresponding to a fragment, we define a function FRAGPOOL(\cdot) that combines node embeddings. Contrasting the combined embedding of nodes against the fragment embedding allows flexibility within the latent representations of individual nodes. Intuitively, the learning task enforces consistency between fragment embeddings and collective node embeddings, incorporating higher-order connectivity information. Through the learning process, the information is optimally distributed among the nodes so that each node only holds a part of this collective knowledge. Arguably, such setup is semantically reasonable because nodes and fragments represent different orders of connectivity. A node does not necessarily hold the same structural information as the corresponding fragment. Instead, it is sufficient for a group of nodes to collectively represent such higher-order information. Essentially, the fragment embeddings and the collective embeddings of atom nodes can be considered different views of molecular fragments.
Extending the notations in 3.1.2, given a pretraining dataset of molecular graphs  $\mathcal{D}_M = \{G_M^{(1)}, G_M^{(2)}, \dots, G_M^{(N)}\}$ , we obtain a set of fragmentations  $\mathcal{F} = \{F^{(1)}, F^{(2)}, \dots, F^{(N)}\}$  and a set of corresponding fragment graphs  $\mathcal{D}_F = \{G_F^{(1)}, G_F^{(2)}, \dots, G_F^{(N)}\}$ , with  $G_M^{(i)} = (V_M^{(i)}, E_M^{(i)})$  and  $G_F^{(i)} = (V_F^{(i)}, E_F^{(i)})$ . More specifically,  $V_M^{(i)}$  is the set of atom nodes in the molecular graph  $G_M^{(i)}$  and  $V_F^{(i)}$  is the set of fragment nodes in the fragment graph  $G_F^{(i)}$ . For the  $i$ -th molecule, let  $H_M^{(i)} \in \mathbb{R}^{|V_M^{(i)}| \times d}$  and  $H_F^{(i)} \in \mathbb{R}^{|V_F^{(i)}| \times d}$  be the final node embeddings and fragment embeddings after applying GNN_  $M$  and GNN_  $F$ , respectively, i.e:
$$
H _ {M} ^ {(i)} = \operatorname {G N N} _ {M} \left(V _ {M} ^ {(i)}, E _ {M} ^ {(i)}\right) \tag {1}
$$
$$
H _ {F} ^ {(i)} = \operatorname {G N N} _ {F} \left(V _ {F} ^ {(i)}, E _ {F} ^ {(i)}\right) \tag {2}
$$
We further compute the fragment-based aggregation of node embeddings:
$$
H _ {A} ^ {(i)} = \operatorname {F R A G P O O L} \left(H _ {M} ^ {(i)}, F ^ {(i)}\right), \tag {3}
$$
where  $H_A^{(i)} \in \mathbb{R}^{|V_F^{(i)}| \times d}$  has the same dimensions as those of  $H_F^{(i)}$ . The  $r$ -th rows,  $h_{A,r}^{(i)} \in H_A^{(i)}$  and  $h_{F,r}^{(i)} \in H_F^{(i)}$ , from both matrices are embeddings of the same fragment from the original molecule. In our contrastive framework,  $h_{F,r}^{(i)}$  is the positive example for the anchor  $h_{A,r}^{(i)}$ . The negative learning examples are sampled from:
$$
\mathcal {X} _ {i, r} ^ {-} = \left\{H _ {F, q} ^ {(j)} \mid j \neq i \vee q \neq r \right\} \tag {4}
$$
Figure 2: Fragment-based predictive pretraining.  $\mathrm{GNN}_M$  processes molecular graphs and produces graph-level embeddings used for prediction. The right upper box shows unique fragments that exist in the input molecule while the right lower box shows the ground-truth structural backbone. A few other backbones are also visualized for comparison.
We minimize the contrastive learning objective based on the InfoNCE loss [25]:
$$
\mathcal {L} _ {C} = - \mathbb {E} _ {i, r} \left[ \frac {\exp \left(\left\langle h _ {A , r} ^ {(i)} , h _ {F , r} ^ {(i)} \right\rangle\right)}{\exp \left(\left\langle h _ {A , r} ^ {(i)} , h _ {F , r} ^ {(i)} \right\rangle\right) + \sum_ {h ^ {-} \sim \mathcal {X} _ {i , r} ^ {-}} \exp \left(\left\langle h _ {A , r} ^ {(i)} , h ^ {-} \right\rangle\right)} \right] \tag {5}
$$
3.3 Fragment-based Predictive Pretraining
We further define two predictive pretraining tasks of which labels are extracted from the properties and the topologies of the fragment graphs. Only the molecular GNN is pretrained in this case. The labels provided by the fragment graphs guide the graph-based pretraining of the molecular GNNs, encouraging the model to learn higher-order structural information.
Fragment Existence Prediction A multi-label prediction task that outputs a vocabulary-size binary vector indicating which fragments exist in the molecular graph. Thanks to the optimized fragmentation procedure [20] that we use, the output dimension is compact without extremely rare classes or fragments, resulting in more robust learning.
Fragment Graph Structure Prediction We predict the structural backbones of fragment graphs. The number of classes is the number of unique structural backbones. Essentially, a backbone is a fragment graph with no node or edge attributes. Graphs share the same backbone if their fragments are arranged similarly. For example, fragment graphs in which three fragments connect in a line correspond to the same structural backbone, a line graph with three nodes. This task also benefits from the optimized fragmentation. As the extracted fragments are sufficiently large, the fragment graphs are small enough that the enumeration of all unique backbones is feasible.
The predictive tasks pretrain the graph embedding  $h_G$ . In particular, one task injects local community information of fragments while the other task injects the arrangements of fragments into  $h_G$ . We train both tasks together, optimizing the objective  $\mathcal{L}_P = \mathcal{L}_{P_1} + \mathcal{L}_{P_1}$ , where  $\mathcal{L}_{P_1}$  and  $\mathcal{L}_{P_2}$  are the predictive objective of each task. The predictive tasks are illustrated in Figure 2.
We can combine the predictive pretraining with the contrastive pretraining presented in Section 3.2. We found parallel pretraining more effective than the sequential pretraining described in [15]. With  $\alpha$  being the weight hyperparameter, the joint pretraining objective is then:
$$
\mathcal {L} = \alpha \mathcal {L} _ {P} + (1 - \alpha) \mathcal {L} _ {C} \tag {6}
$$
3.4 Combining Models based on Fragment Graphs and Molecule Graphs
We further propose to utilize both the molecule encoder  $\mathrm{GNN}_M$  and the fragment encoder  $\mathrm{GNN}_F$  for downstream prediction. Following the procedure described in Section 3.1, we can easily fragment and extract molecular graph and fragment graph representations from any unseen molecule. Such convenience is not always the case. For example, existing works that contrast 2D and 3D molecular graphs [23, 31] can only finetune the 2D encoder since 3D information is expensive to obtain for unseen molecules. Other works that rely on graph augmentation [35, 45, 46] cannot utilize the
augmented views as signals for downstream prediction since they are not faithful representations of the original molecule. In contrast, the representations extracted by our framework are faithful views that do not require expensive information other than the vanilla molecular structure.
Let  $h_M$  and  $h_F$  be the graph embeddings produced by  $\mathrm{GNN}_M$  and  $\mathrm{GNN}_F$ , respectively. We obtain the downstream prediction by applying a fully-connected layer on the concatenation of  $h_M$  and  $h_F$ .