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 paper proposes to leverage 3D information to pretrain graph neural networks for learning molecular representations. To achieve this, contrastive and generative SSL strategies are developed accordingly. The empirical results on several datasets show that the improvements of this method over previous 2D graph pretraining methods are consistent and obvious.