Principal Neighbourhood Aggregation for Graph Nets

Introduction

Graph Neural Networks (GNNs) have been an active research field for the last ten years with significant advancements in graph representation learning . However, it is difficult to understand the effectiveness of new GNNs due to the lack of standardized benchmarks and of theoretical frameworks for their expressive power.

In fact, most work in this domain has focused on improving the GNN architectures on a set of graph benchmarks, without evaluating the capacity of their network to accurately characterize the graphs’ structural properties. Only recently there have been significant studies on the expressive power of various GNN models . However, these have mainly focused on the isomorphism task in domains with countable features spaces, and little work has been done on understanding their capacity to capture and exploit the underlying properties of the graph structure.

We hypothesize that the aggregation layers of current GNNs are unable to extract enough information from the nodes’ neighbourhoods in a single layer, which limits their expressive power and learning abilities.

We first mathematically prove the need for multiple aggregators and propose a solution for the uncountable multiset injectivity problem introduced by . Then, we propose the concept of degree-scalers as a generalization to the sum aggregation, which allow the network to amplify or attenuate signals based on the degree of each node. Combining the above, we design the proposed Principal Neighbourhood Aggregation (PNA) model and demonstrate empirically that multiple aggregation strategies improve the performance of the GNN.

Dehmamy et al. have also empirically found that using multiple aggregators (mean, sum and normalized mean), which extract similar statistics from the input message, improves the performance of GNNs on the task of graph moments. In contrast, our work extends the theoretical framework by deriving the necessity to use complementary aggregators. Accordingly, we propose the use of different statistical aggregations to allow each node to better understand the distribution of the messages it receives, and we generalize the mean as the first of a set of possible n-moment aggregators. In the setting of graph kernels, Cai et al. constructed a simple baseline using multiple aggregators. In the field of computer vision, Lee et al. empirically showed the benefits of combining mean and max pooling. These give us further confidence in the validity of our theoretical analysis.

We present a consistently well-performing and parameter efficient encode-process-decode architecture for GNNs. This differs from traditional GNNs by allowing a variable number of convolutions with shared parameters. Using this model, we compare the performances of some of the most diffused models in the literature (GCN , GAT , GIN and MPNN ) with our PNA.

Previous work on tasks taken from classical graph theory focuses on evaluating the performance of GNN models on a single task such as shortest paths , graph moments or travelling salesman problem . Instead, we took a different approach by developing a multi-task benchmark containing problems both on the node level and the graph level. Many of the tasks are based on dynamic programming algorithms and are, therefore, expected to be well suited for GNNs . We believe this multi-task approach ensures that the GNNs are able to understand multiple properties simultaneously, which is fundamental for solving complex graph problems. Moreover, efficiently sharing parameters between the tasks suggests a deeper understanding of the structural features of the graphs. Furthermore, we explore the generalization ability of the networks by testing on graphs of larger sizes than those present in the training set.

To further demonstrate the performance of our model, we also run tests on recently proposed real-world GNN benchmark datasets with tasks taken from molecular chemistry and computer vision. Results show the PNA outperforms the other models in the literature in most of the tasks hence further supporting our theoretical findings.

The code for all the aggregators, scalers, models (in PyTorch, DGL and PyTorch Geometric frameworks), architectures, multi-task dataset generation and real-world benchmarks is available here.

Principal Neighbourhood Aggregation

In this section, we first explain the motivation behind using multiple aggregators concurrently. We then present the idea of degree-based scalers, linking to prior related work on GNN expressiveness. Finally, we detail the design of graph convolutional layers which leverage the proposed Principal Neighbourhood Aggregation.

Most work in the literature uses only a single aggregation method, with mean, sum and max aggregators being the most used in the state-of-the-art models . In Figure 1, we observe how different aggregators fail to discriminate between different messages when using a single GNN layer.

We formalize our observations in the theorem below:

proposition1 The moments of a multiset (as defined in Equation 4) exhibit a valid example using $n$ aggregators.

