An Image-enhanced Molecular Graph Representation Learning Framework

1 Introduction

The molecular representation learning plays an import role in high-precision drug discovery (such as molecular property prediction, target activity prediction) [Hu et al., 2020a;

∗Correspondding author (xzeng@hnu.edu.cn)

Zeng et al., 2022; Zhang et al., 2023]. As the most direct representation of molecules, graphs treat atoms and bonds in molecule as nodes and edges in graph and have been widely studied and applied [Hu et al., 2020a; Liu et al., 2021; Xia et al., 2023]. In view of the objective existence of molecules in nature, molecules can also be represented as images, allowing us to directly observe the morphology and internal structure of molecules through an electron microscope [Shen et al., 2022] or a rendering tool [Landrum, 2013].

Limited by a single modality [Hu et al., 2020a; Sun et al., 2020], the GraphMVP [Liu et al., 2021] and 3D InfoMax [St¨ark et al., 2021] try to leverage 2 modalities (2D and 3D graphs) to enhance features but the performance improvements are still limited. There are two reasons: (1) similar modalities (2D graphs and 3D graphs) and encoding ways (graph neural network-based encoder [Wu et al., 2020]) and (2) weak feature extraction ability, resulting in insufficient complementary information between modalities [Tumer and Ghosh, 1995]. Image is a modality that is significantly different from graph in terms of modality and encoding way (convolutional neural network-based encoder [Li et al., 2021b]). We empirically prove that 2D graph has high pearson correlation [Cohen et al., 2009] with 3D graph and has significantly low pearson correlation with image, as shown above in Figure 1(a) (See Appendix A for more details 1 ). Unlike graph, images understand molecules from a visual perspective, which contains texture information of molecules and allows for direct visualization of spatial arrangements without introducing any conformation, such as chiral changes of molecules. We further illustrate the advantages of images in basic prior knowledge, as shown below in Figure 1(a) (See Appendix B for more details). Given the popularity of graph neural networks in drug discovery, a meaningful question is raised: can we exploit the rich information in molecular images to facilitate representation learning of molecular graphs?

Considering that the simplest multi-modal learning framework [Song et al., 2022; Wang et al., 2020] requires additional computational costs in the training and inference stages, inspired by knowledge distillation (KD) [Hinton et al., 2015; Hou et al., 2021], we describe the process of information transfer as how to use a knowledgeable teacher (image) to teach an excellent student (graph). The success of KD depends on good teacher and KD strategies (See Section 3.6 for proof). In order to obtain a knowledgeable image-based teacher, we propose a novel multi-view molecular image representation learning method with 3D conformation to enrich the representation with 5 pre-training strategies. Considering that naive feature alignment between cross-modal features will lead to limited gain or even negative transfer [Yang et al., 2021; Yan et al., 2022] in the cross-modal KD (CMKD) stage, we design a novel CMKD framework (called IEM) to enrich the features of graph, which alleviates knowledge enhancer and task enhancer to align graph and image in logit space to avoid modality gaps in feature space [Liang et al., 2022]. As shown in Figure 1(b), IEM distills features from images into graph representations. It is worth noting that compared with multimodal learning methods, the proposed IEM has the following advantages: (1) Universality: IEM can be integrated with any graph-based method. (2) Effectiveness: IEM significantly improves the performance of several graph-based baselines, such as Figure 1 (c). (3) Efficiency: As low as 5% of training images can still improve performance; (4) Compatibility: IEM is compatible with both 2D and 3D molecular images and different rendering strategies.

1 Appendix is available at https://github.com/HongxinXiang/IEM/ blob/main/assets/appendix.pdf

In summary, the main contributions of this paper are summaried as follows:

• To the best of our knowledge, we are the first to reconstruct the graph-based molecular representation learning into image-based cross-modal distillation paradigm.
• We propose the 3D conformational molecular image representation method to train a knowledgeable teacher, achieving a one-time pre-training on 2 million molecules that enables the use of multi-view features to enhance existing graph-based baseline models.
• We propose an image-enhanced molecular graph representation learning framework, called IEM, which equips with knowledge enhancer and task enhancer to improve the performance of a large number of graph-based models. In addition, we also demonstrate the effectiveness of our framework theoretically.
• We show that our method achieves significantly better

