Motif-based Graph Representation Learning with Application to Chemical Molecules

Ideally, the motif vocabulary should be learned in an end-to-end fashion; however, this would incur an extremely high computational complexity. Therefore, we turned to a straightforward method for building a motif vocabulary that represents recurrent spatial patterns in training ARGs. First, we sampled a large number of subgraphs (e.g., $k$-hop neighborhoods) from the dataset. Each subgraph records its own center node. To make the extracted subgraphs cover local contexts as much as possible, we reduced the probability of sampling a subgraph by 50% if the center node of the subgraph already appears in a sampled subgraph. This allows unvisited local contexts to be sampled with greater probability. Highly similar subgraphs (up to 3D rotation+translation transformations) can be represented by one motif. To achieve this, the sampled subgraphs are grouped into a user-specified number of clusters using a hierarchical clustering technique using average linkage [26], implemented in the Orange3 library [27]. A representative subgraph is selected from each cluster as a motif. If the size of the whole subgraph set is too big for the hierarchical clustering algorithm, we can randomly partition the whole subgraph set into many smaller subsets, and apply the above procedure to extract representative subgraphs from each subset. The above procedure is then applied to the representative subgraphs extracted from all subsets to obtain the final motifs. Pair-wise similarity calculations are required to perform hierarchical clustering between subgraphs (each of which are ARGs).