We prove Theorem 1 in Appendix A and Proposition 1 in Appendix B. Note that unlike Xu et al. , we consider a continuous input feature space; this better represents many real-world tasks where the observed values have uncertainty, and better models the latent node features within a neural network’s representations. Continuous features make the space uncountable, and void the injectivity proof of the sum aggregation presented by Xu et al. .

Hence, we redefine aggregators as continuous functions of multisets which compute a statistic on the neighbouring nodes, such as mean, max or standard deviation. The continuity is important with continuous input spaces, as small variations in the input should result in small variations of the aggregators’ output.

Theorem 1 proves that the number of independent aggregators used is a limiting factor of the expressiveness of GNNs. To empirically demonstrate this, we leverage four aggregators, namely mean, maximum, minimum and standard deviation. Furthermore, we note that this can be extended to the normalized moment aggregators, which allow advanced distribution information to be extracted whenever the degree of the nodes is high.

The following paragraphs will describe the aggregators we leveraged in our architectures.

Also often used in literature, they are very useful for discrete tasks, for domains where credit assignment is important and when extrapolating to unseen distributions of graphs . Alternatively, we present the softmax and softmin aggregators in Appendix E, which are differentiable and work for weighted graphs, but don’t perform as well on our benchmarks.

The standard deviation (STD or $\sigma$) is used to quantify the spread of neighbouring nodes features, such that a node can assess the diversity of the signals it receives. Equation 3 presents, on the left, the standard deviation formulation and, on the right, the STD of a graph-neighbourhood. ReLU is the rectified linear unit used to avoid negative values caused by numerical errors and $\epsilon$ is a small positive number to ensure $\sigma$ is differentiable.

The mean and standard deviation are the first and second normalized moments of the multiset ($n=1,n=2$). Additional moments, such as the skewness ($n=3$), the kurtosis ($n=4$), or higher moments, could be useful to better describe the neighbourhood. These become even more important when the degree of a node is high because four aggregators are insufficient to describe the neighbourhood accurately. As described in Appendix D, we choose the nth root normalization, as presented in Equation 4, because it gives a statistic that scales linearly with the size of the individual elements (as the other aggregators); this gives the training adequate numerical stability. Once again we add an $\epsilon$ to the absolute value of the expectation before applying the nth root for numerical stability of the gradient.

2 Degree-based scalers

We introduce scalers as functions of the number of messages being aggregated (usually the node degree), which are multiplied with the aggregated value to perform either an amplification or an attenuation of the incoming messages.

Xu et al. show that the use of mean and max aggregators by themselves fail to distinguish between neighbourhoods with identical features but with differing cardinalities, and the same applies to all the aggregators described above. They propose the use of the sum aggregator to discriminate between such multisets. We generalise their approach by expressing the sum aggregator as the composition of a mean aggregator and a linear-degree amplifying scaler $S_{\text{amp}}(d)=d$.

theorem2 The mean aggregation composed with any scaling linear to an injective function on the neighbourhood size can generate injective functions on bounded multisets of countable elements.

We formalize and prove Theorem 2 in Appendix C; the results proven in about the sum aggregator become then a particular case of this theorem, and we can use any kind of injective scaler to discriminate between multisets of various sizes.

Recent work shows that summation aggregation doesn’t generalize well to unseen graphs , especially when larger. One reason is that a small change of the degree will cause the message and gradients to be amplified/attenuated exponentially (a linear amplification at each layer will cause an exponential amplification after multiple layers). Although there are different strategies to deal with this problem, we propose using a logarithmic amplification $S\propto\log(d+1)$ to reduce this effect. Note that the logarithm is injective for positive values, and $d$ is defined non-negative.

Further motivation for using logarithmic scalers is to better describe the neighbourhood influence of a given node. Suppose we have a social network where nodes A, B and C have respectively 5 million, 1 million and 100 followers: on a linear scale, nodes B and C are closer than A and B; however, this does not accurately model their relative influence. This scenario exhibits how a logarithmic scale can discriminate better between messages received by influencer and follower nodes.

We propose the logarithmic scaler $S_{\text{amp}}$ presented in Equation 5, where $\delta$ is a normalization parameter computed over the training set, and $d$ is the degree of the node receiving the message.

