Building powerful and equivariant graph neural networks with structural message-passing

Introduction

Graph neural networks have recently emerged as a popular way to process and analyze graph-structured data. Among the numerous architectures that have been proposed, the class of message-passing neural networks (MPNNs) has been by far the most widely adopted. In addition to being able to efficiently exploit the sparsity of graphs, MPNNs exhibit an inherent tendency to learn relationships between nearby nodes. This inductive bias is generally considered as a good fit for problems that require relational reasoning , such as tractable relational inference , problems in combinatorial optimization or the simulation of physical interactions between objects .

A second key factor to the success of MPNNs is their equivariance properties. Since neural networks can ultimately only process tensors, in order to use a graph as input, it is necessary to order its nodes and build an adjacency list or matrix. Non-equivariant networks tend to exhibit poor sample efficiency as they need to explicitly learn that all representations of a graph in the (enormous) symmetry group of possible orderings actually correspond to the same object. On the contrary, permutation equivariant networks, such as MPNNs, are better equipped to generalize as they already implement the prior knowledge that any ordering is arbitrary.

Despite their success, equivariant MPNNs possess limited expressive power . For example, they cannot learn whether a graph is connected, what is the local clustering coefficient of a node, or if a given pattern such as a cycle is present in a graph . For tasks where the graph structure is important, such as the prediction of chemical properties of molecules and the solution to combinatorial optimization problems, more powerful graph neural networks are necessary.

Aiming to address these challenges, this work puts forth structural message-passing (SMP)—a new type of graph neural network that is strictly more powerful than MPNNs, while also sharing the attractive inductive bias of message-passing architectures. SMP inherits its power from its ability to manipulate node identifiers. However, in contrast to previous studies that relied on identifiers , it does so in a permutation equivariant way without introducing new sources of randomness. As a result, SMP can be powerful without sacrificing its ability to generalize to unseen data. In particular, we show that if SMP is built out of powerful layers, the resulting model is computationally universal over the space of equivariant functions.

Concretely, SMP maintains at each node a matrix called “local context” (instead of a feature vector as in MPNNs) that is initialized with a one-hot encoding of the nodes and the node features. These local contexts are then propagated in such a way that a permutation of the nodes or a change in the one-hot encoding will reorder the lines of each context without changing their content, which is key to efficient learning and good generalization.

We evaluate SMP on a diverse set of structural tasks that are known to be difficult for message-passing architectures, such as cycle detection, connectivity testing, diameter and shortest path distance computation. In all cases, our approach compares favorably to previous methods: for example, SMP solves cycle detection in all evaluated configurations, whereas other powerful networks struggle when the graphs become larger, and MPNNs do not manage to solve the task completely.

Finally, we evaluate our method on the ZINC chemistry dataset and achieve state-of-the-art performance among methods that do not use expert features. It shows that SMP is able to successfully learn both from the features and from topological information, which is essential in chemistry applications. Overall, our method is able to overcome a major limitation of MPNNs, while retaining their ability to process features with a bias towards locality.

Related work

Originally introduced by Scarselli et al. , MPPNs have progressively been extended to handle edge and graph-level attributes . Despite the flexibility in their parametrization, MPNNs without special node attributes however all have limited expressive power, even in the limit of infinite depth and width. For instance, they are at most as good at isomorphism testing as the Weisfeiler-Lehman (WL) vertex refinment algorithm . The WL test has higher dimensional counterparts ($k$-WL) of increasing power, which has motivated the introduction of the more powerful $k$-WL networks . However, these higher-order networks are global, in the sense that they iteratively update the state of a $k$-tuple of nodes based on all other nodes (and not only neighbours), a procedure which is very costly both in time and memory. While a faster procedure was proposed in concurrently to our work, key differences with SMP remain: we propose to learn richer embeddings for each node instead of one embedding per k-tuple of nodes, and build our theoretical analysis on distributed algorithms rather than vertex refinement methods.

Recent studies have also characterized the expressive power of MPNNs from other perspectives, such as the ability to approximate continuous functions on graphs and solutions to combinatorial problems , highlighting similar limitations of MPNNs — see also .

Beyond higher-order message-passing architectures, there have been efforts to construct more powerful equivariant networks. One way to do so is to incorporate hand-crafted topological features (such as the presence of cliques or cycles) , which requires expert knowledge on what features are relevant for a given task. A more task-agnostic alternative is to build networks by arranging together a set of simple permutation equivariant functions and operators. These building blocks are:

Linear equivariant functions between tensors of arbitrary orders: a basis for these functions was computed by Maron et al. , by solving the linear system imposed by equivariance.

Element-wise functions, applied independently to each feature of a tensor.

Operators that preserve equivariance, such as $+$, $-$, tensor and elementwise products, composition and concatenation along the dimension of the channels.