performance on 12 MoleculeNet benchmarks and can substantially enhance the performance of existing molecular graph representation models.

2 RelatedWork

Graph-based Molecular Representation Learning. Graph neural networks have achieved remarkable success in drug discovery tasks. In view of the high cost of annotating molecules, recent studies mainly learn from large-scale label-free molecular databases by designing pre-training strategies. [Hu et al., 2020a; Li et al., 2021a; Guo et al., 2021] consider both nodelevel and graph-level pre-training strategies to capture the local and the global information in graph. GPT-GNN [Hu et al., 2020b] uses two generative pre-training tasks, including attribute generation and edge generation, to extract fine-grained information of molecules. GROVER [Rong et al., 2020] and MGSSL [Zhang et al., 2021] propose a motif-based prediction or generation tasks to capture the information of molecular motifs. Mole-BERT [Xia et al., 2023] proposed masked atoms modeling and triplet masked contrastive learning to further optimize mask-based GNN. Considering the importance of 3D geometric information in drug discovery tasks, GraphMVP [Rong et al., 2020] and 3D Infomax [St¨ark et al., 2021] pretrained GNN on molecules with 3D geometric information. In the paper, we still follow the graph-based paradigm for drug discovery tasks and try to improve all graph-based models.

Image-based Molecular Representation Learning. Because graphs are discrete and unordered, some researchers consider representing molecules as images and utilizing mature computer vision techniques to extract features. Chemception [Goh et al., 2017], 2DConvNet [Fernandez et al., 2018] and DenseNet121 [Zhong et al., 2021] use molecular images to predict chemical properties, toxicity of compounds and contaminant reactivity, respectively. With the explosion of selfsupervised learning in computer vision [Arnab et al., 2021; Likhosherstov et al., 2021], the first pre-training model based on molecular images (called ImageMol) is proposed for learning representation from 10 million molecules with 5 welldesigned pre-training tasks [Zeng et al., 2022]. CGIP is further proposed to extract fine-grained image representations via carefully designed intra- and inter-modal contrastive learning between graph and image [Xiang et al., 2023]. Different from the images used by previous methods, which are 2D images generated by RDKit, we propose a clearer and more informative 3D conformation-based multi-view image representation.

Cross-Modal Knowledge Distillation. As an important branch of knowledge distillation [Hinton et al., 2015], crossmodal knowledge distillation (CMKD) is still a relatively emerging field, which refers to using a teacher from another modality to supervise the learning model of the current modality and improve the performance of the student during inference. For example, [Gupta et al., 2016] transfer supervision from labeled RGB images to unlabeled depth and optical flow images and [Sun et al., 2021] learns TIR (Thermal Infrared)-specific target representations transferred from the RGB modality. Pri3D [Hou et al., 2021] distills 3D point cloud information into 2D images and improves performance of image encoder in semantic segmentation, object detection, and instance segmentation. PointCMT [Yan et al., 2022] and UniDistill [Zhou et al., 2023] utilize images as teachers to guide point cloud-based students and airbornelidar-based students for improving point cloud-related tasks and 3D object detection in bird's-eye view, respectively. VGSR [Jin et al., 2023] distills the knowledge of face image into audio to improve the performance of speaker recognition. Different from previous works, this is the first image-to-graph cross-modal knowledge distillation framework to our best knowledge.

3 Our Method

3.1 Preliminaries

Background

