PaperlayerView exact revision on arXiv

Source-linked AI summary

Nested Graph Neural Networks

Muhan Zhang, Pan Li

arXiv:2110.13197v1cs.LG

TL;DR

Message-passing GNNs are limited by rooted-subtree representations, especially for cyclic and non-tree structures. NGNN replaces these with pooled rooted-subgraph representations learned by nested base and outer GNNs. It is theoretically more expressive than 1-WL and consistently improves diverse base GNNs without the O(n^3) complexity of other more powerful GNNs.

  • Problem

    Rooted-subtree representations limit message-passing GNNs to 1-WL-level discrimination and poorly represent arbitrary subgraphs, especially cycles.

  • Method

    NGNN extracts a rooted local subgraph around each node, applies a shared base GNN and subgraph pooling, then uses an outer GNN to form the whole-graph representation.

  • Results

    NGNN consistently improves various base GNNs across datasets, while Nested GIN reaches 99.9% accuracy on 1-WL-indistinguishable graph pairs.

  • Takeaways & Limitations

    NGNN provides a plug-and-play way to increase GNN representation power and can discriminate almost all r-regular graphs where 1-WL fails.

  • Takeaways & Limitations

    Standard message-passing GNNs remain limited to rooted-subtree representations and cannot discriminate all n-node r-regular graphs.

Abstract

from arXiv · show

Graph neural network (GNN)'s success in graph classification is closely related to the Weisfeiler-Lehman (1-WL) algorithm. By iteratively aggregating neighboring node features to a center node, both 1-WL and GNN obtain a node representation that encodes a rooted subtree around the center node. These rooted subtree representations are then pooled into a single representation to represent the whole graph. However, rooted subtrees are of limited expressiveness to represent a non-tree graph. To address it, we propose Nested Graph Neural Networks (NGNNs). NGNN represents a graph with rooted subgraphs instead of rooted subtrees, so that two graphs sharing many identical subgraphs (rather than subtrees) tend to have similar representations. The key is to make each node representation encode a subgraph around it more than a subtree. To achieve this, NGNN extracts a local subgraph around each node and applies a base GNN to each subgraph to learn a subgraph representation. The whole-graph representation is then obtained by pooling these subgraph representations. We provide a rigorous theoretical analysis showing that NGNN is strictly more powerful than 1-WL. In particular, we proved that NGNN can discriminate almost all r-regular graphs, where 1-WL always fails. Moreover, unlike other more powerful GNNs, NGNN only introduces a constant-factor higher time complexity than standard GNNs. NGNN is a plug-and-play framework that can be combined with various base GNNs. We test NGNN with different base GNNs on several benchmark datasets. NGNN uniformly improves their performance and shows highly competitive performance on all datasets.

1 Introduction

Message-passing GNNs represent graphs through pooled rooted subtrees, which cannot capture important cyclic and non-tree structures. NGNN instead pools representations learned from rooted local subgraphs, improving expressive power while retaining a flexible framework.

  • Motivation: Graph neural networks propagate neighbor features into node representations and pool them for graph-level classification.Their message-passing scheme encodes local structure and feature information before graph pooling.
  • Motivation: 1-WL and message-passing GNNs can assign identical representations to two non-isomorphic graphs when their rooted subtrees match at every height.In Figure 1, two triangles and a hexagon remain indistinguishable despite having different global structures.
  • Motivation: Rooted subtrees are limited for representing arbitrary subgraphs, especially cyclic structures.This limitation bounds message-passing GNN discrimination by 1-WL and prevents distinguishing all n-node r-regular graphs.
  • NGNN framework: NGNN encodes a rooted local h-hop subgraph around each node instead of a rooted subtree.Rooted subgraphs can distinguish the Figure 1 graphs by comparing their height-1 neighborhoods.
  • NGNN framework: NGNN applies a base GNN independently to each rooted subgraph, pools each subgraph into the root representation, then pools root representations into a graph representation.This nested design uses inner subgraph-level processing followed by outer graph-level aggregation.
  • Advantages: NGNN is plug-and-play, can augment node features with subgraph-specific structure, and improves expressive power with linear time and space complexity in graph size.The framework can use different base GNNs and is evaluated across synthetic and real-world graph tasks.