Similarly to Morris et al. , networks built this way obtain a better expressive power than MPNN by using higher-order tensors . Since $k$-th order tensors can represent any $k$-tuple of nodes, architectures manipulating them can exploit more information to compute topological properties (and be as powerful as the $k$-WL test). Unfortunately, memory requirements are exponential in the tensor order, which makes these methods of little practical interest. More recently, Maron et al. proposed provably powerful graph networks (PPGN) based on the observation that the use of matrix multiplication can make their model more expressive for the same tensor order. This principle was also used in the design of Ring-GNN , which has many similarities with PPGN. Key differences between such methods and ours are that (i) SMP can be parametrized to have a lower time complexity, due to the ability of message-passing to exploit the sparsity of adjacency matrices, (ii) SMP retains the message-passing inductive bias, which is different from PPGN and, as we will show empirically, makes it better suited to practical tasks such as the detection of substructures in a graph.

2 Non-equivariant graph neural networks

In order to better understand the limitations of current graph neural networks, analogies with graph theory and distributed systems have been exploited. In these fields, a large class of problems cannot be solved without using node identifiers . The reasoning is that, in message-passing architectures, each node has access to a local view of the graph created by the reception of messages. Without identifiers, each node can count the number of incoming messages and process them, but cannot tell from how many unique nodes they come from. They are therefore unable to reconstruct the graph structure.

This observation has motivated researchers to provide nodes with randomly selected identifiers . Encouragingly, by showing the equivalence between message-passing and a model in distributed algorithms, Loukas proved that graph neural networks with identifiers and sufficiently expressive message and update functions can be Turing universal, which was also confirmed on small instances of the graph isomorphism problem .

Nevertheless, the main issue with these approaches is sample efficiency. Identifiers introduce a dependency of the network on a random input and the loss of permutation equivariance, causing poor generalization. Although empirical evidence has been presented that the aforementioned dependency can be overcome with large amounts of training data or other augmentations , overfitting and optimization issues can occur. In this work, we propose to overcome this problem by introducing a network which is both powerful and permutation equivariant.

Structural message-passing

We present the structural message-passing neural networks (SMP), as generalization of MPNNs that follows a similar design principle. However, rather than processing vectors with permutation invariant operators, SMP propagates matrices and processes them in a permutation equivariant way. This subtle change greatly improves the network’s ability of to learn information about the graph structure.

Layers

At layer $l+1$, the state of each node is updated as in standard MPNNs : messages are computed on each edge before being aggregated into a single matrix via a symmetric function. The result can then be updated using the local context of previous layer at this node:

Above, $u^{(l)}$, $m^{(l)}$, $\phi$ are the update, message and aggregation functions of the $(l+1)$-th layer, respectively, whereas $c_{l+1}$ denotes the layer’s width.

It might be interesting to observe that, starting from a one-hot encoding and using the update rule ${\bm{U}}_{i}^{(l+1)}=\sum_{v_{j}\in N_{i}}{\bm{U}}_{j}^{(l)}$, SMP iteratively compute powers of ${\bm{A}}$. Since ${\bm{A}}^{l}[i,j]$ corresponds to the count of walks of length $l$ between $v_{i}$ and $v_{j}$, there is a natural connection between the propagation of identifiers and the detection of topological features: even with simple parametrizations, SMP can manipulate polynomials in the adjacency matrix and therefore learn spectral properties that MPNNs cannot detect.

In the following, it will be convenient to express each SMP layer $f^{(l)}$ in a tensor form:

Pooling

whereas a permutation invariant representation is obtained after the application of a pooling function pool. It may be a simple sum or average followed by a soft-max, or a more complex operator :

2 Analysis

The following section characterizes the equivariance properties and representation power of SMP. For the sake of clarity, we defer all proofs to the appendix.

Before providing sufficient conditions for permutation equivariance, we define it formally. A change in the ordering of $n$ nodes can be described by a permutation $\pi$ of the symmetric group $\mathfrak{S}_{n}$. $\pi$ acts on a tensor by permuting the axes indexing nodes (but not the other axes):

We stress that an equivariant SMP network should yield the same output (up to a permutation) for every one-hot encoding used to construct the node identifiers. We can now state some sufficient conditions for equivariance:

Let functions $m$, $\phi$ and $u$ be permutation equivariant, that is, for every permutation $\pi\in\mathfrak{S}_{n}$ we have $u({\pi}.{\bm{U}},\pi.{\bm{U}}^{\prime})={\pi}.u({\bm{U}},{\bm{U}}^{\prime})$, $\phi(\{\pi.{\bm{U}}_{j}\}_{v_{j}\in N_{i}})=\pi.\phi(\{{\bm{U}}_{j}\}_{v_{j}\in N_{i}})$, and $m(\pi.{\bm{U}},\pi.{\bm{U}}^{\prime},{\bm{y}})={\pi}.m({\bm{U}},{\bm{U}}^{\prime},{\bm{y}})$. Then, SMP is permutation equivariant.

The proof is presented in Appendix A. This theorem defines the class of functions that can be used in our model. For example, if the message and update functions are operators applied simultaneously to each row of the local context, the whole layer is guaranteed to be equivariant. However, more general functions can be used: each ${\bm{U}}_{i}$ is a $n\times c$ matrix which can be viewed as the representation of a set of nodes. Hence, any equivariant neural network for sets can be used, which allows the network to have several desirable properties:

