Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD
text: 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.

heading: 3.1 AUGMENTATION ENCODER
text: 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).

heading: 3.2 POLICY
text: 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.

heading: 3.3 AUGMENTATIONS
text: We use five graph augmentations including three structural augmentations: node dropping,edge perturbation, and sub-graph inducing, one feature augmentation: feature masking, and one identity augmentation. These augmentations enforce the priors that the semantic meaning of a graph should not change due to perturbations applied to its features or structure. For efficiency, we do not use compute-intensive augmentations such as graph diffusion. Unlike previous works in which the parameters of the augmentations are chosen either randomly or by heuristics, we opt to learn them end-to-end. For example, rather than dropping nodes randomly(You et al., 2020) or computing the probability proportional to a centrality measure (Zhu et al., 2021), we train a model to predict the distribution over all nodes within a graph and then sample from it to decide which nodes to drop. Unlike the policy module, the augmentations are conditioned on the individual graphs. We use a dedicated head for each augmentation modeled as a two layer MLP that learns a distribution over the augmentation parameters. The inputs to the heads are the original graph G and representations Hv and hG from the augmentation encoder. We sample the learned distribution using Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2016) with temperature t.
Node Dropping Head is conditioned on the node and graph representations to decide which nodes within a graph to drop. It receives the concatenation of the node and graph representations as input and predicts a categorical distribution over the nodes. The distribution is then sampled using GumbelTop-K trick (Kool et al., 2019) with a ratio hyper-parameter. We also tried Bernoulli sampling but we observed that it aggressively drops nodes in the few first epochs and the model cannot recover later. To let the gradients flow back from the augmented graph to the head, we introduce edge weights on the augmented graph where an edge weight wij is computed as p(vi) + p(vj) and p(vi) is the probability assigned to node vi. See Algorithm 2 in Appendix.
Edge Perturbation Head is conditioned on head and tail nodes to decide which edges to add or remove. To achieve this, |E| negative edges (Ē) are first randomly sampled to form a set of negative and positive edges E ∪ Ē with size of 2|E|. Edges represented as [ hvi + hvj ‖ 1E(eij) ] (hvi and hvj are the representations of head and tail nodes of edge eij and 1E(eij) is an indicator function indicating if the edge belongs to positive or negative edges) are fed into the head to learn Bernoulli distributions over the edged. We use the predicted probabilities p(eij) as the edge weights to let the gradients flow back to the head. See Algorithm 3 in Appendix.
Sub-graph Inducing Head is conditioned on the node and graph representations to decide which node to select as the center node. It receives the concatenation of the node and graph representations (i.e., [hv ‖ hg]) as input and learns a categorical distribution over the nodes. The distribution is then sampled to select a central node per graph around which a sub-graph is induced using Breadth-First Search (BFS) with K hops. We use a similar trick to node dropping augmentation to overpass the gradients back to the original graph. See Algorithm 4 in Appendix.
Feature Masking Head is conditioned on the node representation to decide which dimensions of the node feature to mask. The head receives the node representation hv and learns a Bernoulli distribution over each feature dimension of the original node feature. The distribution is then sampled to construct a binary mask m over the initial feature space. Because initial node features can consist of categorical attributes (one-hot or multi-hot), we use a linear layer to project them into a continuous space resulting in x′v . The augmented graph has the same structure as the original graph with initial node features x′v m where is the Hadamard product. See Algorithm 5 in Appendix.

heading: 3.4 BASE ENCODER
text: Base encoder gθ(.) : R|V ′|×d′x × R|V′|×|V′| 7−→ R|V′|×dh × Rdh is a shared graph encoder among the augmentations which receives an augmented graph G′ = (V ′, E ′) from the corresponding augmentation head and learns a set of node representations H′v ∈ R|V
′|×dh and a graph representation h′G ∈ Rdh over the augmented graph G′. The goal of learning the augmentations is to help the base encoder learn invariances to such augmentations and as a result produce robust representations. The base encoder is trained end-to-end with the policy and augmentation heads. At inference time, the input graphs are directly fed to the base encoder to compute the encodings for down-stream tasks.

