🤖 machine learning

Invariant-Based Weight Sharing for Message Passing

This paper introduces ShareGNNs, a novel message-passing neural network architecture that enhances expressivity and structural awareness by sharing weights indexed directly by user-chosen graph invariants, thereby outperforming standard MPNNs in both synthetic and real-world tasks.

Original authors: Florian Seiffarth

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Florian Seiffarth

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Teaching Computers to See Patterns, Not Just Neighbors

Imagine you are trying to teach a computer to understand shapes made of dots and lines (which mathematicians call graphs). These could be molecules, social networks, or road maps.

The Old Way (Standard AI):
Think of a standard AI model like a person walking through a crowded room. They can only talk to the people standing immediately next to them. To learn about the whole room, they have to pass a message from neighbor to neighbor, step by step.

The Problem: If two people are far apart but have the exact same "vibe" (e.g., they are both wearing red hats and standing near a window), the old AI treats them as totally different strangers just because they aren't standing next to each other. It misses the big picture.

The New Way (ShareGNNs):
The authors of this paper invented a new way for the AI to learn. Instead of just listening to immediate neighbors, the AI learns to recognize structural patterns.

Imagine the AI has a special "pattern book." If it sees two people who are both wearing red hats and are exactly 5 steps away from a window, it realizes: "Hey, these two situations are identical!"

Because they are identical in structure, the AI uses the same set of instructions (weights) to process both of them. It doesn't matter if they are in different rooms or different buildings; if the pattern is the same, the rule is the same.

The Core Innovation: "Invariant-Based Weight Sharing"

The paper calls this Invariant-Based Weight Sharing. Let's break that down with an analogy:

The "Invariant": This is a property that doesn't change even if you shuffle things around. Imagine a necklace. If you rotate it or flip it, the beads are still in the same order relative to each other. That order is the "invariant."
The "Weight Sharing": In the old AI, every connection between two dots had its own unique, random instruction. In this new AI, the instructions are indexed by the pattern.
- Analogy: Think of a library. In the old system, every book had a unique, random code. In the new system, books are sorted by genre. If you want to read a "Mystery" book, you grab the "Mystery" instruction. If another "Mystery" book appears in a different library, you use the same "Mystery" instruction.

The AI doesn't need to memorize every single connection in every single graph. It just needs to learn the rules for specific patterns (like "two carbon atoms 3 steps apart"). Once it learns that rule, it can apply it to any molecule or network that has that same pattern.

How It Works: The "ShareGNN"

The authors built a model called ShareGNN to put this idea into action.

The Encoder (The Detective):
Instead of just looking at neighbors, the detective looks at any two dots in the graph. It asks: "What is the label of the first dot? What is the label of the second? How far apart are they?"
- If the answer is "Carbon, Carbon, 3 steps," it pulls out the specific "Carbon-to-Carbon-3-steps" rule from its memory bank.
- This allows information to jump across the whole graph in a single step, rather than walking one step at a time.
The Decoder (The Summarizer):
Once the detective has gathered all the clues, the decoder summarizes the whole graph into a single answer (like "This molecule is toxic" or "This social network is a community"). It does this by grouping nodes based on their patterns, ensuring the final answer doesn't change just because the dots were listed in a different order.

Why Is This Better?

The paper claims three main advantages:

It Sees Further: Because it can jump between any two points based on their pattern, it solves problems that require "long-range thinking" much faster. It's like having a teleporter instead of a walking path.
It's Smarter with Less Data: Since it reuses the same rules for similar patterns, it doesn't need to memorize millions of unique connections. It learns the logic of the structure.
It's Explainable: Because the rules are tied to specific patterns (like "distance 5"), we can look at the AI and say, "Ah, it made this decision because it recognized a specific triangle shape." We know why it decided what it decided.

The Results: Did It Work?

The authors tested their new AI on:

Molecules: Predicting chemical properties.
Social Networks: Classifying groups of people.
Synthetic Puzzles: Made-up graphs designed to trick standard AI.

