PaperlayerView exact revision on arXiv

Source-linked AI summary

Identity-aware Graph Neural Networks

Jiaxuan You, Jonathan Gomes-Selman, Rex Ying, Jure Leskovec

arXiv:2101.10320v2cs.LGcs.AIcs.SI

TL;DR

Existing message-passing GNNs are limited by the 1-WL test and can fail to distinguish nodes and graphs with different structures. ID-GNNs incorporate node identities through heterogeneous message passing, improving corresponding GNNs across diverse prediction tasks, including a 40% average accuracy gain on challenging tasks.

  • Problem

    Existing message-passing GNNs are bounded by the 1-WL test, limiting their ability to distinguish different neighborhood structures, regular graphs, and shortest-path relationships.

  • Method

    ID-GNNs incorporate node identities through inductive identity coloring, heterogeneous message passing, or augmented cycle-count features.

  • Results

    ID-GNNs outperform corresponding GNNs across node, edge, and graph tasks, with a 40% average accuracy improvement on challenging property-prediction tasks.

  • Takeaways & Limitations

    ID-GNNs provide a general extension of message-passing GNNs with expressive power beyond 1-WL and broad empirical gains across tasks.

  • Takeaways & Limitations

    ID-GNNs can require substantially more memory than mini-batch GNNs, increasing memory usage by 34%, 79%, and 162% at batch sizes 16, 32, and 64 on CiteSeer.

Abstract

from arXiv · show

Message passing Graph Neural Networks (GNNs) provide a powerful modeling framework for relational data. However, the expressive power of existing GNNs is upper-bounded by the 1-Weisfeiler-Lehman (1-WL) graph isomorphism test, which means GNNs that are not able to predict node clustering coefficients and shortest path distances, and cannot differentiate between different d-regular graphs. Here we develop a class of message passing GNNs, named Identity-aware Graph Neural Networks (ID-GNNs), with greater expressive power than the 1-WL test. ID-GNN offers a minimal but powerful solution to limitations of existing GNNs. ID-GNN extends existing GNN architectures by inductively considering nodes' identities during message passing. To embed a given node, ID-GNN first extracts the ego network centered at the node, then conducts rounds of heterogeneous message passing, where different sets of parameters are applied to the center node than to other surrounding nodes in the ego network. We further propose a simplified but faster version of ID-GNN that injects node identity information as augmented node features. Altogether, both versions of ID-GNN represent general extensions of message passing GNNs, where experiments show that transforming existing GNNs to ID-GNNs yields on average 40% accuracy improvement on challenging node, edge, and graph property prediction tasks; 3% accuracy improvement on node and graph classification benchmarks; and 15% ROC AUC improvement on real-world link prediction tasks. Additionally, ID-GNNs demonstrate improved or comparable performance over other task-specific graph networks.

Introduction

Message-passing GNNs are efficient and powerful but are bounded by the 1-WL test, limiting their ability to distinguish nodes with different neighborhood structures. ID-GNN addresses this limitation inductively through identity-aware heterogeneous message passing and improves performance across challenging and real-world tasks.

  • Background: Message-passing GNNs learn node embeddings by repeatedly aggregating local-neighborhood information through nonlinear transformations.GCN, GraphSAGE, and GAT are highlighted as widely used because of their simplicity, efficiency, and strong real-world performance.
  • Limitation: Existing GNNs are upper-bounded by the 1-WL test, so nodes with different neighborhood structures can receive identical embeddings.This occurs when their computational graphs—the procedures producing their embeddings—are indistinguishable.
  • Method: ID-GNN inductively incorporates node identity by distinguishing the root from other nodes in its local computational graph, unlike transductive one-hot node labeling.The identity coloring is designed to generalize to unseen graphs.
  • Method: ID-GNN embeds each node by extracting its centered ego network and applying heterogeneous message passing with separate parameters for the center and surrounding nodes.The general formulation also applies to tasks involving node or edge features.
  • Results: 40% average accuracy improvement is obtained on challenging node, edge, and graph property prediction tasks after transforming existing GNNs into ID-GNNs.The reported gains also include 3% accuracy improvement on node and graph classification benchmarks and 15% ROC AUC improvement on real-world link prediction tasks.