heading: 3.5 TRAINING
text: We follow (Sun et al., 2020) and train the framework end-to-end using deep InfoMax (Hjelm et al., 2019) and maximize the MI between the node representations of one view with graph representation of the other view and vice versa with the following objective:
max ω,µφi,φj ,θ
1 |G| ∑ G∈G
[ 1
|V| ∑ v∈V [ I ( hiv, h j G ) + I ( hjv, h i G )]] (1)
where ω, µ, φi, φj , θ are parameters of modules to be learned, hiv, h j G are representations of node v and graph G encoded by augmentation i and j, and I is the mutual information estimator. We use the Jensen-Shannon MI estimator:
I (hv, hG) = Ep [−log (1 + exp (−D(hv, hG)))]− Ep×p̃ [ −log ( 1 + exp ( D(hv, hG̃) ))] (2)
D(., .) : Rdh × Rdh 7−→ R is a discriminator that takes in a node and a graph representation and scores the agreement between them and is implemented as D(hv, hg) = hn.hTg . We provide the positive samples from the joint distribution (p) and the negative samples from the product of marginals p× p̃, and optimize the model parameters with respect to the objective using mini-batch stochastic gradient descent. We found that regularizing the encoders by randomly alternating between training the base and augmentation encoders helps the base encoder to generalize better. For this purpose, we train the policy and the augmentation heads at each step, but we sample from a Bernoulli to decide whether to update the weights of the base or augmentation encoders . The training process is summarized in Algorithm 1.
Algorithm 1: End-to-end training algorithm. Input: Augmentation heads τi, τj ∈ T , policy r, graph encoders gθ and gω , MI estimator I , loss L, and graphs G ∈ G for sampled batch B = {Gk}Nk=1 ∈ G do {(Hvk , hGk)}Nk=1 = gω (Gk) // Compute the augmentation encodings i, j, pi, pj = rµ({(Hvk , hGk)}Nk=1) // Sample the policy for k = 1 to N do
Gik = τφi(Gk,Hvk , hGk) // Sample the first view Hivk , h i Gk = gθ(G i k) // Compute the first view encodings Gjk = τφj (Gk,Hvk , hGk) // Sample the second view Hjvk , h j Gk = gθ(G j k) // Compute the second view encodings hiGk = h i Gk × pi // Scale the encodings by their probabilities hjGk = h j Gk × pj
end for k = 1 to N and k′ = 1 to N do sik,k′ = I ( hiGk ,H j vk′ ) , sjk,k′ = I ( hjGk ,H i vk′ ) // Compute pairwise similarity: end
5ω,µ,φi,φj ,θ 1N2 N∑ k=1 N∑ k′=1 [ L ( sik,k′ ) + L ( sjk,k′ )] // compute gradients and update
end

heading: 4 EXPERIMENTAL RESULTS

