S PARSIFYING THE U PDATE S TEP IN G RAPH N EURAL N ETWORKS
A BSTRACT
Message-Passing Neural Networks (MPNNs), the most prominent Graph Neural Network (GNN) framework, celebrate much success in the analysis of graphstructured data.
Concurrently, the sparsification of Neural Network models attracts a great amount of academic and industrial interest.
In this paper we conduct a structured, empirical study of the effect of sparsification on the trainable part of MPNNs known as the Update step.
To this end, we design a series of models to successively sparsify the linear transform in the Update step.
Specifically, we propose the ExpanderGNN model with a tuneable sparsification rate and the
Activation-Only GNN, which has no linear transform in the Update step.
In agreement with a growing trend in the literature the sparsification paradigm is changed by initialising sparse neural network architectures rather than expensively sparsifying already trained architectures.
Our novel benchmark models enable a better understanding of the influence of the Update step on model performance and outperform existing simplified benchmark models such as the Simple Graph Convolution.
The ExpanderGNNs, and in some cases the Activation-Only models, achieve performance on par with their vanilla counterparts on several downstream tasks, while containing significantly fewer trainable parameters.
Our code is publicly available at: https://github.com/ChangminWu/ExpanderGNN .
I NTRODUCTION
In recent years we have witnessed the blossom of Graph Neural Networks (GNNs).
They have become the standard tools for analysing and learning graph-structured data (Wu et al., 2020) and have demonstrated convincing performance in various application areas, including chemistry (Duvenaud et al., 2015), social networks (Monti et al., 2019), natural language processing (Yao et al., 2019) and neural science (Griffa et al., 2017).
Among various GNN models, Message-Passing Neural Networks (MPNNs, Gilmer et al.
and their variants are considered to be the dominating class.
In MPNNs, the learning procedure can be separated into three major steps: Aggregation , Update and Readout , where Aggregation and
Update are repeated iteratively so that each node’s representation is updated recursively based on the transformed information aggregated over its neighbourhood.
There is thus a division of labour between the Aggregation and the Update step, where the Aggregation utilises local graph structure, while the Update step is only applied to single node representations at a time independent of the local graph structure.
From this a natural question arises: What is the impact of the graph-agnostic Update step on the performance of GNNs?
Since the Update step is the main source of model parameters in MPNNs, understanding its impact is fundamental in the design of parsimonious GNNs.
In this paper we empirically analyse the impact of the Update step and its sparsification in a systematic way.
To this end, we propose two nested model classes, where the Update step is successively sparsified.
In the first model class which we refer to as Expander GNN, the linear transform layers of the Update step are sparsified; while in the second model class, the linear transform layers are removed and only the activation functions remain in the model.
We name the second model
Activation-Only GNN and it contrasts the Simple Graph Convolution model (SGC, Wu et al., 2019) where the activation functions where removed to merge the linear layers.
Through a series of empirical assessments on different graph learning tasks (graph and node classification as well as graph regression), we demonstrate that the Update step can be heavily simplified without inhibiting performance or relevant model expressivity.
Our findings partly agree with the work in Wu et al.
(2019), in that dense Update steps in GNN are expensive and often ineffectual.
In contrast to their proposition, we find that there are many instances in which leaving the Update step out completely significantly harms performance.
In these instances our Activation-Only model shows superior performance while matching the number of parameters and efficiency of the SGC.
R ELATED W ORK
In recent years the idea of utilising expander graphs in the design of neural networks is starting to be explored in the CNN literature.
Most notably, Prabhu et al.
(2018) propose to replace linear fully connected layers in deep networks using an expander graph sampling mechanism and hence, propose a novel, well-performing CNN architecture they call X-nets.
McDonald & Shokoufandeh (2019) and Kepner & Robinett (2019) build on the X-net design and propose alternative expander sampling mechanisms to extend the simplistic design chosen in the X-nets.
The Sparisification and Pruning of neural networks is a very active research topic (Hoefler et al., 2021; Blalock et al., 2020).
In particular, a wealth of algorithms sparsifying neural network architectures at initialisation has recently been proposed (Tanaka et al., 2020; Wang et al., 2020a; Lee et al., 2019).
While these algorithms make pruning decisions on a per-weight basis, Frankle et al.
(2021) find that these algorithms produce equivalent results to a per-layer choice of a fraction of weights to prune, as is directly done in our chosen sparsification scheme.
All of these research efforts are pruning CNNs, typically the VGG and ResNet architectures.
To the best of our knowledge, our work is the first investigating the potential of sparsifying the trainable parameters in GNNs.
Both Wu et al.
and Salha et al.
(2019) observed that simplifications in the Update step of the Graph Convolutional Network (GCN, Kipf & Welling, 2017) model is a promising area of research.
Wu et al.
(2019) proposed the SGC model, where simplification is achieved by removing the non-linear activation functions from the GCN model.
This removal allows them to merge all linear transformations in the Update steps into a single linear transformation without sacrificing expressive power.
Salha et al.
(2019) followed a similar rationale in their simplification of the graph autoencoder and variational graph autoencoder models.
In our work we aim to extend these efforts by providing more simplified benchmark models for GNNs without a specific focus on the GCN.
I NVESTIGATING THE R OLE OF THE U PDATE STEP
In this section, we present the general model structure of MPNNs and our two proposed model classes, where we sparsify or remove the linear transform layer in the Update step.
We define graphs G =
( A , X ) in terms of their adjacency matrix A =
[0 , 1] n × n , which contains the information of the graph’s node set V , and the node features X ∈ R n × s .
The MPNN structure is a prominent paradigm for performing machine learning tasks on graphs such as node or graph classification.
The learning procedure of MPNNs can be divided into the following phases:
Initial (optional).
In this phase, the initial node features X are mapped from the feature space to a hidden space by a parameterised neural network U (0) , usually a fully-connected linear layer;
[Equation]
where h (1) i denotes the hidden representation of node i.
Aggregation.
For each node, information is gathered from the node’s neighbourhood, denoted
N ( i ) for node i.
The gathered pieces of information are called “messages”, denoted by m i .
Formally, if f ( l ) ( · ) denotes the aggregation function at iteration l, then
[Equation]
Due to the isotropic nature of graphs (arbitrary node labelling), this function needs to be permutation equi- or invariant.
It also has to be differentiable so that the framework will be end-to-end trainable.
Update.
The nodes then update their hidden representations based on their current representations and the received “messages”.
Let U ( l ) denote the neural network learning the update function at iteration l.
Then, for node i, we have
[Equation]
Expander Update Step.
If we denote the two sets of computational units, i.e., neurons, being connected in a fully-connected layer of a neural network by S 1 and S 2 , respectively.
Then, following
Prabhu et al.
(2018), we propose to sparsify this layer by uniform randomly sampling only d connections for the computational units in the smaller of the two connected sets S 1 and S 2 .
We refer to the density of such an expander linear layer as the ratio of the number of sampled connections to the number of connections in the fully-connected layer, i.e., d/ max( | S 1 | , | S 2 | ) .
The theoretical computational cost of an expander linear layer is equal to 2 nd min( | S 1 | , | S 2 | )
Floating Point Operations (FLOPs).
The tunable parameter d can therefore lead to significant computational savings as the computational cost of a fully connected linear layer equals 2 n | S 1 || S 2 | FLOPs.
When we replace all linear layers in the Update steps of a GNN with expander linear layers, we get the Expander GNN.
When compared to pruning algorithms which sparsify neural network layers by iteratively removing parameters according to certain metric during training, the expander sparsifiers have two advantages:
The expander design assures that paths exist between consecutive layers, avoiding the risk of layer-collapse that is common in many pruning algorithms, where the algorithm prunes all parameters (weights) in one layer and cuts down the flow between input and output (Tanaka et al., 2020).
The expander sparsifier removes parameters at initialisation and keeps the sparsified structures fixed during training, which avoids the expensive computational cost stemming from adapting the neural network architecture during or after training and then retraining the network as is done in the majority of pruning algorithms (Frankle & Carbin, 2019; Han et al., 2015). Theoretical motivation and implementation details of the Expander Update Step can be found in Appendix Section A.
Activation-Only Models.
In the Activation-Only GNN models, we propose to fully remove the linear transformation in the Update step such that each message-passing step is immediately followed by a pointwise activation function.
The resulting model can be seen as a natural extension of the Expander GNN, where the linear transformation of the Update step is completely forgone.
In Gama et al.
(2020) it is argued that the non-linearity present in GNNs, in form of the activation functions, has the effect of frequency mixing in the sense that “part of the energy associated with large eigenvalues” is brought “towards low eigenvalues where it can be discriminated by stable graph filters.”
The theoretical insight that activation functions help capture information stored in the high energy part of graph signals is strong motivation to consider an alternative simplification to the one made in the SGC.
In this alternative simplification, which we refer to as the ActivationOnly GNN models, we remove linear transformations instead of activation functions such that each message-passing step is immediately followed by a pointwise activation function.
Readout (optional).
After L aggregation and update iterations, depending on the downstream tasks, the MPNN will either output node representations directly or generate a graph representation via a differentiable readout function, [Equation]
E XPERIMENTS AND D ISCUSSION
In order to investigate the influence of the Update step in GNNs, we now observe the performance of the proposed benchmark models on the tasks of graph classification, graph regression and node
[Table] The format follows [Table] classification.
In the Appendix Section B we provide an overview of our experimentation setup.
We have made our experimentation code publicly available online 1 .
Graph Classification.
Table 1 shows the experiment results of the vanilla GCN model and its Expander and Activation-Only variants on the ENZYMES, PROTEINS, MNIST and CIFAR datasets for graph classification.
One direct observation is that the Expander GCN models, even at 10% density, perform on par with the vanilla models.
Surprisingly, the same is true for the Activation-Only model on the ENZYMES and PROTEINS datasets.
The SGC performs, often significantly, worse than the Activation-Only model, especially on the computer vision MNIST and CIFAR datasets.
Additional experiments on the GIN architecture can be found in Appendix Table 6.
They reaffirm the conclusions drawn on the GCN model on graph classification in this section.
Graph Regression.
In Table 2 the Mean Absolute Error (MAE) of our studied and proposed models on the ZINC dataset for graph regression is displayed.
Similar to the graph classification task, the Expander GCN and Expander GraphSage models match the performance of their corresponding vanilla models, regardless of their densities.
The performance of the Expander GIN and Expander PNA models exhibits greater variance across the different densities, especially in the case of the PNA models the performance increases as the network gets denser indicating that the density of the Update step does positively contribute to the model performance of the PNA for the task of graph regression on the ZINC dataset.
The Activation-Only models perform worse than their
Expander counterparts on this task, confirming the insight from the results of the Expander GNNs that the linear transform in the Update step does improve performance in this graph regression task.
Again we see that Activation-Only GCNs outperform the SGC benchmark in this set of experiments.
Hence, for the task of graph regression we observe that both the linear transformation and non-linear activation function in the Update step have a small, positive impact on model performance. [Table]
Node Classification Results from the node classification experiments on four citation graphs (CORA, CITESEER, PUBMED and ogbn-arxiv) can be found in Table 3.
For medium-sized datasets such as CORA, CITESEER and PUBMED, we have the same observation as for the graph classification and graph regression tasks, the Expander models, regardless of their sparsity, are performing on par with the vanilla ones.
The Activation-Only models also perform as well as or even better than (on CITESEER)
the vanilla model. The performance of the GCN Activation-Only model and SGC is equally good across all three datasets. These conclusions remain true for the large-scale dataset ogbn-arxiv with 169,343 nodes and 1,166,243 edges. The ExpanderGCNs are on par with the vanilla GCN while the Activation-Only model and SGC perform slightly worse.
However, the training time of Activation-Only model and SGC is five times faster than that of the Expander and vanilla models.
The Activation-Only model outperforms the SGC.
Additional experiments on the
GIN architecture can be found in Appendix Table 7.
We observe that in the node classification task both the linear transformation and the non-linear activation function offer no benefit for the medium scale datasets.
For the large-scale dataset we find that the linear transformation can be sparsified heavily without a loss in performance, but deleting it entirely does worsen model performance.
C ONCLUSION
With extensive experiments across different GNN models and graph learning tasks, we are able to confirm that the Update step can be sparsified heavily without a significant performance cost.
In fact for six of the nine tested datasets across a variety of tasks we found that the linear transform can be removed entirely without a loss in performance, i.e., the Activation-Only models performed on par with their vanilla counterparts.
The Activation-Only GCN model consistently outperformed the
SGC model and especially in the computer vision datasets we witnessed that the activation functions seem to be crucial for good model performance accounting for an accuracy difference of up to 59%.
These findings partially support the hypothesis by Wu et al.
that the Update step can be simplified significantly without a loss in performance.
Contrary to Wu et al.
we find that the nonlinear activation functions result in a significant accuracy boost and the linear transformation in the Update step can be removed or heavily sparsified.
The Activation-Only GNN is an effective and simple benchmark model framework for any message passing neural network.
It enables practitioners to test whether they can cut the large amount of model parameters used in the linear transform of the Update steps.
If the linear transform does contribute positively to the model’s performance then the Expander GNNs provide a model class of tuneable sparsity which allows efficient parameter usage.
A
CKNOWLEDGEMENTS
The work of Dr. Johannes Lutzeyer and Prof. Michalis Vazirgiannis is supported by the ANR chair AML-HELAS (ANR-CHIA-0020-01). A PPENDIX A M OTIVATION AND I MPLEMENTATION OF THE E XPANDER GNN S
Motivation of the Expander Linear Layer From our random expander sampling scheme discussed in Section 3 we sample bipartite graphs with good expansion properties, which are commonly discussed in the field of error correcting codes under the name “lossless expanders” (Hoory et al., 2006, pp.
Expander graphs can be informally defined to be highly connected and sparse graphs (Lubotzky, 2012).
They are successfully applied in communication networks where communication comes at a certain cost and is to be used such that messages are spread across the network efficiently (Lubotzky, 2012).
Equally, in a neural network each parameter (corresponding to an edge in the neural network architecture) incurs a computational cost and is placed to optimise the overall performance of the neural network architecture.
Therefore, the use of expander graphs in the design of neural network architectures is conceptually well-motivated.
In B¨olcskei et al.
(2019) the connectedness of a sparse neural network architecture was linked to the complexity of a given function class which can be approximated by sparse neural networks.
Hence, utilising neural network parameters to optimise the connectedness of the network maximises the expressivity of the neural network.
In Kepner & Robinett (2019) and Bourely et al.
(2017) the connectedness of the neural network architecture graph was linked – via the path-connectedness and the graph Laplacian eigenvalues – to the performance of neural network architectures.
Therefore, for both the expressivity of the neural network and its performance, the connectedness, which is optimised in expander graphs, is a parameter of interest.
Implementation of Expander Linear Layer
The most straightforward way of implementing the expander linear layer is to store the weight matrix W as a sparse matrix.
Sparse matrix multiplications can be accelerated on several processing units released in 2020 and 2021 such as the
Sparse Linear Algebra Compute (SLAC) cores used in the Cerebras WSE-2 (CEREBRAS SYSTEMS, 2021), the Intelligence Processing Unit (IPU) produced by Graphcore (Strategy, 2020) and the NVIDIA A100 Tensor Core GPU (NVIDIA, 2020).
However, since we ran experiments on a
NVIDIA RTX 2060 GPU, we use masks in our implementation, similar to those of several existing pruning algorithms, to achieve the sparsification.
A mask M ∈ { 0 , 1 } | S 1 |×| S 2 | is of the same dimension as weight matrix and M u,v =
if and only if ( u, v ) ∈
E ′ .
An entrywise multiplication, denoted by ⊙ , is then applied to the mask and the weight matrix so that undesired parameters in the weight matrix are removed.
To illustrate this alteration we use the following matrix representation of the GCN’s model equation,
[Equation]
In our implementation the Expander GCN has the following model equation,
[Equation]
B G ENERAL S ETTINGS AND B ASELINES
Considered GNNs Throughout this section we refer to the standard, already published, architectures as “vanilla” architectures.
We compare the performance of the vanilla GNN models, the
Expander GNN models with different densities ( 10% and 50% ), the Activation-Only GNN model (we report the best result obtained from the ReLU, PReLU and Tanh activation functions), as well as the SGC for the GCN models.
To ensure that our inference is not specific to a certain GNN architecture only, we evaluate the performance across 4 representative GNN models of the literature state-of-the-art.
The considered models are the Graph Convolutional Network (GCN, Kipf &
Welling (2017)), the Graph Isomorphism Network (GIN, Xu et al.
(2019)), the GraphSage Network
Hamilton et al.
(2017), and the Principle Neighborhood Aggregation (PNA, Corso et al.
The precise model equations of our proposed architectures applied to these GNNs can be found in
Table 4.
We make the following additional remarks with regard to this table:
In the model equations of the GCN, d i denotes the degree of node i.
[Table]
h ( l ) j (cid:105) i,s,a h ( l +1)
i =
σ (cid:16) m ( l ) i (cid:17)
[Table]
In the model equations of the GIN, ϵ is a learnable ratio added explicitly to the central node’s own representation.
In the model equations of the PNA, (cid:76) corresponds to an operator formed by taking the tensor product of a vector containing three scaler functions and four aggregator functions, resulting in a tensor indexed by i, s, a, where the index i corresponds to the currently considered node, s corresponds to the scaler dimension and a indexes the aggregator dimension.
For more details see Corso et al.
Datasets We experiment on eleven datasets from areas such as chemistry, social networks, computer vision and academic citation, for three major graph learning tasks.
For graph classification, we have two TU datasets (Kersting et al., 2016) which are chemical graphs, and two Image datasets (MNIST/CIFAR10) that are constructed from original images following the procedure in Knyazev et al.
To perform this conversion they first extract small regions of homogeneous intensity from the images, named “Superpixels” (Dwivedi et al., 2020a), and construct a K -nearest neighbour graph from these superpixels.
The technique we implemented to extract superpixels, the choice of
K and distance kernel for constructing a nearest neighbour graph are the same as in Knyazev et al.
(2019) and Dwivedi et al. (2020a).
For graph regression, we consider molecule graphs from the
ZINC dataset (Irwin et al., 2012).
And for node classification, we use four citation datasets (Sen et al., 2008; Wang et al., 2020b; Hu et al., 2020), where the nodes are academic articles linked by citations.
Details of the used datasets can be found in Table 5, where in the number of nodes and edges column we display average values if the dataset contains multiple graphs.
Experimentation Details Since we aim to observe the performance of our benchmark models independent of the GNN choice we use the model hyperparameters found to yield a fair comparison of GNN models in Dwivedi et al.
(2020a).
Specifically, we follow the same training procedure, such as train/valid/test dataset splits, choice of optimiser, learning rate decay scheme, as well as the same hyperparameters, such as initial learning rate, hidden feature dimensions and number of
GNN layers.
We also implement the same normalisation tricks such as adding batch normalisation after non-linearity of each Update step.
Their setting files (training procedure/hyperparameters) are made public and can be found in Dwivedi et al.
(2020b).
For the node classification task on citation datasets, we follow the settings from Wu et al.
Our experiments found that the [Table] [Table] node classification task on citation graphs of small to medium size can be easily overfit and model performances heavily depend on the choice of hyperparameters.
Using the same parameters with
Wu et al.
(2019), such as learning rate, number of training epochs and number of GNN layers, helps us achieve similar results with the paper on the same model, which allows a fair comparison between the proposed Activation-Only models and the SGC.
Additional Results on the GIN architecture In Tables 6 and 7 we display additional experiments on the GIN architecture on graph classification and node classification, respectively.
These results further support the conclusions drawn in the main paper.
Loss functions After L message-passing iterations, we obtain
[Equation]
as the final node embedding, where we denote p as its feature dimension.
Depending on the downstream task, we either keep working with H ( L ) or construct a graph-level representation g from
H ( L ) ,
[Equation]
which we referred to as the Readout step in Section 3.
g or H ( L ) is then fed into a fully-connected network (MLP) to be transformed into the desired form of output for further assessment, e.g., a scalar value as a prediction score in graph regression.
We denote this network as f ( · ) , which, in our experiments, is fixed to be a three-layer MLP of the form
[Equation]
where W 1 ∈ R p × ( p/ 2) , W 2 ∈ R ( p/ 2) × ( p/ 4) , W 3 ∈ R ( p/ 4)
× k with k being the desired output dimension.
The final output, either f ( g ) or f ( H ( L ) ) , is compared to the ground-truth by a taskspecific loss function.
For graph classification and node classification, we choose cross-entropy loss and for graph regression, we use mean absolute error.