GeomGCL: Geometric Graph Contrastive Learning for
Molecular Property Prediction

Since only the existing bonds in a molecule have the initial edge features (i.e., bond features), the edge embedding should be firstly updated through aggregating the pairwise node embeddings with involving the associated features. To enrich the connection information from different aspects, we use MLP to integrate the chemical bond feature $\bm{e}^{0}_{uv}$ and global distance embeddings $\bm{r}_{uv}$ for 2D- and 3D-edge embeddings respectively:

where $\bm{e}^{2d,t}_{uv}$ and $\bm{e}^{3d,t}_{uv}$ are edge embeddings at $t$-th layer, $\bm{a}_{u}$ is the node (atom) embedding, the superscript $2d$ or $3d$ indicates the view channel, and $\mathbin{\|}$ represents the concatenation operator. Then the dual-channel edge embeddings can contain both geometric and chemical semantic information.

Different from the general graph, both the 2D view and 3D view graphs of a molecule have the unique geometric attributes, which can significantly influence the specific property of the molecule. After the derivation of edge embeddings, GeomMPNN performs an edge$\bm{\rightarrow}$edge message passing process to perceive the geometric distribution in the molecule through 2D and 3D angle-aware aggregations. Considering that the neighbors in 2D view are more sparse by contrast with the neighbors in 3D view, we develop the well-directed layers for both views. Specifically, for 2D view graph learning, the following angle-injected function is employed to update the 2D edge embedding:

where $\mathcal{A}(e_{uv})$ denotes the set of neighboring edges of the edge $e_{uv}$, $\bm{\phi}_{wuv}$ is the 2D angle embedding between the edge $e_{wu}$ and $e_{uv}$, $\odot$ is the element-wise dot operation, $W^{t}_{\phi}$ and $W_{e}^{2d,t}$ are learnable parameters. For the 3D view graph, there are sufficient neighbors around each edge. Inspired by the recent work (Li et al. 2021), we divide the neighboring edges of each target edge into several angle domains $\mathcal{A}_{1},...,\mathcal{A}_{n}$ according to 3D spatial angle $\theta$. Then we apply a hierarchical aggregation process among edges for 3D view graph learning, which consists of local and global stages:

where $\bm{e}^{3d,t}_{uv,i}$ is the aggregated 3D edge embedding at $i$-th angle domain through the local stage, $W_{e,i}$ and $W^{t}_{\theta,i}$ are trainable parameters, $\bm{\theta}_{wuv}$ is the 3D angle embedding between the edge $e_{wu}$ and $e_{uv}$. ${\rm Pool}$ means the max pooling function over all $n$ local edge embeddings, which can generally extract the high-level spatial distribution information to strengthen the geometric structure learning.

After obtaining the angle-aware edge embeddings $\bm{e}^{2d,t}_{uv}$ and $\bm{e}^{3d,t}_{uv}$, we apply the edge$\bm{\rightarrow}$node message passing to fulfil the propagation process from the edge back to the node. The essential distance factor between nodes is well-considered via the similar adaptive scheme with the previous edge$\bm{\rightarrow}$edge component:

where $W^{t}_{l}$, $W_{a}^{t}$, $W^{t}_{r,i}$ and $W_{a,i}^{t}$ are learnable parameters, $\bm{l}_{uv}$ and $\bm{r}_{uv}$ are the geometric embeddings of local and global distances respectively, $\mathcal{D}(a_{u})$ is the set of all neighboring edges for node $a_{u}$ in 2D view, $\mathcal{D}_{i}(a_{u})$ is the set of neighboring edges located in $i$-th distance domain among all $m$ divided domains in 3D view.

Since all geometric factors thoroughly enrich the node and edge embeddings in both 2D and 3D views via $T$ stacked message passing layers, the final representations $\bm{a}^{2d,T}_{v}$ and $\bm{a}^{3d,T}_{v}$ can reflect the graph topology along with the molecular geometry. To further get the graph-level representation with identifying the important nodes, we follow (Xiong et al. 2019) and adopt the node$\bm{\rightarrow}$graph attentive pooling layer. For simplicity, we use $\bm{a}_{v}$ to represent $\bm{a}^{2d,T}_{v}$ or $\bm{a}^{3d,T}_{v}$ and use $\bm{h}$ to represent $\bm{h}^{2d}$ or $\bm{h}^{3d}$. The graph-level embedding $\bm{h}$ is updated iteratively through the attentive propagation process, which starts with the initial embedding $\bm{h}^{0}=\sum_{v}\bm{a}_{v}$.

where $\bm{g}^{t}$ is the global context message at $t$-th pooling layer, which aggregates the valuable nodes from the full set $\mathcal{V}$. After performing $T_{g}$ pooling layers, the final graph representations $\bm{h}^{2d}$ for 2D view and $\bm{h}^{3d}$ for 3D view are acquired.

To evaluate the performance of our proposed model with the existing molecular representation learning methods, we use seven molecular datasets from MoleculeNet (Wu et al. 2018) including ClinTox, Sider, Tox21 and ToxCast four physiology datasets for graph classification tasks, as well as ESOL, FreeSolv and Lipophilicity three physical chemistry datasets for graph regression tasks. The main statistics of datasets are summarized in Table 1.

