Source-linked AI summary
Principal Neighbourhood Aggregation for Graph Nets
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, Petar Veličković
TL;DR
GNN expressive-power analyses have largely focused on isomorphism tasks and countable features, despite continuous features being common in real-world inputs and hidden layers. This paper extends the theory to continuous features, proposes PNA with multiple aggregators and degree-scalers, and finds stronger performance across synthetic and real-world tasks than prior models.
Problem
Existing GNN research lacks standardized benchmarks and theoretical frameworks for expressive power, particularly for continuous features common in real-world inputs and hidden layers.
Method
The paper proves the need for multiple aggregators, introduces degree-scalers that generalize sum aggregation, and combines them in the Principal Neighbourhood Aggregation architecture.
Results
Empirical results from synthetic and real-world domains support the theoretical analysis, with PNA outperforming other literature models on most real-world benchmark tasks.
Takeaways & Limitations
PNA demonstrates that combining multiple aggregators with degree-scalers can improve GNN performance and capture graph structure across benchmark domains.
Takeaways & Limitations
The tests are not exhaustive, and the proofs do not generate optimal aggregators for every task, so the proposed architecture may perform sub-optimally in novel domains.
Abstract
from arXiv · showhide
Graph Neural Networks (GNNs) have been shown to be effective models for different predictive tasks on graph-structured data. Recent work on their expressive power has focused on isomorphism tasks and countable feature spaces. We extend this theoretical framework to include continuous features - which occur regularly in real-world input domains and within the hidden layers of GNNs - and we demonstrate the requirement for multiple aggregation functions in this context. Accordingly, we propose Principal Neighbourhood Aggregation (PNA), a novel architecture combining multiple aggregators with degree-scalers (which generalize the sum aggregator). Finally, we compare the capacity of different models to capture and exploit the graph structure via a novel benchmark containing multiple tasks taken from classical graph theory, alongside existing benchmarks from real-world domains, all of which demonstrate the strength of our model. With this work, we hope to steer some of the GNN research towards new aggregation methods which we believe are essential in the search for powerful and robust models.
1 Introduction
The paper addresses gaps in GNN benchmarks and expressive-power theory, especially for continuous features and neighbourhood structure. It proposes PNA and evaluates it on multi-task, classical graph-theory, and real-world benchmarks.
- GNN research lacks standardized benchmarks and theoretical frameworks for evaluating expressive power.
- Existing expressive-power studies mainly target isomorphism tasks with countable feature spaces, leaving graph-structure properties less understood.
- The authors hypothesize that single aggregation layers extract insufficient neighbourhood information, limiting GNN expressiveness and learning.
- PNA combines multiple aggregators with degree-scalers after proving multiple aggregators are needed for continuous-feature multisets.Degree-scalers generalize sum aggregation by amplifying or attenuating signals according to node degree.
- A new multi-task benchmark evaluates node- and graph-level classical graph-theory problems, complementing comparisons with GCN, GAT, GIN, and MPNN.The study also tests molecular-chemistry and computer-vision datasets.
2 Principal Neighbourhood Aggregation
PNA is motivated by the limited ability of single aggregators to distinguish continuous-feature neighbourhoods. It combines complementary statistics with degree-based scaling to improve structural information and expressiveness.
- Motivation: Single aggregators can fail to distinguish different neighbourhood messages when one GNN layer processes continuous features.
- Theoretical basis: At least n aggregators are needed to discriminate multisets of size n over the real numbers.
- Theoretical basis: Continuous features invalidate the earlier sum-aggregation injectivity argument for countable inputs and require continuous multiset statistics.
- Aggregators: The proposed aggregator set uses mean, maximum, minimum, and standard deviation, with normalized moments available to capture richer neighbourhood distributions.Standard deviation measures feature spread, while higher moments can become useful for high-degree nodes.
- Degree-based scalers: Degree-scalers multiply aggregated messages to amplify or attenuate signals based on node degree, generalizing sum aggregation through injective scaling functions.The authors propose logarithmic amplification to reduce exponential effects associated with repeated linear amplification.
- Combined aggregation: PNA combines four neighbour aggregators with three degree-scalers, producing twelve operations per kernel.The added weights increase parameters modestly, while the analysis identifies single aggregation as a likely expressiveness bottleneck.
3 Architecture
The models are compared within a shared encode-process-decode architecture, differing primarily in their graph convolution layers. Weight sharing, GRUs, and a variable number of layers support parameter-efficient processing.
- Architecture layout: The multi-task architecture uses M convolutions, three fully connected layers for node labels, and set2set readout for graph labels.
- Architecture layout: GRUs follow each layer’s update function and help retain information as the number of convolutional layers increases.
- Architecture layout: Sharing weights in all GNN layers except the first creates an encode-process-decode architecture with variable depth and parameter efficiency.The authors describe this configuration as effective across their experimental tasks.
- Architecture comparison: The comparison replaces the first and final graph convolutions, GC1 and GCm, with different graph convolution types.The remaining architecture is held constant when models are compared.
- Architecture comparison: The architecture was selected for performance and parameter efficiency, while alternative layouts produced similar comparative performance among GNN layers.
4 Multi-task benchmark
The paper introduces a multi-task benchmark of classical graph-theory problems on randomly generated graphs. It combines node-level and graph-level labels while training models to share useful subroutines across tasks.
- Dataset construction: The benchmark uses undirected, unweighted randomly generated graphs spanning diverse graph types, with presented experiments using 15–50 nodes.Random toggling is used to increase graph diversity.
- Task design: The benchmark contains three node labels and three graph labels based on standard graph-theory problems.
- Task design: Node tasks are single-source shortest-path lengths, eccentricity, and Laplacian features, while graph tasks are connectivity, diameter, and spectral radius.
- Input features: Inputs include a one-hot source vector and random node features X sampled independently from U[0, 1].The random features also provide node identifiers in tasks other than Laplacian features.
- Training: Because tasks share subroutines such as graph traversals, the models can potentially share parameters between tasks during training.The authors expect these commonalities to reinforce one another during training.
- Training: Training used Adam for up to 10,000 epochs with early stopping after 1,000 epochs of patience, and reported the five lowest-validation-error runs from ten trials.
5 Results and discussion
PNA is evaluated on artificial and real-world benchmarks using controlled comparisons with existing GNNs. It generally retains an advantage on graph-structure-sensitive settings and extrapolation to larger graphs.
- Evaluation setup: The comparisons use the same architecture and near-optimal hyper-parameters across models, with Figure 6 also marking runs trained with additional patience.
- Artificial benchmark: The parameter-increase test gives competing models 30% more parameters than PNA, yet they remain qualitatively less capable of capturing graph structure.This addresses whether PNA’s gains arise only from its approximately 15% parameter increase over other models.
- Artificial benchmark: PNA consistently outperforms the other models across graph sizes when extrapolating beyond its training range.The models were trained on sizes 15–25, validated on 25–30, and evaluated on 20–50 nodes.
- Artificial benchmark: Performance relative to the baseline gradually worsens on larger graphs, and some models exhibit feature explosions.
- Real-world benchmarks: On real-world datasets, PNA outperforms other literature models on most molecular chemistry and computer-vision tasks.The evaluated datasets are ZINC, MolHIV, CIFAR10, and MNIST.
- Real-world benchmarks: PNA performs strongly on chemical datasets because its design retains neighbourhood information, whereas constant-degree vision graphs make scalers redundant.The no-scaler PNA performs better on the vision datasets, where parameters can instead support larger hidden sizes.
6 Conclusion
The paper extends GNN expressive-power analysis to continuous features, proposes PNA with multiple aggregators and degree-scalers, and evaluates it through new multi-task benchmarks. Synthetic and real-world results support the theoretical analysis and show PNA outperforming prior GNN layer designs.
- Conclusion: The theoretical framework is extended to continuous features, for which the paper proves that multiple aggregators are needed.
- Conclusion: PNA combines multiple aggregators with degree-scalers, including a logarithmic scaling that generalizes sum aggregation.
- Conclusion: The paper introduces a multi-task benchmark and an encode-process-decode architecture to study how GNNs capture graph structures.
- Conclusion: Empirical results from synthetic and real-world domains support the theoretical evidence, with PNA appearing to outperform prior GNN layer designs.
Broader Impact
The paper analyzes GNN expressive power and proposes PNA to address theoretical limitations, while noting risks, non-exhaustive testing, and uncertain performance on novel domains.
- The work studies GNN expressive power and may indirectly affect the impacts of graph representation learning in applied domains.
- Proofs of limitations on continuous feature spaces may clarify GNN behaviour and motivate research toward more reliable models.
- PNA is proposed to overcome some theoretical limitations, but the tests are not exhaustive and the proofs do not generate optimal aggregators for every task.
- Out-of-distribution tests showed consistent PNA superiority alongside a noticeable performance drop on graphs much larger than those used for training.
Funding Disclosure
The authors disclose affiliations with InVivo AI, the University of Cambridge, and DeepMind.
- Dominique Beaini is a Machine Learning Researcher at InVivo AI, Pietro Liò is a Cambridge computer science professor, and Petar Veličković is a DeepMind research scientist.
A Proof for Theorem 1 (Number of aggregators needed)
The proof establishes that continuously aggregating a size-n multiset of real-valued features requires at least n aggregators, and gives a matching construction.
- At least n continuous aggregators are required to discriminate all size-n multisets over the real numbers.
- The proof represents sorted size-n multisets as an n-dimensional subspace S of R^n and aggregators as a continuous map from S to R.
- Assuming n−1 aggregators yield an injective continuous map from S to R^(n−1) creates the contradiction setup.
- Borsuk–Ulam supplies distinct antipodal points with identical outputs, contradicting injectivity.
- Exactly n aggregators suffice when each returns one ordered element of the multiset, uniquely reconstructing the multiset.
B Proof for Proposition 1 (Moments of the multiset)
The moments of a multiset provide n aggregators that uniquely determine its elements through centered power sums, symmetric sums, and polynomial roots.
- The multiset moments defined in Equation 4 form a valid example using n aggregators.
- Centering each element by the mean yields residuals whose power sums are determined by the moments.
- The residual power sums recursively determine the elementary symmetric sums through identities for e1, e2, e3, and higher orders.
- These symmetric sums determine the coefficients of the monic polynomial whose roots are the residual elements.
- Adding the mean back to the uniquely determined residual roots reconstructs the original multiset.
- When multiset size is variable up to n, an additional aggregator can infer the node degree using an injective scaler.
C Proof for Theorem 2 (Injective functions on countable multisets)
Theorem 2 shows that mean aggregation combined with an injective degree scaler can uniquely represent bounded multisets of countable elements. The section then reports empirical evidence that aggregation choices affect expressive power and performance across graph tasks.
- Proof: The proof encodes countable feature values as natural numbers and uses bounded multiset size to construct an injective representation.The construction recovers the multiset from the digits of a base-N representation after determining its cardinality.
- Proof: An injective degree scaler prevents ambiguity in multiset cardinality when mean aggregation produces the same value for differently sized neighbourhoods.The contradiction argument shows that the multiset size can be inferred from the scaled mean.
- Normalized moments aggregation: Around 3 aggregators, the multi-task models reach a task- and graph-dependent performance peak before higher moments reduce performance.The authors attribute the decline to harder optimization and overfitting, while expecting higher moments to help more on higher-degree graphs.
- Normalized moments aggregation: The complete PNA with mean, max, min, and standard-deviation aggregators performs better and more consistently than using higher moments alone in these tasks.The authors note that discrete aggregators may be especially valuable for algorithmic tasks.
- Graph type experiments: PNA improves across graph types but performs worst on high-diameter graphs, suggesting that the available message-passing depth cannot reach the complete graph.The authors identify the message-passing framework as the main limitation to PNA performance.
I Standard architecture
The standard-architecture experiments compare models using fixed convolutions, skip connections, mean readout, and no GRU or S2S readout. PNA remains consistently stronger despite using fewer parameters, while multi-task training benefits models without scalers most.
- Architecture motivations: Skip connections are added because GAT, GIN, and GCN benefit most from the GRU's anti-over-smoothing effect when they lack aggregation skip connections.The same skip-connection treatment is applied across models for a fair comparison.
- Architecture motivations: S2S readout particularly helps models without scalers by providing an alternative counting mechanism.Repeated convolutions are also described as a parameter-saving prior that does not change model ranking.
- Standard architecture: The standard architecture removes the GRU, uses skip connections, replaces S2S with mean readout, and fixes the number of convolutions at 8.This provides the comparison setting described for Figure 9.
- Multi-task training: Multi-task training offers regularization, with models without scalers or a sum aggregator benefiting most from the approach.The authors hypothesize that task-specific supervision helps these models recognize neighbourhood size.
- Parameters comparison: PNA performs consistently better even with fewer parameters, while increasing other models' parameter counts does not improve their performance.The authors interpret this as a qualitative capacity improvement from PNA's multiple aggregators.