We summarized three ways to obtain molecular images: (1) Canvas-based technology, which draws molecular images by creating a sketchpad and using pixels, such as RDKit [Landrum, 2013]; (2) 3D CAD (Computer Aided Design) modeling technology, which uses CAD software to create a virtual three-dimensional space and build a model with threedimensional data or geometric configuration, such as PyMol [DeLano and others, 2002]; (3) Physical microscopy-based technology, including Cryo-EM and single-molecule fluorescence imaging, which utilizes atomic force microscopy (AFM), optical tweezers (OT), magnetic tweezers (MT) or fluorescence microscopy to obtain visual representations of molecules. See Appendix C for the visualization of molecules in three ways. Each of these methods has its own advantages. Canvas technology has low computational complexity and fast processing speed, but is confusing when describing complex, geometric molecules. Existing molecular image-based methods are all based on this imaging method [Zeng et al., 2022; Xiang et al., 2023]. 3D CAD modeling technology retains rich information in molecules but is slow to render. Singlemolecule fluorescence imaging technology can directly reflect the molecular portraits of nature, but obtaining these samples is low-throughput and resource-consuming, making it impossible to expand to large-scale data. Therefore, in this paper, we only consider the first two imaging methods. Moreover, the proposed method can seamlessly support samples collected by the third technology.

Problem Formulation

Let the molecular graphs and corresponding ground-truth labels on downstream tasks are {Gi = (Vi, Ei)}ni=1 and {yi}ni=1 ∈ Rt respectively, where Vi ∈ Rnvi ×dvi with the number of vertices is nvi and the feature dimension of the vertices is dvi and Ei ∈ Rnei ×dei with the number of edges is nei and the feature dimension of the edges is dei represent vertices and edges in i-th molecular graph respectively. t is the number of tasks in downstream tasks. The corresponding single-view 2D images and multi-view 3D images of molecular graphs G can be denoted as V 2D ∈ RH×W×3 and V 3D ∈ RV ×H×W×3 respectively, where V represents the number of views and H and W represent the height and the width of images, respectively. The single-view 2D images can be directly generated by RDKit [Landrum, 2013] and the multi-view 3D images can be obtained by using PyMol [DeLano and others, 2002] to generate snapshots from different viewpoints. The features of graphs F g ∈ Rdg , 2D images F 2D ∈ Rdv and 3D images F 3D ∈ Rdv can be extracted by graph encoder Encg , 2D encoder Enc2D and 3D encoder Enc3D , respectively. In the proposed knowledge distillation framework, after pre-training, both Enc2D and Enc3D from images can be chosen as teacher, while graph encoder is considered as student. This ensures that our proposed knowledge distillation framework remains applicable when it is necessary to perform representation learning on molecular graphs for which 3D conformation cannot be acquired. Considering that 3D will carry more feature information, our experiments choose the Enc3D as the teacher.

3.2 Overview of the Method

Here, we propose the image-enhanced molecular graph representation learning framework (IEM), which equips knowledgeable teachers and distillation strategies to prevent negative transfer. In particular, we pre-train the teacher with 5 pre-tasks to ensure that the features extracted by the teacher are better than those of the students, where 1 pre-task takes into account the intrinsic information of molecular images by aligning 2D and 3D images and 4 pre-tasks consider four priors of molecule (atom S atom ∈ Rnatom , bound S bound ∈ Rnbound , geometry S geom ∈ Rngeom and basic properties S prop ∈ Rnprop ) by using 4 corresponding predictors (Patom, Pbound , Pgeom and Pprop ) to optimize molecular features. These four predictors use the same structure: Full Connection (FC)→Softplus→FC. During the distillation stage, simply aligning features from different modalities dimensionwise can easily lead to negative transfer [Yang et al., 2021; Yan et al., 2022]. Inspired by them, we utilize knowledge enhancer Enhk and task enhancer Enht to align the information of different modalities at the logit level, which preserves both modality-specific semantics from student and logit-specific semantics from teacher guidance. The overview of IEM is illustrated in Figure 2(a). We summarize the main processes in Appendix D. The process of IEM is divided into three steps: (1) Use 5 pre-training tasks to train a knowledgeable teacher (Section 3.3); (2) Exploit image-based teacher to enhance graph-based student by using the knowledge enhancer and task enhancer (Section 3.4); (3) Train IEM and inference in downstream tasks (Section 3.5).