We compare the proposed GeomGCL with a variety of state-of-the-art baseline models for molecular property prediction, which includes molecular message passing-based methods, geometry learning-based GNN methods, and graph contrastive learning methods. The first group consists of three well-designed message passing neural networks. AttentiveFP (Xiong et al. 2019) adopts the graph attention network for molecular representation learning. DMPNN (Yang et al. 2019) and CoMPT (Chen et al. 2021) are message passing models with considering edge features in a node-edge interactive manner. Geometry learning-based models contain several GNN approaches. SGCN (Danel et al. 2020) directly encodes the atomic position information in the aggregation process. MAT (Maziarka et al. 2020) incorporates the local geometric distance into the graph-based transformer model. To comprehensively reflect the superiority of our proposed model, we also compare GeomGCL against with HMGNN (Shui and Karypis 2020) and DimeNet (Klicpera, Groß, and Günnemann 2020), both of which can learn the geometric distance and angle factors in 3D space for quantum property prediction. Besides, the recent graph contrastive models are compared to show the power of our proposed 2D-3D geometric contrastive learning strategy. InfoGraph (Sun et al. 2019) maximizes the mutual information between nodes and graphs, while MoCL (Sun et al. 2021) introduces the multi-level domain knowledge for molecular graphs in a well-designed contrastive learning framework.

Following the previous works, we evaluate all methods through the k-fold cross-validation experiments, and we set k as 10 to report the robustly average experimental results. As recommended by the MoleculeNet benchmarks (Wu et al. 2018), we randomly split each dataset into training, validation, and testing set with a ratio of 0.8/0.1/0.1. The validation set is used for early stop and model selection. We use ROC-AUC and RMSE metrics for graph classification and graph regression tasks respectively.

The 3D structures of molecules are generated for $P=50$ times through the stochastic optimization algorithm of Merck Molecular Force Field (MMFF), which is implemented in the RDKit package (Tosco, Stiefl, and Landrum 2014). For our model, We use Adam optimizer for model training with a learning rate of 1e-3. We set the batch size as 256 for contrastive learning and 32 for finetuning with the scale parameter $\tau=0.5$. The hidden size of all models is set to 128. The cutoff distance $d_{\theta}$ is determined (4 Å or 5 Å) according to the size of the molecule on each dataset. We set the dimension $K$ of geometric embedding as 64. The numbers of 3D angle domains and global distance domains are set to 4. The balancing hyper-parameter $\lambda$ is set to 0.01 according to the performance on validation set. For baseline models, we tune parameters of each method based on recommended settings in the paper to ensure the best performance. As a general setting (Sun et al. 2019, 2021), our proposed GeomGCL and contrastive learning baselines are pretrained on molecular graphs of each dataset, and then we finetune the model for the downstream task on the same dataset.

The performance results of evaluating each model for graph classification and regression tasks are presented in Table 2. As we can see, our model significantly outperforms all the baselines on both types of tasks. On the whole, we can observe that our proposed GeomGCL improves the performance over the best message passing baselines with 2.18% and 10.4% for classification and regression tasks, respectively.

Among all baseline approaches, the well-designed message passing models generally show the best performance, which indicates that the essential bond features of the molecule can provide chemical semantic information for molecular representation learning. Specifically, since DMPNN and CoMPT adopt the node-edge interactive scheme, they perform slightly better than AttentiveFP. As to geometry-based baseline models, MAT can take advantage of the geometric distance from the molecular graph and performs much better than SGCN, which directly encodes the 3D coordinates and can be easily affected by the coordinate systems. Although HMGNN and DimeNet can identify the distance and angle information, they learn the molecular embedding only based on the 3D geometric graph which might be noisy. Besides, these models are designed for quantum property prediction and may not be expert in modeling the larger molecules. By contrast, our GeomMPNN and GeomGCL are capable of learning from the stable 2D graph and the informative 3D geometric structure. For contrastive learning methods, MoCL achieves better results than InfoGraph, showing the significance of domain knowledge for developing the contrastive strategy without changing the chemical semantics. However, the failure of leveraging the critical geometric information limits their ability to model the molecule without labels, while our method can capture geometry-aware structural information by contrasting the 2D-3D geometric views. Therefore, GeomGCL is much effective for molecular representation learning and can accurately predict each targeted property.

To further investigate the factors that influence the performance of the proposed GeomGCL framework, we conduct the ablation study on four benchmark datasets for classification and regression tasks with designing different variants of GeomGCL.

• •

  GeomMPNN-2D only preserves the single-channel 2D geometric message passing layers.

• •

  GeomMPNN-3D only preserves the single-channel 3D geometric message passing layers.

• •

  GeomMPNN uses the dual-channel geometric message passing layers without contrastive learning.

• •

  GeomGCL-NoReg removes the spatial regularizer $\mathcal{L}^{s}$.

As shown in Figure 4, GeomGCL achieves the best performance among all architectures, proving the necessity of learning the 2D-3D geometric structures contrastively and synergistically. To be specific, we can find that learning representations from a single 2D or 3D view can not always perform better than the other view across different datasets, which supports our hypothesis that only a geometric view of molecular graph is not sufficient. Therefore, decoupling 2D or 3D geometric message passing layers from GeomGCL yields a significant drop in performance. Additionally, the use of spatial regularizer when training the model can help GeomGCL to discriminate the relative correlations of different angle domains and then contributes to the performance improvements. What’s more, GeomMPNN significantly performs worse than GeomGCL, which confirms that our geometric contrastive learning scheme is beneficial for molecular representation learning.

Finally, we analyze the performance variation for GeomGCL by varying the coefficient $\lambda$ to look deeper into the impact of the spatial regularizer. As depicted in Figure 5, we observe that the performance first tends to get better with incorporating more 3D angle domain information while training the model, and then begins to drop off slightly. The appropriate trade-off weight can assist the model in identifying the geometric factors and enhancing the representation learning. Overall, the performance of GeomGCL is stable and always better than baseline methods.