Understanding and Extending Subgraph GNNs by Rethinking Their Symmetries
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Subgraph GNNs are a recent class of expressive Graph Neural Networks (GNNs)which model graphs as collections of subgraphs.
So far, the design space ofpossible Subgraph GNN architectures as well as their basic theoretical propertiesare still largely unexplored.
In this paper, we study the most prominent form ofsubgraph methods, which employs node-based subgraph selection policies such asego-networks or node marking and deletion.
We address two central questions: (1) What is the upper-bound of the expressive power of these methods?
and (2) Whatis the family of equivariant message passing layers on these sets of subgraphs? .
Our ﬁrst step in answering these questions is a novel symmetry analysis whichshows that modelling the symmetries of node-based subgraph collections requiresa signiﬁcantly smaller symmetry group than the one adopted in previous works.
This analysis is then used to establish a link between Subgraph GNNs and Invariant
Graph Networks (IGNs).
We answer the questions above by ﬁrst bounding theexpressive power of subgraph methods by 3-WL, and then proposing a generalfamily of message-passing layers for subgraph methods that generalises all previousnode-based Subgraph GNNs.
Finally, we design a novel Subgraph GNN dubbed
SUN , which theoretically uniﬁes previous architectures while providing betterempirical performance on multiple benchmarks.
Introduction
Message Passing Neural Networks (MPNNs) are arguably the most commonly used version of Graph Neural Networks (GNNs).
The limited expressive power of MPNNs [31, 45] has led toa plethora of works aimed at designing expressive GNNs while maintaining the simplicity andscalability of MPNNs [9, 34, 40, 25].
Several recent studies have proposed a new class of sucharchitectures [12, 49, 6, 50, 37, 36], dubbed Subgraph GNNs , which apply MPNNs to collections(‘bags’) of subgraphs extracted from the original input graph and then aggregate the resultingrepresentations.
Subgraphs are selected according to a predeﬁned policy; in the most popular ones,each subgraph is tied to a speciﬁc node in the original graph, for example by deleting it or extractingits local ego-network.
Subgraph GNNs have demonstrated outstanding empirical performance, withstate-of-the-art results on popular benchmarks like the ZINC molecular property prediction [50, 6].
While offering great promise, it is fair to say that we still lack a full understanding of Subgraph
GNNs.
Firstly, on the theoretical side, it is known that subgraph methods are strictly stronger thanthe Weisfeiler-Leman (WL) test [44, 34], but an upper-bound on their expressive power is generallyunknown.
Secondly, on a more practical level, Subgraph GNN architectures differ considerably in the way information is aggregated and shared across the subgraphs, and an understanding of thepossible aggregation and sharing rules is missing.
Both aspects are important: an understanding ofthe former can highlight the limitations of emerging architectures, a study of the latter paves the wayfor improved Subgraph GNNs.
Main contributions.
The goal of this paper is to provide a deeper understanding of node-based Subgraph GNNs in light of the two aforementioned aspects.
The main theoretical tool underpinningour contributions is a novel analysis of the symmetry group that acts on the sets of subgraphs.
Whileseveral previous approaches [37, 12, 6] have (often implicitly) assumed that a subgraph architectureshould be equivariant to independent node and subgraph permutations, we leverage the fact that nodebased policies induce an inherent bijection between the subgraphs and the nodes.
This observationallows us to align the two groups and model the symmetry with a single (smaller) permutation groupthat acts on nodes and subgraphs jointly .
Other works [50, 46, 49] have (again, implicitly) recognisedsuch node-subgraph correspondence but without studying the implications on the symmetry group,and resorting, as a result, to a partial and heuristic choice of equivariant operations.
The use of this stricter symmetry group raises a fruitful connection with k -order Invariant Graph Networks ( k-IGN s) [28, 27], a well studied family of architectures for processing graphs and hypergraphsdesigned to be equivariant to the same symmetry group.
This connection allows us to transfer andreinterpret previous results on IGN s to our Subgraph GNN setup.
As our ﬁrst contribution we showthat the expressive power of Subgraph GNNs with node-based policies is bounded by that of the 3-WLtest.
This is shown by proving that all previous Subgraph GNNs can be implemented by a 3-IGN andby leveraging the fact that the expressive power of these models is bounded by 3-WL [18, 4].
Our second contribution is the proposal of a general layer formulation for Subgraph GNNs, basedon the observation that these methods maintain an n ⇥ n representation of n subgraphs with nnodes, following the same symmetry structure of 2-IGN s (same permutation applied to both rowsand columns of this representation).
We propose a novel extension of 2-IGN s capturing both local(message-passing-like) and global operations.
This extension easily recovers previous methodsfacilitating their comparison.
Also, we present a number of new operations that previous methodsdid not implement.
We build upon these observations to devise a new Subgraph GNN dubbed SUN ,( Subgraph Union Network ).
We prove that SUN generalises all previous node-based Subgraph GNNsand we empirically compare it to these methods, showing it can outperform them.
Previous work
Expressive power of GNNs.
The expressive power of GNNs is a central research focus sinceit was realised that message-passing type GNNs are constrained by the expressivity of the WLisomorphism test [31, 45].
Numerous approaches for more powerful GNNs have been proposed,including positional and structural encodings [1, 38, 9, 15, 24, 26], higher-order message passingschemes [31, 33, 8, 7], and equivariant models [21, 28, 27, 43, 13, 42, 35].
We refer readers to therecent survey by Morris et al.
[34] for additional details.
Invariant graph networks.
IGN s were recently introduced in a series of works by Maron et al.[28, 27, 29] as an alternative to MPNNs for processing graph and hyper-graph data.
For k  2 , k-IGN srepresent hyper-graphs with hyper-edges up to size k with k -order tensor Y 2 R n k , where each entryholds information about a speciﬁc hyper-edge. On these they apply linear S n -equivariant layersinterspersed with pointwise nonlinearities.
These models have been thoroughly studied in terms of:(i) their expressive power; (ii) the space of their equivariant linear layers. As for (i), IGN s were shownto have exactly the same graph separation power as the k-WL graph isomorphism test [27, 4, 18] and,for sufﬁciently large k , to have a universal approximation property w.r.t.
S n -invariant and equivariantfunctions [29, 22, 39].
Concerning (ii), the work in [28] completely characterised the space of linearlayers equivariant to S n from R n k to R n k 0 : the authors derived a basis of bell ( k + k 0 ) linear operatorsconsisting of indicator tensors of equality patterns over the multi-index set { 1 , . . . , n } k + k 0 = [ n ] k + k 0 .
Albooyeh et al.
[2] showed these layers can be (re-)written as sums of pooling-broadcasting operations between elements of Y indexed by the orbits 1 of the action of S n on [ n ] k and [ n ] k {+0 . Take, e.g.,k = k 
 = 2+} .