2 Preliminaries

Graph classification and regression learn mappings from attributed graphs to class labels or target values. The 1-WL test iteratively refines node colors using current colors and neighbor-color multisets to assess graph non-isomorphism.

  • Notation and problem definition: A graph is represented as G = (V, E), with node features x_i and edge features e_ij optionally associated with its elements.The task learns a function mapping the graph to a class or target value y.
  • Notation and problem definition: Graph classification and regression learn a function that maps a graph G to a class or target value y.The graph contains a node set V and edge set E, with optional node and edge feature vectors.
  • Weisfeiler-Lehman test: 1-WL initializes all nodes with one color and repeatedly updates each color from its current color and the multiset of neighbor colors.Two nodes receive the same updated color exactly when both inputs match.
  • Weisfeiler-Lehman test: 1-WL declares graphs non-isomorphic when their node colors differ at some iteration; otherwise, it may fail to determine non-isomorphism.The procedure stops when the number of colors no longer increases between iterations.
  • Weisfeiler-Lehman test: 1-WL colors encode rooted subtrees around nodes at corresponding heights.Nodes share a color at iteration h if and only if their height-h rooted subtrees are identical.

3 Nested Graph Neural Network

NGNN replaces rooted-subtree representations with rooted-subgraph representations computed by nested GNNs, improving expressiveness beyond message passing GNNs and 1-WL. Its theory shows strong discrimination of regular graphs, while its framework remains flexible and computationally practical.

  • The NGNN framework: NGNN extracts each node’s local h-hop rooted subgraph and uses a base GNN plus subgraph pooling to produce that node’s representation.The rooted subgraph is induced by nodes within h hops, and the base GNN is applied independently to each extracted subgraph.
  • The NGNN framework: An outer pooling layer aggregates the root-node representations produced by the inner base GNNs into a whole-graph representation.NGNN is therefore a two-level GNN: inner subgraph-level processing followed by outer graph-level aggregation.
  • The NGNN framework: The same node can receive different representations in different rooted subgraphs because each extracted subgraph is treated as an independent graph.Base GNNs share parameters across rooted subgraphs, while representations remain specific to each root context.
  • The representation power of NGNN: NGNN can discriminate almost all (1 − o(1)) pairs of n-sized r-regular graphs under the stated degree and depth conditions, whereas 1-WL fails on these graphs.The framework is therefore strictly more powerful than 1-WL and message passing GNNs.
  • The representation power of NGNN: Subgraph pooling lets NGNN achieve this theoretical power with few base-GNN layers, and simulations match the theorem’s predictions.The authors attribute the reduced depth requirement to pooling intermediate node representations within each subgraph.
  • The representation power of NGNN: The exact comparison between NGNN and 3-WL remains unresolved, with early analysis indicating both may fail on strongly regular graphs sharing the same parameters.The paper establishes strict superiority over 1-WL and 2-WL but leaves 3-WL for future work.
  • Discussion: NGNN is plug-and-play with higher-order base GNNs and can reduce their complexity from O(n^3) on the full graph to O(nc^3) across subgraphs of at most c nodes.This reduction follows from applying the higher-order model to multiple small rooted subgraphs instead of the whole graph.

4 Related work