Inductivity: as an equivariant neural network for sets can take sets of different size as input, SMP can be trained on graphs with various sizes as well. Furthermore, it can be used in inductive settings on graphs whose size has not been seen during training, which we will confirm experimentally.

Invariance to local isomorphisms: SMP learns structural embeddings, in the sense that it yields the same result on isomorphic subgraphs. More precisely, if the subgraphs $G_{i}^{k}$ and $G_{j}^{k}$ induced by $G$ on the k-hop neighborhoods of $v_{i}$ and $v_{j}$ are isomorphic, then on node classification, any $k$-layer SMP $f$ will yield the same result for $v_{i}$ and $v_{j}$. This is in contrast with several popular methods that learn positional embeddings which do not have this property.

Representation and expressive power

The following theorem characterizes the representation power of SMP when parametrized with powerful layers. Simply put, Theorem 2 asserts that it is possible to parameterize an SMP network such that it maps non-isomorphic graphs to different representations:

Consider the class $S$ of simple graphs with $n$ nodes, diameter at most $\Delta$ and degree at most $d_{\text{max}}$. We assume that these graphs have respectively $c_{X}$ and $c_{Y}$ attributes on the nodes and the edges. Then, there exists a SMP network $f$ of depth at most $\Delta+1$ and width at most $2d_{\text{max}}+c_{X}+n~{}c_{Y}$ such that the full structure of any graph in $S$ (with the attributes) can be recovered from the output of $f$ at any node.

The formal statement and the proof are detailed in Appendix B. We first show the result for the simple case where $f$ can pass messages of size $n\times n$, and then consider the case of $n\times 2d_{\text{max}}$ matrices using the following lemma:

The universality of SMP is a direct corollary: since each node can have the ability to reconstruct the adjacency matrix from its local context, it can also employ a universal network for sets to compute any equivariant function on the graph (cf. Appendix C). Interestingly, this result shows that propagating matrices instead of vectors might be a way to solve the bottleneck problem : while MPNNs need feature maps that exponentially grow with the graph size in order to recover the topology, SMPs can do it with $O(d_{\text{max}}n^{2})$ memory.

Let $G$ be a simple graph of diameter at most $\Delta$ and degree at most $d_{\text{max}}$. Consider an SMP $f=f^{(L)}\circ\cdots\circ f^{(1)}$ of depth $L=\Delta$ and width $2d_{\text{max}}+c_{X}+n~{}c_{Y}$ satisfying the properties of Theorem 3. Then, any equivariant function can be computed as $f_{\textit{eq}}=\sigma\circ f$, where $\sigma$ is a universal function of sets applied simultaneously to each node. Similarly, any permutation invariant function can be computed as $f_{\textit{in}}=\frac{1}{n}\sum_{v_{i}\in V}\sigma\circ f$.

These results show that two components are required to build a universal approximator of functions on graphs. First, one needs an algorithm that breaks symmetry during message passing, which SMP manages to do in an equivariant manner. Second, one needs powerful layers to parametrize the message, aggregation and update functions. Here, we note that the proofs of Theorem 2 and Corollary 1 are not constructive and that deriving practical parametrizations that are universal remains an open question . Nevertheless, we do constructively prove the following more straightforward claim using a simple parametrization:

SMP is strictly more powerful than MPNN: SMP can simulate any MPNN with the same number of layers, but MPNNs cannot simulate all SMPs.

To prove it, we create for any MPNN a corresponding SMP which performs the same operations as the MPNN on the main diagonal of the local context. On the contrary, we can easily create an SMP which is able to distinguish between two small graphs that cannot be distinguished by the Weisfeiler-Lehman test (Appendix D).

Implementation

SMP offers a lot of flexibility in its implementation, as any equivariant function that combines the local context of two nodes and the edge features can be used. We propose two implementations that we found to work well, but our framework can also be implemented differently. In both cases, we split the computation of the messages in two steps. First, the local context of each node is updated using a neural network for sets. Then, a standard message passing network is applied separately on each row of the local contexts. For the first step, we use a subset of the linear equivariant functions computed by Maron et al. :

This architecture corresponds to a standard MPNN, where the message and update functions are two-layer perceptrons. We use a sum aggregator normalized by the average degree $d_{\text{avg}}$ over the graph: it retains the good properties of the sum aggregator , while also avoiding the exploding-norm problem . This network can be written:

Fast SMP

For graphs without edge features, we propose a second implementation with a message function that uses a pointwise multiplication $\odot$:

where ${\bm{W}}_{4}$ and ${\bm{W}}_{5}$ are learnable matrices. The aggregation is the same, and the update is simply a residual connection, so that the $l$-th SMP layer updates each node’s local context as

In this last equation, the arguments of the two sums are only functions of the local context of node $v_{j}$. This allows for a more efficient implementation, where one message is computed per node, instead of one per edge as in default SMP.

One might notice that Fast SMP can be seen as a local version of PPGN (proof in Appendix F):

A Fast SMP with $k$ layers can be approximated by a $2k$-block PPGN.

Despite not being more powerful, Fast SMP has the advantage of being more efficient than PPGN, as it can exploit the sparsity of adjacency matrices. Furthermore, as we will see experimentally, our method manages to learn topological information much more easily than PPGN, a property that we attribute to the inductive bias carried by message-passing.