3.3 Pretraining Teacher

Excellent teacher are one of the keys to knowledge distillation. In this paper, as shown in Figure 2(b), we combine low-computation cost 2D images generated by RDKit and information-rich 3D images rendered by PyMol, both of which can be used to enhance graph representation (See Appendix E for details of image rendering). However, we still recommend using a 3D image encoder as a teacher because it is more discriminative. Here, we choose 4 fixed viewpoints to generate multi-view images. The feature extraction process can be formalized as: F 2D = Enc2D (V 2D ) and F 3D = Enc3D (V 3D ). Due to the uncertainty of downstream tasks, the teacher should have strong knowledge generalization capabilities. Therefore, we designed 5 pre-training tasks based on the principle that generalized knowledge should reflect the basic attribute information of molecules, including (1) contrastive learning between 2D and 3D images (ICL); (2) atom distribution prediction task (ADP); (3) bound distribution prediction task (BDP); (4) geometry distribution prediction task (GDP); (5) property distribution prediction task (PDP).

The main motivation of the ICL is to extract the inherent information in molecular images. Meanwhile, the 2D image encoder has a perception of 3D information and the 3D image encoder has a perception of global information to alleviate the problem of view-occlusion. For convenience, we denote F I = {F 2D , F 3D } ∈ R2×n,dv , where n and dv represent the number of samples and the dimension of visual features. Following InfoNCE [Oord et al., 2018], the contrastive learning loss between 2D and 3D images is formalized as:

where I and τ represent indicator function and temperature, respectively. sim(·) represents cosine similarity, that is, sim(F I i ,F I j ) = F I i · F I j / ∥ F I i ∥2∥ F I j ∥2.

The motivation for the ADP, BDP, GDP and PDP tasks is to allow the teacher to acquire knowledge with strong generalization. Therefore, we extract fundamental prior knowledge of molecules from their atoms (Satom), bonds (Sbound ), geometry (Sgeom), and chemical properties (Sprop ). For details on these priors, see Appendix F. Subsequently, F 2D and F 3D are forward-propagated to 4 predictors Patom, Pbound , Pgeom, Pprop to obtain the corresponding predicted logits {p2D atom, p2D bound , p2D geom, p2D prop } and {p3D atom, p3D bound , p3D geom, p3D prop }. Finally, the losses of 2D image encoder and 3D image encoder in these 4 prediction tasks can be formulated as follows:

Finally, the 2D image encoder and 3D image encoder can be pretrained with the following total loss:

See Appendix G for mode details about pretraining teacher.

3.4 Image-enhanced Distillation Strategy

During CMKD, both 2D image encoder and 3D image encoder can be considered teachers. Here, we take a 3D image encoder as a teacher as an example. We first use frozen teacher and untrained GNN student to extract the corresponding features F v ∈ Rdv and F g ∈ Rdg . Considering that the output dimensions of the teacher and the student may be different, we add a fully connected layer after the student network by default to make dg = dv . Subsequently, instead of using simple feature alignment we use knowledge enhancer Enhk and task enhancer Enht to transfer knowledge from F 3D to F g , which can alleviate the problem of negative transfer.

Knowledge Enhancer: The purpose of the knowledge enhancer is to distill task-irrelevant generalized knowledge from the teacher. As shown in Figure 2(c), Enhk accepts F 3D and F g as input and uses 4 frozen predictors Patom, Pbound , Pgeom, Pprop to predict the distribution of molecules at different knowledge levels, which can be formalized as {pvatom, pvbound , pvgeom, pvprop } and {pgatom, pgbound , pggeom, pgprop }. We use the predicted labels of the teacher as ground-truth to supervise students with L1 loss, which can be formalized as:

It is worth noting that this paper does not use KL-based distillation loss [Hinton et al., 2015] because it is difficult to apply to single-task classification and regression tasks.