Related work seeks graph representations more powerful than message passing GNNs and 1-WL, but higher-order approaches can be difficult to scale. Other methods augment message passing with identity or structural features, while NGNN relates to several rooted-subgraph approaches.

  • Higher-order GNNs: Higher-order GNNs mimic higher-dimensional WL tests but generally represent all node tuples, limiting scalability on large graphs.Universality-oriented approaches may require polynomial(n)-order tensors, which the paper characterizes as more theoretically valuable than practically applicable.
  • Feature augmentation: Node-index and random features enhance message passing when nodes have distinct identities, but their usefulness depends on identities being meaningful independently of graph structure.The paper contrasts these methods with structural feature augmentation intended to preserve generalization.
  • Feature augmentation: Distance-based features augment GNNs using node distances to a target set, and NGNN is naturally compatible with this strategy through independent rooted-subgraph processing.The cited methods include SEAL, IGMC, and DE.
  • Related rooted-subgraph methods: ID-GNN extracts rooted subgraphs and can be viewed as a special case of NGNN with root-specific parameters and no subgraph pooling.The supplied passage also identifies differences in message-passing depth and the use of the root node’s intermediate representation.
  • Related rooted-subgraph methods: MixHop concatenates multi-hop neighbor aggregations, whereas k-hop GNN processes nested neighborhoods sequentially; NGNN differs by retaining connections among other rooted-subgraph nodes.The comparison highlights how rooted-subgraph structure is handled across these methods.

5 Experiments

Experiments evaluate NGNN’s expressive power, performance gains over base GNNs, competitiveness on OGB benchmarks, and computational cost. NGNN distinguishes nearly all tested 1-WL-indistinguishable graphs, improves base models across QM9 and TU, performs competitively on OGB, and incurs higher but comparable training time.

  • Expressive power: NGNN matched Theorem 1 almost perfectly in simulations, demonstrating practical discrimination of r-regular graphs.The simulation tested NGNN’s ability to distinguish r-regular graphs.
  • Expressive power: 99.9% classification accuracy on EXP distinguished almost all 1-WL-indistinguishable graph pairs, versus the 50% expected accuracy of standard message passing GNNs.EXP contains 600 carefully constructed non-isomorphic graph pairs with different labels.
  • Base-GNN improvements: NGNN improved most base-GNN results on QM9 and TU, reducing QM9 MAE by up to 7.9 times and increasing TU accuracy by up to 14.3%.The experiments used QM9 regression and TU graph-classification datasets.
  • OGB benchmarks: Nested GIN achieved highly competitive OGB performance, reaching test scores of 79.86 on ogbg-molhiv and 30.07 test AP on ogbg-molpcba.On ogbg-molpcba, Nested GIN reached 30.07 test AP with ensemble and 28.32 without ensemble, compared with 27.03 for plain GIN.
  • Computational cost: On ogbg-molhiv, training time rose from 54s to 183s per epoch; on ogbg-molpcba, it rose from 10min to 20min, supporting comparable time complexity.The additional cost comes from independently learning representations from rooted subgraphs.
  • Limitation: Current NGNN does not scale to datasets with many nodes or high average degree because rooted subgraphs are copied into GPU memory.Reducing batch size or subgraph height helps memory use but degrades performance.

6 Conclusions

NGNN improves GNN representation power by encoding rooted subgraphs rather than rooted subtrees. It theoretically distinguishes almost all r-regular graphs that 1-WL cannot and empirically improves diverse base GNNs without O(n3) complexity.

  • Conclusion: NGNN encodes rooted subgraphs instead of rooted subtrees to improve GNN representation power.The framework is designed as a general method for improving GNNs.
  • Conclusion: NGNN can discriminate almost all r-regular graphs where 1-WL always fails.This is the paper’s principal theoretical expressive-power result.
  • Conclusion: NGNN consistently improves various base GNNs across datasets without incurring the O(n3) complexity of other more powerful GNNs.The conclusion summarizes both empirical gains and computational positioning.

A Proof of Theorem 1