Related Work

Prior work has extended neural networks beyond 1-WL using additional, often task- or domain-specific components, while this paper incorporates inductive identity information into message-passing GNNs to retain efficiency, simplicity, and broad applicability. Related approaches also use augmented node features to color nodes and improve existing GNNs.

  • Expressive neural networks beyond 1-WL test: Prior neural networks exceeding 1-WL often introduce extra components tailored to specific tasks or domains.The cited examples include Chen et al. (2019), Maron et al. (2019a), Murphy et al. (2019), You, Ying, and Leskovec (2019), and Li et al. (2020).
  • Expressive neural networks beyond 1-WL test: P-GNN embeddings depend on random anchor-sets, limiting their use for node- and graph-level tasks requiring deterministic node embeddings.This limitation is attributed to You, Ying, and Leskovec (2019).
  • Expressive neural networks beyond 1-WL test: ID-GNNs incorporate inductive identity information into message-passing GNNs while maintaining efficiency, simplicity, and broad applicability.The paper presents this as surpassing 1-WL’s expressive power without abandoning the stated benefits of message passing GNNs.
  • Graph Neural Networks with inductive coloring: Related graph models color nodes with augmented features to improve existing GNNs’ performance.The passage identifies this line of work as graph neural networks with inductive coloring and cites Xu et al. (2020) and Veliˇc.

Preliminaries

Message passing GNNs learn node embeddings by iteratively aggregating local neighborhoods, while ID-GNN extends them with inductive identity coloring and heterogeneous message passing. This design increases expressivity beyond existing message passing GNNs and enables cycle counting and distinguishing certain graphs that GIN cannot.

  • Message passing GNNs: Message passing GNNs iteratively aggregate local network neighborhoods to learn node embeddings for node-, edge-, and graph-level prediction tasks.Graphs may include node features and edge features.
  • ID-GNN construction: ID-GNN combines inductive identity coloring with heterogeneous message passing to make any message passing GNN more expressive.It extracts a K-hop ego network, colors its center node, and applies different message functions to colored and uncolored nodes.
  • Expressive power: ID-GNNs are strictly more expressive than existing message passing GNNs, and the ID-GNN version of GIN distinguishes graphs that GIN fails to distinguish.When the two message functions are identical, the construction reduces to the original GNN, while the d-regular graph example establishes added distinguishing power.
  • Cycle counting: ID-GNN embeddings can encode the number of length-j cycles starting and ending at node v in dimension j for j = 1, ..., K.This cycle-counting capability supports computing node clustering coefficients from the learned embedding.
  • ID-GNN-Fast: ID-GNN-Fast injects cycle counts as augmented node features to reduce the ego-network extraction overhead required by heterogeneous message passing.The simplified design targets an important component of ID-GNN’s added expressive power.

Experiments

