Learning Graph Augmentations to Learn Graph Representations

Devising augmentations for graph contrastive learning is challenging due to their irregular structure , drastic distribution shifts , and nonequivalent feature spaces across datasets . We introduce LG2AR , Learning Graph Augmentations to Learn Graph Representations , which is an end-to-end automatic graph augmentation framework that helps encoders learn generalizable representations on both node and graph levels . LG2AR consists of a probabilistic policy that learns a distribution over augmentations and a set of probabilistic augmentation heads that learn distributions over augmentation parameters . We show that LG2AR achieves state-of-the-art results on 18 out of 20 graph-level and node-level benchmarks compared to previous unsupervised models under both linear and semi-supervised evaluation protocols . 1 INTRODUCTION . Graph Neural Networks ( GNNs ) are a class of deep models that learn node representations over order-invariant and variable-size data , structured as graphs , through an iterative process of transferring , transforming , and aggregating the representations from topological neighbors . The learned representations are then summarized into a graph-level representation ( Li et al. , 2015 ; Gilmer et al. , 2017 ; Kipf & Welling , 2017 ; Veličković et al. , 2018 ; Xu et al. , 2019 ; Khasahmadi et al. , 2020a ) . GNNs are applied to non-Euclidean data such as point clouds ( Hassani & Haley , 2019b ) , robot designs ( Wang et al. , 2019b ) , physical processes ( Sanchez-Gonzalez et al. , 2020 ) , molecules ( Duvenaud et al. , 2015 ) , social networks ( Kipf & Welling , 2017 ) , and knowledge graphs ( Vivona & Hassani , 2019 ) . GNNs are mostly trained end-to-end with supervision from task-dependent labels . Nevertheless , annotating graphs is more challenging compared to other common modalities because they usually represent concepts in specialized domains such as biology where labeling through wet-lab experiments is resource-intensive ( You et al. , 2020 ; Hu et al. , 2020 ) and labeling them procedurally using domain knowledge is costly ( Sun et al. , 2020 ) . To address this , unsupervised objectives are coupled with GNNs to learn representations without relying on labels which are transferable to a priori unknown down-stream tasks . Reconstruction-based methods , i.e. , Graph AutoEncoders ( GAE ) , preserve the topological closeness of the nodes in the representations by forcing the model to recover the neighbors from the latent space ( Kipf & Welling , 2016 ; Garcia Duran & Niepert , 2017 ; Pan et al. , 2018 ; Park et al. , 2019 ) . GAEs over-emphasize proximity information at the cost of structural information ( Veličković et al. , 2019 ) . Contrastive methods train graph encoders by maximizing the Mutual Information ( MI ) between node-node , node-graph , or graph-graph representations achieving state-of-the-art results on both node and graph classification benchmarks ( Veličković et al. , 2019 ; Sun et al. , 2020 ; Hassani & Khasahmadi , 2020 ; You et al. , 2020 ; Zhu et al. , 2021 ) . For a review on graph contrastive learning see ( Wu et al. , 2021 ; Xie et al. , 2021 ; Liu et al. , 2021a ) . Contrastive learning is essentially learning invariances to data augmentations which are thoroughly explored for Computer Vision ( CV ) ( Shorten & Khoshgoftaar , 2019 ) and Natural Language Processing ( NLP ) ( Feng et al. , 2021 ) . Learning policies to sample dataset-conditioned augmentations is also studied in CV ( Hataya et al. , 2020 ; Lim et al. , 2019 ; Cubuk et al. , 2019 ; Li et al. , 2020 ) . The irregular structure of graphs complicates both adopting augmentations used on images and also devising new augmentation strategies ( Zhao et al. , 2021 ) . Unlike image datasets where the distribution is mostly from natural images , graph datasets are abstractions diverse in nature and contain shifts on marginal/conditional distributions and nonequivalent feature spaces across datasets . This implies that the effect of augmentations is different across the datasets and hence both augmentations and their selection policy should be learned from the data to adapt to new datasets . Present Work . We introduce LG2AR , Learning Graph Augmentations to Learn Graph Representations , a fully-automated end-to-end contrastive learning framework that helps encoders learn transferable representations . Specifically , our contributions are as follows : ( 1 ) We introduce a probabilistic augmentation selection policy that learns a distribution over the augmentation space conditioned on the dataset to automate the combinatorial augmentation selection process . ( 2 ) We introduce probabilistic augmentation heads where each head learns a distribution over the parameters of a specific augmentation to learn better augmentations for each dataset . ( 3 ) We train the policy and the augmentations end-to-end without requiring an outer-loop optimization and show that unlike other methods , our approach can be used for both node-level and graph-level tasks . Finally , ( 4 ) we exhaustively evaluate our approach under linear and semi-supervised evaluation protocols and show that it achieves state-of-the-art results on 18 out of 20 graph and node level classification benchmarks . 2 RELATED WORK . Graph augmentations are explored in supervised settings to alleviate over-smoothing and over-fitting . DropEdge ( Rong et al. , 2020 ) randomly drops a fraction of the edges during training . ADAEDGE ( Chen et al. , 2020a ) learns to perturb edges between based on the predicted node classes . BGCN ( Zhang et al. , 2019 ) generates an ensemble of deonised graphs by perturbing edges . GAUG ( Zhao et al. , 2021 ) trains a GAE to generate edge probabilities and interpolates them with the original connectivity to sample a denoised graph . FLAG ( Kong et al. , 2020 ) augments node features with adversarial perturbations and GraphMask ( Schlichtkrull et al. , 2021 ) introduces differentiable edge masking to achieve interpretability . These works assume that a specific type of augmentation suffice for all supervised tasks and do not utilize the benefit of mixing the augmentations . Graph augmentations are also studied in contrastive setting to learn transferable graph representations . DGI ( Veličković et al. , 2019 ) and InfoGraph ( Sun et al. , 2020 ) adopt DeepInfoMax ( Hjelm et al. , 2019 ) and enforce the consistency between local ( node ) and global ( graph ) representation . MVGRL ( Hassani & Khasahmadi , 2020 ) augments a graph via graph diffusion and constructs two views by randomly sampling sub-graphs from the adjacency and diffusion matrices . GCC ( Qiu et al. , 2020 ) uses sub-graph instance discrimination and contrasts sub-graphs from a multi-hop ego network . GraphCL ( You et al. , 2020 ) uses trial-and-error to hand-pick graph augmentations and the corresponding parameters of each augmentation . JOAO ( You et al. , 2021 ) extends the GraphCL using a bi-level min-max optimization that learns to select the augmentations . Nevertheless , it does not show much improvement over GraphCL . GRACE ( Zhu et al. , 2020 ) uses a similar approach to GraphCL learn node representations . GCA ( Zhu et al. , 2021 ) uses a set of heuristics to adaptively pick the augmentation parameters . BGRL ( Thakoor et al. , 2021 ) adopts BYOL ( Grill et al. , 2020 ) and uses random augmentations to learn node representations . ADGCL ( Suresh et al. , 2021 ) introduces adversarial graph augmentation strategies to avoid capturing redundant information . Different from these methods , LG2AR emphasizes the importance of conditional augmentations and learns a distribution over the augmentation space along with a set of distributions over the augmentation parameters endto-end without requiring a bi-level optimization and outperforms the previous contrastive methods on both node and graph level benchmarks under linear and semi-supervised evaluation protocols . 3 METHOD . Given a dataset of graphs G = { Gk } Nk=1 where each graph Gk = ( V , E , X ) consists of |V| nodes , |E| edges ( E ⊆ V × V ) , and initial node features X ∈ R|V|×dx , and assuming that the semantic labels are not available during the training , the goal is to learn node-level representations Hv ∈ R|V|×dh and graph-level representation hG ∈ Rdh such that the learned representations are transferable to the down-stream tasks unknown a priori . Assuming a set of possible rational augmentations T over G where each augmentation τi ∈ T , i ∈ { 1 , ... , |T | } is defined as a function over graph Gk that generates an identity-preserving view of the graph : Gik = τi ( Gk ) , a contrastive framework with negative sampling strategy uses T to draw positive samples from the joint distribution p ( τi ( Gk ) , τj ( Gk ) ) in order to maximize the agreement between different views of the same graph Gk and to draw negative samples from the product of marginals p ( τi ( Gk ) ) × p ( τj ( Gk′ ) ) to minimize it for views Augmentation Encoder ( Deterministic ) Policy ( Probabilistic ) Augmentation Heads ( Probabilistic & Dedicated ) Main Encoder ( Deterministic & Shared ) from two distinct graphs Gk and Gk′ , k 6= k′ . Most works , sample the augmentations uniformly and use trail-and-error to determine a single parameter for each augmentation , e.g. , the probability of dropping nodes . LG2AR , on the other hand , learns the sampling distribution over T and also learns parametric augmentation τφi end-to-end along with the representations to generate robust representations . The architecture of LG2AR ( Figure 1 ) achieves this using an augmentation encoder , a probabilistic policy , a set of probabilistic augmentation heads , and a shared base encoder . The details are as follows . 3.1 AUGMENTATION ENCODER . Augmentation encoder gω ( . ) : R|V|×dx × R|E| 7−→ R|V|×dh × Rdh learns a set of node encoding Hv ∈ R|V|×dh and a graph encoding hg ∈ Rdh over the input graph Gk to provide the subsequent modules with expressive encodings so that they can condition their predictions on the input graphs . The augmentation encoder consists of a GNN producing node representations , a read-out function ( summation ) aggregating the representations into a graph representation , and two dedicated projection heads ( three layer MLPs ) applied to the learned representations . To encode graphs , we opted for expressive power and adopted graph isomorphism network ( GIN ) ( Xu et al. , 2019 ) . 3.2 POLICY . Policy rµ ( . ) : R|B|×dh 7−→ R|τ | is a probabilistic module that receives a batch of graph-level representations Hg ∈ R|B|×dh from the augmentation encoder , constructs a categorical distribution over the possible augmentations T , and then samples two augmentations , τφi and τφj from that distribution for each batch with temperature t. It is shown that conditioning the augmentation sampling on the dataset helps achieve better performance ( You et al. , 2020 ) . However , feeding the whole dataset to the policy module is computationally expensive and hence we approximate it by conditioning the policy on mini-batches . Moreover , the policy must be invariant to the order of representations within the batch . To enforce this , we tried two strategies : ( 1 ) a policy instantiated as a deep set ( Zaheer et al. , 2017 ) where representations are first projected and then aggregated into a batch representation , and ( 2 ) a policy instantiated as an RNN where we impose an order on the representations by sorting them based on their L2-norm and then feeding them into a GRU ( Cho et al. , 2014 ) . We use the last hidden state as the batch representation . We observed that GRU based policy performed better . The policy module automates the ad-hoc trial-and-error augmentation selection process . To let the gradients flow back to the policy module , we use a skip-connection and multiply the final graph representations by their associated augmentation probabilities predicted by the policy .

Summary: This paper proposes an end-to-end graph augmentation framework for graph representation learning with contrastive learning. Specifically, the authors introduced a policy sampler and node/edge sampler for each policy. The authors showed empirically that the LG2AR pipeline outperforms previous SOTA on various benchmarks.