heading: 4.1 UNSUPERVISED REPRESENTATION LEARNING EVALUATION
text: We evaluate LG2AR under the linear evaluation protocol on both node-level and graph-level classification benchmarks where unsupervised models first encode the graphs, and then the encoding are fed into a down-stream linear classifier without fine-tuning the encoder. For graph classification benchmarks, we follow GraphCL and use eight datasets from TUDataset (Morris et al., 2020) and
for the node classification, we follow GCA, and use seven datasets from (Mernyei & Cangea, 2020; Shchur et al., 2018). For fair comparisons, we closely follow the linear evaluation protocol from previous unsupervised works. For evaluating graph classification, we report the mean 10-fold cross validation accuracy with standard deviation after ten runs with a down-stream linear SVM classifier, and for node classification evaluation, we report the mean accuracy of twenty runs over different data splits with a down-stream single layer linear classifier. For details on the evaluation protocol see Appendix A.2 and for implementation details and hyper-parameter settings see Appendix A.3.
For both tasks, we train three variants of our framework denoted as LG2AR-GRU, LG2AR-DeepSet, and LG2AR-Random where each variant indicates its policy instantiation, i.e., GRU, deep set, and random (sampling views from uniform distribution) policies. For graph classification benchmarks, we compare the LG2AR with two supervised baselines: GIN and Graph Attention Network (GAT) (Veličković et al., 2018), five graph kernel methods including Shortest Path kernel (SP) (Borgwardt & Kriegel, 2005), Graphlet Kernel (GK) (Shervashidze et al., 2009), Weisfeiler-Lehman sub-tree kernel (WL) (Shervashidze et al., 2011), Deep Graph Kernels (DGK) (Yanardag & Vishwana, 2015), and Multi-scale Laplacian Graph kernel (MLG) (Kondor & Pan, 2016), and four random walk methods including Random Walk (Gärtner et al., 2003), Node2Vec (Grover & Leskovec, 2016), Sub2Vec (Adhikari et al., 2018), Graph2Vec (Narayanan et al., 2017). We also compare the results with stateof-the-art deep contrastive models including InfoGraph, GraphCL, JOAO, and ADGCL. For node classification benchmarks, we compare our results with random walk methods including DeepWalk with and without concatenating node features, and Node2Vec. We also compare the results with deep learning methods including Graph Autoencoders (GAE, VGAE) (Kipf & Welling, 2016), Graphical Mutual Information Maximization (GMI) (Peng et al., 2020), MVGRL, DGI, and GCA.
The results for graph classification are reported in Tables 1 and 7 (Appendix) and for the node classification are reported in Tables 2 and 7 (Appendix). The reported performance for other models are from the corresponding papers under the same experiment setting. As shown for graph classification, LG2AR achieves state-of-the-art results with respect to unsupervised models across all eight benchmarks. For example, on IMDB-Multi and COLLAB datasets we achieve 2.2% and 6.4% absolute improvement over previous state-of-the-art. We also observe that LG2AR narrows the gap with the best performing supervised baselines even surpassing them on MUTAG dataset. Moreover, for node classification, we observe that we achieve state-of-the-art results with respect to previous unsupervised models in five out of seven benchmarks. For example, we achieve 4.7%, 1.6%, and 1.7 % absolute improvement on the PubMed, Amazon-Photo, and Amazon-Computer datasets. When compared to supervised baselines, we outperform or perform equally good across the benchmarks.

heading: 4.2 SEMI-SUPERVISED LEARNING EVALUATION
text: Furthermore, we evaluated LG2AR in a semi-supervised learning setting on graph classification benchmarks. We follow the experimental protocol introduced in GraphCL and pre-train the encoder in an unsupervised fashion and fine-tune it only on 10% of the labeled data. The results reported in Table 3 show that LG2AR achieves state-of-the-art results compared to previous unsupervised models across all five benchmarks. Most notably, LG2AR achieves an absolute accuracy gain of 3.6% and 2.6% over Collab and Reddit-Multi benchmarks.

heading: 4.3 ANALYSIS OF THE OPTIMIZATION FRAMEWORK
text: In this section, we discuss how LG2AR cannot fall into trivial solutions and compare its optimization with a few notable works. SimSiam (Chen & He, 2021) states that collapsing, i.e. minimum possible loss with constant outputs, cannot be prevented by solely relying on architecture designs such as batch normalization. By designing multiple experiments they concluded that the non-collapsing behaviour of SimSiam still remains an empirical observation. Inspired by MoCo (He et al., 2020) and BYOL (Grill et al., 2020), AutoMix (Liu et al., 2021b) avoids a nested optimization by decoupling its momentum pipeline. LG2AR does not require a momentum pipeline because it is based on contrasting local-global information. Even if the augmentation distribution ends being a Dirac peaking on any of the five augmentations, which we did not observe, a collapse cannot happen and a single level optimization is sufficient. If solely any of the node, sub-graph or edge augmentations are sampled, the inductive biases that we injected and discussed, force the model to at least select a sub part of the graph as an augmented view, hence avoiding the trivial solution of a graph with only one node and no edges. If only identity or feature augmentation is sampled, LG2AR would reduce to a single level optimization method such as InfoGraph where there are no augmentations. JOAO,
re-frames the auto-augmentation on graphs to a min-max optimization. Inspired from robustness and adversarial learning literature, they employ Alternating Gradient Descent (Wang et al., 2019a) to design an approach for learning the augmentation distribution and the encoder parameters in a bi-level optimization setting. We are using the gumbel-softmax trick to sample from the augmentation distribution and let the gradient flow through discrete parameters. Our algorithm alternates between updating the encoder and augmentation parameters by tossing a fair coin in each iteration. Alternating gradients between modules in unsupervised learning is shown to be efficient in avoiding the trivial solutions (Caron et al., 2018; Hassani & Haley, 2019a; Khasahmadi et al., 2020b). With these bag of tricks, LG2AR avoid collapsing to trivial solutions and solving a min-max problem that needs an inefficient bi-level optimization. Moreover, Figure 3 in the Appendix is showing a stable training trajectory for multiple datasets.

