SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks
Abstract
In the past few years, graph neural networks (GNNs) have become the de facto model of choice for graph classification.
While, from the theoretical viewpoint, most GNNs can operate on graphs of any size, it is empirically observed that their classification performance degrades when they are applied on graphs with sizes that differ from those in the training data.
Previous works have tried to tackle this issue in graph classification by providing the model with inductive biases derived from assumptions on the generative process of the graphs, or by requiring access to graphs from the test domain.
The first strategy is tied to the quality of the assumptions made for the generative process, and requires the use of specific models designed after the explicit definition of the generative process of the data, leaving open the question of how to improve the performance of generic GNN models in general settings.
On the other hand, the second strategy can be applied to any GNN, but requires access to information that is not always easy to obtain.
In this work we consider the scenario in which we only have access to the training data, and we propose a regularization strategy that can be applied to any GNN to improve its generalization capabilities from smaller to larger graphs without requiring access to the test data.
Our regularization is based on the idea of simulating a shift in the size of the training graphs using coarsening techniques, and enforcing the model to be robust to such a shift.
Experimental results on standard datasets show that popular GNN models, trained on the 50% smallest graphs in the dataset and tested on the 10% largest graphs, obtain performance improvements of up to 30% when trained with our regularization strategy.
Introduction
Graph structured data are found in a large variety of domains, ranging from social networks, to molecules and transportation networks.
When dealing with graph data in machine learning settings, it is common to resort to graph representation learning [29].
Graph representation learning is the task of learning to generate vector representations of the nodes, or of entire graphs, that encode structural and feature-related information that can be used for downstream tasks.
In this area, Graph Neural
[Table] [Figure] Networks (GNNs) are a deep learning model for graph representation learning that has become the de facto standard for many tasks, including graph classification, which is the focus of this work.
GNNs are designed to be able to work on graphs of any size.
However, empirical results show that they struggle at generalizing to graphs of sizes that differ from those in the training data [33, 25, 6, 68].
Obtaining good size-generalization performance from smaller to larger graphs is crucial for several reasons.
First , it is common for graphs in the same domain to vary in size.
For example, social networks can range from tens to millions of nodes, and molecular graphs can go from few atoms, to large compounds with thousands of atoms.
Second , in many scenarios, obtaining labels for smaller graphs is cheaper than for larger graphs.
For example, recently deep learning is being used in combinatorial optimization problems [5], and GNNs are heavily used in this area as many problems can be formulated as graph classification instances [53, 24, 11].
Combinatorial optimization often deals with NP-hard problems, and obtaining ground truth labels can be extremely expensive for large scale graphs.
Third , training on smaller graphs requires less computational resources.
Two approaches have been proposed in literature to tackle the poor size-generalization of GNNs.
The first approach makes assumptions on the causal graph describing the generative process behind the graphs in the dataset, and then designs a model that is tailored towards such causal graph [6]. This strategy requires an explicit definition of the generative process behind the data, which, in the case of
Bevilacqua et al.
[6], leads to a model requiring the computation of induced homomorphism densities over all connected k -vertex subgraphs in order to make a prediction.
While this strategy shows great results on synthetic graphs produced by the assumed generative process, its benefits are reduced on real-world graphs, for which the underlying generative process is unknown.
The second approach assumes access to graphs coming from different domains, or from the test distribution, and applies domain adaptation techniques to transfer knowledge across domains [68, 17].
These methods do not require ad-hoc models, but require knowledge about each graph’s domain, or access to the test distribution, both of which are not always available or easily obtainable in advance.
In this work we consider the graph classification scenario in which we only have access to the training data (as in the first approach described above).
However, in contrast to prior work, we do not try to improve the size-generalization capabilities (from smaller to larger graphs) of GNNs by handcrafting models on the base of specific knowledge or assumptions, but we study a general form of regularization that can be applied to any GNN, and that aims at letting the model learn to be robust to size-shifts.
The idea behind our regularization is to simulate a shift in the size of the training graphs, and to regularize the model to be robust towards this shift.
The size-shift simulation is obtained using graph coarsening techniques, while the regularization is performed by minimizing the discrepancy between the distributions of the node embeddings generated by the GNN model for the original training graphs and their coarsened versions (an overview is shown in Figure 1).
In the experimental section we evaluate GNN models on popular benchmark datasets using a particular train/test split designed to test for size-generalizaiton.
In particular, following [68, 6], we train on the 50% smallest graphs and test on the 10% largest.
This split leads to an average size of the test graphs which is 3 to 9 times larger than the average size of the training graphs.
Our results show that our regularization method applied to standard GNN models (GCN [36], PNA [15], GIN [65]) leads to an average performance improvement on the test set of up to 30% across datasets.
Furthermore, these performance improvements show that standard GNN models trained with our regularization can achieve better size-generalization performance than more expensive models.
The latter more expensive models were designed after making explicit assumptions on the generative process behind the graphs in the dataset [6], while our method relies on learning to be robust towards size-shifts.
Our contributions.
Our contributions are threefold.
(1) We propose a regularization strategy that can be applied to any GNN to improve its size-generalization capabilities on the task of graph classification.
(2) We analyze the impact that our strategy has on the embeddings generated by a
GNN.
(3) We test our strategy on popular GNN models, and compare against previously proposed methods, showing that “standard” GNN models augmented with our regularization strategy achieve comparable or better size-generalization performance than more complex models designed after explicit assumptions on the generative process behind the data.
Preliminaries
Notation.
Let G =
( V, E ) be a graph, where V is the set of nodes with size n , and
E ⊆
V × V is the set of edges.
For a node v ∈ V , we use N v to indicate the set of its one-hop neighbours (i.e., the nodes connected to v by an edge).
We represent an attributed graph G with a tuple ( A , X ) , where
A ∈ { 0 , 1 } n × n is the adjacency matrix (i.e., the element in position i, j is equal to 1 if there is an edge between nodes i and j , and zero otherwise), and X ∈ R n × d is the feature matrix, where the i -th row, X i , contains the d -dimensional feature vector for node i .
2.1 Graph Neural Networks
Most popular Graph Neural Networks (GNNs) [12, 64] models follow the message-passing paradigm [26].
Let H ( ℓ ) ∈ R n × d ′ be the matrix containing the node representations at layer ℓ (i.e., the representation for node i is contained in the i -th row H ( ℓ ) i ), with H (0) =
X .
Given an aggregation function Φ , which is permutation invariant, and a learnable update function Ψ (usually a neural network), a message passing layer updates the representation of every node v as follows:
[Equation]
where m v represents the aggregated message received by node v from its neighbours.
After L message-passing layers, the final node embeddings H ( L ) are used to perform a given task (e.g., they are the input to a network performing the task), and the whole system is trained end-to-end.
2.2 Graph Coarsening
Given a graph G =
( V, E ) , a coarsened version G c =
( V c , E c ) is obtained by partitioning
V and grouping together the nodes that belong to the same subset of the partition.
More formally, given a partition P = { S 1 , S 2 , . . .
, S n ′ } where S i ⊆ V , S i ∩ S
j = ∅ , (cid:83)
i S i
= V , we associate a node s i to each subset S
i
and let V c
= { s 1 , s 2 , . . .
, s n ′ } and E c = { ( s i , s j ) | there is a node in S i connected to at least one node in S j } .
The goal of graph coarsening is to generate a coarsened version of a graph (i.e., n ′ < n ) such that specific properties are preserved.
For instance, Loukas and Vandergheynst [45] aim at preserving the principal eigenvalues and eigenspaces, while Jin et al.
[32] aim at minimizing a distance function that measures the changes in the structure and connectivity that occur in the coarsening process.
Jin et al.
[32] further prove that their algorithm bounds the spectral distance [35] between the original and a “lifted” version of the coarsened graph.
2.3 Central Moment Discrepancy
The Central Moment Discrepancy (CMD) [70] is a metric used to measure the discrepancy between the distribution of high-dimensional random variables.
Let p and q be two probability distributions with support in the interval [ a, b ] d and let c k be the k -th order moment, then the CMD between p and q is defined as:
[Equation]
At a practical level, we follow Zhu et al.
[72], and limit the number of considered moments to 5 .
To align with the definition of CMD it is possible to treat node embeddings as realizations of bounded high dimensional distributions by applying bounded activation functions like tanh or sigmoid at the embeddings layer (as done in [70, 72]).
Our Method
The high level idea behind our method is to “simulate” a size-shift in the training dataset, and to regularize the GNN to be robust to this shift (an overview is shown in Figure 1).
While previous works rely on models that incorporate domain-specific inductive biases and/or are tied to assumptions on the generative process of the data, our method aims at learning to be robust towards size-shifts.
Simulating the shift.
Let G c : ( G, r ) → G cr be a coarsening function , that takes as input a graph
G with n nodes and a ratio r ∈ (0 , 1) , and returns a coarsened version of the graph, G cr , which has ⌊ r × n ⌋ nodes.
Our method is not bound to a specific coarsening function G c , any existing method can be used, e.g. [45, 43, 32].
Our method then proceeds in the following way: given a dataset of graphs D
= { G 1 , G 2 , . .
. , G ℓ } , a coarsening function G c , and a set of k coarsening ratios
C = { r 1 , r 2 , . . .
, r k } , we create a coarsened dataset D r j for each ratio by applying the coarsening function to each graph in the dataset: D r j = { G cr j , 1 = G c
( G 1 , r j ) , G cr j , 2 =
G c ( G 2 , r j ) , · · · , G cr j ,ℓ =
G c ( G ℓ , r j ) } , ∀ j = 1 , 2 , . . . , k .
When the graph is attributed, we obtain the features for each node of the coarsened version of the graph by aggregating the features of the corresponding nodes in the original graph using simple aggregations (e.g., mean, max, sum), as commonly done in readout functions for GNNs [64].
Regularizing the GNN.
Given the graph dataset D , its coarsened versions D r 1 , D r 2 , · · · , D r k , and a GNN f θ , where θ are the parameters, we aim at regularizing the GNN so that the distribution of the node embeddings generated by the model for the graphs in the original dataset and their coarsened version is similar.
In other words, for a given graph, we want the model to generate a distribution of node embeddings that is robust across different coarsened versions of the graph.
In more detail, we regularize a GNN by minimizing the CMD [70] between the distribution of the node embeddings generated by the GNN for the original graph and the distribution of the node embeddings generated by the GNN for the coarsened version(s) of the graph.
We choose CMD as it has proven to be successful and stable as a regularization term for non-linear models [42, 72].
Formally, let L be the supervised loss function that is used to train the model (e.g. cross entropy for classification), and let λ be a term measuring the strength of the regularization.
Our optimization problem is
[Equation]
Our method is model-agnostic, and can hence be applied to any GNN.
Pseudocode and practical aspect.
Algorithm 1 presents the pseudocode for computing our regularization loss for a batch of graphs during training.
The pseudocode is presented in an extended manner for clarity, but at a practical level, the coarsened versions of the training graphs are pre-computed before training, and the computation of the loss is done in a vectored manner so it is computed concurrently for all graphs in the batch in one pass (i.e., we do not iterate through the graphs in the batch).
At a practical level, in our experiments we notice that training a model with our regularization introduces a 50% overhead in the running time for a training epoch w.r.t.
training the same model without regularization.
3.1 Limitations
Similarly to previous works [6, 68] we are assuming that there are some properties that determine the label of a graph and that do not depend on the size of the graph.
In a scenario in which small graphs do not carry information that is relevant for solving the task on larger graphs, we do not expect our regularization to provide substantial benefits, and the best option would be to include larger graphs in the training set.
In our experiments, the ratio between the average size of the graphs in the test set and the average size of the graphs in the training set is between 3 and 9.
While our method shows significant performance improvements in this setting, it may show lower performance improvements when this ratio reaches much higher values
[Table]
[Table]
Analysis of Node Embeddings
[Table]
Before evaluating how our regularization impacts the size-generalization performance of a model, we analyze the effects that our regularization has on the embeddings generated by the model.
In order to do this, we consider two identical GIN [65] models and we train one with our regularization and one without.
We then use these models to generate node embeddings and use the Central Kernel Alignment (CKA) [38] to study the generated representations, similarly to [34].
CKA takes as input two matrices A ∈ R m × d ′ , B ∈ R m × d ′′ of representations and provides a value between 0 and 1 quantifying how aligned the representations are (allowing for d ′ ̸ = d ′′ ).
CKA quantifies the similarity of representations learned by (possibly) different models and gives us a way to study the effects of our regularization.
We report results averaged over 10 different random seeds.
We show only results for the DD dataset for space limitations, but the same trends are observed for the other datasets (see Appendix).
First, we ask if a model trained with our regularization and a model trained without our regularization produce similar node embeddings.
To compare the node embeddings between the two models we obtain a representation for each graph by concatenating the node embeddings for that graph.
We then compute the CKA between the representations generated by the model trained with regularization and the representations from model trained without regularization.
We do this for the original graphs and the coarsened versions of the graphs, and report the average CKA.
Results shown in Tablehighlight that there is low alignment between the embeddings generated by a model trained with regularization and one trained without regularization (the average CKA across datasets is 0 .
48 for the original graphs).
This is even more apparent for the coarsened graphs, showing that our regularization [Figure] is impacting the way a model generates embeddings, especially for size-shifted versions of the graphs.
To analyze the trend more in detail, we also plot the CKA values across coarsening ratios in Figure(a).
Here we notice that the CKA between the embeddings generated by a model trained with our regularization and a model trained without, decreases significantly with ratios lower than 0.8.
This shows that, while the alignment between the representations learned by the models trained with and without our regularization is low even for a coarsening ratio of 0.9 (average CKA value is 0.47), it becomes 15% lower (on average) for ratios smaller than 0.7, indicating that when the size shift is strong, there is a stronger misalignment between models trained with and without regularization.
Second, we ask if a model trained with our regularization is in fact being more robust to size shifts with respect to a model trained without our regularization.
To answer this question we analyze the alignment between the representations generated by the same model, trained with or without regularization, on original graphs and coarsened versions of the graphs.
In Figure 2 (b) we plot the average CKA between the node embeddings for the original graphs and the node embeddings for the coarsened versions of the graphs, both generated by the same model, for different ratios.
The plot shows that a model trained with our regularization is learning to be more robust to size shifts, as its representation for coarsened versions of the graphs are much more aligned to the representations for the original graphs compared to a model trained without regularization.
This is even more apparent for low coarsening ratios (i.e., when the shift is stronger), indicating that our regularization is in fact making the model more robust to size shifts.
Evaluation
Datasets.
Following prior works [6, 68], we evaluate on datasets from the TUDataset library [48] with a train/test split explicitly designed to test for size-generalization.
In particular, the training split contains the graphs with size smaller or equal than the 50-th percentile, and the test split contains graphs with a size larger or equal than the 90-th percentile.
With this split, the average size of the test graphs is 3 to 9 times larger than the average size of the training graphs (in more detail, it isfor NCI1 and NCI109, 9 for PROTEINS, and 5 for DD).
We use 10% of the training graphs as a validation set.
More details on the datasets can be found in the Appendix.
Models.
As in Bevilacqua et al.
[6], we consider three standard GNN models: GCN [36], GIN [65], and PNA [15], a more expressive GNN: RPGNN [49], and two graph kernels: the Graphlet Counting kernel (GC Kernel)
[58], and the WL Kernel [59].
We further apply the Invariant Risk Minimization (IRM) [1] penalty to the standard GNN models, and compare against the E-Invariant models ( Γ 1-hot , Γ GIN , Γ RPGIN ) introduced in Bevilacqua et al. For IRM, we follow [6], and [Table]
[Table] artificially create two environments: one with graphs with size smaller than the median size of the graphs in the training set [Table] and one with graphs with size larger than the median size in the training set.
Hyperparameters and Evaluation Protocol.
For the standard GNN models we adopt the same hyperparameters of Bevilacqua et al.
[6], which were obtained from a tuning procedure involving number of layers, learning rate, batch size, hidden layers dimension [Table] and regularization terms.
All models are trained with early stopping (i [Table] e., taking the weights at the epoch with best performance on the validation set).
We do not modify the original hyperparameters when we apply our regularization.
We tune the regularization coefficient λ and the coarse ratios C for GIN on the PROTEINS dataset (we found λ = 0 .
1 and C = { 0 . 8 , 0 . 9 } to be the best on the validation set), and we apply these settings to all models and datasets when using our regularization, to show that our method can work without extensive (and expensive) hyperparameter tuning.
We use the SGC coarsening algorithm [32] to obtain the coarsened versions of the graphs (chosen for its theoretical properties).
The only hyperparameter that is tuned on a per-dataset basis is the aggregation strategy used to obtain the features for the nodes in the coarsened versions of the graphs (in particular we take the best performing between ‘sum’, ‘max’, and ‘mean’).
For all other models we use the hyperparameters introduced by their respective papers.
Bevilacqua et al.
[6] presented results by averaging overruns (each time with a different random seed), but given the high variance observed, we re-run the experiments for all the models and present the results averaged over 50 runs.
More details on the hyperparameters can be found in the Appendix and source code is publicly available 1 .
5.1 Results
Table 2 shows how the three considered standard GNN models perform on the test set, containing the 10% largest graphs in the dataset, when trained with and without our regularization on the 50% smallest graphs in the dataset.
As this split leads to an imbalanced dataset, we follow previous works and report the results in terms of Matthews correlation coefficient (MCC), which has been shown to be more reliable in imbalanced settings with respect to other common metrics [14].
MCC gives a value between − 1 and 1, where − 1 indicates perfect disagreement, 0 is the value for a random guesser, and 1 indicates perfect agreement between the predictions and the true labels.
The results show that the use of the proposed regularization is always beneficial to the performance of the models, and that it leads to an average improvement across datasets of 21 to 30% .
[Table] In Table 3, we compare the standard GNNs trained with our regularization strategy against more expressive models (RPGNN), graph kernels, models trained with the IRM strategy, and the E-invariant models.
We notice that on 3 out of 4 datasets, a “standard” GNN trained with our regularization obtains the best performance on the test set composed of graphs with size larger than those present in the training set.
Furthermore on all datasets there is at least one model trained with our regularization in the top-3 best performing models.
We also confirm previous results [6, 17] showing that IRM is not effective in the graph domain, and that using theoretically more expressive models, like RPGNN, does not necessarily lead to good size-generalization performance.
Graph kernels are highly dataset-dependent, and, while they can perform well for some datasets, in many cases they fail to perform better than a random classifier.
Perhaps more surprisingly, on 3 out ofdatasets there is at least one “standard” GNN trained with our regularization that performs comparably or better than the best E-Invariant model.
In fact, E-invariant models are tied to the assumed causal model for the generation of graphs, which is not guaranteed to hold reliably for real-world datasets for which we do not have this kind of information.
Futhermore, E-invariant models require the computation of induced homomorphism densities over all possible connected k-vertex subgraphs both at training time and at test time.
Our method instead tries to learn to be robust to size-shifts, does not require any additional computation at inference time, and yet can lead even simple models like GCN to have size-generalization performance comparable to E-invariant models.
Finally, in addition to the benchmarks considered by Bevilacqua et al.
[6], we also experiment on a dataset from a different domain.
In particular, we consider a social network dataset obtained from
Deezer [55].
We follow the same evaluation strategy as before: train on the 50% smallest and test on the 10% largest, and we apply our regularization with λ = 0 .
1 and C = { 0 . 8 , 0 . 9 } without any additional hyperparameter tuning. Results shown in Table 4 confirm the effectiveness of our method. 5.2 Ablation Study
In this section we use a PNA model to study the contribution and importance of the different components of our regularization.
We choose PNA because of its performance in the previous analysis, but similar conclusions are observed for GIN and GCN (as we report in the Appendix).
[Table]
Changing Size of Coarsened Graphs.
In Table 5 we train a PNA model with our regularization strategy using different coarsening ratios C .
We notice an overall trend in which the performance decreases as the coarsening ratio decreases.
This follows intuitively as very low coarsening ratios may lead to uninformative graphs.
Furthermore we notice that using all ratios (from 0.1 to 0.9) is usually not effective and setting C = { 0 . 8 , 0 . 9 } seems a good default option, leading to the best performance on 3 out of 4 datasets.
Changing Coarsening Method.
In Figurewe show how the performance of a PNA model change when it is regularized with our method using graphs coming from different coarsening functions.
We consider Spectral Graph Coarsen-
[Figure] ing (SGC) and Multi-level Graph Coarsening (MLGC) introduced by Jin et al. [32], and two baselines, referred to as SC and KMEANS, based on standard spectral clustering [28] and K-Means clustering [40].
For the baselines we apply the clustering algorithm, merge the nodes belonging to the same cluster into a single super-node, and create an edge between two super-nodes if there is at least one edge in the original graph going from a node in one cluster to the other.
For all coarsening methods we use the same coarsening ratios as above (i.e., C = { 0 . 8 , 0 . 9 } ), and we apply the same training procedure and regularization weight λ , keeping all hyperparameters equal across methods.
We notice that our regularization strategy is quite robust towards the choice of the coarsening method, but, on average, specialised methods like SGC and MLGC tend to provide the best results.
Related Work
Graph neural networks and size-generalization.
Since the first instances of graph neural networks have appeared [60, 57], the field has exploded in recent years with a plethora of different architectures (e.g., [36, 65, 15, 62]), applications (e.g., [69, 67, 37, 23, 8]), and studies on their expressivity and properties (e.g., [65, 44, 47, 9]).
For a complete review of the field, we refer to the excellent surveys by Chami et al.
[12] and by Wu et al.
Many works have noticed poor size-generalization in GNNs [33, 25, 6, 68].
For instance, Joshi et al.
[33] consider the task of learning solutions for the traveling salesman problem, and show that GNNs struggle at generalizing to graphs larger than those appearing in the training set.
Similarly, Gasteiger et al.
[25] observe that GNNs applied to the molecular domain can be very sensible to several shifts in the data, including shifts in the size of the graphs.
Some works have reported good size-generalization by creating ad-hoc models that leverage the properties of the domain [2, 3, 56, 61, 22].
For example, Tsubaki and Mizoguchi [61] impose constraints dictated by the physical properties of their data, and report good size-generalization in predicting the energy of a molecule.
While these works show that placing the right inductive biases can help the size-generalization capabilities of GNNs, the current literature is lacking a general method that can be applied on generic GNNs.
Bevilacqua et al.
[6] tackle the issue of poor size-generalization by assuming a causal model describing the generative process for the graphs in the dataset, and designing a specific model which can be invariant to the size of the graphs obtained through this causal model.
While this strategy shows great size-generalization capabilities on synthetic graphs generated according to their causal model, its benefits decrease when applied to real-world graphs where there are no guarantees that the causal model is correct.
This method requires an ad-hoc model and additional computations at both train and test time, while our method can be applied on any GNN and does not require pre-computations at test time.
Yehudai et al.
[68] provide a theoretical and empirical analysis showing that, even for simple tasks, it is not trivial to obtain a GNN with good size-generalization properties.
Yehudai et al.
[68] consider a different scenario from ours, as, together with the labelled training graphs, they assume access to graphs from the test distribution, and treat the task as a domain-adaptation scenario.
Graph coarsening.
While there is no consensus on what is the best metric to consider for coarsening a graph, many methods and metrics have been proposed in recent years [45, 43, 18, 7, 32].
There have also been some attempts at using GNNs for the task of graph coarsening [10, 46].
Orthogonally, graph coarsening has been used to reduce the computational resources needed for training GNNs on large graphs [71, 31].
For a detailed presentation of graph coarsening we refer to Chen et al.
Invariant risk minimization domain adaptation, and regularization
with discrepancy measures. Our method is conceptually related to Invariant Risk Minimization (IRM) [1], which aims at learning representations that are invariant across training environments (where different training environments are intended as sets of data points collected under different conditions).
IRM has been shown to have several shortcomings, specially when the data comes from a single environment (e.g. when there is access only to small graphs) [54, 6].
Our method does not require access to different training environments, which may not be easy to obtain in practical scenarios, and our results show that, as also observed in [6, 17], IRM seems ineffective for improving size-generalization in GNNs.
Similarly to IRM, the field of domain adaptation [4] subsumes access to (unlabelled) external data, in addition to labelled training data, and aims at transferring knowledge between domains.
Domain adaptation has been used in a large variety of fields, and many surveys are available [16, 63, 27].
Ding et al.
[17] analyze how existing domain adaptation algorithms perform on graph data, and observe that these tend to not be effective, indicating that new methods specific for graph data are needed.
Discrepancy measures that have been used for regularizing deep learning models are the Maximum mean discrepancy (MMD) [19, 41, 66, 42] and the central moment discrepancy (CMD) [52, 70, 72], which has been shown to be empirically more effective.
We further mention the work by Zhu et al. [72] which, similarly to ours, uses CMD to obtain a GNN that is robust to biases in the sampling process that is used to select the nodes for the training set in node classification scenarios. Finally, we mention the work by Joshi et al. [34] which uses CMD to study the representations of GNNs.
Conclusions
GNNs are heavily used models for the task of graph classification.
In many domains it is typical for graphs to vary in size, and while GNNs are designed to be able to process graphs of any size, empirical results across literature have highlighted that GNNs struggle at generalizing to sizes unseen during training.
In this work we introduce a regularization strategy that can be applied on any GNN, and that improves size-generalization performance from smaller to larger graphs by up to 30%.
Acknowledgments and Disclosure of Funding
The authors thank Chaitanya Joshi and Iulia Duta for the great feedback provided on earlier versions of the paper.
{+This work is supported, in part, by the Italian Ministry of Education, University and Research (MIUR), under PRIN Project n. 20174LF3T
+}
“AHeAD” and the initiative “Departments of Excellence” (Law 232/2016) and by University of Padova under project “SID 2020:
RATED-X”. Appendix
This Appendix contains the following additional material. In Section A we report more statistics regarding the considered datasets. In Section B we report the results for the CKA analysis on the other datasets not shown in the main paper for space limitations. In Section C we report more details on the hyperparameters of the models used for our experiments.
In Section D we report the results for the ablation study on the other models not shown in the main paper for space limitations. In Section
E we show the overhead (in terms of training time) incurred by our method. In Section F we report information on the computing resources used to perform our experimental results, links to access the datasets, and the licence of the publicly available libraries used in our code. A Dataset Information
[Table]
In Table 6 (taken from [68, 6]) we report the information about the graphs in the considered datasets. B Embedding Analysis Results for Other Datasets
In Figure 4 we show the continuation of Figure 2 (b) from the main paper, including plots on all datasets. In more detail, we report plots for the average CKA [38] values between the node embeddings for the original graphs and the node embeddings for the coarsened versions of the graphs, both generated by the same model, for different ratios. The plot shows that a model trained with our regularization is learning to be more robust to size shifts, as its representation for coarsened versions of the graphs are much more aligned (higher CKA values) to the representations for the original graphs compared to a model trained without regularization. This is even more apparent for low coarsening ratios (i.e., when the shift is stronger), indicating that our regularization is in fact making the model more robust to size shifts.
The reason why the CKA values for a model trained with and without regularization tend to become similar for small ratios (which is instead not apparent on the DD dataset shown in Figure 2 (b) in the main paper) is that the graphs in NCI1, NCI109,
[Figure]
and PROTEINS are much smaller than the graphs in DD (see Section A), and hence graphs tend to become completely uninformative with ratios lower than 0 . C Hyperparameters and Evaluation Procedure
Hyperparameters. As we follow exactly the same procedures as previous work [68, 6], we refer the reader to those papers for an in depth discussion, and we report below the description of the standard GNN models used in our experiments (which follows from the above mentioned papers). The GCN[36], GIN [65], and PNA [15] models used in our experiments have 3 graph convolution layers and a final multi-layer perceptron with a softmax activation function to obtain the predictions. Batch norm is used in between graph convolution layers, and ReLU is used as activation function.
The networks are trained with a dropout of 0 . 3 , and were tuned using the validation set. In particular the batch size was chosen between 64 and 128 , the learning rate between 0 . 01 , 0 . 005 , and 0 . 001 , and the network width between 32 and 64 . All models are trained with early stopping (taking the wights related to the epoch with the highest MCC on the validation set), and the different classes were weighted in the supervised loss function according to the frequency of a class in the training data (to deal with the imbalance in the data, as explained in the main paper).
When applying our method in Section 3 and in Table 2 and Table 3 in the main paper, we use λ = 0 .and C = (0 . 8 , 0 . 9) for all models and datasets. These parameters were taken from a tuning procedure on the validation set described below. Tuning and Evaluation Procedure.
To identify the values of λ and C to use for our method, we tried different values ( λ = { 1 . 0 , 0 . 1 , 0 . 01 , 0 . 001 } and C = { (0 . 5) , (0 . 8) , (0 . 5 , 0 . 9) , (0 . 8 , 0 . 9) } ) on the validation set using a GIN model on the PROTEINS dataset.
We found that λ = 0 . 1 and C = (0 . 8 , 0 . 9) performed best on the validation set, and we then used those values on the validation set of other datasets and with other models. As the results on the validation sets were good (i.e., they were leading to better performance with respect to a model trained without regularization), we adopted the same values of λ and C for all datasets and models to show that our method can work without extensive (and expensive) hyperparameter tuning.
It is possible that dataset-specific and model-specific tuning can lead to higher results.
We first obtained the results for Table 1, Figure 2, Table 2, and Table 3, in the main paper.
Then, only after having those results, we have performed the ablation study.
We remark that the results shown in the ablation study are not coming from an hyperparameter tuning procedure, but were performed after all the tuning (done on the validation sets) and evaluation results were obtained.
The ablation is used to understand the impact of the components of our method only after having evaluated it, as is the standard procedure for ablation studies.
D Full Ablation Study Results
We report the results of the ablation study also for the GIN and GCN models.
Matthews correlation coefficient (MCC) for
GCN model trained with proposed regularization with the Spectral Graph Coarsening strategy applied with different coarsening ratios
[Table] [Table] Matthews correlation coefficient (MCC) for a GCN model trained with our proposed regularization and different coarsening strategies: Spectral Clustering (
), Spectral Graph Coarsening ( Graph K-Means [Table]
[Table] Changing Size of Coarsened Graphs.
We report in Table 7 and Table 8 the performance of a GCN [36] and GIN [65] model when trained with our regularization strategy with different coarsening ratios C r .
All other hyperparameters are kept the same as for the main results in our paper. As for the
PNA model shown in the main paper, we notice an overall trend in which the performance decreases as the coarsening ratio decreases. This follows intuitively as very low coarsening ratios may lead to uninformative graphs.
Furthermore we notice that using all ratios (from 0.1 to 0.9) is usually not effective and setting C = { 0 . 8 , 0 .
9 } seems a good default option.
Changing Coarsening Method. In Table 9 and Table 10 we report the performance of a GCN [36] and a GIN [65] model trained with our regularization using different coarsening functions.
All other parameters are kept the same.
As for the PNA shown in the main paper we notice that our method is quite robust to the choice of the coarsening function. However, in most cases, specialised methods like SGC and MLGC provide the best results.
Table 11 shows the time in seconds for training a model with and without our regularization.
We notice that our method introduces an average 50% overhead in training time.
F Compute and Licence Information
All our experiments were performed on a machine with a Nvidia 1080Ti GPU and a CPU cluster equipped with 8 CPUs 12-core Intel Xeon Gold 5118@2.30GHz with 1.5Tb of RAM.
[Table]
The TUDataset [48] we used for our experiments are publicly available online 2 .
Our code makes use of the following publicly available libraries: PyTorch [50] (BSD-Style License), PyTorch Geometric [21] (MIT License), PyTorch Lightning [20] (Apache-2.0 license), scikit-learn [51] (BSD-3 license), numpy [30] (BSD-3 license), ray[tune] [39] (Apache 2.0).