We need to measure the similarity between two ARGs when building the motif vocabulary (Section 2.1 and performing motif convolutions (Section 2.3). Such a similarity measurement should be invariant to the permutation of nodes, which requires node-to-node matching between two graphs. In addition, the similarity measurement should not be sensitive to graph sizes. Otherwise, a larger graph could have a higher chance to be more similar to a motif than a smaller graph. Assuming we have the node-to-node matching, which is represented by a matching matrix M, between two ARGs $G_{1}$ and $G_{2}$. Each element $\textbf{M}_{ui}\in\{0,1\}$ indicates whether node $u$ in $G_{1}$ matches with node $i$ in $G_{2}$. Inspired by [28, 29], we define the normalized similarity between $G_{1}$ and $G_{2}$ defined as:

where $n_{1}$ and $n_{2}$ are the numbers of nodes in $G_{1}$ and $G_{2}$, respectively. $l_{1}$ and $l_{2}$ are the numbers of edges in $G_{1}$ and $G_{2}$, respectively, $s_{1}(e_{uv}^{(1)},e_{ij}^{(2)})$ is the relation compatibility function measuring the similarity between $e_{uv}^{(1)}\in G_{1}$ and $e_{ij}^{(2)}\in G_{2}$, $s_{2}(u,i)$ is the node compatibility function measuring the similarity between node $u\in G_{1}$ and node $i\in G_{2}$. $\alpha$ is the trade-off parameter to balance the contributions from edge similarities and node similarities. Theorem 2.1 shows that $S(G_{1},G_{2})$ is independent of graph sizes. A matching matrix M is required to compute $S(G_{1},G_{2})$. Finding an optimal matching between two ARGs is an NP problem and has been widely studied. We leave the details of problem definition and the efficient algorithm for finding a sub-optimal M in Appendix C.

The motif convolution layer (MCL) computes the similarity (see Section 2.2) between every motif and the subgraph centered at each node in an input graph. A motif representation of each input node is obtained by concatenating the similarity scores between the subgraph of the node and all motifs. This representation can be fed into a trainable multi-layer perceptron (MLP) with non-linear activation functions (e.g., ReLU) to produce a further embedding that encodes interactions among motif features. We denote this operation as:

where $G$ is an input ARG, $u$ is a node in $G$, $\{\mathcal{M}_{i}\}_{i=0}^{N}$ represents the motif vocabulary of size $N$, and $\mathbf{a}_{u}^{(0)}$ is the MCM-embedding of $u$. Figure A.1 in Appendix shows an expanded illustration of the MCM computation flow.

The MCM can serve as a preceding module for any GNN to form MCM+GNN. The output of the MCM is still an ARG, which can be fed into any GNN that accepts an ARG as input. A readout function of an MCM+GNN model is needed to obtain the final representation of an input $G$:

where $L$ is the number of GNN layers. The READOUT function should be invariant to permutation, and thus average, sum, and max pooling functions are widely used. The final representation $\mathbf{h}_{G}$ can then be fed into a trainable component (e.g., a fully connected layer or a linear regressor) to generate the desired predictions.

This experiment shows the advantage of motif convolution in capturing local structural context over GNNs. We designed 5 ARG templates (Figure 2), and one synthetic dataset of 5 classes, which share similar node attributes but have different structures. This template can only be well distinguished by their overall structures. For example, templates 2 and 5 are very similar to each other except for two edges have different attributes. Sample ARGs were produced from these 5 ARG templates by randomly adding nodes to templates and adding Gaussian noises of $\mathcal{N}(0,0.1)$ to node attributes. The number of added nodes in each sample ARG was sampled from a binomial distribution $B(4,0.1)$. Each sample ARG is labeled by the ID of its ARG template. The task is to predict the template ID of any given synthetic ARG. We synthesized two datasets of sizes 500 and 10,000, respectively. Each template contributed to 20% of each dataset.

We only used the MCL of the MCM as it was already sufficient. The readout of the MCL is fed to a logistic regressor (LR) to output the classification result. Standardization was applied to the readout by removing the mean and scaling to unit variance. We named this model MCL-LR. Two readout functions (average pooling and max pooling) were tried, and max pooling always outperformed average pooling. A motif vocabulary of size 5 was constructed. We tried using more than 5 motifs, and found no significant advantage. We compared MCL-LR with several baseline models built from GNN variants with edge weight normalization implemented by [30], including GCN [31], GIN [32] and GAT [33] (detailed model configurations in Appendix D.2).

We ran each model 20 times on both datasets. In each run, each dataset was randomly split into 8:1:1 for training, validation and test. The average prediction accuracy, as well as the standard deviation, are reported in Table 1. The MCL-LR models significantly outperform other models by an average of 20%. In addition, MCL-LR requires substantially smaller training data as it is able to achieve near-perfect results on the 500 datasets. We observed that the learned motifs were quite similar to the underlying templates (Appendix Figure E.1) and contains necessary local structures for discriminant purpose, which explain the superior performance of MCL-LR. The performance by template categories (Appendix Table E.1) suggests that MCL-LR is able to discriminate highly similar templates, as in the case of templates 2 and 5 in the appendix table. In addition, we observed that training of GNNs on the larger dataset took more time and computational resources than MCL-LR.

We conducted an experiment using several small and medium sized molecular benchmark datasets in MoleculeNet [34] and OGB [35]. The results demonstrate that MCM can be integrated with GNNs in a broad way. An MCM+GNN model uses the MCM component to preprocess input graphs. We used the open-source package RDKit [36] to parse the SMILES formula of molecules and performed scaffold-split [37, 38] to get the train-validation-test split as 8:1:1. Following the suggestions in OGB [35], both baseline models (GIN and GCN) have 5-layer with hidden dimension of 300. Mean pooling is used as the readout function after convolutional layers. Both MCM+GCN and MCM+GIN use a motif vocabulary of size 100. Smaller baseline models (3 conv layers and 64 hidden dim in GCN/GIN) are used in MCM-GCN/GIN on datasets bace, bbbp, clintox, sider and tox21. On five MoleculeNet benchmarks, we compared our model with MICRO-Graph [39] and MGSSL [40] with different generation orders (BFS and DFS), which are also pre-training frameworks for GNNs with a motif-aware fashion. For each dataset, we carried out 5 independent runs and reported means and standard deviations. Table 2 shows that GNNs integrated with MCM consistently perform better than the base models. Figure 3 compares the training and test curves of MCM+GIN and GIN, and shows that MCM significantly speeds up and stabilizes training, suggesting MCM+GIN is fundamentally more expressive than GIN. We believe this is because MCM encodes local structural information that is not sufficiently captured with traditional message passing in GNNs. The details of training settings and data preprocessing are provided in Appendix D.3.

The QM9 dataset [19] is a widely used benchmark for evaluating models that predict quantum molecular properties. It consists of about 130k organic molecules with up to 9 heavy atoms (C, O, N and F). The mean absolute error (MAE) of target properties is the commonly used evaluation metric. We adopted the data-splitting setting used in [23, 24, 25]. More specifically, following [41], we removed about 3k molecules that failed the geometric consistency check or were hard to converge. We applied random splitting to the dataset, which takes 110,000 molecules for training, 10,000 for validation, and the rest for test. We only used the atomization energy for $U_{0}$, $U$, $H$ and $G$, by subtracting the atomic reference energies as in [23]. For property $\Delta\epsilon$, we followed the DFT calculation and calculate it by simply taking $\epsilon_{LUMO}-\epsilon_{HOMO}$.

We designed the MCM to be MCL + 2-layer MLP (MLP: input $\rightarrow$ 128 $\rightarrow$ ReLU $\rightarrow$ 128 $\rightarrow$ output). Due to limited computational resources, we only tried two motif vocabulary sizes (100 and 600). The number of motifs is represented as a hyper-parameter. We formed our model MCM+MXMNet by connecting the above MCM to an MXMNet. Two options (5Å and 10Å) were tested for the distance cut-off hyper-parameter $d_{g}$ of MXMNet. A separate model was trained for each target property and used grid search on learning rate, batch size, motif number, and cut-off distance $d_{g}$. Edges in molecules are defined by connecting atoms that lie within the cut-off distance $d_{g}$. Following [23], we did not include auxiliary features such as bond types or electronegativity of atoms. Detailed training settings are provided in Appendix D.4, and details of motif vocabulary construction and efficiency in implementation are well discussed in Appendix D.5.

We compared our model MCM+MXMNet with several other state-of-the-art models including SchNet [21], DimeNet [23], DimeNet++ [24] and MXMNet [25]. For other models, we use the results reported in their original works. All experiments were run on one NVIDIA Tesla V100 GPU (32 GB). Table 3 summarizes the comparison results, and shows that our model MCM+MXMNet outperforms others on eight molecule property prediction tasks. For two MXMNet settings, a larger cut-off distance (i.e., $d_{g}=10$Å) can lead to better results for some tasks, but not all of them. This is because larger $d_{g}$ leads to a larger receptive field and thus helps to capture longer range interactions. However, higher $d_{g}$ might cause redundancy or oversmoothing in message passing and will also increase computation cost. We observed a similar phenomenon for MCM+MXMNet. We also observed that under the same $d_{g}$ setting, MCM+MXMNet tends to perform better than MXMNet. We believe that this is because MCM helps to produce more informative node representations that better encode local chemical context.

The embeddings that MCM learns are highly explainable and encode domain semantics. We visualize the representations of carbons produced by a MCM with 600 motifs in the QM9 experiment. The visualization is done using the T-distributed Stochastic Neighbor Embedding (t-SNE) algorithm  [42]. We randomly sampled 15,000 molecules from the QM9 dataset, and then randomly selected 2 carbons from each chosen molecule. Figure 4 shows the t-SNE visualization of these 30,000 atoms’ representations learned by MCM. To better understand our representation, we manually labelled 300 carbons randomly sampled from the above 30,000 carbons according to their 1-hop local structures. We observe that carbons in the same local context tend to cluster together and are separated from those in different local structures.

More interestingly, we observe that node representations learned by MCM encode meaningful chemical properties. For example, the carbons (red in Figure 4A) in the Trifluoromethyl (-CF₃) groups are tightly clustered together, actually stacked into one point. It is known that the more fluorines are connected to a carbon, the shorter the bonds from this carbon [43], which makes the Trifluoromethyl groups very different from other substructures. Moreover, Methylene (-CH₂-) is the most common ‘bridge’ in organic chemistry, connecting all kinds of functional groups (R, R’). Hence, the carbons (pink in Figure 4A) in the Methylene groups are scattered apart because of their diverse contexts. The carbons in the alcohol functional groups (-CH₂OH, green in Figure 4A) are clustered into two separate sub-groups. This is because they are connected to two very different chemical structures (Figure 4B): cyclic functional groups and linear functional groups.

The highest workload in MCM comes from matching motifs with subgraphs, which can be sped up tremendously using parallel computing in GPUs. We developed a CUDA-enabled graph matching kernel (Appendix C.4) for matching multiple Motif-ARG pairs concurrently, which offer an essential boost to this work. We tested the efficiency of our graph matching kernel under various settings. All experiments were run on NVIDIA GeForce RTX 2080 11GB GPUs. We created 4 test datasets with graph sizes of 10, 15, 20, and 25, respectively. Each set contains 500 molecules sampled from the QM9 dataset. We ran our CUDA-enabled graph matching kernel using up to 8 GPUs to compute pair-wise matching within each dataset. In total, there are 124,750 pairs. The execution times (including loading data from hard disks) of different settings are compared in Figure 5. In general, as expected, it took longer to match larger ARGs. More GPUs help to accelerate the computation. When using # GPUs $\leq 4$, doubling GPU devices approximately reduced the execution time by half, which indicates that our kernel achieved a balanced workload in parallel. Using more than 5 GPUs only offered marginal speed improvements because GPUs spent significant amounts of time waiting for data to be loaded.