Adversarial Graph Augmentation to Improve Graph Contrastive Learning

1 Introduction . Graph representation learning ( GRL ) aims to encode graph-structured data into low-dimensional vector representations , which has recently shown great potential in many applications in biochemistry , physics and social science [ 1–3 ] . Graph neural networks ( GNNs ) , inheriting the power of neural networks [ 4 , 5 ] , have become the almost de facto encoders for GRL [ 6–9 ] . GNNs have been mostly studied in cases with supervised end-to-end training [ 10–16 ] , where a large number of task-specific labels are needed . However , in many applications , annotating labels of graph data takes a lot of time and resources [ 17 , 18 ] , e.g. , identifying pharmacological effect of drug molecule graphs requires living animal experiments [ 19 ] . Therefore , recent research efforts are directed towards studying self-supervised learning for GNNs , where only limited or even no labels are needed [ 18 , 20–31 ] . Designing proper self-supervised-learning principles for GNNs is crucial , as they drive what information of graph-structured data will be captured by GNNs and may heavily impact their performance in downstream tasks . Many previous works adopt the edge-reconstruction principle to match traditional network-embedding requirement [ 32–35 ] , where the edges of the input graph are expected to be reconstructed based on the output of GNNs [ 20 , 21 , 36 ] . Experiments showed that these GNN models learn to over-emphasize node proximity [ 23 ] and may lose subtle but crucial structural information , thus failing in many tasks including node-role classification [ 16 , 35 , 37 , 38 ] and graph classification [ 17 ] . ∗Pan Li and Jennifer Neville co-correspond this work . 2https : //github.com/susheels/adgcl 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To avoid the above issue , graph contrastive learning ( GCL ) has attracted more attention recently [ 18 , 22 , 23 , 25–31 ] . GCL leverages the mutual information maximization principle ( InfoMax ) [ 39 ] that aims to maximize the correspondence between the representations of a graph ( or a node ) in its different augmented forms [ 18 , 24 , 25 , 28–31 ] . Perfect correspondence indicates that a representation precisely identifies its corresponding graph ( or node ) and thus the encoding procedure does not decrease the mutual information between them . However , researchers have found that the InfoMax principle may be risky because it may push encoders to capture redundant information that is irrelevant to the downstream tasks : Redundant information suffices to identify each graph to achieve InfoMax , but encoding it yields brittle representations and may severely deteriorate the performance of the encoder in the downstream tasks [ 40 ] . This observation reminds us of another principle , termed information bottleneck ( IB ) [ 41–46 ] . As opposed to InfoMax , IB asks the encoder to capture the minimal sufficient information for the downstream tasks . Specifically , IB minimizes the information from the original data while maximizing the information that is relevant to the downstream tasks . As the redundant information gets removed , the encoder learnt by IB tends to be more robust and transferable . Recently , IB has been applied to GNNs [ 47 , 48 ] . But IB needs the knowledge of the downstream tasks that may not be available . Hence , a natural question emerges : When the knowledge of downstream tasks are unavailable , how to train GNNs that may remove redundant information ? Previous works highlight some solutions by designing data augmentation strategies for GCL but those strategies are typically task-related and sub-optimal . They either leverage domain knowledge [ 25 , 28 , 30 ] , e.g. , node centralities in network science or molecule motifs in bio-chemistry , or depend on extensive evaluation on the downstream tasks , where the best strategy is selected based on validation performance [ 24 , 30 ] . In this paper , we approach this question by proposing a novel principle that pairs GCL with adversarial training , termed AD-GCL , as shown in Fig.1 . We particularly focus on training self-supervised GNNs for graph-level tasks , though the idea may be generalized for node-level tasks . AD-GCL consists of two components : The first component contains a GNN encoder , which adopts InfoMax to maximize the correspondence/mutual information between the representations of the original graph and its augmented graphs . The second component contains a GNN-based augmenter , which aims to optimize the augmentation strategy to decrease redundant information from the original graph as much as possible . AD-GCL essentially allows the encoder capturing the minimal sufficient information to distinguish graphs in the dataset . We further provide theoretical explanations of AD-GCL . We show that with certain regularization on the search space of the augmenter , AD-GCL can yield a lower bound guarantee of the information related to the downstream tasks , while simultaneously holding an upper bound guarantee of the redundant information from the original graphs , which matches the aim of the IB principle . We further give an instantiation of AD-GCL : The GNN augmenter adopts a task-agnostic augmentation strategy and will learn an input-graph-dependent non-uniform-edge-drop probability to perform graph augmentation . Finally , we extensively evaluate AD-GCL on 18 different benchmark datasets for molecule property classification and regression , and social network classification tasks in different setting viz . unsuper- vised learning ( Sec . 5.1 ) , transfer learning ( Sec . 5.3 ) and semi-supervised learning ( Sec . 5.4 ) learning . AD-GCL achieves significant performance gains in relative improvement and high mean ranks over the datasets compared to state-of-the-art baselines . We also study the theoretical aspects of AD-GCL with apt experiments and analyze the results to offer fresh perspectives ( Sec . 5.2 ) : Interestingly , we observe that AD-GCL outperforms traditional GCL based on non-optimizable augmentation across almost the entire range of perturbation levels . 2 Notations and Preliminaries . We first introduce some preliminary concepts and notations for further exposition . In this work , we consider attributed graphs G = ( V , E ) where V is a node set and E is an edge set . G may have node attributes { Xv ∈ RF | v ∈ V } and edge attributes { Xe ∈ RF | e ∈ E } of dimension F . We denote the set of the neighbors of a node v as Nv . Learning Graph Representations . Given a set of graphs Gi , i = 1 , 2 , ... , n , in some universe G , the aim is to learn an encoder f : G → Rd , where f ( Gi ) can be further used in some downstream task . We also assume that Gi ’ s are all IID sampled from an unknown distribution PG defined over G. In a downstream task , each Gi is associated with a label yi ∈ Y . Another model q : Rd → Y will be learnt to predict Yi based on q ( f ( Gi ) ) . We assume ( Gi , Yi ) ’ s are IID sampled from a distribution PG×Y = PY|GPG , where PY|G is the conditional distribution of the graph label in the downstream task given the graph . Graph Neural Networks ( GNNs ) . In this work , we focus on using GNNs , message passing GNNs in particular [ 49 ] , as the encoder f . For a graph G = ( V , E ) , every node v ∈ V will be paired with a node representation hv initialized as h ( 0 ) v = Xv . These representations will be updated by a GNN . During the kth iteration , each h ( k−1 ) v is updated using v′s neighbourhood information expressed as , h ( k ) v = UPDATE ( k ) ( h ( k−1 ) v , AGGREGATE ( k ) ( { ( h ( k−1 ) u , Xuv ) | u ∈ Nv } ) ) ( 1 ) where AGGREGATE ( · ) is a trainable function that maps the set of node representations and edge attributes Xuv to an aggregated vector , UPDATE ( · ) is another trainable function that maps both v ’ s current representation and the aggregated vector to v ’ s updated representation . After K iterations of Eq . 1 , the graph representation is obtained by pooling the final set of node representations as , f ( G ) : ≜ hG = POOL ( { h ( K ) v | v ∈ V } ) ( 2 ) For design choices regarding aggregation , update and pooling functions we refer the reader to [ 3,7,8 ] . The Mutual Information Maximization Principle . GCL is built upon the InfoMax principle [ 39 ] , which prescribes to learn an encoder f that maximizes the mutual information or the correspondence between the graph and its representation . The rationale behind GCL is that a graph representation f ( G ) should capture the features of the graph G so that representation can distinguish this graph from other graphs . Specifically , the objective of GCL follows InfoMax : max f I ( G ; f ( G ) ) , where G ∼ PG . ( 3 ) where I ( X1 ; X2 ) denotes the mutual information between two random variables X1 and X2 [ 50 ] . Note that the encoder f ( · ) given by GNNs is not injective in the graph space G due to its limited expressive power [ 14 , 15 ] . Specifically , for the graphs that can not be distinguished by 1-WL test [ 51 ] , GNNs will associate them with the same representations . We leave more discussion on 1-WL test in Appendix C. In contrast to using CNNs as encoders , one can never expect GNNs to identify all the graphs in G based their representations , which introduces a unique challenge for GCL . 3 Adversarial Graph Contrastive Learning . In this section , we introduce our adversarial graph contrastive learning ( AD-GCL ) framework and one of its instantiations based on edge perturbation . 3.1 Theoretical Motivation and Formulation of AD-GCL . The InfoMax principle in Eq . 3 could be problematic in practice for general representation learning . Tschannen et al . have shown that for image classification , representations capturing the information that is entirely irrelevant to the image labels are also able to maximize the mutual information but such representations are definitely not useful for image classification [ 40 ] . A similar issue can also be observed in graph representation learning , as illustrated by Fig.2 : We consider a binary graph classification problem with graphs in the dataset ogbg-molbace [ 52 ] . Two GNN encoders with exactly the same architecture are trained to keep mutual information maximization between graph representations and the input graphs , but one of the GNN encoders in the same time is further supervised by random graph labels . Although the GNN encoder supervised by random labels still keeps one-to-one correspondance between every input graph and its representation ( i.e. , mutual information maximization ) , we may observe significant performance degeneration of this GNN encoder when evaluating it over the downstream ground-truth labels . More detailed experiment setup is left in Appendix G.1 . This observation inspires us to rethink what a good graph representation is . Recently , the information bottleneck has applied to learn graph representations [ 47 , 48 ] . Specifically , the objective of graph information bottleneck ( GIB ) follows GIB : max f I ( f ( G ) ; Y ) − βI ( G ; f ( G ) ) , ( 4 ) where ( G , Y ) ∼ PG×Y , β is a positive constant . Comparing Eq . 3 and Eq . 4 , we may observe the different requirements between InfoMax and GIB : InfoMax asks for maximizing the information from the original graph , while GIB asks for minimizing such information but simultaneously maximizing the information that is relevant to the downstream tasks . As GIB asks to remove redundant information , GIB naturally avoids the issue encountered in Fig.2 . Removing extra information also makes GNNs trained w.r.t . GIB robust to adverserial attack and strongly transferrable [ 47 , 48 ] . Unfortunately , GIB requires the knowledge of the class labels Y from the downstream task and thus does not apply to self-supervised training of GNNs where there are few or no labels . Then , the question is how to learn robust and transferable GNNs in a self-supervised way . To address this , we will develop a GCL approach that uses adversarial learning to avoid capturing redundant information during the representation learning . In general , GCL methods use graph data augmentation ( GDA ) processes to perturb the original observed graphs and decrease the amount of information they encode . Then , the methods apply InfoMax over perturbed graph pairs ( using different GDAs ) to train an encoder f to capture the remaining information . Definition 1 ( Graph Data Augmentation ( GDA ) ) . For a graph G ∈ G , T ( G ) denotes a graph data augmentation of G , which is a distribution defined over G conditioned on G. We use t ( G ) ∈ G to denote a sample of T ( G ) . Specifically , given two ways of GDA T1 and T2 , the objective of GCL becomes GDA-GCL : max f I ( f ( t1 ( G ) ) ; f ( t2 ( G ) ) ) , where G ∼ PG , ti ( G ) ∼ Ti ( G ) , i ∈ { 1 , 2 } . ( 5 ) In practice , GDA processes are often pre-designed based on either domain knowledge or extensive evaluation , and improper choice of GDA may severely impact the downstream performance [ 17 , 24 ] . We will review a few GDAs adopted in existing works in Sec.4 . In contrast to previous predefined GDAs , our idea , inspired by GIB , is to learn the GDA process ( over a parameterized family ) , so that the encoder f can capture the minimal information that is sufficient to identify each graph . AD-GCL : We optimize the following objective , over a GDA family T ( defined below ) . AD-GCL : min T∈T max f I ( f ( G ) ; f ( t ( G ) ) ) , where G ∼ PG , t ( G ) ∼ T ( G ) , ( 6 ) Definition 2 ( Graph Data Augmentation Family ) . Let T denote a family of different GDAs TΦ ( · ) , where Φ is the parameter in some universe . A TΦ ( · ) ∈ T is a specific GDA with parameter Φ . The min-max principle in AD-GCL aims to train the encoder such that even with a very aggressive GDA ( i.e. , where t ( G ) is very different from G ) , the mutual information / the correspondence between the perturbed graph and the original graph can be maximized . Compared with the two GDAs adopted in GDA-GCL ( Eq.5 ) , AD-GCL views the original graph G as the anchor while pushing its perturbation T ( G ) as far from the anchor as it can . The automatic search over T ∈ T saves a great deal of effort evaluating different combinations of GDA as adopted in [ 24 ] . Relating AD-GCL to the downstream task . Next , we will theoretically characterize the property of the encoder trained via AD-GCL . The analysis here not only further illustrates the rationale of AD-GCL but helps design practical T when some knowledge of Y is accessible . But note that our analysis does not make any assumption on the availability of Y . Note that GNNs learning graph representations is very different from CNNs learning image representations because GNNs are never injective mappings between the graph universe G and the representation space Rd , because the expressive power of GNNs is limited by the 1-WL test [ 14 , 15 , 51 ] . So , we need to define a quotient space of G based on the equivalence given by the 1-WL test . Definition 3 ( Graph Quotient Space ) . Define the equivalence ∼= between two graphs G1 ∼= G2 if G1 , G2 can not be distinguished by the 1-WL test . Define the quotient space G′ = G/ ∼= . So every element in the quotient space , i.e. , G′ ∈ G′ , is a representative graph from a family of graphs that can not be distinguished by the 1-WL test . Note that our definition also allows attributed graphs . Definition 4 ( Probability Measures in G′ ) . Define PG′ over the space G′ such that PG′ ( G′ ) = PG ( G ∼= G′ ) for any G′ ∈ G′ . Further define PG′×Y ( G′ , Y ′ ) = PG×Y ( G ∼= G′ , Y = Y ′ ) . Given a GDA T ( · ) defined over G , define a distribution on G′ , T ′ ( G′ ) = EG∼PG [ T ( G ) |G ∼= G′ ] for G′ ∈ G′ . Now , we provide our theoretical results and give their implication . The proof is in the Appendix B. Theorem 1 . Suppose the encoder f is implemented by a GNN as powerful as the 1-WL test . Suppose G is a countable space and thus G′ is a countable space . Then , the optimal solution ( f∗ , T ∗ ) to AD-GCL satisfies , letting T ′∗ ( G′ ) = EG∼PG [ T ∗ ( G ) |G ∼= G′ ] ,

Summary: In this paper, the authors try to improve the graph contrastive learning. They point out that the representation learned by maximizing mutual information does not necessarily help for downstream tasks. And they propose adversarial graph augmentation to enhance the InfoMax objective for self-supervised learning of GNNs. They argue that AD-GCL will make GNNs avoid learning the redundant information during the training, which are not beneficial to GraphCL. The author also gave a relatively sufficient theoretical explanation. Extensive experiments are conduct to demonstrate the effectiveness of the proposed method.