Task Enhancer: The purpose of the task enhancer is to distill knowledge relevant to downstream tasks from teacher. As shown in Figure 2(d), task predictor accepts F 3D and F g and generates task-related prediction logits ¯ yv and ¯ yg . We use ¯ yv from the teacher as ground-truth to supervise ¯ yg from the student with smooth L1 loss [Girshick, 2015], which can be formalized as:

3.5 Training and Inference

For efficiency, we usually hope to use a single modality rather than multiple modalities when inferring. We further use the true ground-truth y of the downstream task to train student, and the loss is formalized as:

where φ represents cross-entropy loss for classification and MSE loss for regression. The final loss is formalized as:

During inference, predicted results can be obtained by inputting graph data G to the student network Encg .

3.6 Theoretical Justification of IEMEffectiveness

We utilize CMKD to transfer knowledge from image-based teachers to graph-based students to obtain image-enhanced graph features F IE . For a useful knowledge distillation, F IE with knowledge γ from teacher should be more informative than F g from student only. We denote the information enhancement useful to the task after knowledge distillation as Idiff = I IE - I g , where I IE = I(F IE |V, y;Encg , γ) and I g = I(F g |G, y;Encg ) represent the information amount of features enhanced by teacher and features extracted only by student, respectively. I(·|⋆) represents the amount of information produced by · under the conditions of given ⋆. We prove that when negative transfer does not occur, the lower bound Ω of the information increment Idiff depends on the information difference between the teacher and the student, which can be formulated as Ω = I(F 3D |V, y;Enc3D )-I(F g |G, y;Encg ). Therefore, a knowledgeable teacher and a effective distillation strategy to prevent negative transfer are crucial for CMKD. Please see Appendix H for detailed proof. To ensure effectiveness of IEM, we design teacher with 5 pre-training tasks and 2 enhancers for alleviating negative transfer.

4 Experiments and Results

4.1 Experimental Settings

Datasets and evaluation protocol. For pre-training teacher model, we sample 2 millions unlabeled molecules with 3D conformations from PCQM4Mv2 database [Hu et al., 2017]. In evaluation stage, we use the widely-used 8 binary classification datasets fromMoleculeNet [Wu et al., 2018] with ROC-AUC metric. For broader evaluation, we also test the effectiveness of IEM on 4 regression datasets included in GraphMVP [Liu et al., 2021] with root mean square error (RMSE) metric. See Appendix I for more dataset details. Notably, we use strict scaffold splitting [Hu et al., 2020a] to divide all datasets into training set, validation set and test set according to 8:1:1.

Implementation details. In pre-training teacher model stage, we use 2 independent ResNet-18 [He et al., 2016] as the architecture for 2D and 3D image encoders. The 2D image encoder extracts 512-dimensional visual features directly from a single 2D molecule image, while the 3D image encoder obtains 512-dimensional visual features by applying view-wise mean-pooling to the features of multi-view 3D images. In addition, the atom, binding, geometry, property predictors are multi-layer perceptrons (MLPs), including a 512-dimensional input layer, a 256-dimensional hidden layer, a softplus activation function and an output layer related to prediction. The teacher model is pre-trained for more than 30 epochs (about

450k steps) with temperature of 0.1, batch-size of 128 and learning rate of 0.01 (see Appendix J for details of training losses). Note that the pre-trained teacher model is able to distill any graph-based model. In order to align the experimental settings and fair comparison, we uniformly use those methods with 5-layer Graph Isomorphism Networks (GIN) [Xu et al., 2018] and 300 hidden layer dimensions as our baseline models for evaluation on downstream tasks. Following [Hu et al., 2020a], we train for 100 epochs with batch-size of 32 and learning rate of 0.001. We select hyper-parameters λKE and λTE from {0.001, 0.01, 0.1, 1, 5} and report test scores corresponding to the best validation performance. Notice that the results of some baselines may differ from their original papers because inconsistent evaluation settings and we reproduced them with the same evaluation.

4.2 Main Results