The Outcome:
The ShareGNN consistently beat the standard AI models. In some tricky puzzles where the old AI failed completely (because it couldn't see the long-range patterns), the ShareGNN solved them with near-perfect accuracy. It also performed very well on real-world data, matching or beating the most advanced models currently available, but often using a much simpler, "shallower" structure.

Summary

The paper introduces a new way for computers to learn from graphs. Instead of treating every connection as unique, it groups connections by their structural shape. By sharing the same "brain rules" for identical patterns, the AI becomes better at seeing the big picture, learns faster, and is easier to understand. It's like teaching a child to recognize a "face" by the arrangement of eyes and nose, rather than memorizing the face of every single person they meet.

Technical Summary: Invariant-Based Weight Sharing for Message Passing

Problem Statement
Message-passing neural networks (MPNNs) are a dominant framework for learning representations of graph-structured data. However, standard MPNNs and graph transformers typically rely on weight sharing based solely on node or edge features, or they learn instance-specific attention coefficients. This approach limits the model's ability to capture inherent structural patterns because weights are not explicitly tied to the graph's topology. Furthermore, standard MPNNs are often restricted to local neighborhood aggregation, making them susceptible to over-smoothing and over-squashing in deep architectures, while graph transformers often lack systematic parameter reuse across structurally equivalent subgraphs.

Methodology: ShareGNNs
The authors introduce ShareGNNs, a novel architecture grounded in invariant-based weight sharing. The core principle is to index learnable parameters directly by graph invariants—functions preserved under node permutations—rather than by fixed feature dimensions or instance-specific positions.

Invariant-Based Weight Sharing:
- Encoder (Message Passing): Instead of fixed weight matrices, the model dynamically constructs message-passing weights for every pair of nodes $(v, w)$ based on an invariant triple: $(l(v), l(w), d(v, w))$ , where $l$ is a labeling function (e.g., atomic number, Weisfeiler-Leman label, pattern count) and $d$ is the shortest-path distance. All node pairs sharing the same invariant triple share a single learnable scalar weight $\omega_\tau$ . This allows information to flow between arbitrary nodes in a single layer, not just immediate neighbors.
- Decoder (Pooling): Similarly, bias terms and pooling weights for graph-level aggregation are indexed by the invariant labels of individual nodes.
- Dynamic Construction: For any input graph, the relevant weight matrices are assembled on-the-fly from a shared parameter pool. This ensures the model is permutation equivariant in the encoder and permutation invariant in the decoder by construction.
Labeling Functions:
The expressivity of the model is controlled by the choice of labeling functions $l$ . The paper explores:
- Molecular Labeling: Atomic numbers.
- Weisfeiler-Leman (WL) Labeling: Iterative refinement of node labels based on neighbor multisets.
- Pattern Labeling: Counts of specific subgraph patterns (e.g., triangles, cycles) involving a node.
  These can be combined to capture multi-level structural information.
Architecture:
ShareGNNs utilize a simple encoder-decoder structure. The encoder propagates information across arbitrary distances in a single layer, and the decoder aggregates node embeddings into a graph representation using invariant-indexed pooling. The architecture supports multi-head extensions where different heads utilize different invariant sets.

Key Contributions

Invariant-Based Weight Sharing Principle: A new mechanism for weight sharing that explicitly incorporates graph structure, enabling systematic parameter reuse across structurally equivalent node pairs within and across different graphs.
ShareGNN Architecture: A practical instantiation of this principle that combines message passing with transformer-like connectivity (all-to-all interactions) but with parameters tied to structural invariants rather than learned attention scores.
Theoretical and Empirical Validation:
- Expressivity: The authors prove that ShareGNNs are at least as expressive as the chosen labeling function. If two graphs are distinguishable by the labeling function (e.g., WL labels), the ShareGNN can distinguish them.
- Control: The model offers explicit control over complexity and expressivity through the selection of invariants.
- Performance: Experiments demonstrate that ShareGNNs outperform standard MPNNs (e.g., GIN, GAT) and are competitive with graph transformers on various benchmarks, including molecular classification, social network analysis, and substructure counting.

Results

Graph Classification: On diverse datasets (NCI1, NCI109, MUTAG, IMDB-B/M), ShareGNNs achieve state-of-the-art or highly competitive accuracy. Notably, they achieve near-perfect accuracy (100%) on synthetic benchmarks (RingTransfer, Snowflakes) designed to challenge standard MPNNs, demonstrating their ability to capture long-range dependencies and distinguish graphs indistinguishable by the 1-WL test.
Substructure Counting: On the Substructure Counting Benchmark, ShareGNNs achieve state-of-the-art results in single-task settings and remain competitive in multi-task settings, validating that expressive invariants lead to expressive models even with shallow architectures (single encoder layer).
Scalability: The model scales to large datasets like ZINC (12k and 250k graphs) while maintaining a shallow architecture, avoiding the over-smoothing issues common in deep MPNNs.
Efficiency: Ablation studies show that increasing the maximum message-passing distance $D$ has a more significant impact on performance than adding more encoder layers. Furthermore, pruning rare invariant triples significantly reduces the parameter count without sacrificing accuracy.

Significance and Claims
The paper claims that invariant-based weight sharing provides a principled alternative to feature-based weight sharing and instance-specific attention. By tying weights to structural invariants, ShareGNNs offer:

Improved Interpretability: Learned weights can be directly interpreted in terms of structural patterns (e.g., specific distances between specific atom types).
Transferability: Parameters are naturally transferable across graphs with similar structural patterns, suggesting potential for graph foundation models.
Expressivity Control: The model allows researchers to explicitly control the discriminative power of the network by selecting appropriate invariants, bridging the gap between theoretical expressivity (e.g., k-WL tests) and practical implementation.

The authors acknowledge limitations, including the preprocessing cost of computing invariants and pairwise distances, and the quadratic memory complexity in the worst case, though they note this is comparable to graph transformers and can be mitigated by pruning. They also note that the current approach requires user-defined invariants, leaving room for future work in automated invariant selection.