We further generalize this scaler in Equation 6, where $\alpha$ is a variable parameter that is negative for attenuation, positive for amplification or zero for no scaling. Other definitions of $S(d)$ can be used—such as a linear scaling—as long as the function is injective for $d>0$.

3 Combined aggregation

We combine the aggregators and scalers presented in previous sections obtaining the Principal Neighbourhood Aggregation (PNA). This is a general and flexible architecture, which in our tests we used with four neighbour-aggregations with three degree-scalers each, as summarized in Equation 7. The aggregators are defined in Equations 1–3, while the scalers are defined in Equation 6, with $\otimes$ being the tensor product.

As mentioned earlier, higher degree graphs such as social networks could benefit from further aggregators (e.g. using the moments proposed in Equation 4). We insert the PNA operator within the framework of a message passing neural network , obtaining the following GNN layer:

Using twelve operations per kernel will require the usage of additional weights per input feature in the $U$ function, which could seem to be just quantitatively—not qualitatively—more powerful than an ordinary MPNN with a single aggregator . However, the overall increase in parameters in the GNN model is modest and, as per our theoretical analysis above, a limiting factor of GNNs is likely their usage of a single aggregation.

This is comparable to convolutional neural networks (CNN) where a simple $3\times 3$ convolutional kernel requires 9 weights per feature (1 weight per neighbour). Using a CNN with a single weight per $3\times 3$ kernel will reduce the computational capacity since the feedforward network won’t be able to compute derivatives or the Laplacian operator. Hence, it is intuitive that the GNNs should also require multiple weights per node, as previously demonstrated in Theorem 1. In Appendix K, we will demonstrate this observation empirically, by running experiments on baseline models with larger dimensions of the hidden features (and, therefore, more parameters).

Architecture

We compare the performance of the PNA layer against some of the most popular models in the literature, namely GCN , GAT , GIN and MPNN on a common architecture. In Appendix F, we present the details of these graph convolutional layers.

For the multi-task experiments, we used an architecture, represented in Figure 3, with $\mathcal{M}$ convolutions followed by three fully-connected layers for node labels and a set2set (S2S) readout function for graph labels. In particular, we want to highlight:

Gated Recurrent Units (GRU) applied after the update function of each layer, as in . Their ability to retain information from previous layers proved effective when increasing the number of convolutional layers $\mathcal{M}$.

Weight sharing in all the GNN layers but the first makes the architecture follow an encode-process-decode configuration . This is a strong prior which works well on all our experimental tasks, yields a parameter-efficient architecture, and allows the model to have a variable number $\mathcal{M}$ of layers.

Variable depth $\mathcal{M}$, decided at inference time (based on the size of the input graph and/or other heuristics), is important when using models over high variance graph distributions. In our experiments we have only used heuristics dependant on the number of nodes $N$ $(\mathcal{M}=f(N))$ and, for the architectures in the results below, we settled with $\mathcal{M}=\lfloor N/2\rfloor$. It would be interesting to test heuristics based on properties of the graph, such as the diameter, or an adaptive computation time heuristic based on, for example, the convergence of the nodes features . We leave these analyses to future work.

This architecture layout was chosen for its performance and parameter efficiency. We note that all architectural attempts yielded similar comparative performance of GNN layers and in Appendix I we provide the results for a more standard architecture.

Multi-task benchmark

The benchmark consists of classical graph theory tasks on artificially generated graphs.

Following previous work , the benchmark contains undirected unweighted randomly generated graphs of a wide variety of types. In Appendix G, we detail these types, and we describe the random toggling used to increase graph diversity. For the presented multi-task results, we used graphs of small sizes (15 to 50 nodes) as they were already sufficient to demonstrate clear differences between the models.

Multi-task graph properties

In the multi-task benchmark, we consider three node labels and three graph labels based on standard graph theory problems. The node properties tasks are the single-source shortest-path lengths, the eccentricity and the Laplacian features ($LX$ where $L=(D-A)$ is the Laplacian matrix and $X$ the node feature vector). The graph properties tasks are whether the graph is connected, the diameter and the spectral radius.

Input features

