Invariant-Based Weight Sharing for Message Passing

2026-05-25Machine Learning

Machine Learning
AI summary

The authors present a new way to improve message-passing neural networks (MPNNs) by making their weights depend on the graph's structure, not just node features. They do this by indexing weights using special graph properties that stay the same even if the graph nodes are rearranged, allowing the model to recognize and reuse patterns in similar parts of the graph. Their new model, ShareGNN, uses this idea in a simple encoder-decoder setup with learnable connections that act like transformers. Tests show that ShareGNN better captures graph structures, works well on various tasks including counting subgraphs, and scales nicely to big datasets.

Message-Passing Neural NetworksGraph InvariantsWeight SharingGraph Neural NetworksEncoder-Decoder ArchitectureTransformerSubgraph Counting1-WL TestGraph Representations
Authors
Florian Seiffarth
Abstract
Message-passing neural networks (MPNNs) are a powerful framework for learning representations of graph-structured domains. However, weights in MPNNs act on features only, limiting their ability to capture structural patterns. We introduce a novel structure-aware weight sharing principle that explicitly incorporates information inherent to the graph structure. Weights are indexed directly by user-chosen graph invariants, i.e., functions preserved under node permutations, enabling systematic reuse across structurally equivalent subgraphs. We present ShareGNNs, which instantiate this principle within a simple encoder-decoder architecture, resulting in an MPNN with learnable adjacency and transformer-like connectivity. We show that their expressivity is at least as strong as the discriminative power of the chosen invariants, providing explicit control over the model complexity. Experiments on synthetic and real-world data, as well as subgraph counting tasks, demonstrate consistent improvements over standard MPNNs, competitive expressivity beyond the 1-WL test, and scalability to large datasets.