In this case there are only two orbits: { i, i } , i 2 [ n ] corresponding to on-diagonal terms,and { i, j } , i 6 =
j 2 [ n ] , off-diagonal terms.
According to Albooyeh et al. [2] any equivariant linearlayer L : R n 2 ! R n 2 can be represented as a composition of pooling and broadcasting operationson the elements indexed by these orbits. One example is the linear map that sums the on-diagonalelements and broadcasts the result to the off-diagonal ones: L ( Y ) ij =
{+P k Y kk for i 
+}
= j , 0 otherwise. See Appendix B, for additional details. These results particularly important as they underpin most ofour theoretical derivations. Lastly, a more comprehensive coverage of IGN s can be found in [34].
Subgraph GNNs
Notation.
Let G =
( A, X ) be a member of the family G of node-attributed , undirected, ﬁnite, simplegraphs 2 .
The adjacency matrix A 2 R n ⇥ n represents G ’s edge set E over its set of n nodes V .
Thefeature matrix X 2 R n ⇥ d gathers the node features; we denote by x j 2 R d ⇥ 1 the features of nodej corresponding to the j -th row of X .
B G is used to denote a multiset (bag) of m subgraphs of G .
Adjacency and feature matrices for subgraphs in B G are arranged in tensors A 2 R m ⇥ n ⇥ n and X 2 R m ⇥ n ⇥ d .
Superscript i, ( t ) refers to representations on subgraph i at the t -th layer of a stacking,as in x i, ( t ) j .
Finally, we denote [ n ] = { 1 , . . . , n } .
All proofs are deferred to Appendix B and D.
Formalising Subgraph GNNs.
Subgraph GNNs compute a representation of G 2 G as(
[Equation]
Here, ⇡ : G 7! { G 1 , ..., G m } =
{ ( A 1 , X 1 ) , ..., ( A m , X m ) }
=
B (0) G is a selection policy generatinga bag of subgraphs from G ; S =
L T  . . .
L 1 : B (0) G 7! B ( T ) G is a stacking of T (node- andsubgraph-) permutation equivariant layers; ⇢ : ( G, B ( T ) G ) 7! x G is a permutation invariant poolingfunction , µ is an MLP.
Subgraph GNNs differ in the implementation of ⇡ , S and, in some cases,⇢ .
For example, in 1-Reconstruction GNNs [12], ⇡ selects node-deleted subgraphs and S appliesa Siamese MPNN to each subgraph independently.
To exemplify the variability in S , ESAN [6]extends this method with cross-subgraph node and connectivity aggregation.
Layers in S comprise abase-encoder in the form of a GNN applied to subgraphs; throughout this paper, we assume it to be a1-WL maximally expressive MPNN such as the one in [31].
Node-based selection policies.
In this work, we focus on a speciﬁc family of node-based subgraphselection policies, wherein every subgraph is associated with a unique node in the graph.
Formally,we call a subgraph selection policy node-based if it is of the form ⇡ ( G ) =
{ f ( G, v ) } v 2 V , for someselection function f
( G, v ) that takes a graph G and a node v as inputs and outputs a subgraph G v .
In the following, we refer to v as the root of subgraph G v .
We require f to be a bijection and wenote that such policies produce m = n different subgraphs.
Amongst the most common examplesare node-deletion (ND), node-marking (NM), and ego-networks (EGO) policies.
For input graph G ,f ND ( G, v ) removes node v and the associated connectivity; f NM ( G, v ) adds a special ‘mark’ attributeto v ’s features (with no connectivity alterations), and f EGO( h ) ( G, v ) returns the subgraph induced bythe h -hop-neighbourhood around the root v .
EGO policies can be ‘marked’:
f EGO+( h ) ( G, v ) extractsthe h -hop ego-net around v and marks this node as done by f NM .
For convenience, we denote theclass of such node-based selection policies by ⇧ :
Deﬁnition 1 (Known node-based selection policies ⇧ ) .
Let ⌃ be the set of all node-basedsubgraph selection policies.
Class ⇧ ⇢ ⌃ collects the node-based policies node-deletion(ND), node-marking (NM), ego-nets (EGO) and marked ego-nets (EGO+) of any depth: ⇧ ={ ⇡ ND , ⇡ NM , ⇡ EGO( h ) , ⇡ EGO+( h ) | h > 0 } .
Appendix A presents the different formulations of currently known node-based Subgraph GNNs indetail, which we group in the the following family of models:
In ESAN [6] symmetries are modelled as a direct product of node and subgraph permutation groups; however, node-based policies enable the use of one single permutation group, the same as in 3-IGN s.
3-IGN s are less constrained, thus more expressive than ESAN and other Subgraph GNNs.
See diagram on the right and formal statement in Section 5.
Deﬁnition 2 (Known Subgraph GNNs ⌥ ) .
Let ⌅ be the set of all node-based Subgraph GNNs.
Class ⌥ ⇢ ⌅ collects known Subgraph GNNs when equipped with 1-WL base-encoders:
⌥ = { 1Reconstr. GNN , DS-GNN , DSS-GNN , GNN-AK , GNN-AK-ctx , NGNN , ID-GNN } .
Importantly, all these methods apply MPNNs to subgraphs of the original graph, but differ in the wayinformation is shared between subgraphs/nodes.
In all cases, their expressive power is strictly largerthan 1-WL , but an upper -bound is currently unknown.
Symmetries of node-based subgraph selection policies
In an effort to characterise the representational power of node-based Subgraph GNNs, we ﬁrst studythe symmetry group of the objects they process: ‘bags of subgraphs’ represented as tensors ( A , X ) {+
+}
R m ⇥ n ⇥ n ⇥ R m ⇥ n ⇥ d , assuming n nodes across m subgraphs.
Previous approaches [12, 6, 37] usedtwo permutation groups: one copy of the symmetric group S n models node permutations , whileanother copy S m models subgraph permutations in the bag.
These two were combined by a groupproduct 3 acting independently on the nodes and subgraphs in ( A , X ) .
For example, Bevilacqua et al.[6] model the symmetry as:
[Equation]
Our contributions stem from the following crucial observation: When using node-based policies,the subgraphs in ( A , X ) can be ordered consistently with the nodes by leveraging the bijectionf : v 7! G v characterising this policy class.
In other words, f suggests a node-subgraph alignmentinducing a new structure on ( A , X ) , whereby the subgraph order is not independent of that of nodesanymore.
Importantly, this new structure is preserved only by those permutations operating identicallyon both nodes and subgraphs.
Following this observation, the symmetry of a node-based bag ofsubgraphs is modelled more accurately using only one single permutation group S n jointly acting onboth nodes and subgraphs:
[Equation]
It should be noted that S n is signiﬁcantly smaller than S n ⇥ S n 4 .
Informally, the latter group containsmany permutations which are not in the former: those acting differently on nodes and subgraphsand, thus, not preserving the new structure of ( A , X ) .
Since they are restricted by a smaller set ofequivariance constraints, we expect GNNs designed to be equivariant to S n to be be more expressivethan those equivariant to the larger groups considered by previous works [29] (see Figure 1).
The insight we obtain from Equation (3) is profound: it reveals that the symmetry structure of
A exactly matches the symmetries of third-order tensors used by 3-IGN s, and similarly, that the symmetry structure for X matches the symmetries of second-order tensors used by 2-IGN s. In thefollowing, we will make use of this insight and the fact that IGN s are well-studied objects to provean upper-bound on the expressive power of Subgraph GNNs and to design principled extensionsto these models. We remark that bags of node-based subgraphs can also be represented as tensors Y 2 R n 3 ⇥ d , the same objects on which 3-IGN s operate. Here, X is embedded in the main diagonalplane of Y , A in its remaining entries.
Within this context, it is informative to study the semantics ofthe 5 orbits induced by the action of S n on Y ’s multi-index set [ n ] 3 : each of these uniquely identifyroot nodes, non-root nodes, links from and to root nodes as well as links between non-root nodes (see Figure 4 and additional details in Appendix B.1.1). We build upon this observation, along with thelayer construction by Albooyeh et al. [2], to prove many of the results presented in the following.
A representational bound for Subgraph GNNs
In this section we prove that the expressive power of known node-based Subgraph GNNs is boundedby 3-WL by showing that they can be implemented by 3-IGN s, which have the same expressive poweras 3-WL .
Underpinning the possibility of IGN s to upper-bound a certain Subgraph GNN N in itsexpressive power is the ability of IGN s to (i) implement N ’s subgraph selection policy ( ⇡ ) and (ii)implement N ’s (generalised) message-passing and pooling equations ( µ  ⇢  S ).
This would ensurethat whenever N assigns distinct representations to two non-isomorphic graphs, an IGN implementing N would do the same.
We start by introducing a recurring, useful concept.
Deﬁnition 3 (“implements”) .
Let f : D f ! C f , g : D g !
C g be two functions and such that
D g ✓ D f , C g ✓ C f .
We say f implements g (and write f ⇠ = g ) when 8 x 2 D g , f ( x ) =
g ( x ) .
Our ﬁrst result shows that 3-IGN s can implement the selection policies in class ⇧ (Deﬁnition 1),which, to the best of our knowledge, represent all known node-based policies utilised by previouslyproposed Subgraph GNNs.
Lemma 4 ( 3-IGN s implement known selection policies) .
For any ⇡ 2 ⇧ there exists a stacking of3-IGN layers M ⇡ s.t.
M ⇡ ⇠ = ⇡ .
Intuitively, 3-IGN s start from a R n 2 representation of G and, ﬁrst, move to a R n 3 tensor ‘copying’this latter along its ﬁrst (subgraph) dimension.
This is realised via an appropriate broadcast operation.
Then, they proceed by adding a ‘mark’ to the features of some nodes and/or by nullifying elementscorresponding to some edges.
Next, we show that 3-IGN s can implement layers of any model in ⌥ .
Lemma 5 ( 3-IGN s implement Subgraph GNN layers) .
Let G 1 , G 2 be two graphs in G and N amodel in family ⌥ equipped with Morris et al.
[31] message-passing base-encoders.
Let B ( t ) 1 , B ( t ) 2be bags of subgraphs in the input of some intermediate layer L in N .
Then there exists a stacking of3-IGN layers M L for which M L ( B ( t ) i ) =
B ( t +1) i
=
L ( B ( t ) i ) for i = 1 , 2 .
Lemmas 4 and 5 allow us to upper-bound the expressive power of all known instances of node-based Subgraph GNNs by that of 3-IGN s:
Theorem 6 ( 3-IGN s upper-bound Subgraph GNNs) .
For any pair of non-isomorphic graphs G 1 , G 2in family G and Subgraph GNN model N 2 ⌥ equipped with Morris et al.
[31] message-passingbase-encoders, if there exists weights ⇥ such that G 1 , G 2 are distinguished by instance N ⇥ , thenthere exist weights ⌦ for a 3 -IGN instance M ⌦ such that G 1 , G 2 are distinguished by M ⌦ as well.
Theorem 6 has profound consequences in the characterisation of the expressive power of node-based Subgraph GNNs, as we show in the following
Corollary 7 ( 3-WL upper-bounds Subgraph GNNs) .
Let G 1 , G 2 2 G be two non-isomorphic graphsand N ⇥ 2 ⌥ one instance of model N with weights ⇥ .
If N ⇥ distinguishes G 1 , G 2 , then the 3-WLalgorithm does so as well.
Proof idea : If there is a pair of graphs undistinguishable by 3-WL , but for which there exists a Subgraph GNN separating them, there must exists a 3-IGN separating these (Theorem 6).
This is incontradiction with the result in Geerts [18], Azizian and Lelarge [4].
A design space for Subgraph GNNs
As discussed, different formulations of Subgraph GNNs differ primarily in the speciﬁc rules for updating node representations across subgraphs.
However, until now it is not clear whether existing rulesexhaust all the possible equivariant options.
We devote this section to a systematic characterisation ofthe ‘layer space’ of Subgraph GNNs.
In the spirit of the previous Section 5, where we “embedded” Subgraph GNNs in 3-IGN s, one optionwould be to consider all bell (6) = 203 linear equivariant operations prescribed by this formalism.
However, this choice would be problematic for three main reasons: (i) This layer space is too vast tobe conveniently explored; (ii) It includes operations involving O ( n 3 ) space complexity, impracticalin most applications; (iii) The linear IGN basis does not directly support local message passing, akey operation in subgraph methods.
Following previous Subgraph GNN variants, which use O ( n 2 )storage for the representation of n nodes in n subgraphs, we set the desideratum of O ( n 2 ) memorycomplexity as our main constraint, and use this restriction to reduce the design space.
Precisely, weare interested in modelling S n -equivariant transformations on the subgraph-node tensor X .
6.1 Extended 2-IGNs
As we have already observed in Equation 3 in Section 4, such a second order tensor X abides by thesame symmetry structure of 2-IGN s. We therefore gain intuition from the characterisation of linearequivariant mappings as introduced by Maron et al.
[28], and propose an extension of this formalism.2-IGN layer space.
A 2-IGN layer L ⇥ updates X 2 R n ⇥ n ⇥ d as X ( t +1) =
L ⇥  X ( t )  by applying aspeciﬁc transformation to on- ( x ii ) and off-diagonal terms ( x ij , i 6 = j ):x i, ( t +1) i =  ✓
[Equation] j x j, ( t ) j
[Equation]
⇤
[Equation]
[Equation]
i x i
[Equation]
( t )
j
[Equation]
⇤
[Equation]
[Equation]
h
[Equation]
( t )
[Equation]
⇤
[Equation] h [Equation] ( t )
[Equation] (4) x k, ( t +1) i =  ✓ [Equation] k
[Equation] ( t ) i [Equation]
x i [Equation] [Equation] [Equation]
Here, ⇤ indicates a permutation invariant aggregation function,  ✓ 1 ,  ✓ 2 apply a speciﬁc d ⇥ d 0 lineartransformation to each input term and sum the outputs including bias terms.
ReIGN(2) layer space.
As 2-IGN layers are linear, the authors advocate setting ⇤ ⌘ P , performing pooling as global summation.
Here, we extend this formulation to additionally includedifferent local aggregation schemes.
In this new extended formalism, entry x ki represents nodei in subgraph k ; accordingly, each aggregation in Equation 4 can be also performed locally, i.e.extending only over i ’s neighbours, as prescribed by the connectivity of subgraph k or of the original input graph. As an example, when updating entry x k, ( t ) i , term ⇤ j 6 = k x k, ( t ) j is expanded as ⇤ j
6 = k x k, ( t ) j , ⇤ j ⇠ k i x k, ( t ) j , ⇤ j ⇠ i x k, ( t ) j  , with ⇠ k denoting adjacency in subgraph k , and ⇠ that inthe original graph connectivity.
Each term in the expansion is associated with a speciﬁc learnablelinear transformation.
We report a full list of pooling operations in Appendix D, Table 3.
Theselocal pooling operations allow to readily recover sparse message passing, which constitutes the maincomputational primitive of all popular (Subgraph) GNNs.
We also note that other characteristic Subgraph GNN operations are also recovered by this formalism: for example, ⇤ h
6 = i x h, ( t ) i operatesglobal pooling of node i ’s representations across subgraphs, as previously introduced in Bevilacquaet al.
[6], Zhao et al.
We generally refer to this framework as ReIGN(2)
(“Rethought 2-IGN ”).
ReIGN(2) architectures.
ReIGN(2) induces (linear) layers in the same form of Equation 4, butwhere ⇤ terms are expanded to both local and global operations, as explained.
These layers canoperate on any bag generated by a node-based selection policy ¯ ⇡ , and can be combined together in ReIGN(2) stacks of the form
S R =
L ( T )    L ( T  1)    . . .
L (1) , where  ’s are pointwisenonlinearities and L ’s are ReIGN(2) layers.
This allows us to deﬁne ReIGN(2) models as Subgraph
GNNs in the form of Equation 1, where S is a ReIGN(2) layer stacking and ⇡ is node-based: [Equation]
More generally, ReIGN(2) induces a ‘layer space’ for node-based Subgraph GNNs: the expandedterms in its update equations represent a pool of atomic operations that can be selected and combinedto deﬁne new equivariant layers.
Compared to that of 3-IGN s, this space is of tractable size, yet itrecovers previously proposed Subgraph GNNs and allows to deﬁne novel interesting variants.
Recovering previous Subgraph GNNs.
The following result states that the ReIGN(2) generalisesall known subgraph methods in ⌥ , as their layers are captured by a ReIGN(2) stacking.
Theorem 8.
Let N be a model in family ⌥ equipped with Morris et al.
[31] message-passingbase-encoders.
For any instance N ⇥ , there exists ReIGN(2) instance R ⌦ such that R ⌦ ⇠ = N ⇥ .
This shows that known methods are generalised without resorting to the O ( n 3 ) computationalcomplexity of 3-IGN s. Figure 2 illustrates the aggregation and sharing rules used by previoussubgraph GNNs, visually demonstrating how ReIGN(2) generalises previous methods.
It is importantto notice that, when employing policies in ⇧ and 3-IGN -computable invariant pooling functions ⇢(as those used by models in ⌥ ), ReIGN(2) s are upper-bounded by 3-IGN s:
Proposition 9.
For any pair of non-isomorphic graphs G 1 , G 2 in family G , if there exist policy¯ ⇡ 2 ⇧ , parameters ⇥ and 3-IGN -computable invariant pooling function ⇢ such that the ReIGN(2)instance R ⇢ , ⇥ , ¯ ⇡ distinguishes G 1 , G 2 , then there exist weights ⌦ for a 3-IGN instance M ⌦ suchthat G 1 , G 2 are distinguished by M ⌦ as well.
This proposition entails an upper-bound on the expressive power of ReIGN(2) .
Corollary 10.
The expressive power of a ReIGN(2) model with policy ⇡ 2 ⇧ and 3-IGN -computableinvariant pooling function ⇢ is upper-bounded by 3-WL .
We note that there may be layers equivariant to S n over R n 2 not captured by ReIGN(2) .
Yet,previously proposed Subgraph GNN layers do not exhaust the ReIGN(2) design space, whichremains largely unexplored.
One, amongst possible novel constructions, is introduced next.
6.2 A unifying architecture: Subgraph Union Networks
We now show how the ReIGN(2) layer space can guide the design of novel, expressive, Subgraph GNNs.
Our present endeavour is to conceive a computationally tractable architecture subsumingknown node-based models: in virtue of this latter desideratum, we will dub this architecture “Subgraph Union Network” ( SUN ).
To design the base equivariant layer for SUN , we select and combine speciﬁcaggregation terms suggested by the ReIGN(2) framework:
[Equation]
[Equation]
[Table]
where  ’s sum their inputs after applying a speciﬁc linear transformations to each term.
One ofthe novel features of SUN is that roots are transformed by a different set of parameters ( ✓ 1 ) than theother nodes 5 ( ✓ 2 , see Figure 2).
In practice, the ﬁrst and last two terms in each one of Equations (5)and (6) can be processed by maximally expressive MPNNs [31, 45], the remaining terms by MLPs.
We test these variants in our experiments, with their formulations in Appendix G. SUN remains aninstantiation of the ReIGN(2) framework:
Proposition 11 (A ReIGN(2) stacking implements SUN layers) .
For any SUN layer L deﬁned according to Equations 5 and 6, there exists a ReIGN(2) layer stacking S L , such that S L ⇠ = L .
Finally, we show that a stacking of SUN layers can implement any layer of known node-based
Subgraph Networks, making this model a principled generalisation thereof.
Proposition 12 (A SUN stacking implements known Subgraph GNN layers) .
Let N be a model infamily ⌥ employing Morris et al.
[31] as a message-passing base-encoder.
Then, for any layer L in N , there exists a stacking of SUN layers S L such that S L ⇠ =
L .
Beyond SUN .
As it can be seen in Figure 2, SUN does not use all possible operations in the ReIGN(2)framework.
Notably, two interesting operations that are not a part of SUN are: (i) The ‘transpose’:x ki =
✓ ( x ik ) , which shares information between the i -th node in the k -th subgraph and the k -th nodein the i -th subgraph; (ii) Local vertical pooling x ki =
✓ (
P h ⇠ i x hi ) .
The exploration of these andother operations is left to future work.
Experiments
We experimentally validate the effectiveness of one ReIGN(2) instantiation, comparing SUN topreviously proposed Subgraph GNNs.
We seek to verify whether its theoretical representationalpower practically enables superior accuracy in expressiveness tasks and real-world benchmarks.
Concurrently, we pay attention to the generalisation ability of models in comparison.
SUN layersare less constrained in their weight sharing pattern, resulting in a more complex model.
As this istraditionally associated with inferior generalisation abilities in low data regimes, we deem it importantto additionally assess this aspect.
Synthetic.
Counting substructures and regressing graph topological features are notoriously hardtasks for GNNs [10, 15, 11].
We test the representational ability of SUN on common benchmarks ofthis kind [10, 11].
Table 1 reports results on the substructure counting suite, on which SUN attainsstate-of-the-art results in 3 out of 4 tasks.
Additional results on the regression of global, structuralproperties are reported in Appendix G. As result [Figure]
Real-world.
On the molecular ZINC-12k benchmark (constrained solubility regression) [41, 19, 14],
SUN exhibits best performance amongst all domain-agnostic GNNs under the 500 k parameter budget,including other Subgraph GNNs (see Table 1).
A similar trend is observed on the large-scale Molhivdataset from the OGB [20] (inhibition of HIV replication).
Results are in Table 2.
Remarkably, onboth datasets, SUN either outperforms or approaches HIMP [16], GSN [9] and CIN [7], GNNs whichexplicitly model rings.
We experiment on smaller-scale TUDatasets [32] in Appendix G, where wealso compare selection policies.
[Table]
Generalisation from limited data.
In this set of experimentswe compare the test performance of Subgraph GNNs whentrained on increasing fractions of the available training data.
Each architecture is selected by tuning the hyperparameterswith the entire training and validation sets.
We run this experiment on the 4 -cycle counting task and the real-world ZINC-12k.
We illustrate results in Figures 3a to 3c.
Except for a short initial phase in the EGO policy, SUN generalises better than other Subgraph GNNs on cycle-counting.
On ZINC-12k, SUN outruns DSS-, DS-GNN and GNN-AK variants from, respectively,20 , 30 and 40% of the samples.
These results demonstrate that SUN ’s expressiveness is not at the expense of sample efﬁciency,suggesting that its modelled symmetries guarantee strong representational power while retaining important inductive biasesfor learning on graphs.
Conclusions
Our work uniﬁes, extends, and analyses the emerging class of Subgraph GNNs.
Notably, wedemonstrated that the expressive power of these methods is bounded by 3-WL .
Towards a systematicstudy of models whose expressivity lies between 1- and 3-WL , we proposed a new family of layersfor the class of Subgraph GNNs and, unlike most previous works on the expressive power of GNNs,we also investigated the generalisation abilities of these models, for which SUN shows considerableimprovement.
Appendix E lists several directions for future work, including an extension of our workto higher-order node-based policies.
Societal impact.
We do not envision any negative, immediate societal impact originating from ourtheoretical results, which represent most of our contribution.
Experimentally, our model has shownpromising results on molecular property prediction tasks and strong generalisation ability in low-dataregimes.
This leads us to believe our work may contribute to positively impactful pharmaceuticalresearch, such as drug discovery [17, 3].