heading: 4.4 ABLATION STUDY
text: Effect of the Policy. To investigate the effect of policy, we trained the models with three variants of the policy including GRU, deep set, and random policies. As shown in Tables 1-2, the GRU-based policy outperforms or performs equally well on 12 out of 15 benchmarks whereas the random policy outperforms the other variants in only 1 out of 15 datasets, indicating the importance of learning the view sampling policy. Also, in order to probe what the policy module is learning, we computed the normalized frequency of the sampled augmentations by the GRU-based policy during the training. The frequencies for two graph classification benchmarks (MUTAG and Reddit-binary) and two node classification benchmarks (Coauthor-CS and PubMed) are shown in Figure 2. We observe the following: (1) The policy is learning different distributions over augmentations for each benchmark suggesting that it is adapting to the given datasets. (2) In node classification benchmarks, because we already induce sub-graphs to transform them to inductive tasks, we observe that the policy inclines towards sampling the identity augmentation more frequently which is essentially a sub-graph of the original graph. (3) We observe that regardless of the task, edge perturbation and sub-graph inducing are the two most commonly sampled augmentations. This confirms the observation that sub-graphs are generally beneficial across datasets (You et al., 2020). (4) We observe that for node classification benchmarks, the probability of sampling feature masking augmentation is positively correlated with the initial node feature dimension.
Effect of the Augmentations. To investigate the effect of the augmentations, we run the experiments with single augmentation, and structural vs feature space augmentations. The results shown in Table 4 indicate that: (1) using our single augmentations performs on par or better than baselines, (2) generally structural augmentations contribute more than feature space augmentations, and (3) all augmentations are contributing to the final performance which suggests that the model should use all augmentations while learning the sampling frequencies.
Effect of the Augmentation Heads. To investigate the effect of the augmentations heads without benefiting from the policy, we compare our framework when trained with a random policy with GraphCL for graph classification benchmarks. Both LG2AR-Random and GraphCL sample the augmentations from a uniform distribution, where the former learns distributions over the augmentation parameters and the latter randomly samples those. As shown in Tables 1 and 7, in 8 out of 8 benchmarks LG2AR-Random outperforms GraphCL. For instance, we see an absolute improvement of 2.7% accuracy on Reddit-Binary. This implies that learning the augmentations contributes to the performance boost on the graph classification benchmarks, and when combined with the policy learning, results are further improved. We see less of this effect in transductive tasks suggesting that the policy learning is playing a more important rule in node classification benchmarks. One reason for this may be the fact that GCA unlike GraphCL uses strong topological inductive biases to select the augmentation parameters.
Effect of the Mutual Information Estimator and Discriminator. We investigated four MI estimators including: noise-contrastive estimation (NCE) Gutmann & Hyvärinen (2010); Oord et al. (2018), Jensen-Shannon (JSD) estimator following formulation in Nowozin et al. (2016), normalized temperature-scaled cross-entropy (NT-Xent) Chen et al. (2020b), and Donsker-Varadhan (DV) representation of the KL-divergence Donsker & Varadhan (1975). The results shown in Table 4 suggests that JSD and NT-Xent perform better compared to the other estimators. We also investigated the effect of discriminator by training the model using four variants including dot product, cosine
distance, Bilinear, and MLP discriminators. The results shown in Table 4 suggests that discriminator instantiated as an MLP performs better across the datasets.

