Source-linked AI summary
Design Space for Graph Neural Networks
Jiaxuan You, Rex Ying, Jure Leskovec
TL;DR
Existing GNN research largely studies specific architectures on limited task sets, leaving systematic design exploration and rapid architecture selection for novel tasks underdeveloped. The paper defines and evaluates a 315,000-design GNN space across 32 tasks, adds a task-similarity metric and efficient evaluation, and reports transferable designs, design guidelines, and state-of-the-art performance.
Problem
GNN research often evaluates specific architectures on limited tasks, while the best design for a novel task or dataset remains unclear.
Method
The paper defines a 12-dimension GNN design space, a task space using Kendall rank correlation of anchor-model performance, and an efficient design-space evaluation procedure.
Results
The study provides design guidelines, transfers designs across tasks, and finds state-of-the-art models from the design space.
Takeaways & Limitations
Systematically studying GNN design and task spaces supports cross-task design transfer and reproducible, standardized experimentation through GraphGym.
Takeaways & Limitations
The evaluation focuses on message-passing GNNs and controls experiments by matching trainable parameter counts.
Abstract
from arXiv · showhide
The rapid evolution of Graph Neural Networks (GNNs) has led to a growing number of new architectures as well as novel applications. However, current research focuses on proposing and evaluating specific architectural designs of GNNs, as opposed to studying the more general design space of GNNs that consists of a Cartesian product of different design dimensions, such as the number of layers or the type of the aggregation function. Additionally, GNN designs are often specialized to a single task, yet few efforts have been made to understand how to quickly find the best GNN design for a novel task or a novel dataset. Here we define and systematically study the architectural design space for GNNs which consists of 315,000 different designs over 32 different predictive tasks. Our approach features three key innovations: (1) A general GNN design space; (2) a GNN task space with a similarity metric, so that for a given novel task/dataset, we can quickly identify/transfer the best performing architecture; (3) an efficient and effective design space evaluation method which allows insights to be distilled from a huge number of model-task combinations. Our key results include: (1) A comprehensive set of guidelines for designing well-performing GNNs; (2) while best GNN designs for different tasks vary significantly, the GNN task space allows for transferring the best designs across different tasks; (3) models discovered using our design space achieve state-of-the-art performance. Overall, our work offers a principled and scalable approach to transition from studying individual GNN designs for specific tasks, to systematically studying the GNN design space and the task space. Finally, we release GraphGym, a powerful platform for exploring different GNN designs and tasks. GraphGym features modularized GNN implementation, standardized GNN evaluation, and reproducible and scalable experiment management.
1 Introduction
GNN research has largely studied specific architectures on limited task sets, leaving the broader design space, cross-task transfer, and unified evaluation insufficiently addressed. This work defines general design and task spaces, an efficient evaluation procedure, and GraphGym to study them systematically.
- Issues in GNN architecture design: Specific GNN architectures are instances within a larger cross-product design space, so focusing on individual designs can limit discovery of successful models.Changing aggregation or adding skip connections can produce variants that outperform named baseline architectures on some tasks.
- Issues in GNN evaluation: GNN models are commonly evaluated on limited task sets, while novel tasks may not resemble existing benchmarks, making effective architecture design unclear.Examples include circuit design, SAT generation, data imputation, and subgraph matching.
- Issues in GNN implementation: A unified platform for exploring GNN designs across node-, edge-, and graph-level tasks is lacking.The paper identifies this implementation gap as a major contributor to the preceding architecture and evaluation issues.
- Present work: The study combines a general GNN design space, a task-similarity space, and efficient design-space evaluation, and releases GraphGym for modular, standardized, reproducible experiments.The task space uses Kendall rank correlation of fixed anchor-model performances to quantify task similarity.
- Contribution: GraphGym standardizes model and task descriptions and evaluation, making reproduction and fair comparison require minimal effort.The platform provides reproducible and scalable experiment management.
2 Related Work
Prior work has explored limited GNN architecture subsets, fairness of comparisons, and architecture transfer under narrower task assumptions. This paper differs by studying a broader, modular GNN design space across diverse tasks.
- Graph Architecture Search: GNN architecture-search methods have focused on within-layer design and small numbers of node-classification tasks.
- Evaluation of GNN Models: Fair-comparison studies generally evaluate specific architectures such as GCN, GAT, and GraphSAGE rather than the general GNN design space.
- Other graph learning models: This work focuses on message-passing GNNs because of their performance and efficient implementation across graph tasks.Alternative graph-learning models have different, less-modular design spaces.
- Transferable Architecture Search: Prior transferable architecture approaches often assume tasks share a distribution or that one architecture performs well across all tasks.
3 Preliminaries
The paper formalizes GNN designs, design spaces, task spaces, and experiment spaces as combinations of architectural choices and predictive tasks.
- Terminology: A design is a concrete GNN instantiation, while design choices are selected values within dimensions such as layer count or aggregation.For example, L = 2 is a choice within the layer-count dimension.
- Terminology: A design space is the Cartesian product of its design dimensions.Layer counts with four options and aggregation functions with three options yield 12 designs.
- Terminology: A task space contains tasks such as node classification or graph classification, and an experiment applies one GNN design to one task.
- Terminology: An experiment space covers all combinations of designs and tasks.
4 Proposed Design Space for GNNs
The proposed GNN design space organizes architectural and training choices into intra-layer, inter-layer, and learning-configuration dimensions. It uses modest, literature-informed option ranges to support systematic exploration.
- Design-space principles: The design-space principles prioritize important development choices, few dimensions, and modest option ranges informed by prior literature.Model-specific dimensions are omitted to keep the space compact and general.
- Intra-layer design: Intra-layer design orders batch normalization, dropout, activation, and aggregation after a linear layer.Candidate activations are RELU, PRELU, and SWISH; aggregations are MEAN, MAX, and SUM.
- Intra-layer design: The adopted intra-layer space varies batch normalization, dropout, activation, and aggregation choices.Dropout options include False, 0.3, and 0.6, while batch normalization is True or False.
- Inter-layer design: Inter-layer design studies how message-passing layers are organized, including residual SKIP-SUM and dense concatenating SKIP-CAT connections.Skip connections are examined systematically in combination with other design dimensions.
- Training configurations: Learning configuration varies batch size, learning rate, optimizer, and training epochs.
5 Proposed Task Space for GNNs
The paper defines a task space that measures relationships among diverse GNN tasks through architecture-performance rankings, enabling design transfer and discovery beyond fixed task taxonomies.
- Proposed task similarity metric: Task similarity is computed by ranking fixed anchor GNNs on each task and applying Kendall rank correlation to those rankings.The metric uses anchor-model performance rather than dataset-domain or prediction-type labels.
- Proposed task similarity metric: The task space supports transferring promising GNN designs between similar tasks and identifying novel tasks that may inspire new designs.The authors present these as the two intended uses of quantitative task similarity.
- Anchor models: The anchor models are selected to represent diverse performance levels after random designs are evaluated across tasks.Random designs are ranked by average performance and evenly divided into groups, from which representative models are selected.
- Metric generality: The proposed metric is intended to generalize across predictive and non-predictive tasks by ranking models with task-appropriate scores.Examples include ROC AUC for binary classification, mean square error for regression, and drug-likeness for molecule generation.
6 Evaluation of GNN Design Space
The evaluation framework studies individual GNN design dimensions with controlled random comparisons across the large model-task space, while matching computational budgets for fairness.
- Evaluation framework: Over 10M model-task combinations make exhaustive grid search impractical, motivating controlled random search with equalized model budgets.The method is designed to distill design insights while ensuring fair comparisons.
- Controlled random search: For each design choice, 96 sampled setups are paired while all other dimensions remain fixed.For BatchNorm, each BN = TRUE setup is altered to BN = FALSE under the same remaining configuration.
- Ranking analysis: Design choices are ranked within each paired setup, treating performance differences within ϵ = 0.02 as ties.The ranking distribution across setups provides the basis for design-space insights.
- Computational budget: The framework controls trainable parameter counts by adjusting hidden dimensions around a reference architecture.The reference uses 1 pre-processing layer, 3 message-passing layers, 1 post-processing layer, and 256 hidden dimensions.
7 Experiments
Experiments use GraphGym and ranking analysis to evaluate GNN designs across tasks, derive condensed search spaces, test task-based transfer, and assess performance on ogbg-molhiv.
- 7.1 GraphGym: Platform for GNN Design: GraphGym provides modularized implementation, standardized evaluation, and reproducible, scalable experiment management for GNN design studies.Configurations specify model and evaluation choices, supporting reproducible comparisons.
- 7.3 Results on GNN Design Space Evaluation: Design-space analysis finds BatchNorm, PRELU, SUM aggregation, skip connections, ADAM, and longer training generally favorable, while several depth choices remain task specific.The reported findings distinguish broadly useful choices from choices whose performance depends on the task.
- 7.3 Results on GNN Design Space Evaluation: 96 setups per design dimension and controlled random search over 10M combinations provide broader evidence than evaluating one architectural variant on a few benchmarks.The authors report that 7 of 12 design dimensions significantly influence performance under Bonferroni-corrected one-way ANOVA.
- 7.4 Results on the Efficacy of GNN Task Space: The condensed design space reduces 315,000 designs to 96, enabling full-grid experiments for testing task-space transfer.The condensed space is derived from the design guidelines.
- 7.5 Case Study: Applying to a Challenging New Task ogbg-molhiv: On ogbg-molhiv, the best discovered design reaches ROC AUC 0.792 versus 0.771 for existing SOTA.A design transferred from a highly similar task reaches ROC AUC 0.785 versus 0.771, whereas one from a dissimilar task reaches AUC 0.736.
8 Conclusion
The conclusion presents joint study of GNN design and task spaces as a principled route to new design insights, lower development costs, and empirical performance gains.
- 8 Conclusion: Jointly studying GNN design and task spaces with tractable evaluation techniques yields new understanding of models and tasks.The conclusion connects this approach with reduced algorithm development costs and empirical performance gains.
- 8 Conclusion: The paper recommends moving beyond individual GNN designs and tasks toward systematic study of their broader spaces.This is the paper’s stated overall direction for GNN research.
Broader Impact
The framework extends beyond individual GNN architectures by supporting fair evaluation, task comparison, and transfer across applications. It also aims to lower barriers for applying GNNs in other domains while remaining an intentionally incomplete framework.
- Impact on GNN research: The framework emphasizes design principles and task spaces rather than particular GNN instantiations, helping analyze where algorithmic advancements are useful.It can also identify novel tasks that are dissimilar to existing ones and may require new algorithmic development.
- Impact on machine learning research: Controlled random search can evaluate algorithmic advancements across random model-task combinations to identify where they improve performance.The approach is presented as applicable beyond GNNs for fair evaluation of novel algorithmic advancements.
- Impact on machine learning research: The ranking-based task similarity metric can compare general machine learning tasks whenever model designs can be ranked by performance.The paper gives MNIST and CIFAR-10 classification as examples of potential use beyond GNN tasks.
- Impact on other research domains: Domain experts can provide formatted datasets and receive recommended GNN designs selected through similarity-based transfer or design-space search.The fastest mode applies anchor models to estimate similarity, then transfers top designs from similar tasks.
- Scope: The framework is intended as a systematic, extensible example rather than a complete account of all design and evaluation aspects.The authors state that new design dimensions and tasks can be added to the framework.
B.2 Case Study: Link Prediction as a New Type of Tasks
The link-prediction case study extends the task space beyond node and graph classification. Most link-prediction tasks cluster separately, and their strongest discovered designs differ from those for the other task types.
- Task-space extension: The framework is extended by adding link-prediction tasks to the GNN task space.The authors re-plot the task-space analysis after including these tasks.
- Task-space structure: Most link-prediction tasks form a cluster distant from node- and graph-classification tasks in the extended task space.The task-space visualization is presented in Figure B.2.
- Best designs: The best discovered GNN designs for link prediction differ from the best designs for node or graph classification.The comparison is reported between Tables B.1 and C.1.
C Best GNN Designs for Each Task
The design-space search identifies task-specific best GNN designs and outperforms the best tested GCN configuration on most tasks. Task similarity is also evaluated as an indicator for transferring designs across tasks.
- Per-task designs: The best model from the design space outperforms the best GCN model on 24 out of 32 tasks.The GCN comparison uses message-passing depths of 2, 4, 6, and 8 with other optimal hyperparameters.
- Design transfer: The quantitative task-similarity metric is reported as a strong indicator for transferring the best designs between tasks.Additional visualizations show task similarity and performance ranking after task transfer.
- Design transfer: Figure D.3 ranks the performance of a design selected on one task after transfer to another task across all task pairs.The ranking is normalized, with higher values indicating better performance.