As input features, the network is provided with two vectors of size $N$, a one-hot vector (representing the source for the shortest-path task) and a feature vector $X$ where each element is i.i.d. sampled as $X_{i}\sim\mathcal{U}$. Apart from taking part in the Laplacian features task, this random feature vector also provides a unique identifier for the nodes in other tasks. Similar strengthening via random features was also concurrently discovered by . This allows for addressing some of the problems highlighted in ; e.g. the task of whether a graph is connected could be performed by continually aggregating the maximum feature of the neighbourhood and then checking whether they are all equal in the readout.

Model training

While having clear differences, these tasks also share related subroutines (such as graph traversals). While we do not take this sharing of subroutines as prior as in , we expect models to pick up on these commonalities and efficiently share parameters between the tasks, which reinforce each other during the training.

We trained the models using the Adam optimizer for a maximum of 10,000 epochs, using early stopping with a patience of 1,000 epochs. Learning rates, weight decay, dropout and other hyper-parameters were tuned on the validation set. For each model, we run 10 training runs with different seeds and different hyper-parameters (but close to the tuned values) and report the five with least validation error.

Results and discussion

The multi-task results are presented in Figure 4a, where we observe that the proposed PNA model consistently outperforms state-of-the-art models, and in Figure 4b, where we note that the PNA performs better on all tasks. The baseline represents the MSE from predicting the average of the training set for all tasks.

The trend of these multi-task results follows and amplifies the difference in the average performances of the models when trained separately on the individual tasks. This suggests that the PNA model can better capture and exploit the common sub-units of these tasks. Appendix J provides the average results of the models when trained on individual tasks. Moreover, PNA showed to perform the best on all architecture layouts that we attempted (see Appendix I) and on all the various types of graphs (see Appendix H).

To demonstrate that the performance improvements of the PNA model are not due to the (relatively small) number of additional parameters it has compared to the other models (about 15%), we ran tests on all the other models with latent size increased from 16 to 20 features. The results, presented in Appendix K, suggest that even when these models are given 30% more parameters than the PNA, they are qualitatively less capable of capturing the graph structure.

Finally, we explored the extrapolation of the models to larger graphs, in particular, we trained models on graphs of sizes between 15 and 25, validated between 25 and 30 and evaluate between 20 and 50. This task presents many challenges, two of the most significant are: firstly, unlike in the models are not given any step-wise supervision or trained on easily extendable subroutines; secondly, the models have to cope with their architectures being augmented with further hidden layers than trained on, which can sometimes cause problems with rapidly increasing feature scales.

Due to the aforementioned challenges, as expected, the performance of the models (as a proportion of the baseline performance) gradually worsens, with some of them having feature explosions. However, the PNA model keeps consistently outperforming all the other models on all graph sizes. Our results also follow the findings in , i.e. that between single aggregators the max tends to perform best when extrapolating to larger graphs.

2 Real-world benchmarks

The recent works by Dwivedi et al. and Hu et al. have shown problems with many benchmarks used for GNNs in recent years and proposed a new range of datasets across different artificial and real-world tasks. To test the capacity of the PNA model in real-world domains, we assessed it on their chemical (ZINC and MolHIV) and computer vision (CIFAR10 and MNIST) datasets.

To ensure a fair comparison of the different convolutional layers, we followed their method for training procedure (data splits, optimizer, etc.) and GNN structure (layers, normalization and approximate number of parameters). For the MolHIV dataset, we used the same GNN structure as in .

To better understand the results in the table, we need to take into account how graphs differ among the four datasets. In the chemical benchmarks, graphs are diverse and individual edges (bonds) can significantly impact the properties of the graphs (molecules). This contrasts with computer vision datasets made of graphs with a regular topology (every node has 8 edges) and where the graph structure of the representation is not crucial (the good performance of the MLP is evidence).

With this and our theoretical analysis in mind, it is understandable why the PNA has a strong performance in the chemical datasets, as it was designed to understand the graph structure and better retain neighbourhood information. At the same time, the version without scalers suffers from the fact it cannot distinguish between neighbourhoods of different size. Instead, in the computer vision datasets the average improvement of the PNA on SOTA was lower due to the smaller importance of the graph structure and the version of the PNA without scalers performs better as the constant degree of these graphs makes scalers redundant (and it is better to ’spend’ parameters for larger hidden sizes).

