NodeNAS: Node-Specific Graph Neural Architecture Search for Out-of-Distribution Generalization

TL;DR

NodeNAS introduces node specific graph neural architecture search to tailor aggregation methods for individual nodes, addressing OOD generalization with limited data. Building on this, MNNAS extends to multi-dimension architecture search with adaptive aggregation attention and a disentangled mapping encoder that leverages power-law degree distributions and assortativity to guide architecture customization. The framework achieves superior OOD generalization on both supervised and unsupervised tasks, including graph classification and community detection, often with single graph training. This work provides a scalable, interpretable approach to tailor node representations across diverse graphs, with potential impact on real-world graph learning under distribution shifts.

Abstract

Graph neural architecture search (GraphNAS) has demonstrated advantages in mitigating performance degradation of graph neural networks (GNNs) due to distribution shifts. Recent approaches introduce weight sharing across tailored architectures, generating unique GNN architectures for each graph end-to-end. However, existing GraphNAS methods do not account for distribution patterns across different graphs and heavily rely on extensive training data. With sparse or single training graphs, these methods struggle to discover optimal mappings between graphs and architectures, failing to generalize to out-of-distribution (OOD) data. In this paper, we propose node-specific graph neural architecture search(NodeNAS), which aims to tailor distinct aggregation methods for different nodes through disentangling node topology and graph distribution with limited datasets. We further propose adaptive aggregation attention based Multi-dim NodeNAS method(MNNAS), which learns an node-specific architecture customizer with good generalizability. Specifically, we extend the vertical depth of the search space, supporting simultaneous node-specific architecture customization across multiple dimensions. Moreover, we model the power-law distribution of node degrees under varying assortativity, encoding structure invariant information to guide architecture customization across each dimension. Extensive experiments across supervised and unsupervised tasks demonstrate that MNNAS surpasses state-of-the-art algorithms and achieves excellent OOD generalization.

Figures (4)

• Figure 1: An overview of NodeNAS
• Figure 2: An overview of our proposed MNNAS model.
• Figure 3: Degree distribution of datasets.
• Figure 4: (a)-(d) illustrate the statistically derived architecture customization patterns of MNNAS for different degree distributions across datasets