Experiments across synthetic and real-world node, edge, and graph prediction tasks show that ID-GNNs generally improve over same-budget message-passing GNNs. The gains are strongest for graph-property and shortest-path prediction, while runtime and task-specific comparisons clarify the trade-offs and competitiveness of the two ID-GNN variants.

  • Graph property prediction: 36.8% average absolute gain separates the best ID-GNN from the best GNN on node clustering coefficient prediction.Both ID-GNN formulations achieve near 100% accuracy for at least one architecture on each dataset, with ID-GNN-Fast showing the most consistent improvements.
  • Graph property prediction: 39.9% average improvement gives ID-GNN-Full a substantial advantage over GNNs for pairwise shortest path distance prediction.ID-GNN-Full reaches 100% or near 100% accuracy across architectures and datasets, whereas ID-GNN-Fast and GNNs perform comparatively poorly.
  • Real-world prediction: 9.2% and 20.6% ROC AUC improvements occur for ID-GNNs over GNNs on synthetic and real-world link prediction, respectively.For node classification, improvements are smaller: 1% on Cora and 1.6% on CiteSeer.
  • Efficiency and expressive-network comparisons: 3.8x slower runtime makes ID-GNN-Full costlier than its GNN equivalent for forward and backward passes, while ID-GNN-Fast has essentially zero overhead.ID-GNNs outperform other graph networks on MUTAG and PROTEINS and remain comparable to two of four powerful models on PTC.
  • Overall results: Fewer than 10 of 160 model-task combinations fail to improve accuracy with ID-GNNs over GNNs.The experiments use matched computational budgets, and the authors attribute rare ID-GNN-Full non-improvements partly to underfitting caused by controlled model complexity.

Conclusion

ID-GNNs extend existing GNNs with expressive power beyond the 1-WL test. A feature-augmented variant preserves theoretical guarantees and empirical success while requiring only one-time feature preprocessing when runtime efficiency matters.

  • ID-GNNs are presented as a general extension to existing GNNs with expressive power beyond the 1-WL test.
  • The feature-augmented ID-GNN preserves theoretical guarantees and empirical success of heterogeneous message passing while requiring only one-time feature preprocessing.This variant is intended for settings where runtime efficiency is the primary concern.

Ethics Statement

GNNs are presented as promising models for relational data analysis across diverse application domains. The simplicity of ID-GNNs motivates further research into more expressive message passing, coloring schemes, and feature augmentation.

  • GNNs support relational-data applications including fraud detection, molecular drug structure discovery, recommender systems, and network analysis.
  • ID-GNNs’ simplicity motivates further exploration of GNN expressiveness.
  • The work encourages research into heterogeneous message passing, coloring schemes, and generic feature augmentation.

Proof of Proposition 2

The proof constructs heterogeneous message-passing weights so that each embedding coordinate counts paths to an identity-colored node. Applying this result to the node itself yields counts of cycles rooted at that node.

  • Weight construction: The construction assigns separate trainable weights to identity-colored and uncolored nodes, with layer-specific matrices designed to propagate path-count coordinates.The proof assumes constant initial features and specifies weight assignments across the first and subsequent layers.
  • Inductive lemma: After n heterogeneous message-passing layers, coordinate j of node u’s embedding equals the number of length-j paths from u to identity node v.The claim is established by induction, beginning with the one-layer case and extending each coordinate by one hop.
  • Inductive step: The inductive step sums neighbors’ length-(j−1) path counts, thereby counting length-j paths from u to v after one additional message-passing layer.Each path from a neighboring node is extended by the edge to u, preserving the desired coordinate-wise interpretation.
  • Proposition 2: Choosing v itself as the identity node makes coordinate j count length-j paths from v back to v, equivalently cycles rooted at v.This conclusion directly derives Proposition 2 from Lemma 1 for j = 1, ..., k.

Memory consumption of ID-GNN

ID-GNN-Full has no memory overhead compared with GraphSAGE-style mini-batch GNNs when computational complexity is controlled. Compared with overlap-based memory-efficient mini-batching, ID-GNNs use moderately more memory because that approach reuses embeddings across ego-nets.

  • Memory consumption of ID-GNN: ID-GNN-Full has no memory overhead compared with GraphSAGE-style mini-batch GNNs when computational complexity is controlled.ID-GNN uses disjoint extracted ego-nets in a GraphSAGE-style mini-batch implementation.
  • Memory consumption of ID-GNN: Moderate memory increases arise relative to mini-batch GNNs that exploit overlap among extracted ego-nets and compute shared-node embeddings only once.The overlap-based approach saves embedding computation but increases mini-batch processing time because nodes from different ego-nets must be aligned and deduplicated.
Loading 2101.10320v2…