Conclusion

We have extended the theoretical framework in which GNNs are analyzed to continuous features and proven the need for multiple aggregators in such circumstances. We also have generalized the $sum$ aggregation by presenting degree-scalers and proposed the use of a logarithmic scaling. Taking the above into consideration, we have presented a method, Principal Neighbourhood Aggregation, consisting of the composition of multiple aggregators and degree-scalers. With the goal of understanding the ability of GNNs to capture graph structures, we have proposed a novel multi-task benchmark and an encode-process-decode architecture for approaching it. Empirical results from synthetic and real-world domains support our theoretical evidence. We believe that our findings constitute a step towards establishing a hierarchy of models w.r.t. their expressive power, where the PNA model appears to outperform the prior art in GNN layer design.

Broader Impact

Our work focuses mainly on theoretically analyzing the expressive power of Graph Neural Networks and can, therefore, play an indirect role in the (positive or negative) impacts that the field of graph representation learning might have on the domains where it will be applied.

More directly, our contribution in proving the limitations of existing GNNs on continuous feature spaces should help to provide an insight into their behaviour. We believe this is a significant result which might motivate future research aimed at overcoming such limitations, yielding more reliable models. However, we also recognize that, in the short-term, proofs of such weaknesses might spark mistrust against applications of these systems or steer adversarial attacks towards existing GNN architectures.

In an effort to overcome some of these short-term negative impacts and contribute to the search for more reliable models, we propose the Principal Neighbourhood Aggregation, a method that overcomes some of these theoretical limitations. Our tests demonstrate the higher capacity of the PNA compared to the prior art on both synthetic and real-world tasks; however, we recognize that our tests are not exhaustive and that our proofs do not allow for generating “optimal” aggregators for any task. As such, we do not rule out sub-optimal performance when applying the exact architecture proposed here to novel domains.

We propose the usage of aggregation functions, such as standard deviation and higher-order moments, and logarithmic scalers. To the best of our knowledge, these have not been used before in GNN literature. To further test their behaviour, we conducted out-of-distribution experiments, testing our models on graphs much larger than those in the training set. While the PNA model consistently outperformed other models and baselines, there was still a noticeable drop in performance. We therefore strongly encourage future work on analyzing the stability and efficacy of these novel aggregation methods on new domains and, in general, on finding GNN architectures that better generalize to graphs from unseen distributions, as this will be essential for the transition to industrial applications.

Acknowledgements

The authors thank Saro Passaro for the valuable insights and discussion for the mathematical proofs.

Funding Disclosure

Dominique Beaini is currently a Machine Learning Researcher at InVivo AI. Pietro Liò is a Full Professor at the Department of Computer Science and Technology of the University of Cambridge. Petar Veličković is a Research Scientist at DeepMind.

References

Appendix A Proof for Theorem 1 ( (Number of aggregators needed).)

Assume by contradiction that it is possible to discriminate between all the multisets of size $n$ using only $n-1$ aggregators, viz. $g_{1},g_{2},\ldots,g_{n-1}$.

As $S$ is a $n$-dimensional Euclidean subspace, it is possible to define a $(n-1)$-sphere $C^{n-1}$ entirely contained within it, i.e. $C^{n-1}\subseteq S$. According to Borsuk–Ulam theorem , there are two distinct (in particular, non-zero and antipodal) points $\vec{x}_{1},\vec{x}_{2}\in C^{n-1}$ satisfying $f(\vec{x}_{1})=f(\vec{x}_{2})$, showing $f$ not to be injective; hence the required contradiction. ∎

Note: $n$ aggregators are actually sufficient. A simple example is to use $g_{1},g_{2},\ldots,g_{n}$ where $g_{k}(X)=$ the $k$-th smallest item in $X$. It’s clear to see that the multiset whose elements are $g_{1}(X),\,g_{2}(X),\,\ldots\,,g_{n}(X)$ is $X$, which can hence be uniquely determined by the aggregators.