We first evaluate the performance of IEM on the 8 molecular property prediction datasets with 5 baselines (GIN, EdgePred, GraphMVP, GraphMVP-C and Mole-BERT) and Table 1 reports the main results (See Appendix K for more comparison methods). We observe that the baselines equipped with IEM achieve the state-of-the-art performance. It is worth noting that all baselines obtain consistent performance improvement after being equipped with IEM with an absolute improvement ranging from 1.3% to 3.5% on average ROC-AUC. To verify that the performance improvement of IEM is not caused by the standard deviation of the baseline, we counted the results that are higher than the α = mean+standard deviation of the baseline. We find that 70% (28 out of 40) results outperform the α of the baseline. Especially for GIN, except for the Tox21, other improvements are better than the α of the baseline.

We further evaluate on a wider range of drug discovery tasks, which includes 4 regression benchmarks. As shown in Table 2, we find the same conclusion as the classification task, that is, IEM comprehensively improved all baselines with a maximum 8.56% relative improvement and achieved the state-of-the-art performance. Therefore, IEM is promising as a universal plug-in to improve any graph-based model.

4.3 Different GNNArchitectures

To verify that IEM is effective on different GNN architectures, we used 4 different GNN architectures, including GCN, GIN, GAT, and GraphSAGE. The Table 3 shows the average ROC-AUC results of different GNN architectures on 8 classification datasets. We find that IEM can significantly improve different GNN architectures with a relative performance improvement of 3.92% to 5.23%. In particular, we observe that GIN equipped with IEM is competitive with pre-trained models such as InfoGraph, EdgePred, and 3D InfoMax. This shows the strong generalization ability of the IEM, which is able to achieve performance comparable to pre-training even without any pre-training.

4.4 Different Image Rendering Strategies

To explore the performance of the IEM on conformation-free molecules, we discuss the impact of different image rendering strategies, including conformation-free 2D images rendered by RDKit and PyMol. As shown in Table 4. We find that IEM can improve the performance of EdgePred and GraphMVP with an average performance improvement of more than 2% in all rendering strategies, indicating that IEM can successfully be compatible with conformation-free 2D images. We also find that the performance of IEM on 2D images rendered by RDKit and PyMol is similar (72.21% v.s. 72.00% and 73.34% v.s. 73.41%), indicating that we can use the more economical RDKit to improve graph-based models under limited computing resources.

4.5 Image Efficiency

To verify demonstrating the image-efficiency of IEM, we use different numbers of images to distill GraphMVP. Table 5 shows the average ROC-AUC with different number of images. We find that as the number of image samples increases, the performance of the IEM continues to improve with a range

Table 5: The average ROC-AUC (%) performance on 8 classification datasets with different number of images. The image size represents the proportion of image samples used. We use GraphMVP as baseline model. ∆ represents the relative improvement percentage.

4.6 Ablation Study

Table 6 shows the ablation results on knowledge enhancer and task enhancer using 4 different GNNs. We find that a single enhancer can improve the performance of baselines with an absolute improvement of up to 2.06%, which shows the effectiveness of two enhancers. By combining the two enhancers, more performance improvements can be achieved, with absolute performance improvements ranging from 2.62% to 3.48%. We also note that the performance improvement of KE is always better than that of TE, indicating that the fundamental prior knowledge (atom, bound, geometry and property) in the teacher transfers well to the student.

5 Conclusion

In this work, we propose a novel image-enhanced molecular graph representation learning framework (called IEM), which is the first attempt to use images to improve the performance of graphs. Equiped with a knowledgeable image-based teacher and 2 enhancers (knowledge enhancer and task enhancer), our IEM can significantly improve the performance of graphbased methods without any architectural modifications on a large number of drug discovery benchmarks. Therefore, IEM has the potential to leverage images to empower a wider range of graph representation learning fields, such as grid representation learning and skeleton representation learning. In particular, we experimently demonstrate that performance can be improved by cheap image rendering for microscopic entities where image data is difficult to obtain, which will encourage us to use IEM for more biological entities in life sciences, such as protein and ribonucleic acid.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (grant nos. 62122025, U22A2037, 62250028), Postgraduate Scientific Research Innovation Project of Hunan Province (grant no. CX20220380).