Complexity

Table 1 compares the per-layer space and time complexity induced by the forward pass of SMP with that of other standard graph networks. Whereas local order-$3$ Weisfeiler-Lehman networks need to store all triplets of nodes, both PPGN and SMP only store information for pairs of nodes. However, message-passing architectures (such as SMP) can leverage the sparsity of the adjacency matrix and hence benefit from a more favorable time complexity than architectures which perform global updates (as PPGN).

An apparent drawback of SMP (shared by all equivariant powerful architectures we are aware of) is the need for more memory than MPNN. This difference is partially misleading since it is known that the width of any MPNN needs to grow at least linearly with $n$ (for any constant depth) for it to be able to solve many graph-theoretic problems . However, for graphs with a large diameter, the memory requirements of SMP can be relaxed by using the following observation: if each node is colored differently from all nodes in its $2k$-hop neighborhood, then no node will see the same color twice in its $k$-hop neighborhood. It implies that nodes which are far apart can use the same identifier without conflict. We propose in Appendix E a procedure (Fast SMP with coloring) based on greedy coloring which can replace the initial one-hot encoding, so that each node can manipulate smaller matrices ${\bm{U}}_{i}$. This method allows to theoretically improve both the time and space complexity of SMP, although the number of colors needed usually grows fast with the number of layers in the network.

Experiments

We first evaluate different architectures on the detection of cycles of length 4, 6 and 8 (for several graph sizes), implemented as a graph classification problemOur implentation with Pytorch Geometric is available at github.com/cvignac/SMP.. Models are retrained for each cycle length and graph size on $10$k samples with balanced classes, and evaluated on $10,000$ samples as well. The same architecture (detailed in Figure 2) is used for all models, as we found it to perform better than the original implementation of each method: the methods under comparison thus only differ in the definition of the convolution, making comparison easy. We use the fast implementation of SMP, as we find its expressivity to be sufficient for this task.

Results are shown in Tab 2(c). For a given cycle length, the task becomes harder as the number of nodes in the graph grows: the bigger the graph, the more candidate paths that the network needs to verify as being cycles. SMP is able to solve the task almost perfectly for all graph and cycle sizes. For standard message-passing models, we observe a correlation between accuracy and the presence of identifiers: random identifiers and weak identifiers (a one-hot encoding of the degree) tend to perform better than the baseline GIN and MPNN. PPGN and RING-GNN solve the task well for small graphs, but fail when $n$ grows. Perhaps due to a miss-aligned inductive bias, we encountered difficulties with training them, whereas message-passing architectures could be trained more easily. We provide a more detailed comparison between SMP, PPGN and Ring-GNN in Appendix G. We also compare the generalization ability of the different networks that can be used in inductive settings. GIN generalizes well, but SMP is the only one that achieves good performance among the powerful networks. This may be imputable to the inductive bias of message passing architectures, shared by GIN and SMP. Finally, we compare SMP and GIN with random identifiers in settings with less training data: SMP requires much fewer samples to achieve good performance, which confirms that equivariance is important for good sample efficiency.

2 Multi-task detection of graph properties

We further benchmark SMP on the multi-task detection of graph properties proposed in Corso et al. . The goal is to estimate three node-defined targets: geodesic distance from a given node (Dist.), node eccentricity (Ecc.), and computation of Laplacian features ${\bm{L}}{\bm{x}}$ given a vector ${\bm{x}}$ (Lap.), as well as three graph-defined targets: connectivity (Conn.), graph diameter (Diam.), and spectral radius (Rad.). The training set is composed of 5120 graphs with up to 24 nodes, while graphs in the test set have up to 19 nodes. Several MPNNs are evaluated as well as PNA , a message-passing model based on the combination of several aggregators. Importantly, random identifiers are used for all these models, so that all baseline methods are theoretically poweful , but not equivariant.

All models are benchmarked using the same architecture, apart from the fact that SMP manipulates local contexts. In order to pool these contexts into node features and use them as input to the Gated Recurrent Unit , we use an extractor described in Figure 2. As an ablation study, we also consider for each model a corresponding MPNN with the same architecture.

The results are summarized in Table 3. We find that both SMPs are able to exploit the local contexts, as they perform much better than the corresponding MPNN. SMP also outperforms other methods by a significant margin. Lastly, standard SMP tends to achieve better results than fast SMP on tasks that require graph traversals (shortest path computations, excentricity, checking connectivity), which may be due to a better representation power.

3 Constrained solubility regression on ZINC

The ZINC database is a large scale dataset containing molecules with up to 37 atoms. The task is to predict the constrained solubility of each molecule, which can be seen as a graph regression problem. We follow the setting of : we train SMP on the same subset of 10,000 molecules with a parameter budget of around 100k, reduce the learning rate when validation loss stagnates, and stop training when it reaches a predefined value. We use an expressive parametrization of SMP, with 12 layers and 2-layer MLPs both in the message and the update functions. In order to reduce the number of parameters, we share the same feature extractor after each layer (cf Fig. 2). Results are presented in Table 4. They show that in both cases (with or without edge features, which are a one-hot encoding of the bond type), SMP is able to achieve state of the art performance. Note however than even better results (0.108 MAE using a MPNN with edge features ) can be achieved by augmenting the input with expert features. We did not use them in order to compare fairly with the baseline results.

