Pre-training Molecular Graph Representation with 3D Geometry

1 INTRODUCTION . In recent years , drug discovery has drawn increasing interest in the machine learning community . Among many challenges therein , how to discriminatively represent a molecule with a vectorized embedding remains a fundamental yet open challenge . The underlying problem can be decomposed into two components : how to design a common latent space for molecule graphs ( i.e. , designing a suitable encoder ) and how to construct an objective function to supervise the training ( i.e. , defining a learning target ) . Falling broadly into the second category , our paper studies self-supervised molecular representation learning by leveraging the consistency between 3D geometry and 2D topology . Motivated by the prominent success of the pretraining-finetuning pipeline [ 16 ] , unsupervisedly pretrained graph neural networks for molecules yields promising performance on downstream tasks and becomes increasingly popular [ 39 , 48 , 72 , 80 , 91 , 92 ] . The key to pre-training lies in finding an effective proxy task ( i.e. , training objective ) to leverage the power of large unlabeled datasets . Inspired by [ 50 , 69 ] that molecular properties [ 27 , 48 ] can be better predicted by 3D geometry due to its encoded energy knowledge , we aim to make use of the 3D geometry of molecules in pre-training . However , the stereochemical structures are often very expensive to obtain , making such 3D geometric information scarce in downstream tasks . To address this problem , we propose the GraphMulti-View Pre-training ( GraphMVP ) framework , where a 2D molecule encoder is pre-trained with the knowledge of 3D geometry and then fine-tuned on downstream tasks without 3D information . Our learning paradigm , during pre-training , injects the knowledge of 3D molecular geometry to a 2D molecular graph encoder such that the downstream tasks can benefit from the implicit 3D geometry even if there is no 3D information available . We attain the aforementioned goal by leveraging two pretext tasks on the 3D and 2D molecular graphs : one contrastive and one generative SSL . Contrastive SSL creates the supervised signal at an inter-molecule level : the 3D and 2D graph pairs are positive if they are from the same molecule , and negative otherwise ; Then contrastive SSL [ 83 ] will align the positive pairs and contrast the negative pairs simultaneously . Generative SSL [ 36 , 45 , 81 ] , on the other hand , obtains the supervised signal in an intra-molecule way : it learns a 2D/3D representation that can reconstruct its 3D/2D counterpart view for each molecule itself . To cope with the challenge of measuring the quality of reconstruction on molecule 3D and 2D space , we further propose a novel surrogate objective function called variation representation reconstruction ( VRR ) for the generative SSL task , which can effectively compute such quality in the continuous representation space . The knowledge acquired by these two SSL tasks is complementary , so our GraphMVP framework integrates them to form more discriminative 2D molecular graph representation . Consistent performance improvements empirically validate the effectiveness of GraphMVP . We give additional insights to justify the effectiveness of GraphMVP . First , GraphMVP is a selfsupervised learning approach based on maximizing mutual information ( MI ) between 3D and 2D views , enabling the learnt representation to capture high-level factors [ 5 , 6 , 76 ] in molecule data . Second , we find that 3D molecular geometry is a form of privileged information [ 78 , 79 ] . It has been proven that using privileged information in training can accelerate the speed of learning . We are aware that privileged information is only used in training , while it is not available in testing . This perfectly matches our intuition of pre-training molecular representation with 3D geometry . Our major contributions include ( 1 ) To our best knowledge , we are the first to incorporate the 3D geometric information into graph SSL ; ( 2 ) We propose one contrastive and one generative SSL tasks for pre-training . Then we elaborate their differences and empirically validate that combining both can lead to a better representation ; ( 3 ) We provide theoretical insights and case studies justify why adding 3D geometry is beneficial ; ( 4 ) We achieve the SOTA performance among all the SSL baselines . Related work . We briefly review the most related works here and include a more detailed summarization in Appendix A. Self-supervised learning ( SSL ) methods have attracted massive attention to graph applications [ 49 , 51 , 85 , 87 ] . In general , there are roughly two categories of graph SSL : contrastive and generative , where they differ on the design of the supervised signals . Contrastive graph SSL [ 39 , 72 , 80 , 91 , 92 ] constructs the supervised signals at the inter-graph level and learns the representation by contrasting with other graphs , while generative graph SSL [ 32 , 39 , 40 , 48 ] focuses on reconstructing the original graph at the intra-graph level . One of the most significant differences that separate our work from existing methods is that all previous methods merely focus on 2D molecular topology . However , for scientific tasks such as molecular property prediction , 3D geometry should be incorporated as it provides complementary and comprehensive information [ 50 , 69 ] . To fill this gap , we propose GraphMVP to leverage the 3D geometry in graph self-supervised pre-training . 2 PRELIMINARIES . We first outline the key concepts and notations used in this work . Self-supervised learning ( SSL ) is based on the view design , where each view provides a specific aspect and modality of the data . Each molecule has two natural views : the 2D graph incorporates the topological structure defined by the adjacency , while the 3D graph can better reflect the geometry and spatial relation . From a chemical perspective , 3D geometric graphs focus on the energy while 2D graphs emphasize the topological information ; thus they can be composed for learning more informative representation in GraphMVP . Transformation is an atomic operation in SSL that can extract specific information from each view . Next , we will briefly introduce how to represent these two views of molecular graphs . 2D Molecular Graph represents molecules as 2D graphs , with atoms as nodes and bonds as edges . We denote it as g2D = ( X , E ) , where X is the atom attribute matrix and E is the bond attribute matrix . Notice that hereE also includes the bond connectivity . Then we will apply one transformation function T2D on the topological graph . Given a 2D molecular graph g2D , its representation h2D can be obtained from a 2D graph neural network ( GNN ) model : h2D = GNN-2D ( T2D ( g2D ) ) = GNN-2D ( T2D ( X , E ) ) . ( 1 ) 3D Molecular Graph additionally includes spatial positions of the atoms , which needless to be static since atoms are in continual motion on a potential energy surface [ 3 ] . 1 The 3D structures at the local minima on this surface are named conformer . As the molecular properties are conformers ensembled [ 34 ] , GraphMVP provides a novel perspective on adopting 3D conformers for learning better representation . Given a conformer g3D = ( X , R ) , its representation via a 3D GNN model is : h3D = GNN-3D ( T3D ( g3D ) ) = GNN-3D ( T3D ( X , R ) ) , ( 2 ) 1A more rigorous way of defining conformer is in [ 56 ] : a conformer is an isomer of a molecule that differs from another isomer by the rotation of a single bond in the molecule . where R is the 3D-coordinate matrix and T3D is the 3D transformation . In what follows , for notation simplicity , we use x and y for the 2D and 3D graphs , i.e. , x , g2D and y , g3D . Then the latent representations are denoted as hx and hy . 3 GRAPHMVP : GRAPH MULTI-VIEW PRE-TRAINING . Our model , termed as Graph Multi-View Pre-training ( GraphMVP ) , conducts self-supervised learning ( SSL ) pre-training with 3D information . The 3D conformers encode rich information about the molecule energy and spatial structure , which are complementary to the 2D topology . Thus , applying SSL between the 3D and 2D views will provide a better 2D representation , which implicitly embeds the ensembles of energies and geometric information for molecules . In the following , we first present an overview of proposed GraphMVP , and then introduce two pretext tasks specialized concerning 3D conformation structures . Finally , we summarize a broader graph SSL family that prevails the 2D molecular graph representation learning with 3D geometry . 3.1 OVERVIEW OF GRAPHMVP 3 As aforementioned , GraphMVP exerts 2D topology and 3D geometry as two complementary views for each molecule . By proceeding SSL between these views , it is expected to learn a 2D representation enhanced with 3D conformation , which can better reflect certain molecular properties . As generic SSL pre-training pipelines , GraphMVP has two stages : pre-training then fine-tuning . In the pre-training , we conduct SSL via auxiliary tasks on data collections providing both 3D and 2D molecular structures . During fine-tuning , the pre-trained 2D GNN models are subsequently fine tuned on specific downstream tasks , where usually only 2D structures are available . At the SSL pre-training stage , we design two pretext tasks : one contrastive and one generative . We conjecture then empirically prove that these two tasks are focusing on different learning aspects , which are concluded into following two points . ( 1 ) From the perspective of representation learning , contrastive SSL is learning from inter-data and generative SSL is learning by intra-data . For contrastive SSL , one key step is to obtain the negative view pairs from inter-data for contrasting ; while generative SSL focuses on each data point itself , by reconstructing the key features at the intra-data level . ( 2 ) From the perspective of distribution learning , contrastive SSL and generative SSL are learning the data distribution from local and global manner , respectively . Contrastive SSL learns the distribution locally by contrasting the pairwise distance at the inter-data level . Thus , with sufficient number of data , the local contrastive operation can iteratively recover the data distribution . Generative SSL , on the other hand , learns the global data density function directly . Therefore , contrastive and generative SSL are essentially conducting representation and distribution learning with different intuitions and disciplines , and we expect that combining these two can lead to better representation . We later carry out an ablation study ( Section 4.4 ) to verify this empirically . In addition , to make the pretext tasks more challenging , we take views for each molecule by randomly masking M nodes ( and corresponding edges ) as the transformation function , i.e. , T2D = T3D = mask . This smart trick has been widely used in graph SSL [ 39 , 91 , 92 ] and has shown robust improvements . 3.2 CONTRASTIVE SELF-SUPERVISED LEARNING BETWEEN 3D AND 2D VIEWS . The main idea of contrastive self-supervised learning ( SSL ) [ 9 , 60 ] is first to define positive and negative pairs of views from the inter-data level , and then to align the positive pairs and contrast the negative pairs simultaneously [ 83 ] . For each molecule , we first extract representations from 2D and 3D views , i.e. , hx and hy . Then we create positive and negative pairs for contrastive learning : the 2D-3D pairs ( x , y ) for the same molecule are treated as positive , and negative otherwise . Finally , we align the positive pairs and contrast the negative ones . The pipeline is shown in Figure 1 . In the following , we discuss two common objective functions on contrastive graph SSL . InfoNCE is first proposed in [ 60 ] , and its effectiveness has been validated both empirically [ 9 , 35 ] and theoretically [ 2 ] . Its formulation is given as follows : LInfoNCE = − 1 2 Ep ( x , y ) [ log exp ( fx ( x , y ) ) exp ( fx ( x , y ) ) + ∑ j exp ( fx ( xj , y ) ) + log exp ( fy ( y , x ) ) exp ( fy ( y , x ) ) + ∑ j exp ( fy ( yj , x ) ) ] , ( 3 ) where x ( j ) , y ( j ) are randomly sampled 2D and 3D views regarding to the anchored pair ( x , y ) . fx ( x , y ) and fy ( y , x ) are scoring functions for the two corresponding views , whose formulation can be flexible . Here we use fx ( x , y ) = fy ( y , x ) = 〈hx , hy〉 . More details are in Appendix D. Energy-Based Model with Noise Contrastive Estimation ( EBM-NCE ) is an alternative that has been widely used in the research line of graph contrastive SSL [ 39 , 72 , 91 , 92 ] . Its intention is essentially the same as InfoNCE , to align positive pairs and contrast negative pairs , while the main difference is the usage of binary cross-entropy and extra noise distribution for negative sampling : LEBM-NCE = − 1 2 Ep ( y ) [ Epn ( x|y ) log ( 1− σ ( fx ( x , y ) ) ) + Ep ( x|y ) log σ ( fx ( x , y ) ) ] − 1 2 Ep ( x ) [ Epn ( y|x ) log ( 1− σ ( fy ( y , x ) ) ) + Ep ( y , x ) log σ ( fy ( y , x ) ) ] , ( 4 ) where pn is the noise distribution and σ is the sigmoid function . We also notice that the final formulation of EBM-NCE shares certain similarities with Jensen-Shannon estimation ( JSE ) [ 59 ] . However , the derivation process and underlying intuition are different : EBM-NCE models the conditional distribution in MI lower bound ( Equation ( 9 ) ) with EBM , while JSE is a special case of variational estimation of f-divergence . Since this is not the main focus of GraphMVP , we expand the a more comprehensive comparison in Appendix D , plus some potential benefits with EBM-NCE . Few works [ 33 ] have witnessed the effect on the choice of objectives in graph contrastive SSL , therefore in GraphMVP , we treat it as a hyper-parameter and further run ablations on them , i.e. , to solely use either InfoNCE ( LC = LInfoNCE ) or EMB-NCE ( LC = LEBM-NCE ) .

Summary: This work aims to leverage additional 3D geometric information for molecular graph representation learning and proposes a multi-view pre-training framework, GraphMVP. Specifically, in GraphMVP, both 2D and 3D information are used, and a combined loss function (i.e., the combination of contrastive self-supervised learning loss and generative self-supervised learning loss ) is adopted to enhance the quality of representations. To verify the model's effectiveness, they conducted experiments and comparisons on a number of datasets. Ablation studies are also carried out to inspect the inner workings of GraphMVP.