Ordered Subgraph Aggregation Networks
Abstract
Numerous subgraph-enhanced graph neural networks (GNNs) have emerged recently, provably boosting the expressive power of standard (message-passing) GNNs.
However, there is a limited understanding of how these approaches relate to each other and to the Weisfeiler–Leman hierarchy.
Moreover, current approaches either use all subgraphs of a given size, sample them uniformly at random, or use hand-crafted heuristics instead of learning to select subgraphs in a data-driven manner.
Here, we offer a unified way to study such architectures by introducing a theoretical framework and extending the known expressivity results of subgraph-enhanced GNNs.
Concretely, we show that increasing subgraph size always increases the expressive power and develop a better understanding of their limitations by relating them to the established k - WL hierarchy.
In addition, we explore different approaches for learning to sample subgraphs using recent methods for backpropagating through complex discrete probability distributions.
Empirically, we study the predictive performance of different subgraph-enhanced GNNs, showing that our data-driven architectures increase prediction accuracy on standard benchmark datasets compared to non-data-driven subgraph-enhanced graph neural networks while reducing computation time.
Introduction
Graph-structured data are ubiquitous across application domains ranging from chemo- and bioinformatics [Barabasi and Oltvai, 2004, Jumper et al., 2021,
Stokes et al., 2020] to image [Simonovsky and Komodakis, 2017] and social-network analysis [Easley and Kleinberg, 2010].
Numerous approaches for graph–based machine learning have been proposed, most notably those based on graph kernels [Borgwardt et al., 2020, Kriege et al., 2020] or using graph neural networks (GNNs)
[Chami et al., 2020, Gilmer et al., 2017, Morris et al., 2021].
Here, graph kernels based on the 1 -dimensional Weisfeiler–Leman algorithm ( 1 - WL )
[Weisfeiler and Leman, 1968], a simple heuristic for the graph isomorphism problem, and corresponding GNNs [Morris et al., 2019, Xu et al., 2019], have recently advanced the state-of-the-art in supervised vertex- and graph-level learning.
However, the 1 - WL and GNNs operate via local neighborhood aggregation, missing crucial patterns in the given data while more expressive architectures based on the k -dimensional Weisfeiler–Leman algorithm ( k - WL )
[Azizian and Lelarge, 2020, Maron et al., 2019, Morris et al., 2020b, 2021, 2022] may not scale to larger graphs.
Hence, several approaches such as Bevilacqua et al.
[2021], Cotta et al.
[2021], Li et al.
[2020], Papp et al. [2021], Thiede et al.
[2021], You et al.
[2021] and Zhao et al.
[2021] have enhanced the expressive power of GNNs, by removing, extracting, or marking (small) subgraphs, so as to allow GNNs to leverage more structural patterns within the given graph, essentially breaking symmetries induced by the GNNs’ local aggregation function.
We henceforth refer to these approaches as subgraph-enhanced GNNs .
Present work First, to bring some order to the multitude of recently proposed subgraph-enhanced GNNs, we introduce a theoretical framework to study these approaches’ expressive power in a unified setting.
Concretely,
• we introduce k -ordered subgraph aggregation networks ( k - OSAN s) and show that they capture most of the recently proposed subgraph-enhanced GNNs.
• We show that any k - OSAN is upper bounded by ( k + 1) - WL in terms of expressive power and show that k - OSAN s and k - WL are incomparable in terms of expressive power.
Consequently, we obtain new upper bounds on the expressive power of recently proposed subgraph-enhanced GNNs.
• We show that increasing k , i.e., using larger subgraphs, always leads to an increase in expressive power, effectively showing that k - OSAN s form a hierarchy.
Second, most approaches consider all subgraphs or use hand-crafted heuristics to select them, e.g., by deleting vertices or edges.
Instead, we leverage recent progress in back-propagating through discrete structures using perturbation-based differentiation [Domke, 2010, Niepert et al., 2021] to sample subgraphs in a data-driven fashion, automatically adapting to the given data distribution.
Concretely,
• we explore different strategies to sample subgraphs leveraging the I-MLE framework [Niepert et al., 2021], resulting in the data-driven k - OSAN architecture.
• We show, empirically, that data-driven k - OSAN s increase prediction accuracy on standard benchmark datasets compared to non-data-driven subgraph-enhanced GNNs while vastly reducing computation time.
In the following, we discuss related work relevant to the present work; see Appendix A for an extended discussion.
GNNs Recently, GNNs [Gilmer et al., 2017, Scarselli et al., 2009] emerged as the most prominent graph representation learning architecture.
Notable instances of this architecture include, e.g., Duvenaud et al.
[2015], Hamilton et al.
[2017] and Veliˇckovi´c et al.
[2018], which can be subsumed under the message-passing framework introduced in Gilmer et al.
In parallel, approaches based on spectral information were introduced in, e.g., Defferrard et al.
[2016], Bruna et al.
[2014], Kipf and Welling [2017] and Monti et al.
[2017]—all of which descend from early work in Baskin et al.
[1997], Kireev [1995], Micheli and Sestito [2005], Merkwirth and Lengauer [2005], Micheli [2009], Scarselli et al. [2009] and Sperduti and Starita [1997].
Limits of GNNs and more expressive architectures Recently, connections between GNNs and
Weisfeiler–Leman type algorithms have been shown [Azizian and Lelarge, 2020, Barceló et al., 2020, Chen et al., 2019b, Geerts et al., 2020, Geerts, 2020, Geerts and Reutter, 2022, Maehara and NT, 2019, Maron et al., 2019, Morris et al., 2019, 2022, Xu et al., 2019].
Specifically, Morris et al.
[2019] and Xu et al.
[2019] showed that the expressive power of any possible GNN architecture is limited by the 1 - WL in terms of distinguishing non-isomorphic graphs.
Recent works have extended the expressive power of GNNs, e.g., by encoding vertex identifiers [Murphy et al., 2019, Vignac et al., 2020], using random features [Abboud et al., 2020, Dasoulas et al., 2020, Sato et al., 2020], homomorphism and subgraph counts [Barceló et al., 2021, Bouritsas et al., 2020, NT and
Maehara, 2020], spectral information [Balcilar et al., 2021], simplicial and cellular complexes [Bodnar et al., 2021b,a], persistent homology [Horn et al., 2021], random walks [Tönshoff et al., 2021], graph decompositions [Talak et al., 2021], or distance [Li et al., 2020] and directional information [Beaini et al., 2020].
See Morris et al.
[2021] for an in-depth survey on this topic.
Subgraph-enhanced GNNs Most relevant to the present work are subgraph-enhanced GNNs .
Cotta et al.
[2021] and Papp et al.
[2021] showed how to make GNNs more expressive by removing one or more vertices from a given graph and using standard GNN architectures to learn vectorial representations of the resulting subgraphs.
The approaches either consider all possible subgraphs or utilize random sampling to arrive at more scalable architectures.
Cotta et al.
[2021] showed that by removing one or two vertices, such architectures can distinguish graphs the 1 - WL and 2 - WL , respectively, are not able to distinguish.
Extensions and refinements of the above were proposed in Bevilacqua et al.
[2021], Papp and Wattenhofer [2022], Thiede et al.
[2021], You et al.
[2021], Zhang and Li [2021] and Zhao et al.
[2021], see Papp and Wattenhofer [2022] for an overview.
For example, Bevilacqua et al.
[2021] generalized several ideas discussed above and proposed the ESAN framework in which each graph is represented as a multiset of its subgraphs and processed them using an equivariant architecture based on the DSS architecture [Maron et al., 2020] and GNNs.
The authors proposed several simple subgraph selection policies, e.g., edge removal, ego networks, or vertex removal, and showed that the architecture surpasses the expressive power of the 1 - WL . Moreover, Frasca et al.
[2022] presented a novel symmetry analysis unifying a series of subgraph-enhanced GNNs, allowing them to upper-bound their expressive power and to define a systematic framework to conceive novel architectures in this family.
We note here that the above works, unlike the present one, mostly do not study the approaches’ expressive power beyond (folklore or non-oblivious) 2 - WL and do not compare at all to the (folklore or non-oblivious) 3 - WL , while our analysis works for the whole k - WL hierarchy.
See Appendix A for a detailed overview of recent progress in differentiating through discrete structures.
Preliminaries
As usual, for n ≥ 1 , let [ n ] := { 1 , . . . , n } ⊂ N .
We use {{ . . . }} to denote multisets, i.e., the generalization of sets allowing for multiple instances of each of its elements.
A graph G is a pair ( V ( G ) , E ( G )) with finite sets of vertices V ( G ) and edges E ( G ) ⊆ {{ u, v } ⊆
V ( G ) | u ̸ = v } .
If not otherwise stated, we set n :=
| V ( G ) | .
For ease of notation, we denote the edge { u, v } in E ( G ) by ( u, v ) or ( v, u ) .
In the case of directed graphs , E ⊆ { ( u, v ) ∈ V × V | u ̸ = v } .
A labeled graph G is a triple ( V, E, l ) with (vertex) coloring or label function l : V ( G ) → N .
Then l ( v ) is a label of v for v in V ( G ) .
The neighborhood of v in G is denoted by N G ( v ) :
= { u ∈ V ( G ) | { v, u } ∈ E ( G ) } and the degree of a vertex v is | N G ( v ) | .
For S ⊆ V ( G ) , the graph G [ S ]
= ( S, E S ) is the subgraph induced by S , where E S :
= { ( u, v )
∈ E ( G ) | u, v ∈ S } .
Two graphs G and H are isomorphic and we write G ≃ H if there exists a bijection φ : V ( G )
→ V ( H ) that preserves the adjacency relation, i.e., ( u, v ) is in E ( G ) if and only if ( φ ( u ) , φ ( v )) is in E ( H ) .
Then φ is an isomorphism between G and H .
Moreover, we call the equivalence classes induced by the relation ≃ isomorphism types .
In the case of labeled graphs, we additionally require that l ( v ) =
l ( φ ( v )) for v in V ( G ) .
We further define the atomic type atp : V ( G ) k
→ N such that atp ( v ) =
atp ( w ) for v , w ∈ V ( G ) k
if and only if the mapping φ : V ( G ) k
→ V ( G ) k where v i
(cid:55)→ w i induces a partial isomorphism, i.e., we have v i =
v j ⇐⇒ w i =
w j and ( v i , v j
) ∈ E ( G )
⇐⇒ ( φ ( v i ) , φ ( v j ))
∈ E ( G ) .
Let v be a tuple in V ( G ) k for k > 0 , then G [ v ] is the ordered k -vertex subgraph induced by the multiset of elements of v , where the vertices are labeled with integers from [ k ] corresponding to their positions in v .
Moreover, let t ( G [ v ]) :
= v , i.e., the k -tuple v underlying the ordered k -vertex subgraph G [ v ] .
We denote the set of all ordered k -vertex subgraphs of a graph G by G k .
Finally, let G be the set of all vertex-labeled graphs.
The 1 - WL and the k - WL
The 1 - WL or color refinement is a simple heuristic for the graph isomorphism problem, originally proposed by Weisfeiler and Leman [1968]. 2 Intuitively, the algorithm determines if two graphs are non-isomorphic by iteratively coloring or labeling vertices.
Given an initial coloring or labeling of the vertices of both graphs, e.g., their degree or application-specific information, in each iteration, two vertices with the same label get different labels if the number of identically labeled neighbors is not equal.
If, after some iteration, the number of vertices annotated with a specific label is different in both graphs, the algorithm terminates and a stable coloring (partition) is obtained.
We can then conclude that the two graphs are not isomorphic.
It is easy to see that the algorithm cannot distinguish all non-isomorphic graphs [Cai et al., 1992].
Nonetheless, it is a powerful heuristic that can successfully test isomorphism for a broad class of graphs [Babai and Kucera, 1979].
Formally, let G =
( V, E, l ) be a labeled graph.
In each iteration, i > 0 , the 1 - WL computes a vertex coloring C 1 i : V ( G ) → N , which depends on the coloring of the neighbors.
That is, in iteration i > 0 , we set
[Equation]
where RELABEL injectively maps the above pair to a unique natural number, which has not been used in previous iterations.
In iteration 0 , the coloring C 10 := l .
To test if two graphs G and H are non-isomorphic, we run the above algorithm in “parallel” on both graphs.
If the two graphs have a different number of vertices colored c in N at some iteration, the 1 - WL distinguishes the graphs as non-isomorphic.
Moreover, if the number of colors between two iterations, i and ( i + 1) , does not change, i.e., the cardinalities of the images of C 1 i and C 1 i +1 are equal, the algorithm terminates.
For such i , we define the stable coloring
C 1 ∞ ( v ) = C 1 i ( v ) for v in V ( G ) . The stable coloring is reached after at most max {| V ( G ) | , | V ( H ) |} iterations [Grohe, 2017]. Due to the shortcomings of the 1 - WL or color refinement in distinguishing non-isomorphic graphs, several researchers [Babai, 1979, 2016, Immerman and Lander, 1990], devised a more powerful generalization of the former, today known as the k -dimensional Weisfeiler-Leman algorithm ( k - WL ); see Appendix B for a detailed description.
Graph Neural Networks Intuitively, GNNs learn a vectorial representation, i.e., a d -dimensional vector, representing each vertex in a graph by aggregating information from neighboring vertices.
Formally, let
G =
( V, E, l ) be a labeled graph with initial vertex features h (0) v ∈ R d that are consistent with l .
That is, each vertex v is annotated with a feature h (0) v ∈ R d such that h (0) u
= h (0) v if l ( u )
= l ( v ) , e.g., a one-hot encoding of the labels l ( u ) and l ( v ) .
Alternatively, h (0) v can be an arbitrary real-valued feature vector or attribute of the vertex v , e.g., physical measurements in the case of chemical molecules.
A GNN architecture consists of a stack of neural network layers, i.e., a composition of parameterized functions.
Similarly to 1 - WL , each layer aggregates local neighborhood information, i.e., the neighbors’ features, around each vertex and then passes this aggregated information on to the next layer.
Following, Gilmer et al.
[2017] and Scarselli et al.
[2009], in each layer, i > 0 , we compute vertex features
[Equation]
where UPD ( i +1) and AGG ( i +1) may be differentiable parameterized functions, e.g., neural networks.
In the case of graph-level tasks, e.g., graph classification, one uses [Equation] to compute a single vectorial representation based on learned vertex features after iteration T . Again,
READOUT may be a differentiable parameterized function. To adapt the parameters of the above three functions, they are optimized end-to-end, usually through a variant of stochastic gradient descent, e.g., [Kingma and Ba, 2015], together with the parameters of a neural network used for classification or regression. The Weisfeiler–Leman hierarchy and permutation-invariant function approximation The Weisfeiler– Leman hierarchy is a purely combinatorial algorithm for testing graph isomorphism.
However, the graph isomorphism function, mapping non-isomorphic graphs to different values, is the hardest to approximate permutation-invariant function.
Hence, the Weisfeiler–Leman hierarchy has strong ties to
GNNs’ capabilities to approximate permutation-invariant or equivariant functions over graphs.
For example, Morris et al.
[2019], Xu et al.
[2019] showed that the expressive power of any possible GNN architecture is limited by 1 - WL in terms of distinguishing non-isomorphic graphs.
Azizian and Lelarge [2020] refined these results by showing that if an architecture is capable of simulating k - WL and allows the application of universal neural networks on vertex features, it will be able to approximate any permutation-equivariant function below the expressive power of k - WL ; see also Chen et al.
[2019b].
Hence, if one shows that one architecture distinguishes more graphs than another, it follows that the corresponding GNN can approximate more functions.
These results were refined in Geerts and Reutter [2022] for color refinement and taking into account the number of iterations of k - WL .
Ordered subgraph Weisfeiler–Leman and MPNNs
In the following, we introduce a variant of 1 - WL , denoted k -ordered subgraph WL ( k - OSWL ).
Essentially, the k - OSWL labels or marks ordered subgraphs and then executes 1 - WL on top of the marked graphs, followed by an aggregation phase.
Although unordered subgraphs are also possible, ordered ones lead to more expressive architectures and also encompass the unordered case; see Appendix F.1 for a discussion.
To make the procedure permutation-invariant, we consider all possible ordered subgraphs.
Based on the ideas of k - OSWL , we then introduce k - OSAN s, which can be seen as a neural variant of the former, allowing us to analyze various subgraph-enhanced GNNs.
3.1 Ordered subgraph WL
We now describe the algorithm formally.
Let G be a graph, and let g ∈ G k be an ordered k -vertex subgraph.
Then k - OSWL computes a vertex coloring, similarly to 1 - WL , with the main distinction that it can use structural graph information related to the ordered subgraph G [( v, t ( g ))] , where v is a vertex in V ( G ) .
More precisely, at each iteration i ≥ 0 , k - OSWL computes a coloring C i : V ( G )
× G k
→ N where we interpret elements ( v, g
) ∈ V ( G ) ×
G k as a vertex v along with an ordered k -vertex subgraph.
Given an ordered k -vertex subgraph g ∈ G k , initially, for i = 0 , we set C {+
+}
( v, g ) :=
atp ( v, t ( g )) , and for i > 0 , we set
[Equation]
where □ is either N G ( v ) or V ( G ) .
We compute the stable partition analogously to 1 - WL .
Finally, to compute a single color for a vertex v , we aggregate all ordered k -vertex subgraphs, i.e., we compute
[Equation]
In other words, one can regard the k - OSWL as running 1 - WL in parallel over n k graphs, one for each ordered k -vertex subgraph g ∈ G k , followed by combining the colors of each vertex in all these graphs.
Furthermore, by restricting the number of considered subgraphs, the algorithm allows for more fine-grained control over the trade-off between scalability and expressivity.
Note that 0 - OSWL is equal to 1 - WL .
We also define a variation of the k - OSWL , denoted vertex-subgraph k - OSWL , which, unlike Equation (2), first computes a color C ( g ) for each ordered k -vertex subgraph g by aggregating over vertices; see Appendix E for details.
We remark that in contrast to k - WL , which has to update the coloring of all n k ordered k -vertex subgraphs in a complicated manner, the computation of k - OSWL ’s coloring relies on the simple and easy-to-implement 1 - WL .
Furthermore, k - OSWL ’s computation can be either done in parallel or sequentially across all n vertices and n k graphs.
Despite its simplicity, in Section 3.2, we will show that the k - OSWL has high expressivity.
3.2 Ordered subgraph MPNNs
In the following, to study the expressivity of subgraph-enhanced GNNs, we introduce k -ordered subgraph MPNNs ( k - OSAN s), which can be viewed as neural variants of the k - OSWL .
At initialization, k - OSAN s learn two features for each element in G k and each vertex v
[Equation]
and
[Equation]
where UPD and UPD πππ are differentiable, parameterized function, e.g., a neural network.
Additional vertex features can be concatenated to the first feature.
We use the second feature πππ v, g to select admissible ordered subgraphs for the vertex v ; see below.
Now in each layer ( i + 1) , we update the feature of a vertex v with regard to the k -ordered subgraph g as
[Equation]
where □ is either N G ( v ) or V ( G ) .
After T such layers, for each vertex v , we then learn a joint feature over all k -ordered subgraphs, i.e., we apply subgraph aggregation
[Equation]
Here, we leverage πππ v, g ̸ = 0 to select a subset of the set of k -ordered subgraphs.
Finally, analogous to GNNs, we use a READOUT layer to compute a single graph feature.
Again, AGG ( i +1) , UPD ( i +1) , READOUT , and SAGG are differentiable, parameterized functions, e.g., neural networks.
We also define a variation of k - OSAN s, denoted vertex-subgraph k - OSAN s , which, unlike Equation (3), first compute a color for each ordered k -vertex subgraph; see Appendix E for details.
Expressive power of k - OSAN s
In the following, we study the expressive power of k - OSAN s. The first result shows that any possible k - OSAN has at most the expressive power of k - OSWL in terms of distinguishing non-isomorphic graphs.
Further, k - OSAN s are in principle capable of reaching k - OSWL ’s expressive power.
Hence, the k - OSWL upper bounds k - OSAN s ability to represent permutation-invariant functions.
Proposition 1.
For all k ≥ 1 , it holds that k - OSAN s are upper bounded by k - OSWL in terms of distinguishing non-isomorphic graphs.
Further, there exists a k - OSAN instance that has exactly the same expressive power as the k - OSWL .
The following result shows that any possible k - OSAN is upper-bounded by the ( k + 1) - WL in terms of distinguishing non-isomorphic graphs while the expressive power of k - OSAN s and the k - WL are incomparable.
That is, there exist non-isomorphic graphs k - WL cannot distinguish while k - OSAN s can and vice versa.
Proposition 2.
For all k ≥ 1 , it holds that ( k + 1) - WL is stricly more expressive than k - OSAN s and there exist non-isomorphic graphs k - WL cannot distinguish while k - OSAN s can and vice versa.
Finally, the following results shows that increasing the size of the subgraphs always leads to a strictly more expressive k - OSAN s.
Theorem 3.
For all k ≥ 1 , it holds that ( k + 1) - OSAN s is strictly more expressive than k - OSAN s.
Subgraph-enhanced GNNs captured by k - OSAN s To exemplify the power and generality of k - OSAN s, we show how k - OSAN s cover most subgraph-enhanced GNNs; see Appendix F for a thorough overview.
We say that k - OSAN s capture a subgraph-enhanced GNN G if there exists a k - OSAN instance that is at least as expressive as G .
The first results shows that k - OSAN s capture k -marked GNNs ( k - mGNN s) [Papp and Wattenhofer, 2022] and k -reconstruction GNNs ( k - recGNN s) [Cotta et al., 2021].
For both approaches, the sets of k vertices to be marked or deleted correspond to unordered k -vertex subgraphs.
It then suffices to ensure that the update and aggregation functions in the k - OSAN s treat the vertices in the selected k -vertex subgraphs as being marked or deleted.
Proposition 4.
For k ≥ 1 , k - OSAN s capture k - mGNN s and vertex-subgraph k - OSAN s capture k - recGNN s.
Further, 1 - OSAN s capture identity-aware GNNs ( idGNN s) [You et al., 2021], GNN As Kernel ( kernelGNN s) [Zhao et al., 2021], and nested GNNs ( nestedGNN s) [Zhang and Li, 2021].
Intuitively, in these approaches GNNs are used locally around each vertex.
It thus suffices to ensure that the update and aggregation functions in the 1 - OSAN s use the selected single vertex subgraph to only pass messages locally.
Proposition 5.
1 - OSAN s capture idGNN s, kernelGNN s, and nestedGNN s.
Finally, k - OSAN s capture the DS - GNN s with the vertex-deleted policy [Bevilacqua et al., 2021].
Proposition 6.
Vertex-subgraph k - OSAN s capture DS - GNN s with the k -vertex-deleted policy.
We note that the above result can be further extended to accommodate the edge-deleted and ego-networks policy from Bevilacqua et al.
[2021]; see Appendix F.
Importantly, by viewing existing subgraph-enhanced GNNs as k - OSAN s we immediately gain insights into their expressive power.
Previous results primarily focused on showing more expressivity than 1 - WL .
Data-driven Subgraph-enhanced GNNs
In the above section, we thoroughly investigated the expressive power of subgraph-enhanced GNNs.
Specifically, we showed that they are strictly limited by the ( k + 1) - WL and that they can distinguish graphs which are not distinguishable by k - WL .
As indicated by Proposition 1 and to reach maximal expressive power, however, we need to consider all possible ordered subgraphs, resulting in an exponential running time.
Hence, in this section, we leverage the I-MLE framework [Niepert et al., 2021], to sample ordered subgraphs in a data-driven fashion.
We first address the problem of learning the parameters of a probability distribution over ordered subgraphs using a GNN. Secondly, we show how to approximately sample from this intractable distribution.
Subsequently, these subgraphs are used within a k - OSAN to compute a graph representation.
Finally, we propose a gradient estimation scheme that allows us to use backpropagation in the resulting discrete-continuous architecture.
Parameterizing probability distributions over subgraphs Contrary to existing approaches, which often consider all possible subgraphs or sample a fraction of subgraphs uniformly at random, our method maintains a probability distribution over (ordered) subgraphs.
Let G be a graph where each vertex v has an initial feature h (0) v , which we stack row-wise over all vertices into the feature matrix H ∈ R n × d .
Further, let h W 1 : G × R n × d → R m × n be a permutation-equivariant function, e.g., a message-passing GNN, parameterized by W 1 , called upstream model , mapping a graph G and its initial features H to a parameter matrix
[Equation]
Intuitively, each parameter θ ij is an unnormalized prior probability of vertex j being part of the i th sampled subgraph of G .
Let θ i :=
( θ i 1 , . .
, θ in ) for i ∈ [ m ] . We use these to parameterize m probability distributions p ( z ; θ i ) , for i ∈ [ m ] , over vector encodings of ordered k -vertex subgraphs of G , i.e., p
[Equation]
θ
[Equation]
where ⟨· , ·⟩ is the standard inner product and A ( θ i ) is the log-partition function defined as A ( θ i ) :
= log (cid:0)(cid:80)
z ∈Z exp ( ⟨ z , θ i ⟩ ) (cid:1) .
Furthermore, for a distribution over unordered k -vertex subgraphs, Z is the set of all binary n -component vectors with exactly k non-zero entries indicating which vertices are part of a subgraph of G .
Hence, there is a bijection between Z and the set of unordered k -vertex subgraphs of G .
For a distribution over ordered k -vertex subgraphs, the set Z is the set of all vectors in [ k ] n with k non-zero entries.
For each non-zero entry z i for z ∈ Z it holds that z i
=
j if and only if vertex i has rank k + 1 − j in the ordered subgraph, encoding the position in the ordered graph.
This encoding is required for the gradient computation we perform later.
For instance, in an ordered 5 -vertex subgraph, if a node has rank 1 but should have rank 5 to obtain a lower loss, then the gradient of the downstream loss is proportional to 5 − 1 = 4 .
{+Similarly, if a node i is not part of the ordered subgraph, that is, z i = 
+}
but should be in position 1 of the ordered subgraph, then the gradient of a downstream loss is proportional to 0 − 5 .
For any i, j ∈ [ k ] with i ̸
=
j we have that z i ̸ =
z j .
Hence, again, there is a bijection between Z and the set of ordered k -vertex subgraphs of G .
Efficient approximate sampling of subgraphs Computing the log-partition function and sampling exactly from the probability distribution in Equation (4) is intractable for both ordered and unordered graphs.
Since it is tractable, however, to compute a configuration with a highest probability, a maximum a posteriori (MAP) configuration z ∗ ( θ i ) , we can use perturb-and-MAP
[Papandreou and Yuille, 2011, Niepert et al., 2021] to sample approximately.
For unordered graphs, determining the top- k values in θ i suffices, while for ordered graphs, we additionally require their rank.
Therefore, the worst-case running time of computing z ∗ ( θ i ) for ordered graphs of size k is in O ( n + k log k ) .
That is, we first use a selection algorithm to find the k th largest element E in the list of weights, taking time O ( n ) , e.g., using the Quickselect algorithm.
Now, we go through the list again and select all entries larger to E , taking time O ( n ) .
Finally, we sort the k values.
Now, to use perturb-and-MAP to approximately sample the i -th ordered k -vertex
subgraph g i from the above probability distributions, we compute
[Equation] with [Equation] where ρ ( ϵ ) is a noise distribution such as the Gumbel distribution and adj converts the above vector encoding of the (ordered) subgraph to an n × n adjacency matrix as follows.
The j th row or column encodes the vertex of the ordered subgraph with rank j and its incident edges within the ordered subgraph.
All other entries are set to 0 , i.e., they are masked out.
We therefore sample a multiset of (ordered) subgraphs S :
= {{ g 1 , ..., g m }} ⊆
G k , which act as the input to a k - OSAN instance f W 2 , called downstream model , where πππ v, g ̸ = 0 for v ∈ V ( G ) if g ∈ S , to compute the target outputs [Equation]
Backpropagating through the subgraph distribution Now that we have outlined a way to approximately sample subgraphs, we still need to learn the parameters ω
= ( W 1 , W 2 ) of the upstream and downstream model.
Hence, given a set of examples { ( G j , H j , ˆ y j ) } Nj =1 , we are concerned with finding approximate solutions to min ω 1 / N (cid:80) j L ( G, H , ˆ y j ; ω ) , where L is the expected training error
[Equation]
with θ := h W 1 ( G, H ) and ℓ : Y ×
Y → R + is a point-wise loss function.
The challenge of training a model as defined in Equation (5) is to compute ∇ θ L ( G, H , ˆ y j ; ω ) , i.e., the gradient with respect to the parameters θ of the probability distribution for the expected loss.
In this work, we utilize implicit maximum likelihood learning, a recent framework that allows us to use algorithmic solvers of combinatorial optimization problems as black-box components [Rolinek et al., 2020, Niepert et al., 2021].
A particular instance of the framework uses implicit differentation via perturbation [Domke, 2010].
We derive the gradient computation for a single ordered subgraph g i to simplify the notation.
We compute the gradients of the downstream loss with respect to parameters θ i as
[Equation]
where λ > 0 and (cid:98) ∇ v , the approximated gradient for a single vertex v , is defined as
[Equation]
with v ∈ V ( G ) .
Here, agg can be any aggregation function such as the element-wise sum or mean.
Hence, to approximate the gradients with respect to the parameter θ iv , which corresponds to vertex v of the input graph, we aggregate the gradients of the downstream loss with respect to all edges ( v, w ) incident to vertex v in the original graph.
Hence, the above techniques allow us to efficiently learn to sample subgraphs, which are then fed into a k - OSAN to learn a graph representation while optimizing the parameters of all components in an end-to-end fashion.
Experimental evaluation
Here, we aim to empirically investigate the learning performance and efficiency of data-driven subgraphenhanced GNNs, instances of the k - OSAN framework, compared to non-data-driven ones.
Specifically, we aim to answer the following questions.
Q1 Do data-driven subgraph-enhanced GNNs exhibit better predictive performance than non-data-driven ones?
Q2 Does the graph structure of the subgraphs sampled affect predictive performance?
Q3 Does data-driven sampling have an advantage in efficiency and predictive performance when used within state-of-the-art subgraph-enhanced GNNs?
All experimental results are fully reproducible from the source code provided at https://github.com/
Spazierganger/OSAN Datasets To compare our data-driven, subgraph-enhanced GNNs to non-data-driven ones and standard
GNN baselines, we used the A LCHEMY
[Chen et al., 2019a], the Q M 9 [Ramakrishnan et al., 2014, Wu et al., 2018], O GBG -M OLESOL
[Hu et al., 2020], and the Z INC [Dwivedi et al., 2020, Jin et al., 2017] graph-level regression datasets; see Table 16 in Appendix C for dataset statistics and properties.
In addition, we used the E XP dataset [Abboud et al., 2020] to investigate the additional expressive power of subgraph-enhanced GNNs over standard ones.
Following Morris et al.
[2020b], we opted not to use the 3D-coordinates of the A LCHEMY dataset to solely show the benefits of the data-driven subgraph-enhanced [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
GNNs regarding graph structure.
All datasets, excluding E XP and O GBG -M OLESOL , are available from Morris et al.
[2020a].
Neural architectures and experimental protocol For all datasets and architectures, we used the competitive GIN layers [Xu et al., 2019] for the baselines and the downstream models.
For data with (continuous) edge features, we used a 2 -layer MLP to map them to the same number of components as the vertex features and combined them using summation.
We describe the upstream and downstream models’ architecture used for each dataset in the following.
We stress here that we always used the same hyperparameters for the downstream model and the baselines.
Sampling subgraphs Since the number of unordered k -vertex subgraphs is considerably smaller than the number of ordered k -vertex subgraphs, we opted to consider unordered k -vertex subgraphs; see also Appendix F.1.
Further, since vertex-subgraph k - OSAN s, see Appendix E, are easier to implement efficiently and are closer to DS - GNN s variant of ESAN [Bevilacqua et al., 2021], we opted to use them for the empirical evaluation.
In addition, we used a simple GNN architecture for the upstream model to compute initial features for the subgraphs for ease of implementation.
We experimented with selecting and deleting a various number of vertices, edges, and subgraphs induced by k -hop neighborhoods ( k -Ego) for all datasets; see Appendix C for details.
Upstream models For all datasets and experiments, we used a GCN model [Kipf and Welling, 2017] consisting of three GCN layers, with batch norm and ReLU activation after each layer.
We set the hidden dimensions to that of the downstream model one.
The model either outputs the vertex or edge embeddings according to the task.
We computed edge embeddings based on the vertex features of the incident vertices after the last layers and the edge attributes provided by the dataset.
When sampling multiple subgraphs with I-MLE, they tend to have similar structures.
In other words,
I-MLE learns similar distributions in different channels of the neural network.
This phenomenon is not in our favor, as we need to cover the original full graph as much as possible.
To mitigate this issue, we propose an auxiliary loss for the diversity of subgraphs. We calculate the cosine similarity between the selected vertex or edge masks of every two subgraphs and try to minimize the average similarity value. We tune the weight for the auxiliary loss on the log scale, e.g., 0 . 1 , 1 , 10 , and so on. Downstream and baseline models See Appendix A for a detailed description of the architecture used for the downstream and baseline models, and how we processed subgraphs.
For processing the subgraphs, we performed similar steps like ESAN .
We first applied intra-subgraph aggregation for the vertices within each subgraph and obtained graph embeddings for each subgraph.
After that, we performed inter-subgraph mean pooling to obtain a single embedding vector for the original graph.
It is worth noting that ESAN does not exclude the vertices deleted during graph pooling but removes the adjacent edges of those nodes.
In our experiments, we masked out the deleted or unselected nodes.
See Appendix C for further details on the experiments.
5.1 Results and discussion
In the following, we answer the research questions Q1 to Q3 .
A1 See Tables 1, 2a, 4 and 5 (in the appendix).
On all five datasets, the subgraph-enhanced GNN models based on I-MLE beat the random baseline, excluding edge sampling configurations on the A LCHEMY and the Q M 9 dataset; see Tables 1b and 2a.
For example, on the O GBG -M OLESOL the average gain over the random baseline is over 11%.
Similar improvements can be observed over the other four datasets.
Moreover, the results on the E XP dataset, see Appendix C, clearly indicate that the added expressivity of the (data-driven) subgraph-enhanced GNNs translates into improved predictive performance.
The data-driven subgraph-enhanced GNNs improve the accuracy of the non-subgraph-enhanced GNN by almost 50% in all configurations while improving over the random subgraph-enhanced GNN baseline by almost 6%.
The data-driven subgraph-enhanced GNNs also clearly improve over the (non-subgraph-enhanced)
GNN baseline on four out of five datasets.
A2 See Tables 1 and 2a (in the appendix).
Deleting or selecting subgraphs leads to a clear boost in predictive performance across datasets over the random baseline while also improving over the non-subgraph-enhanced
GNN baseline.
Further, on all datasets, learning to delete or select k -hop neighborhood subgraphs for k ∈ { 2 , 3 } leads to a clear boost over the random as well as non-subgraph-enhanced baselines.
However, the number of deleted vertices seems to affect the predictive performance.
For example, on the O GBG -M OLESOL dataset, going from deleting one 2-vertex subgraph to one 10-vertex subgraph leads to a drop in performance.
Hence, the drop in performance of the latter is in contrast to our theoretical findings, i.e., larger subgraphs lead to improved expressivity, indicating that more work should be done to understand subgraph-enhanced GNNs’ generalization ability.
Interestingly, deleting edges did not perform as well as deleting vertices or other subgraphs.
We speculate that a more powerful edge embedding method is needed here, which computes edge features directly instead of learning them from vertex features.
A3 See Table 2b (in the appendix).
The I-MLE-based ESAN severely speeds up the computation time.
That is, across all configurations, we achieve a significant speed-up.
For some configurations, e.g., sampling three vertices, the I-MLE based ESAN is more than 3.5 times faster than the non-data-driven ESAN while taking about the same time as the simple random baseline.
We stress here that the ESAN implementation provided by Bevilacqua et al.
[2021] precomputes subgraphs in a preprocessing step, which is not possible when learning to sample subgraphs using ESAN .
Regarding predictive performance, the I-MLE based ESAN is slightly behind the non-data-driven one, although always better than the non-subgraph enhanced GNN baseline; see Table 5 in Appendix C.
Conclusion
We introduced the k - OSAN framework to study the expressive power of recently introduced subgraphenhanced GNNs.
We showed that any such architecture is strictly less powerful than the ( k + 1) - WL while being incomparable to the k - WL in representing permutation-invariant functions over graphs.
Further, to circumvent random or heuristic subgraph selection, we devised a data-driven variant of k - OSAN s which learn to select subgraph for a given data distribution.
Empirically, we verified that such data-driven subgraph selection is superior to previously used random sampling in predictive performance.
Further, when compared to state-of-the-art models, we showed promising performance in terms of computation time while still providing good predictive performance.
We believe that our paper provides a first step in unifying combinatorial insights on the expressive power of GNNs with data-driven insights.
Acknowledgments and Disclosure of Funding
CM is partially funded a DFG Emmy Noether grant (468502433) and RWTH Junior Principal Investigator
Fellowship under the Excellence Strategy of the Federal Government and the Länder.
GR is funded by the
DFG Research Grants Program–RA 3242/1-1–411032549.
MN acknowledges funding by the German
Research Foundation under Germany’s Excellence Strategy–