Conclusion

We introduced structural message-passing (SMP), a new architecture that is both powerful and permutation equivariant, solving a major weakness of previous message-passing networks. Empirically, SMP significantly outperforms previous models in learning graph topological properties, but retains the inductive bias of MPNNs and their good ability to process node features. We believe that our work paves the way to graph neural networks that efficiently manipulate both node and topological features, with potential applications to chemistry, computational biology and neural algorithmic reasoning.

Broader Impact

This paper introduced a new methodology for building graph neural networks, conceived independently of a specific application. As graphs constitute a very abstract way to represent data, they have found a lot of different applications . The wide applicability of graph neural networks makes it challenging to foresee how our method will be used and the ethical problems which might occur.

Nevertheless, as we propose to overcome limitations of previous work in learning topological information, our method is likely to be used first and foremost in fields were graph topology is believed to be important. We hope in particular that it can contribute to the fields of quantum chemistry and drug discovery. The good performance obtained on the ZINC dataset is an encouraging sign of the potential of SMP in these fields. Other applications come to mind: material science , computational biology , combinatorial optimization or code generation .

Acknowledgments and Disclosure of Funding

Clément Vignac would like to thank the Swiss Data Science Center for supporting him through the PhD fellowship program (grant P18-11). Andreas Loukas would like to thank the Swiss National Science Foundation for supporting him in the context of the project “Deep Learning for Graph-Structured Data” (grant number PZ00P2 179981).

References

Appendix A Proof of Theorem 1

The action of a permutation $\pi$ on the inputs is defined as $f(\pi.({\bm{\mathsfit{U}}},{\bm{\mathsfit{Y}}},{\bm{A}}))=f(\pi.{\bm{\mathsfit{U}}},\pi.{\bm{\mathsfit{Y}}},\pi.{\bm{A}})$. In order to simplify notation, we will consider $\pi^{-1}$ instead of $\pi$. We have for example $(\pi^{-1}.~{}{\bm{A}})[{i,j}]={A}[{\pi_{i},\pi_{j}}]$ and $(\pi^{-1}.{\bm{\mathsfit{U}}})[i,j,k]={\bm{\mathsfit{U}}}[\pi_{i},\pi_{j},k]$, which can be written as

As shown next, the theorem’s conditions suffice to render SMP equivariant:

which matches the definition of equivariance.

Appendix B Proof of Theorem 2

We first present the formal version of the theorem:

If $G$ and $G^{\prime}$ are not isomorphic, then for all $\pi\in\mathfrak{S_{n}}$,

If $G$ and $G^{\prime}$ are isomorphic, then for some $\pi\in\mathfrak{S}_{n}$ independent of $i$ and $j$,

The fact that embeddings produced by isomorphic graphs are permutations one of another is a consequence of equivariance, so we are left to prove the first point. To do so, we will first ignore the features and prove that there is an SMP that maps the initial one-hot encoding of each node to an embedding that allows to reconstruct the adjacency matrix. The case of attributed graphs and the statement of the theorem will then follow easily.

These edges correspond to the receptive field of node $v_{i}$ after $l$ layers of message-passing. We denote by ${\bm{A}}_{i}^{(l)}$ the adjacency matrix of $G_{i}^{(l)}$.

To build intuition, it is useful to first consider the case where $\bm{U}_{i}$ are $n\times n$ matrices (rather than $n\times c$ as in SMP). In this setting, messages are $n\times n$ matrices as well. If the initial state of each node $v_{i}$ is its one-hop neighbourhood (${\bm{U}}_{i}^{(1)}={\bm{A}}_{i}^{(1)}$), then each node can easily recover the full adjacency matrix by updating its internal state as follows:

Recursion 2 yields ${\bm{U}}_{i}^{(l)}={\bm{A}}_{i}^{(l)}$.

We prove the claim by induction. It is true by construction for $l=1$. For the inductive step, suppose that ${\bm{U}}_{i}^{(l)}={\bm{A}}_{i}^{(l)}$. Then,

Conversely, if ${\bm{A}}_{i}^{(l+1)}[p,q]=1$, then there exists either a path of length $l$ of the form $(v_{i},v_{j},\ldots,v_{p})$ or $(v_{i},v_{j},\ldots,v_{q})$. This node $v_{j}$ will satisfy ${\bm{U}}_{j}^{(l)}[p,q]=1$ and thus ${\bm{U}}_{i}^{(l+1)}[p,q]=1$. ∎

It is an immediate consequence that, for every connected graph of diameter $\Delta$, we have ${\bm{U}}_{i}^{(\Delta)}={\bm{A}}$.

B.2 SMP: nodes manipulate node embeddings

We now shift to the case of SMP. We will start by proving that we can find an $n\times 2d_{\text{max}}$ embedding matrix (rather than $n\times n$) that still allows to reconstruct ${\bm{A}}_{i}^{(l)}$. For this purpose, we will use the following result:

There exists a sequence $(f_{l})_{l\geq 1}$ of permutation equivariant SMP layers defining $\bm{U}_{i}^{(l+1)}=f^{(l+1)}(\bm{U}_{i}^{(l)},\{\bm{U}_{j}^{(l)}\}_{j\in N_{i}})$ such that $\bm{u}_{j}^{(l)}\perp\bm{u}_{k}^{(l)}\iff(v_{j},v_{k})\in E_{i}^{(l)}$ for every layer $l$ and nodes $v_{j},v_{k}\in V$. These functions do not depend on the choice of $v_{i}\in V$.

We use an inductive argument. An initialization (layer $l=1$), we have $\bm{U}_{j}^{(0)}=\bm{\delta}_{j}$ for every $v_{j}$. We need to prove that there exists $\bm{U}_{i}^{(1)}=f^{(1)}({\bm{U}}_{i}^{(0)},\{\bm{U}_{j}^{(0)}\}_{v_{j}\in N_{i}})$ which satisfies

Rewritten in matrix form, it is sufficient to show that there exists $\bm{U}_{i}^{(1)}$ such that $\bm{U}_{i}^{(1)}(\bm{U}_{i}^{(1)})^{\top}=\bm{1}\bm{1}^{\top}-\bm{A}_{i}^{(1)}$, with $\bm{1}$ being the all-ones vector. $\bm{A}_{i}^{(1)}$ is the adjacency matrix of a star consisting of $v_{i}$ at the center and all its $d_{i}$ neighbors at the spokes. Further, it can be constructed in an equivariant manner from the layer’s input as follows:

Since the rank of $\bm{A}_{i}^{(1)}$ is at most $d_{i}$ (there are $d_{i}$ non-zero rows), the rank of $\bm{1}\bm{1}^{\top}-\bm{A}_{i}^{(1)}$ is at most $d_{i}+1\leq 2d_{i}\leq 2d_{\textit{max}}$. It directly follows that there exists a matrix $\bm{U}_{i}^{(1)}$ of dimension $n\times 2d_{\textit{max}}$ which satisfies $\bm{U}_{i}^{(1)}(\bm{U}_{i}^{(1)})^{\top}=\bm{1}\bm{1}^{\top}-\bm{A}_{i}^{(1)}$. Further, as the construction of this matrix is based on the eigendecomposition of $\bm{A}_{i}^{(1)}$, it is permutation equivariant as desired.

Inductive step. According to the inductive hypothesis, we suppose that:

The function $f^{(l+1)}$ builds the embedding $\bm{U}_{i}^{(l+1)}$ from $(\bm{U}_{i}^{(l)},\{\bm{U}_{j}^{(l)},~{}v_{j}\in N_{i}\})$ in three steps:

Each node $v_{j}\in N_{i}$ sends its embedding ${\bm{U}}_{j}^{(l)}$ to node $v_{i}$. This is done using the message function $m^{(l)}$.

The aggregation function $\phi$ reconstructs the adjacency matrix ${\bm{A}}_{j}^{(l)}$ of $G_{j}^{(l)}$ from $\bm{U}_{j}^{(l)}$ for each $v_{j}\in N_{i}\cup\{v_{i}\}$. This is done by testing orthogonality conditions, which is a permutation equivariant operation. Then, it computes ${\bm{A}}_{i}^{(l+1)}$ as in Lemma 2 using ${\bm{A}}_{i}^{(l+1)}=\max(\{{\bm{A}}_{j}^{(l+1)}\}_{v_{j}\in N_{i}\cup\{v_{i}\}})$, with the maximum taken entry-wise.

Therefore, the constructed embedding matrix $\bm{U}_{i}^{(l+1)}$ satisfies

and the function $f^{(l+1)}$ is permutation equivariant (as a composition of equivariant functions). ∎

It is a direct corollary of Lemma 1 that, when the depth is at least as large as the graph diameter, such that $E_{i}^{(l)}=E$ for all $v_{i}$ and the width is at least as large as $2d_{\text{max}}$, then there exist a permutation equivariant SMP $f=f^{(L)}\circ\ldots\circ f^{(1)}$ that induces an injective mapping from the adjacency matrix ${\bm{A}}$ to the local context ${\bm{U}}_{i}$ of each node $v_{i}$. As a result, given two graphs $G$ and $G^{\prime}$, if there are two nodes $v_{i}\in V$ and $v^{\prime}_{j}\in V^{\prime}$ and a permutation $\pi\in\mathfrak{S}_{n}$ such that ${\bm{U}}_{i}^{(L)}={\bm{\Pi}}^{T}~{}{\bm{U}}^{\prime(L)}_{j}$, then the orthogonality conditions will yield ${\bm{A}}={\bm{\Pi}}^{T}~{}{\bm{A}}^{\prime}~{}{\bm{\Pi}}$. The contraposition is that if two nodes belong to graphs that are not isomorphic, their embedding will belong to two different equivalence classes (i.e. they will be different even up to permutations).

B.3 Extension for attributed graphs

For attributed graphs, the reasoning is very similar: we are looking for a SMP network that maps the attributes to a set of local context matrices such that all the attributes of the graph can be recovered from the context matrix at any node. We treat the case of node and edge attributes separately:

Using $c_{X}$ extra channels in SMP is sufficient to create the desired embedding. We recall that the input to the SMP is a local context such that the $i$-th row of $v_{i}$ contains $[1,{\bm{x}}_{i}]$, where ${\bm{x}}_{i}$ is the vector of attributes of $v_{i}$, while the other rows are zero. Ignoring the first entry of this vector (which was used to reconstruct the topology), we propose the following update rule:

where the max is taken element-wise on each entry of the matrix. This function is simply an extension of the max aggregator that allows to replace the zeros of the local context by non zero values, even if they are negative. Using it, each node can progressively fill the rows corresponding to nodes that are more and more distant. With the assumption that the graph is connected, each node will eventually have access to all node features.

Edge attributes

As each edge attribute can be seen as a $n\times n$ matrix, edges attributes are handled in a very similar way as the adjacency matrix of unattributed graphs. If nodes could send $n\times n$ matrices as messages, they would be able to recover all the edge features using the previous update rule (equation 3). However, SMP manipulates embeddings that transform under the action of a permutation as $\pi~{}.~{}{\bm{U}}_{i}={\bm{\Pi}}^{T}{\bm{U}}_{i}$, whereas a $n\times n$ matrix ${\bm{M}}$ transforms as $\pi~{}.~{}{\bm{M}}={\bm{\Pi}}^{T}{\bm{M}}~{}{\bm{\Pi}}$. As a result, we cannot directly pass the incomplete edge features as messages, and we need to embed them into a matrix that permutes in the right way.

The construction of a SMP that embeds the input to local contexts that allow to reconstruct an edge feature matrix ${\bm{E}}$ is the same as for the adjacency matrix, except for one difference: lemma 1, which was used to embed adjacency matrices into a smaller matrix cannot be used anymore, as it is specific to unweighted graphs. Therefore, we propose another way to embed each matrix ${\bm{E}}_{i}^{(l)}$ obtained at node $v_{i}$ after $l$ message passing layers:

For undirected graphs, ${\bm{E}}_{i}^{(l)}$ is symmetric. We can therefore compute its eigendecomposition ${\bm{E}}_{i}^{(l)}={\bm{V}}{\bm{\Lambda}}{\bm{V}}^{T}$.

We add a given value $\lambda$ to the diagonal of ${\bm{\Lambda}}$ to make sure that all coefficients are non-negative.

We compute the square root matrix ${\bm{U}}={\bm{V}}({\bm{\Lambda}}+\lambda{\bm{I}})^{1/2}$. This matrix permutes as desired under the action of a permutation: $\pi~{}.~{}{\bm{U}}={\bm{\Pi}}^{T}{\bm{U}}$. In addition, it allows to reconstruct the matrix ${\bm{E}}_{i}^{(l)}={\bm{U}}{\bm{U}}^{T}$, so that it constitutes a valid embedding for the rest of the proof.

Note that the square root matrix permutes as desired, but that it does not compress the representation of ${\bm{E}}_{i}^{(l)}$: for each edge features, $n$ additional channels are needed, so that a SMP should have $n\times c_{Y}$ more channels to be able to reconstruct all edge features.

B.4 Conclusion

We have shown that there exists an SMP that satisfies the conditions of the theorem, and specifically, we demonstrated that each layer can be decomposed in a message, aggregation and update functions that should be able to internally manipulate $n\times n$ matrices in order to produce embeddings of size $n\times 2d_{\text{max}}+c_{X}+n~{}c_{Y}$.

The main assumption of our proof is that the aggregation and update functions can exactly compute any function of their input — this is impossible in practice. An extension of our argument to a universal approximation statement would entail substituting the aggregation and update functions by appropriate universal approximators. In particular, the aggregation function manipulates a set of $n\times c$ matrices, which can be represented as a $n\times n\times c$ tensor with some lines zeroed out. Some universal approximators of equivariant functions for these tensors are known , but they have large memory requirements. Therefore, proving that a given parametrization of an SMP can be used to approximately reconstruct the adjacency matrix hinges on the identification of a simple universal approximator of equivariant functions on $n\times n\times c$ tensors.

Appendix C Proof of Corollary 1

Lemma 1 proves the existence of an injective mapping from adjacency matrices of simple graphs to features for a set of nodes. Therefore, any permutation equivariant function $h_{\text{eq}}(\bm{A})$ on adjacency matrices can be expressed by an equivariant function on sets

as long as the node embeddings $\bm{u}_{1},\ldots,\bm{u}_{n}$ allow the reconstruction of $\bm{A}$, e.g., through orthogonality conditions. It was proven in Theorem 3 that, under the corollary’s conditions, the local context ${\bm{U}}_{i}^{(L)}$ of any node $v_{i}$ yields an appropriate matrix ${\bm{U}}$. In order to compute $h_{\text{eq}}$, each node can then rely on the universal $\sigma$ to compute the invariant function:

Appendix D Proof of Proposition 1

We will show by induction that any MPNN can be simulated by an SMP:

For any MPNN mapping initial node features $({\bm{x}}_{i}^{(0)})_{v_{i}\in V}$ to $({\bm{x}}_{i}^{(L)})_{v_{i}\in V}$, there is an SMP with the same number of layers such that

Consider a graph with node features $(\bm{x}_{i}^{(0)})_{v_{i}\in V}$ and edge features $({{\bm{y}}_{ij}})_{(v_{i},v_{j})\in E}$.

Initialization: The context tensor is initialized by mapping the node features on the diagonal of ${\bm{\mathsfit{U}}}$: ${\bm{\mathsfit{U}}}^{(0)}[i,i,\colon]={\bm{x}}_{i}^{(0)}$. The desired property is then true by construction.

Inductive step: Denote by $({\bm{x}}_{i}^{(l)})_{v_{i}\in V}$ the features obtained after $l$ layers of the MPNN. Assume that there is a k-layer SMP such that the local context after $l$ layers contains the same features in its diagonal elements: ${\bm{\mathsfit{U}}}^{(l)}[i,i,:]={\bm{x}}_{i}^{(l)}$ and 0 in the other entries. Consider one additional layer of MPNN:

Finally, the function diag extracts the main diagonal of the tensor ${\bm{\mathsfit{U}}}$ along the two first axes. Let $\delta_{i,j}$ be the function that is equal to 1 if $i=j$ and , otherwise. We have: $\textit{diag}({\bm{\mathsfit{U}}})[i,j,:]={\bm{\mathsfit{U}}}[i,j,:]~{}\delta_{i,j}$. Note that this function can equivalently be written as an update function applied separately to each node: $\textit{diag}({\bm{U}}_{i})[j,:]={\bm{U}}_{i}[j,:]\delta_{i,j}$. We now have ${\bm{\mathsfit{U}}}^{(l+1)}[i,i;:]={\bm{x}}_{i}^{l+1}$ and ${\bm{\mathsfit{U}}}$ equal to on all the other entries, so that the induction hypothesis is verified at layer $l+1$. ∎

As any MPNN can be computed by an SMP, we conclude that SMPs are at least as powerful as MPNNs.

SMP are strictly more powerful

To prove that SMPs are strictly more powerful than MPNNs, we use a similar argument to :

There is an SMP network which yields different outputs for the two graphs of Fig. 3, while any MPNN will view these graphs are isomorphic.

The two graphs of Fig. 3 are regular, which implies that they cannot be distinguished by the Weisfeiler-Lehman test or by MPNNs without special node features . On the contrary, consider an SMP $f$ made of three layers computing ${\bm{U}}_{i}^{(l+1)}=\sum_{v_{i}\in N_{i}}{\bm{U}}_{i}^{(l)}$, followed by the trace of ${\bm{U}}^{(3)}$ as a a pooling function. As each layer can be written ${\bm{U}}^{(l+1)}={\bm{A}}{\bm{U}}^{(l)}$ and ${\bm{U}}^{(0)}={\bm{I}}_{n}$, we have $f({\bm{A}})=\textit{tr}({\bm{A}}^{3})$. In particular $f({\bm{A}})=2$ for the graph on the left, while $f({\bm{A}})=0$ on the right.∎

Appendix E A more compact representation with graph coloring

Appendix F Proof of Proposition 2

We will prove by induction that any Fast SMP layer can be approximated by two blocks of PPGN. It implies that the expressive power of Fast SMP is bounded by that of PPGN.

Recall that a block of PPGN is parameterized as:

To simplify the presentation, we assume that:

At each layer $l$, one of the channels of ${\bm{\mathsfit{T}}}^{(l)}$ corresponds to the adjacency matrix ${\bm{A}}$, another contains a matrix full of ones $\bm{1}_{n}\bm{1}_{n}^{\top}$ and a third the identity matrix ${\bm{I}}_{n}$, so that each PPGN layer has access at all times to these quantities. These matrices can be computed by the first PPGN layer and then kept throughout the computations using residual connections.

The neural network can compute entry-wise multiplications $\odot$. This computation is not possible in the original model, but it can be approximated by a neural network.

${\bm{\mathsfit{U}}}$ and ${\bm{\mathsfit{T}}}$ have only one channel (so that we write them ${\bm{U}}$ and ${\bm{T}}$). This hypothesis is not necessary, but it will allow us to manipulate matrices instead of tensors.

Initially, we simply use the same input for PPGN as for SMP (${\bm{U}}^{(0)}={\bm{T}}^{(0)}={\bm{I}}_{n}$).

Induction

Assume that at layer $l$ we have ${\bm{U}}^{(l)}={\bm{T}}^{(l)}$. Consider a layer of Fast SMP:

A first PPGN block can be used to compute $\hat{\bm{U}}_{i}^{(l)}$ for each node. This block is parametrized by:

The output of this block exactly corresponds to $\hat{\bm{U}}^{(l)}$. Then, a second PPGN block can be used to compute the rest of the Fast SMP layer. It should be parametrized as:

By plugging these expressions into the definition of a PPGN block, we obtain that the output of this block corresponds to ${\bm{U}}^{(l+1)}$ as desired.

Appendix G Comparison between SMP, Provably powerful graph networks and Ring-GNN