Source-linked AI summary
Provably Powerful Graph Networks
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, Yaron Lipman
TL;DR
Message passing GNNs are limited to the distinguishing power of 1-WL, motivating more expressive yet scalable graph models. The paper proves k-order networks match k-WL and proposes a 2-order MLP-plus-matrix-multiplication model with 3-WL expressiveness. Experiments report performance on par with or better than state of the art across graph classification and regression tasks.
Problem
Message passing GNNs cannot distinguish graphs indistinguishable by 1-WL, motivating more expressive graph learning models that retain scalability.
Method
The paper analyzes k-order invariant GNNs and proposes a reduced 2-order network interleaving feature-wise MLPs with matrix multiplication.
Results
The proposed model achieves 3-WL discrimination power and performs on par with or better than state-of-the-art methods on graph classification and regression tasks.
Takeaways & Limitations
The model offers a practical combination of guaranteed 3-WL expressiveness and k = 2 tensor scalability, while operating with quadratic complexity.
Abstract
from arXiv · showhide
Recently, the Weisfeiler-Lehman (WL) graph isomorphism test was used to measure the expressive power of graph neural networks (GNN). It was shown that the popular message passing GNN cannot distinguish between graphs that are indistinguishable by the 1-WL test (Morris et al. 2018; Xu et al. 2019). Unfortunately, many simple instances of graphs are indistinguishable by the 1-WL test. In search for more expressive graph learning models we build upon the recent k-order invariant and equivariant graph neural networks (Maron et al. 2019a,b) and present two results: First, we show that such k-order networks can distinguish between non-isomorphic graphs as good as the k-WL tests, which are provably stronger than the 1-WL test for k>2. This makes these models strictly stronger than message passing models. Unfortunately, the higher expressiveness of these models comes with a computational cost of processing high order tensors. Second, setting our goal at building a provably stronger, simple and scalable model we show that a reduced 2-order network containing just scaled identity operator, augmented with a single quadratic operation (matrix multiplication) has a provable 3-WL expressive power. Differently put, we suggest a simple model that interleaves applications of standard Multilayer-Perceptron (MLP) applied to the feature dimension and matrix multiplication. We validate this model by presenting state of the art results on popular graph classification and regression tasks. To the best of our knowledge, this is the first practical invariant/equivariant model with guaranteed 3-WL expressiveness, strictly stronger than message passing models.
1 Introduction
Message passing GNNs are limited by 1-WL expressiveness, motivating models that improve graph distinction while retaining scalability. The paper proves k-order networks match k-WL and introduces a scalable 3-WL-expressive model, with competitive empirical results.
- Motivation: Message passing models cannot distinguish graphs that 1-WL cannot distinguish.Figure 1 provides an example of such graphs.
- Contributions: The paper proves k-order invariant GNNs are as expressive as k-WL tests for k ≥2.This establishes a hierarchy of expressive graph networks.
- Contributions: A reduced k = 2 model combines feature-wise MLPs with matrix multiplication and achieves 3-WL expressiveness.The model uses tensors with the same order as the graph input.
- Contributions: The proposed model is provably more powerful than message passing models because 3-WL is strictly stronger than 1-WL.The paper notes that 1-WL and 2-WL are equivalent, whereas 3-WL is stronger.
- Experiments: Experiments on graph classification and regression datasets found the model on par with or better than state-of-the-art methods.The datasets include social networks, molecules, and chemical compounds.
2 Previous work
Previous work developed deep graph models, WL-based graph methods, and higher-order approaches, but higher-order expressiveness can remain computationally expensive or require relaxed guarantees.
- Deep learning on graph data: Graph learning research has generalized neural convolutions and recurrent architectures to graph-structured data.These approaches include spectral graph convolutions and recurrent node-representation models.
- Weisfeiler-Lehman tests: WL tests form a hierarchy used in graph kernels, GNN analysis, link prediction, and higher-order graph features.Related work also studies sparse higher-order WL representations.
3 Preliminaries
The preliminaries formalize permutation actions, invariant and equivariant k-order networks, and the WL hierarchy, including tuple-color refinement and its discrimination-power relations.
- Notation: A k-tuple is an ordered collection of k vertex indices, and permutations act entrywise on tuple indices and tensor entries.The tensor action reindexes tuple positions according to the inverse permutation.
- 3.1 k-order graph networks: A k-order invariant graph network composes equivariant linear layers, entrywise nonlinearities, an invariant layer, and an MLP.Invariance is achieved by propagating permutation actions through equivariant layers.
- 3.1 k-order graph networks: For k = 2, these networks can approximate any message passing neural network, while universality at k = poly(n) has limited practical value.The higher-order universal construction uses very high-order tensors.
- 3.2 The Weisfeiler-Lehman graph isomorphism test: The WL family tests graph isomorphism by iteratively coloring k-tuples and comparing the resulting color histograms.Initial colors encode tuple isomorphism types, including equality and adjacency relations.
- 3.2 The Weisfeiler-Lehman graph isomorphism test: WL refinement repeatedly aggregates colors from tuple neighborhoods until the coloring stabilizes.The WL and FWL variants differ in how neighboring k-tuples are aggregated.
4 Colors and multisets in networks
The section develops vector-based color and multiset representations for networks, using Power-sum Multi-symmetric Polynomials (PMP) to encode multisets invariantly and uniquely.
- Implementation challenges: The section frames three implementation challenges: representing colors, representing multisets, and implementing the encoding function.Color-pair encoding can be implemented by concatenating the corresponding vector representations.
- Color representation: Network colors are represented as vectors, replacing discrete WL color tensors with tensors in R^(n^k×a).Each k-tuple receives a vector color in R^a.
- Multiset representation: A tuple of vector colors forms a matrix whose row order must not affect its multiset representation.Permuting the rows corresponds to the same multiset and must produce the same color.
- Multiset representation: The method encodes multisets with Power-sum Multi-symmetric Polynomials, which are invariant under permutations of the rows.PMP generalize power-sum symmetric polynomials to vector-valued multisets.
- Multiset representation: PMP with total degree at most n generate the ring of multisymmetric polynomials, allowing every such invariant polynomial to be expressed through u(X).For any multisymmetric polynomial q, there is a polynomial r such that q(X)=r(u(X)).
- Multiset representation: The vector u(X) uniquely represents a multiset: u(X)=u(X′) exactly when X′ is a row permutation of X.This gives the network an injective, permutation-invariant multiset encoding.
5 k-order graph networks are as powerful as k-WL
The section shows that k-order graph networks can match the distinguishing power of k-WL by encoding the neighborhood multisets used in WL updates.
- Expressiveness result: For every 2 ≤ k ≤ n, the construction aims to show that k-order graph networks are at least as powerful as the k-WL test.The proof constructs a k-order network and learnable weights that implement k-WL.
- Network construction: The WL update is equivariant, motivating its implementation with equivariant linear layers, pointwise nonlinearities, and feature-wise MLP approximation.The polynomial function can be replaced by an approximating MLP using universal approximation.
- Expressiveness result: Theorem 1 states that k-WL-distinguishable graphs can be separated by some k-order network, while isomorphic graphs receive the same output.Thus the network preserves isomorphism invariance and achieves the stated k-WL distinguishing guarantee.
- Network construction: Graphs are encoded as tensors whose final channel stores adjacency and whose diagonal feature channels store vertex colors.The representation has shape B ∈ R^(n^2×(e+1)).
- Network construction: The key construction encodes the neighborhood multiset of each k-tuple, the only nontrivial component of the WL update rule.The relevant colors are arranged as rows of a matrix extracted from the input tensor.
- Network construction: The network computes the multiset encoding by applying a polynomial map entrywise and then summing over one coordinate with a linear equivariant operator.This produces the desired encoded neighborhood information at each k-tuple.
6 A simple network with 3-WL discrimination power
The proposed second-order GNN combines feature-wise MLPs with matrix multiplication and achieves 3-WL discrimination power while using only second-order tensors.
- Architecture: The model applies feature-wise MLPs, multiplies matching feature matrices, and outputs both transformed features and their matrix products.For input X, three MLPs produce m1(X), m2(X), and m3(X); matrix multiplication forms W, and the block outputs (m3(X), W).
- Architecture: Matrix multiplication is equivariant under simultaneous node permutations, making each block equivariant and the complete network invariant.The proof uses (g · A) · (g · B) = g · (A · B).
- Expressiveness: Two matrix-multiplication blocks can compute A^3 and use tr(A^3) to count length-3 cycles, distinguishing the Figure 1 graphs with 0 versus 12 such cycles.The adjacency matrix is stored as an input channel, and the invariant layer takes the trace.
- Expressiveness: The model can implement 2-FWL updates, which establishes 3-WL expressive power for graph discrimination.The construction represents the required multiset of paired colors using second-order tensors and matrix multiplication.
- Extensions: A generalized matrix multiplication is proposed as a possible extension for making k-order networks as expressive as k-FWL and hence (k + 1)-WL.This extension is presented as a possible direction rather than as the main model evaluated here.
- Complexity: The model uses O(n^2) space and O(n^3) time per block for dense edge data and constant feature depth.The cited comparison gives Morris et al. (2018) O(n^3) space and O(n^4) time, while noting possible improvements for both methods.
7 Experiments
Experiments evaluate the model on graph classification, graph regression, and an equivariant-layer comparison, with results competitive with existing methods.
- Datasets: The experiments cover eight real-world graph classification datasets spanning social networks, chemical compounds, and protein structures.Graphs use adjacency matrices and, for bioinformatics datasets, possibly categorical node features.
- Graph classification: The classification protocol uses 10-fold cross-validation and reports the best averaged accuracy across folds.Hyperparameters are selected using a fixed random 90%-10% split.
- Graph classification: The model ranked among the top three methods on 6 of 8 classification datasets.The reported ranking compares the method with all previous methods in Table 1.
- Graph regression: Graph regression uses an 80%-10%-10% train-validation-test split and evaluates both joint and quantity-specific prediction networks.The experiment uses an absolute error loss and selects test results according to validation error.
- Equivariant layer evaluation: On NCI1, the matrix-product model and matrix-product-plus-linear-basis model are comparable, while the matrix-product model is considerably more efficient.All methods except the MLP achieved zero training error, but the matrix-product variants had better generalization than the linear basis model.
8 Conclusions
The paper presents a theoretically stronger hierarchy of k-order networks and a simpler model with 3-WL expressiveness, while identifying quadratic complexity as a remaining drawback.
- Contributions: k-order invariant networks parallel the distinction power of k-WL tests but have limited practical interest because they use high-dimensional tensors.The conclusion presents this hierarchy as the first of two explored models.
- Contributions: The proposed simple GNN combines MLPs with matrix multiplication and achieves 3-WL expressiveness using input tensors of size n^2.The model is described as useful for problems with dense edge data.
- Limitations and future work: The model’s complexity remains quadratic, which is worse than message passing methods.The conclusion identifies more efficient expressive GNNs and generalization analysis as future work.
A Proof of Proposition 1
The proof establishes that the power-sum representation separates permutation orbits by constructing an invariant polynomial contradiction.
- Orbit separation: If X′ is a permutation of X, all power-sum features agree and therefore u(X) = u(X′).The proof first establishes invariance under the permutation action.
- Orbit separation: For distinct permutation orbits, Stone–Weierstrass supplies a polynomial that is at least 1 on one orbit and at most 0 on the other.The construction uses a compact set containing both disjoint orbits.
- Contradiction: The constructed invariant polynomial factors through u, contradicting the assumption that u(X) = u(X′) for distinct orbits.The contradiction completes the separation argument.
B Proof of equivairance of WL update step
The k-WL update is shown to preserve equivariance under vertex permutations by transforming the input coloring consistently and recovering the correspondingly permuted output coloring.
- The proof represents k-WL neighborhood information using a formal tensor Bj whose entries are multisets.
- The k-WL update rule can be expressed through this multiset tensor representation.
- Each entry of the update tuple is checked for equivariance under the permutation action.
- Feeding the k-WL update rule with the permuted coloring g · Cl−1 produces the permuted output Cl.
C Proof of Theorem 1
Theorem 1 constructs a k-order equivariant network that implements k-WL initialization and updates, then uses a color histogram to distinguish every graph pair that k-WL distinguishes in a finite collection.
- The construction targets a finite graph set and assigns different outputs to non-isomorphic graphs distinguishable by k-WL.
- The network implements k-WL through initialization, coloring updates, and a final histogram of tuple colors.
- Input: The input tensor encodes adjacency in its last channel and vertex colors on its diagonal feature channels.
- Initialization: A linear equivariant operator extracts each k-tuple’s subtensor and equality pattern, thereby encoding its tuple isomorphism type.
- Initialization: The operator is equivariant because applying a vertex permutation before or after extraction yields the correspondingly permuted subtensor.
- Coloring update: Each k-WL update applies an entrywise polynomial feature map, an equivariant linear operator, and concatenation of the current and aggregated colors.
- Coloring update: On the finite set of possible colors, MLPs can approximate the polynomial map closely enough to preserve injective multiset representations.
- Graph-level output: After sufficiently many update steps, an entrywise one-hot MLP and summation produce a graph-level histogram of k-tuple colors.
D Proof of Theorem 2
Theorem 2 implements the 2-FWL update with a reduced second-order network: initialization encodes pair types, while two feature maps and matrix multiplication construct an injective multiset representation.
- The construction considers finite input tensors and builds a network distinguishing graphs that are 3-WL distinguishable.
- Input: The input tensor includes vertex features, adjacency information, and an identity-matrix channel.
- Initialization: Initialization encodes each ordered vertex pair using endpoint colors, adjacency or complement adjacency, and an equality indicator.
- 2-FWL update step: The 2-FWL update requires constructing a tensor whose coloring represents the multiset neighborhood information for each pair.
- 2-FWL update step: A matrix X ∈ Rn×2a is formed to represent the relevant pairwise neighborhood features, and the output tensor applies an injective multiset map u(X).
- 2-FWL update step: Two polynomial maps transform the input features before their products are used to construct the multiset representation.
- 2-FWL update step: MLPs replace the polynomial maps while preserving the required injective property on the finite set of possible colors.
- 2-FWL update step: Taking the third map as the identity and concatenating the original tensor with the matrix-product features completes the 2-FWL update implementation.