heading: 5 CONCLUSION
text: We introduced LG2AR, and end-to-end framework to automate graph contrastive learning. The proposed framework learns the augmentations, view selection policy, and the encoders end-to-end without requiring ad-hoc trial-and-error processes for devising the augmentations for each and every dataset. Experimental results showed that LG2AR achieves state-of-the-art results on 8 out of 8
graph classification benchmarks, and 6 out of 7 node classification benchmarks compared to the previous unsupervised methods. The results also suggest that LG2AR narrows the gap with its supervised counterparts. Furthermore, the results suggest that both learning the policy and learning the augmentations contributes to the performance. In future work, we are planning to investigate large pre-training and transfer learning capabilities of the proposed method.

heading: A APPENDIX
text: A.1 BENCHMARKS
We use seven node classification and eight graph classification benchmarks reported by previous state-of-the-art methods. For node classification benchmarks, we follow GCA (Zhu et al., 2021) and use Wiki-CS (Mernyei & Cangea, 2020) which is a computer science subset of the Wikipedia, Amazon-Computers and Amazon-Photo (Shchur et al., 2018) which are networks of co-purchase relationships constructed from Amazon, Coauthor-CS and Coauthor-Physics (Shchur et al., 2018) which are two academic networks containing co-authorship graphs, and two other citation networks, Cora and Pubmed Sen et al. (2008). For graph classification benchmarks, we follow GraphCL (You et al., 2020) and use benchmarks from TUDatasets (Morris et al., 2020). We use Proteins and DD (Dobson & Doig, 2003) modeling neighborhoods in the amino-acid sequences and protein structures, respectively, MUTAG Kriege & Mutzel (2012) modeling compounds tested for carcinogenicity, COLLAB (Yanardag & Vishwana, 2015) derived from 3 public physics collaboration, Reddit-Binary and Reddit-Multi-5K (Yanardag & Vishwana, 2015) connecting users through responses in Reddit online discussions, and IMDB-Binary and IMDB-Multi Yanardag & Vishwana (2015) connecting actors/actresses based on movie appearances. The statistics of the graph and graph classification benchmarks are summarized in Tables 5 and 6, respectively.
A.2 EVALUATION PROTOCOL DETAILS
For node classification, we follow (Veličković et al., 2019; Zhu et al., 2021) where we train the model with the contrastive method, and then use the resulting embeddings to train and test a simple logistic regression classifier. We train the model for twenty runs over different data splits and report the mean accuracy with standard deviation. For fair evaluation and following GCA, we use a twolayer Graph Convolution Network (GCN) (Kipf & Welling, 2017) for the base encoder across all node classification benchmarks. In order to make the transductive node classification benchmarks compatible with our inductive framework, we sample |B| sub-graphs from the input graph around randomly selected nodes to emulate a batch of graphs, and then feed the batch to our framework. Because sub-graph augmentation occurs before our framework, we remove this augmentation from the policy and also remove its corresponding head. For graph classification, we follow (You et al., 2020; Sun et al., 2020) where we first use the contrastive loss to train the model and then report the best mean 10-fold cross validation accuracy with standard deviation after five runs. The classifier is a linear SVM trained using cross-validation on the training folds of the learned embeddings. Following (Sun et al., 2020), we use GIN layers for the base encoder and treat the number of layers as a hyper-parameter. We observed that contrasting graph-level representation achieves better results in graph classification benchmarks. Finally, we report more node and graph level evaluation results under linear evaluation protocol in Table 7.
A.3 IMPLEMENTATION & HYPER-PARAMETER SELECTION
We implemented the experiments using PyTorch and used Pytorch Geometric library to implement the graph encoders. Each experiments was run on a single RTX 6000 GPU. We initialize the parameters using Xavier initialization and train the model using Adam optimizer. All our graph implementations are sparse and in Pytorch Geometric format. Therefore, in order to let the gradients back-propagate, we use edge weights computed from augmentation heads as a way to pass the gradients. Also, in order to let the gradients back-propagate to the policy module, we multiply the final graph encodings from the two views with the associated probability of each view computed by the policy.
For node graph classification benchmarks, following GCA, we fix the base encoder to a tow-layer GCN model with mean-pooling as the read-out function. We select the number of augmentation encoder layers from [1, 2, 3, 4, 5, 6], number of sub-graphs per batch from [4, 8, 12, 16, 32], hidden dimension from [128, 256, 512], learning rate from [1e-4, 1e-1], number of hops from [1, 2, 3, 4, 5, 6], temperature from [0.7, 1.4], node dropping ration from [0.6, 0.9], and the dropout from [0.0, 0.2]. The augmentation consists of GIN layers with three layer projection heads and a summation read-out function. Following DGI, we use a early-stopping with patience of 50 steps. Following GCA, we train the linear model for 300 epochs with the learning rate of 1e-2.
For graph classification benchmarks, following InfoGraph, we design the both base and augmentation encoders with GIN layers, dedicated three-layer projection heads for node and graph encodings, and a summation read-out function. We share the learning rate and the number of layers between the two encoders and select them from [1e-4, 1e-1] and [1, 2, 3, 4, 5, 6], respectively. We select the batch size from [32, 64, 128], hidden dimension from [128, 256, 512], number of epochs from [10, 20, 40, 60, 100, 200], learning rate from [1e-4, 1e-1], number of hops from [1, 2, 3, 4, 5, 6], temperature from [0.7, 1.4], node dropping ration from [0.6, 0.9], and the dropout from [0.0, 0.2]. We also follow InfoGraph for graph classification and choose the C parameter of the SVM from [10−3, 10−2, ..., 102, 103]. The selected hyper-parameters are shown in Table 9.
Algorithm 2: Node dropping head. Input: Node and graph encodings Hv and hg , graph G = (V, E ,X), Ratio µ p(V) = MLP ([Hv ‖ hg]) V ′ ← Sample-Top-K(p(V), µ) E ′ ← E ⊆ V ′ × V ′ X′ ← X[V ′] WE ← [p(vi) + p(vj) ∀eij ∈ E ′] G′ ← (V ′, E ′,X′,WE) return G′
Algorithm 3: Edge perturbation head. Input: Node and graph encodings Hv and hg , graph G = (V, E ,X), temperature t V ′, E ′,X′,WE ← ∅ Ē = Sample-Negative-Edges(E) for eij to Ē ∪ E do
heij = [ hvi + hvj ‖ 1E(eij) ] p(eij) = MLP(heij ) if Bernoulli-Sample(t, p(eij)) then V ′ ← V ′ ∪ {vi, vj} E ′ ← E ′ ∪ {eij} X′ ← X′ ∪ {hvi , hvj} WE ←WE ∪ {p(eij)} end G′ ← (V ′, E ′,X′,WE)
end return G′
Algorithm 4: Sub-graph inducing head. Input: Node and graph encodings Hv and hg , graph G = (V, E ,X), Number of hops K p(V) = MLP ([Hv ‖ hg]) vcenter ← Sample-Categorical(p(V)) V ′, E ′,X′ ← k-Hop-BFS(vcenter,K) E ′ ← E ⊆ V ′ × V ′ WE ← [p(vi) + p(vj) ∀eij ∈ E ′] G′ ← (V ′, E ′,X′,WE) return G′
Algorithm 5: Feature masking head. Input: Node and graph encodings Hv and hg , graph G = (V, E ,X), temperature t X′ ← Linear(X) M← Bernoulli-Sample(MLP(Hv), t) X′ ← X′ M G′ ← (V, E ,X′,1) return G′

Review summaries:
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.

In this article, the authors investigate the problem of data augmentation in graph contrastive learning. The authors propose an end-to-end learning framework for learning the distribution of data augmentation to achieve an automated learning for graph representation. Experiments are also conducted to demonstrate the effectiveness of the method.

This paper proposes an end-to-end graph augmentation framework, namely LG2AR, that helps encoders learn generalizable representations, to achieve automation and adaptive augmentation. Specifically, L2GAR introduces a probabilistic policy that learns a distribution and then samples two congruent graph views through probabilistic augmentation heads. Finally, L2GAR achieves state-of-the-art results compared to existing methods on various node/graph classification datasets.