P-Tensors: a General Formalism for Constructing Higher Order Message Passing Networks

TL;DR

This work introduces $P$-tensors to unify higher-order permutation equivariant message passing with subgraph neural networks. It defines invariant subgraph selection and two-level equivariance, then derives the complete set of permutation-equivariant linear maps between $P$-tensors and proves that networks built from these maps are globally equivariant. The approach extends traditional MPNNs to operate over subgraphs, edges, and cycles, enabling richer interactions while preserving symmetry through $P$-tensor transformations. Empirically, the framework achieves state-of-the-art results on ZINC-12K and competitive performance on other molecular benchmarks, demonstrating that increased expressivity translates to practical gains. Overall, $P$-tensors offer a reusable, principled paradigm for designing expressive, symmetry-preserving higher-order GNNs and motivate a software library for such models.

Abstract

Several recent papers have proposed increasing the expressive power of graph neural networks by exploiting subgraphs or other topological structures. In parallel, researchers have investigated higher order permutation equivariant networks. In this paper we tie these two threads together by providing a general framework for higher order permutation equivariant message passing in subgraph neural networks. In this paper we introduce a new type of mathematical object called $P$-tensors, which provide a simple way to define the most general form of permutation equivariant message passing in both the above two categories of networks. We show that the P-Tensors paradigm can achieve state-of-the-art performance on benchmark molecular datasets.

Figures (3)

• Figure 1: A subgraph neural network must be equivariant to two different ways that permutations act on it: changing the set of vertices assigned to a given subgraph, and reordering the vertices of the subgraph internally. This is especially important when the subgraph neurons produce matrix/tensor valued outputs indexed by the vertices of the subgraph itself. The $P$-tensors formalism allows us to handle this situation in a simple way, defining the most general form of equivariant linear messages between such tensors, without making reference to the global ordering.
• Figure 2: Given two $P$-tensors $T^{\textrm{in}}$ and $T^{\textrm{out}}$ whose reference domains have $d^{\cap}$ atoms in common, without loss of generality we can rearrange the two tensors so that the indices corresponding to the common atoms appear first. Mapping the corresponding $d^{\cap}\!\times d^{\cap}\!\times \ldots \times d^{\cap}$ subtensor of $T^{\textrm{in}}$ to the analogous $d^{\cap}\!\times d^{\cap}\!\times \ldots\times d^{\cap}$ subtensor of $T^{\textrm{out}}$ with any of the $B(k_1\space+\space k_2)$ linear maps described in Section \ref{['subsec: Maron']} is an equivariant operation. The additional equivariant operations correspond to similar maps except with the summations and broadcast operations extending over not just the overlapping part of the tensors but the entirety of $T^{\textrm{in}}$ or $T^{\textrm{out}}$.
• Figure 3: The number of independent equivariant linear maps from a $k_1$'th order $P$-tensor to a $k_2$ 'th order $P$-tensor when the reference domains are the same vs. when they overlap only partially.

Theorems & Definitions (11)

• Definition 1
• Definition 2
• Definition 3
• Definition 4
• Definition 5: $P$-tensors
• Definition 6: Permutation equivariant maps between $P$-tensors
• Definition 7: Relabeling invariant maps between $P$-tensors
• Theorem 1
• Proposition 2: Maron et al.
• Theorem 3