Appendix B Proof for Proposition 1 ( (Moments of the multiset).)

Since $n\geq 1$, and the first aggregator is mean, we know $\mu$. Let $X=\{x_{1},x_{2},\ldots,x_{n}\}$ be the multiset to be found, and define $R=\{r_{1}=x_{1}-\mu,\;r_{2}=x_{2}-\mu,\;\ldots,\;r_{n}=x_{n}-\mu\}$.

Notice how $\sum{r_{i}}^{1}=0$, and for $1<k\leq n$ we have $\sum{r_{i}}^{k}=n\;M_{k}(X)^{k}$, i.e. all the symmetric power sums $p_{k}=\sum{r_{i}}^{k}$ ($k\leq n$) are uniquely determined by the moments.

Additionally, $e_{k}$, the elementary symmetric sums of $R$, i.e. the sum of the products of all the sub-multisets of size $k$ ($1\leq k\leq n$), are determined as follow:

$e_{1}$, the sum of all elements, is equal to $p_{1}$; $e_{2}$, the sum of the products of all pairs in $R$, is $\left(e_{1}p_{1}-p_{2}\right)/2$; $e_{3}$, the sum of the products of all triplets, is $\left(e_{2}p_{1}-e_{1}p_{2}+p_{3}\right)/3$, and so on. Notice how $e_{1},e_{2},\ldots,e_{n}$ can be computed using the following recursive formula :

Consider polynomial $P(x)=\Pi(x-r_{i})$, i.e. the unique polynomial of degree $n$ with leading coefficient 1 whose roots are $R$. This defines $A$, the coefficients of $P$, i.e. the real numbers $a_{0},a_{1},\ldots,a_{n-1}$ for which $P(x)=x^{n}+a_{n-1}x^{n-1}+\ldots+a_{1}x+a_{0}$. Using Vieta’s formulas :

Hence $A$ is uniquely determined, and so is $P$, being its coefficients a valid definition of it. By the fundamental theorem of algebra, $P$ has $n$ (possibly repeated) roots, which are the elements of $R$, hence uniquely determining the latter.

Finally, $X$ can be easily determined adding $\mu$ to each element of $R$. ∎

Note: the proof above assumes the knowledge of $n$. In the case that $n$ is variable (as in GNNs), and so we have multisets of up to $n$ elements, an extra aggregator will be needed. An example of such aggregator is the mean multiplied by any injective scaler which would allow the degree of the node to be inferred.

Appendix C Proof for Theorem 2 ( (Injective functions on countable multisets).)

Let’s define an injective function $s$, and without loss of generality, assume $s(0),s(1),\ldots,s(N)>0$ (otherwise for the rest of the proof consider $s$ as $s^{\prime}(i)=s(i)-\min_{j\in[0,N]}s(j)+\epsilon$ which is positive for all $i\in[0,N]$). $s(|X|)$ can only take value in $\{s(0),s(1),\ldots,s(N)\}$, therefore let us define $\gamma=\min\left\{\frac{s(i)}{s(j)}\,\mid\,i,j\in[0,N],\>s(i)\geq s(j)\right\}$. Since $s$ is injective, $s(i)\neq s(j)$ for $i\neq j$, which implies $\gamma>1$.

Let $K>\frac{1}{\gamma-1}$ be a positive real number and consider $f(x)=N^{-Z(x)}+K$.

We proceed to show that the cardinality of $X$ can be uniquely determined, and $X$ itself can be determined as well, by showing that exist an injection $h$ over the multisets.

Let us $h$ as a function that scales the mean of $f$ by an injective function of the cardinality:

We want show that the value of $|X|$ can be uniquely inferred from the value of $h(X)$. Assume by contradiction $\exists\,X^{\prime},X^{\prime\prime}$ multisets of size at most $N$ such that $|X^{\prime}|\neq|X^{\prime\prime}|$ but $h(X^{\prime})=h(X^{\prime\prime})$; since $s$ is injective $s(|X^{\prime}|)\neq s(|X^{\prime\prime}|)$, without loss of generality let $s(|X^{\prime}|)>s(|X^{\prime\prime}|)$, then:

which is a contradiction. So it is impossible for the size of a multiset $X$ to be ambiguous from the value of $h(X)$.

Let us define $d$ as the function mapping $h(X)$ to $|X|$.

Considering the $Z(j)$-th digit $i$ after the decimal point in the base $N$ representation of $h^{\prime}(X)$, it can be inferred that $X$ contains $i$ elements $j$, and, so, all the elements in $X$ can be determined; hence $h$ is injective over the multisets in $X$. ∎

Note: this proof is a generalization of the one by Xu et al. on the sum aggregator.

Appendix D Normalized moments aggregation

The main motivation for choosing the nth root normalization for the moments is numerical stability. In fact, one property of our version is that it scales linearly with $L$, for uniformly distributed random variables $U[0,L]$, as do other aggregators such as mean, max and min (std is a particular case). Other common formulations of the moments such as those in Equation 9 scale respectively as the nth power and constantly with $L$. This difference causes numerical instability when combined in the same layer.

To demonstrate the usefulness of higher moments aggregation and further motivate the need for multiple aggregation functions, we ran an ablation study showing how different moments affect the performance of the model. We conduct this by testing five different models, each taking a different number of moments, on our multi-task benchmark.

The results in Figure 7 demonstrate that with the increase of the number of aggregators the models reach a higher expressive power, but at a certain point (dependent on the graphs and tasks, in this case around 3) the increase in expressiveness given by higher moments reduces the performance since the model becomes harder to optimize and prone to overfitting. We expect that higher moments will be more beneficial on graphs with a higher average degree since they will better characterize the neighbourhood distributions.

Finally, we note how the addition of the max and min aggregators in the PNA (rightmost column) gives a better and more consistent performance in these tasks than higher moments. We believe this is task-dependent, and, for algorithmic tasks, discrete aggregators can be valuable. As a side note, we point out how the max and min aggregators of positive values can be considered as the nth-root of the nth (non-centralized) moment as n tends to, respectively, $+\infty$ and $-\infty$.

Appendix E Alternative aggregators

Besides those described above, we have experimented with additional aggregators. We detail some examples below. Domain-specific metrics can also be an effective choice.

As an alternative to max and min, softmax and softmin are differentiable and can be weighted in the case of edge features or attention networks. They also allow an asymmetric message passing in the direction of the strongest signal. Equation 10 presents their direct neighbour formulations, where $X^{l}$ are the nodes features at layer $l$ with respect to node $i$ and $N(i)$ is the neighbourhood of node $i$:

Appendix F Alternative graph convolutions

In this section, we present the details of the four graph convolutional layers from existing models that we used to compare the performance of the PNA in the multi-task benchmark.

Graph Attention Networks (GAT)

perform a linear transformation of the input features followed by an aggregation of the neighbourhood as a weighted sum of the transformed features, where the weights are set by an attention mechanism $a$. We define it in Equation 12, where $W$ is a trainable projection matrix. As in the original paper, we employ the use of multi-head attention.

Graph Isomorphism Networks (GIN)

perform a sum aggregation over the neighbourhood, followed by an update function $U$ consisting of a multi-layer perceptron. We define it in Equation 13, where $\epsilon$ is a learnable parameter. As in the original paper, we use a 2-layer MLP for $U$.

Message Passing Neural Networks (MPNN)

perform a transformation before and after an arbitrary aggregator. We define it in Equation 14, where $M$ and $U$ are neural networks and $\bigoplus$ is a single aggregator. In particular, we test models with sum and max aggregators, as they are the most used in literature. As with PNA layers, we found that linear transformations are sufficient for $M$ and $U$ and, as in the original paper , we employ multiple towers.

Appendix G Random graph generation

In this section, we present the details of the random generation of the graphs we used in the multi-task benchmark. Following previous work , we opted for undirected unweighted graphs from a wide variety of types (we provide, in parentheses, the approximate proportion of such graphs in the benchmark). Letting $N$ be the total number of nodes per graph:

Erdős-Rényi (20%): with probability of presence for each edge equal to $p$, where $p$ is independently generated for each graph from $\mathcal{U}$

Barabási-Albert (20%): the number of edges for a new node is $k$, which is taken randomly from $\{1,2,...,N-1\}$ for each graph

Grid (5%): $m\times k$ 2d grid graph with $N=mk$ and $m$ and $k$ as close as possible

Caveman (5%): with $m$ cliques of size $k$, with $m$ and $k$ as close as possible

Tree (15%): generated with a power-law degree distribution with exponent 3

Caterpillar graphs (10%): with a backbone of size $b$ (drawn from $\mathcal{U}[1,N)\,$), and $N-b$ pendent vertices uniformly connected to the backbone

Lobster graphs (10%): with a backbone of size $b$ (drawn from $\mathcal{U}[1,N)\,$), $p$ (drawn from $\mathcal{U}[1,N-b\,]\,$) pendent vertices uniformly connected to the backbone, and additional $N-b-p$ pendent vertices uniformly connected to the previous pendent vertices.

Additional randomness was introduced to the generated graphs by randomly toggling arcs, without strongly impacting the average degree and main structure. If $e$ is the number of edges and $m$ the number of ’missing edges’ ($2e+2m=N(N-1)$), the probabilities of toggling an existing and missing edge, respectively $P_{e}$ and $P_{m}$, are:

After performing the random toggling, we discarded graphs containing singleton nodes, as they are in no way affected by the choice of aggregation.

Appendix H Graph type experiments

In order to better interpret the improvements in performance that the PNA brings, we tested the models against the various types of graphs in the multi-task benchmark. In particular, in these experiments, we trained the models on the whole dataset with the proportions described above and then tested them against datasets composed by just one category of graphs.

The results, presented in Figure 8, show that the PNA improves across all types. However, it performs the worst on the graphs with a higher diameter (especially graphs close to lines), suggesting that the number of layers is not enough to reach the complete graph. Therefore, the main limitation to the PNA performance seems to be the message passing framework; this could motivate future research to try to improve the framework itself.

Appendix I Standard architecture

In this section we will provide more intuition on the motivation behind our choice of architecture, presented in Section 3, which we will refer to as recurrent,Note that this was only used in the synthetic benchmarks, while in the real-world benchmarks, we kept the same architecture from Dwivedi et al. and present the results on a more standard architecture.

The main motivations behind the choice of the architecture were: (1) provide a fairer comparison between the models (2) showcase a parameter-efficient recurrent architecture with a prior This prior corresponds to the knowledge that these tasks can be solved by the convergence of an aggregation function in the message passing context, potentially with an additional readout/function. that works very well with the tasks at hand. In particular:

The GRU helps to avoid over-smoothing, and the models that do not have a skip connection across the aggregation (GAT, GIN and GCN) are those benefiting the most from it; therefore, to still provide a fair comparison in the results below, we added skip connections from every convolutional layer to the readout, in all the models.

The S2S (as opposed to a mean readout used below) helps the most architectures without scalers as it can provide an alternative counting mechanism.

The repeated convolutions are a parameter-saving prior which works well in these tasks but does not change the rank between the various models.

For completeness, we present in Figure 9 the comparison of the average results of the recurrent architecture and standard one which uses no GRU but skip connections, mean readout rather than S2S and a fixed number of convolutions (8).

Appendix J Single task experiments

Apart from a good method to evaluate the performance on a variety of different problems, the multi-task approach offers a regularization opportunity that some models capture more than others. In particular, we found that models without scalers (or sum aggregator) are those benefiting the most from the approach; we hypothesise that the reason for this lies in some supervision that specific tasks give to recognise the size of a model neighbourhood. Moreover, more complex models are more prone to overfitting when trained on a single task. Figure 10 shows the average performance on the individual tasks of the various models.

Appendix K Parameters comparison

Figure 11 shows the results of testing all the other models on the multi-task benchmark with increased latent size.

We observe that, even with fewer parameters, PNA performs consistently better and an increased number of parameters does not boost the performance of the other models. This suggests that the multiple aggregators in the PNA produce a qualitative improvement to the capacity of the model.