The proof shows that rooted subgraph extraction creates distance-sensitive boundary representations, which message passing propagates inward to distinguish subgraphs and, with injective pooling, whole graphs.

  • Proof strategy: Height-k rooted subgraphs are induced by nodes within k hops of the center, and extraction is equivalent to injecting distance features indicating distances below k + 1.These implicit distance features underpin the proof that NGNN can exceed 1-WL.
  • Proof strategy: For sampled r-regular graphs, edge configurations between successive distance layers capture structural differences around selected nodes.The proof defines edge configurations as lists describing how nodes in one layer connect to the preceding layer.
  • Node discrimination: With at most ϵ⌈log n / log(r−1−ϵ)⌉ message-passing layers, injective pooling distinguishes two extracted subgraphs with probability at least 1 − o(n−1).This is the conclusion of Lemma 2 for randomly sampled n-node r-regular graphs under the stated range of r.
  • Node discrimination: Subgraph extraction removes boundary edges, causing non-default representations to appear at the outer layer and propagate toward the center through message passing.Nodes away from the boundary initially retain default representations because they continue to have degree r within the extracted subgraphs.
  • Graph discrimination: A union bound shows that, with probability 1 − o(1), one node representation from one graph differs from every node representation in the other graph.Injective final graph pooling then yields different graph representations.

B Design choices of NGNN

NGNN supports alternative nesting depths, pooling functions, and receptive-field settings, while the paper identifies practical choices for subgraph height and base-GNN depth.

  • High-order NGNN: High-order NGNN recursively uses an NGNN as the base model, allowing arbitrary nesting orders in principle.The paper leaves the effect of higher-order architectures on representation power and practical performance for future work.
  • Pooling functions: Mean pooling works well for both subgraph and graph representations, alongside alternatives such as sum, max, top-K, and hierarchical pooling.The paper also reports another pooling function as sometimes useful for subgraph representation.
  • Subgraph height and base-GNN layers: Theorem 1 guides choices of subgraph height h and base-GNN layers l, with h = 3 and l = 4 generally performing well across tasks.Small h restricts the receptive field, whereas overly large h can make rooted subgraphs excessively large.

C More details about the experimental settings

The experiments use dataset-specific NGNN configurations, including different rooted-subgraph heights, GIN depths, and pooling choices.

  • Experimental configurations: For ogbg-molhiv, NGNN uses rooted subgraph height h = 4, six GIN layers, and mean pooling at both levels.The experiments run on a Linux server with 64GB memory, two RTX 2080S GPUs, and an Intel i9-9900 CPU.
  • Experimental configurations: For ogbg-molpcba, NGNN uses rooted subgraph height h = 3, four GIN layers, center pooling for subgraphs, and mean pooling for graphs.

D Simulation experiments to verify Theorem 1

The simulation tests NGNN’s theoretical discrimination guarantees on random regular graphs and reports close agreement between empirical behavior and the theory.

  • Experimental setup: 100 uniformly sampled n-node 3-regular graphs are evaluated for n from 10 to 1280 and rooted-subgraph heights h from 1 to 10.An untrained Nested GIN with one message-passing layer and uniform node features is used.
  • Node-level results: The node-level results match Lemma 2 almost perfectly, with the required h tracking its theoretical lower bound.Node representations are compared after subgraph pooling.
  • Graph-level results: Graph-level pooling distinguishes almost all sampled regular graphs even when h is small and does not always exceed 0.5 log(n) / log(r−1).A single node representation differing from all representations in the other graph is sufficient for injective graph pooling to separate the graphs.

E Ablation study on DE

The ablation evaluates Distance Encoding (DE) in Nested GNNs on QM9. DE generally strengthens NGNN improvements over base GNNs while adding negligible time.

  • DE features: DE encodes each node’s distance to the root as an embedding concatenated with its raw features.This augmented vector becomes the node’s input feature within the rooted subgraph.
  • Experimental setup: The QM9 ablation compares each base GNN with Nested GNN variants without and with DE features.The compared models use the same base GNNs as Table 3.
  • Results: NGNNs without DE generally outperform their corresponding base GNNs, showing that the framework helps even without feature augmentation.The comparison uses Table 6’s color coding to identify improvements over the base GNN.
  • Results: When NGNN improves over the base GNN, adding DE can further improve performance by achieving the smallest MAE among the three models.The three-way comparison is between the base GNN, NGNN without DE, and NGNN with DE.
  • Efficiency: DE features add negligible time because they can be computed simultaneously with rooted subgraph extraction.This motivates using DE as NGNN’s default augmentation.
Loading 2110.13197v1…