Model-Agnostic Graph Dataset Compression with the Tree Mover’s Distance

Authors: mika  sarkin jain
stefanie  jegelka
ishani  karmarkar
luana  ruiz
ellen  vitercik
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-03 01:27:09

Abstract: Graph neural networks have demonstrated remarkable success across a variety of domains. However, the acquisition and management of largescale graph datasets poses several challenges. Acquiring graph-level labels can be prohibitively costly, especially for applications in the biosciences and combinatorial optimization. Storage and privacy constraints can pose additional challenges. In this work, we propose an approach for data subset selection for graph datasets, which downsamples graphs and nodes based on the Tree Mover’s Distance. We provide new efficient methods for computing the TMD in our setting; empirical results showing our approach outperforms other node and graph sampling methods; and theoretical results bounding the decrease in accuracy caused by training on the downsampled graphs. Surprisingly, we find that with our method, we can subsample down to 1% of the number of graphs and 10% of the number of nodes on some datasets, with minimal degradation in model accuracy.

Main paper: 
 introduction: 
 Graph neural networks (GNNs) are a popular architecture for learning over graph-structured data. Despite their widespread popularity and demonstrated success (Bongini et al., 2022; Zhou et al., 2020; Peng et al., 2021; Fan et al., 2019; Peng et al., 2021; Zhang & Chen, 2018) , training GNNs remains challenging because the datasets can be massive, both in terms of the size (number of nodes) of each graph, as well as the number of graphs in the dataset.Large-scale datasets can pose many challenges to training GNNs. For instance, training may be computationally intractable (Hashemi et al., 2024; Yan et al., 2020; Zeng et al., 2021) . First of all, obtaining supervision signals can be Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024). expensive on large graphs or on datasets consisting of large graphs. This is the case, for example, when GNNs are used to model molecular properties for computational drug discovery (Igarashi et al., 2024; Shen et al., 2020; Jiang et al., 2021) . Likewise, storing and loading massive graphs into memory can be slow and/or expensive (Huang et al., 2024; Hamilton et al., 2017; Da San Martino et al., 2012) or may raise privacy concerns (Wu et al., 2022) . GNNs have also been employed to circumvent expensive simulation, e.g., in molecular dynamics (Gasteiger et al., 2021; Park et al., 2021; Li et al., 2022) , where the cost of labeling larger instances grows quickly with graph size. Similarly, GNNs have been applied for disease prediction (Sun et al., 2020; Zheng et al., 2022) , in which case labeling a single graph may require expensive clinical or diagnostic procedures. Additionally, there is growing interest in using GNNs to solve challenging or costly optimization problems, in which case the runtime to collect training labels grows with the graph size (e.g., Bengio et al., 2021; Cappart et al., 2023; Vesselinova et al., 2020) .A natural approach to tackle these multifaceted challenges is subsampling. In node or graph subsampling, we aim to reduce the number of nodes or graphs, respectively, in a graph dataset while preserving performance of a downstream GNN trained on the dataset. Existing work on node and graph subsampling focuses primarily on node-level tasks (e.g., node classification), and often makes strong assumptions about the GNN model architecture-and can even assume access to the final model or the ability to approximate the training trajectories on the original, unsampled dataset.In this paper, we propose and theoretically motivate new approaches for (1) subsampling nodes from each training graph and (2) subsampling entire graphs in the training set for graph-level learning tasks (e.g., graph classification and Given a graph budget k, the goal is to find a sample of k graphs that minimizes the total distance (TMD) from every original graph to its closest sample. regression). We must ensure that training a GNN on the subsampled data has minimal accuracy reduction compared to training on the original training set. Moreover, in practice, we often do not know a priori (i.e., before collecting, compressing, and labeling our graph dataset) the architecture or hyperparameters of the GNN. Consequently, it is crucial that the original and subsampled datasets maintain model-agnostic similarity, i.e., they produce similar model behavior regardless of the downstream model.For many GNNs, specific structural similarities-tied to the GNN's computations-are essential to successfully transfer a GNN out of distribution (Yehudai et al., 2021; Chuang & Jegelka, 2022; Ruiz et al., 2021; Levie et al., 2021; Le & Jegelka, 2023) . A GNN uses message-passing to compute predictions about graphs. Initially, each node has an embedding, which it updates by aggregating its neighbors' embeddings and processing the result through a neural network. This process repeats over L iterations-or layersand the nodes' final embeddings are used to compute the GNN's predictions. The L-hop neighbors involved in a node's final embedding can be visualized via the node's computation tree, illustrated in Figure 3 . Computation trees play a critical role in the transfer and stability of GNNs (Yehudai et al., 2021; Chuang & Jegelka, 2022; Georgiev et al., 2022) . Two nodes with similar computation trees will likely have similar embeddings. To capture this intuition, Chuang & Jegelka (2022) introduced a graph distance-the Tree Mover's Distance (TMD)-as an optimal transport distance between sets of computation trees, and demonstrated its correlation with the performance and stability of GNNs under data perturbations, including generalization bounds.Based on the strong theoretical foundations provided by TMD, our primary goal is to ensure that the TMD between the original and subsampled graphs is small. However, achieving this goal requires computing the TMD between many graphs while subsampling, which poses a significant challenge. TMD is a combinatorially complex function: the dynamic programming algorithm developed in prior work for computing the TMD between graphs with n nodes has a runtime of O(Ln 4 ) (Chuang & Jegelka, 2022) . Given our objective to reduce the size of large graphs, this runtime is prohibitively large and presents a computational challenge that we overcome in our methodological contributions.Contributions. In this work, we make the following contributions: We present an approach for node subsampling that, given a budget k and a training graph, returns a subgraph on k nodes with low TMD from the original graph (Figure 1 ). Moreover, we present an approach for graph subsampling that, given a training set of graphs, returns k graphs such that the total TMD between each original graph and its closest sampled graph is small (Figure 2 ).We show that if h is a GNN with minimal loss on the training set subsampled using our methods, then h also has minimal loss on the original training set. When subsampling nodes in each training graph, we bound the resulting increase in loss by the average TMD between the original and subsampled graphs. When subsampling entire graphs, we bound the increase in loss by the average TMD between each training graph and the nearest graph in the subsampled set.Given a set of candidate subgraphs, we develop a faster algorithm for selecting a subgraph with minimum TMD from a given graph. It computes the depth-L TMD between a graph G = (V, E) and any of its subgraphs in just O(L|E|) time. This algorithm leverages structural properties of the TMD that we uncover, which may be of independent interest.Our experiments demonstrate our methods outperform or perform competitively with other standard approaches on real-world graph learning benchmarks. For graph sampling, our method consistently outperforms the baselines on four benchmark datasets, achieving over 90% of the test accuracy of the full training set for three of the datasets using 1% of the graphs. For node subsampling, it is consistently optimal or near-optimal across all datasets, unlike other competitive baselines, whose performance varies significantly with the fraction of deleted data. 
related work: 
 Graph reduction, graph sparsification, and graph condensation approaches aim to to reduce the size of a graph while minimizing the impact on test loss. Existing approaches focus on node subsampling (e.g., Hashemi et al., 2024; Tanaka et al., 2020, and references therein) , primarily for node classification and graph signal processing. These approaches involve subsampling nodes, and do not extend to graph subsampling, a critical focus of this work, and a setting which remains largely unstudied in the graph reduction literature. Two existing approaches to model-agnostic graph reduction for graph classification are Herding (Welling, 2009) and k-Centers (Farahani & Hekmatfar, 2009) , which subsample k nodes by clustering the nodes' features into k clusters and selecting the centers. A key advantage of our approach is that, unlike these approaches, we leverage for the node features and the graph structure.Gradient-matching approaches (Jin et al., 2022; Zhang et al., 2024; Fang et al., 2024) have been used for graph dataset subsampling and ensure that the loss's gradients remain similar when trained on the small and large graphs. While these are powerful methods that can aid in developing reduced graph datasets that are smaller and more interpretable, they are fundamentally incomparable in our setting because (1) they assume knowledge of a downstream GNN architecture and therefore are not model-agnostic; and (2) they require training on the complete, unreduced graphs to generate the gradient trajectories. This obviates several of the advantages of graph reduction with respect to graph labeling and storage costs-and these are the primary motivations for our model agnostic setting. Georgiev et al. (2022) study a normalized variant of TMD to select training datasets for neural algorithmic reasoning tasks. However, this is not agnostic to the downstream model/task and is not directly applicable to subsampling. 
notation and background: 
 Notation. An undirected, unweighted graph G = (V, E, f ) has a node set V , edges E, and node features f ∈ R p×|V | , where f v is the feature vector of v ∈ V . We use N G (v) to denote the neighborhood of v ∈ V and G to denote a set of graphs. An r-rooted tree is denoted by T = (V, E, f, r). We use 1, 0, I to denote the all ones and zeros vectors and identity matrix, respectively. For norms, ∥•∥ denotes any norm over R p and ∥•∥ 1 is the ℓ 1 norm.Message-passing neural networks (GNNs). We analyze message-passing graph neural networks (GNNs), which operate on graphs G = (V, E, f ). For simplicity, we define Graph Isomorphism Networks (GINs) (Xu et al., 2019) here; however, the results of Chuang & Jegelka (2022) and, therefore, our results also apply to other GNNs. A GIN is defined by L (learnable) functions ϕ (ℓ) : R d → R d , such as (shallow) feed-forward neural networks, and a fixed weight η > 0. Each node begins with an initial embedding z (0) v = f v and updates its embedding using the following message-passing routine:Message passing layers: for ℓ ∈ [L -1], z (ℓ) v = ϕ (ℓ) z (ℓ-1) v + η u∈N (v) z (ℓ-1) u , Graph readout: h(G) = ϕ (L) v∈V z (L-1) u .Other GNN architectures may, for instance, replace the summations with an average (as in the case of Graph Convolutional Neural Networks (GCNs) (Kipf & Welling, 2016) ) or other aggregation functions (Hamilton, 2020) .Tree Mover's Distance (TMD). The TMD was introduced by Chuang & Jegelka (2022) and has found applica- tions for graph learning under distribution shifts (Yu et al., 2023; Georgiev et al., 2022) and adversarially robust graph learning (Schuchardt et al., 2024) . TMD is a pseudometric on graphs based on the optimal transport distance. Like a message passing GNN, it views a graph as a set of computation trees. A node's computation tree is constructed by adding the node's neighbors to the tree level by level, as in Figures 1 and 3 .Definition 3.1 (Computation tree). Given G = (V, E, f ) and a node v ∈ V , let T 1 v (G) = (v, ∅, {f v }, v) be a singleton v-rooted tree that consists of only the node v, no edges, and the node feature vector f v . Inductively, we defineT L v (G) to be the depth-L computation tree of node v. That is, T L v (G) is a tree rooted at v obtained as follows. For each leaf ℓ of T L-1 v (G) and each neighbor u ∈ N G (ℓ), we add a new node ℓ u to T L v (G)and add an edge from ℓ to ℓ u . The multiset of depth-L computation trees defined by G is denoted byT L G := {T L v (G)} v∈V .We can compare graphs by comparing their computation trees. TMD does this via a hierarchical optimal transport. Definition 3.2 (OT distance). Let X = {x i } q i=1 and Y = {y i } q i=1 be two multisets of q elements each. Given a distance metric d : X × Y → R, let C ∈ R q×q be a matrix where C i,j = d(x i , y j ) is the transportation cost between x i and y j . Moreover, let the set of all transportation plans between X and Y be Γ(X,Y ) := {γ ∈ R q×q + | γ1 = γ ⊤ 1 = 1} . The transportation cost of γ ∈ Γ is ⟨γ, C⟩ := i,j γ i,j C i,j . The (unnormalized) Wasserstein distance OT d is defined as OT d (X, Y ) := q min γ∈Γ(X,Y ) ⟨C, γ⟩. We say γ ⋆ is an optimal transport (OT) plan if γ ⋆ ∈ argmin γ∈Γ(X,Y ) ⟨C, γ⟩.We need a metric between trees to compute the OT distance between sets of computation trees. TMD compares two trees T v , T u by comparing their roots f v , f u and recursively comparing their subtrees. Definition 3.3 (Tree multisets). Given an r-rooted tree T = (V, E, f, r), let ℓ = depth(T ). We use T r (T ) to denote the multiset of depth (ℓ -1) computation trees of all neighbors u ∈ N T (r), illustrated in Figure 3 . (In other words, if T u is the subtree of T rooted at u, thenT r (T ) := ∪ u∈N T (r) T (ℓ-1) u (T u ), where the union denotes a multiset union).TMD compares T u (T ) and T v (T ′ ) via an OT distance. However, the number of subtrees could differ in T u (T ) and T v (T ′ ). Chuang & Jegelka (2022) therefore augment the smaller set with blank trees. A blank tree T 0 consists of a single node with node features 0 ∈ R d and no edges. Given two tree multisets T v (T ′ ) and T u (T ), the function ρ(T v (T ′ ), T u (T )) returns two multisets of the same size, where the smaller is padded with blank trees (see Appendix A.3 for more details). We are now ready to define a distance metric between trees.Definition 3.4 (Tree distance). Let T a = (V a , E a , f a , r a ) and T b = (V b , E b , f b , r b ) be two trees with ℓ := max(depth(T a ), depth(T b )). Moreover, let w : Z ≥0 → R be a depth-dependent weighting function. The tree distance (TD) between T a and T b is defined asTD w (T a , T b ) := ∥f ra a -f r b b ∥ + w(ℓ) • OT TDw (ρ(T ra (T a ), T r b (T b ))), if ℓ > 1 0, otherwise.Here, OT TDw is the OT distance with metric TD w . We define the TMD, which calculates the cost of the optimal transport plan between two graphs' computation trees.Definition 3.5 (Tree mover's distance (TMD) (Chuang & Jegelka, 2022) ). Given graphs G and G ′ , weight function w : Z ≥0 → R, and a depth parameter L > 0, we defineTMD L w (G, G ′ ) := OT TDw (ρ(T L G , T L G ′ )). The tree norm is ∥G∥ TN L w := TMD L w (G, ∅)where we overload ∅ to denote the empty graph.The tree distance (Definition 3.4) and TMD (Definition 3.5) is defined recursively, so is intractable to compute naively. However, as noted by Chuang & Jegelka (2022) , it can be computed in polynomial time using dynamic programming. 
node subsampling: 
 In this section, we present an approach to subsampling the nodes of a graph dataset while preserving the performance of a downstream GNN trained on the subsampled dataset. Full proofs for all results in this section can be found in Appendix C. To this end, suppose we are given a graph datasetX = {G 1 , ..., G n } where G i = (V i , E i ) together with a node budget k. Given a subset of k nodes S i ⊆ V i , we use G i [S i ] to denote the subgraph of G i induced by S i .Our goal is to select these subsets so that a GNN obtains similar readouts on the original training dataset X and on the induced subgraphsX ′ = {G 1 [S 1 ], ..., G n [S n ]}.Although quantifying the stability of the GNN readouts with respect to node deletion directly is analytically intractable, it is possible to bound the Lipschitz constant of a messagepassing GNN's readouts in terms of the TMD.1 Theorem 4.1 (Informal, Theorem 8 in (Chuang & Jegelka, 2022 ), restated). There exists a weight function w such that given an (L -1)-layer message-passing GNN with readout function h : G → R d and ℓ-th layer Lipschitz constants ϕ ℓ , for any two graphsG a , G b , ∥h(G a ) -h(G b )∥ ≤ L-1 ℓ=1 ϕ ℓ • TMD L w (G a , G b ).Theorem 4.1 implies that if the Lipschitz constants of the GNN layers and the distances TMD(G i [S i ], G i ) are small, then minimizing a Lipschitz loss over the subsampled training set X ′ yields a nearly optimal hypothesis over X. 2 The following corollary summarizes this observation.Corollary 4.2. Let H be a hypothesis class of (L -1)layer GNNs h : G → R d with ℓ-th layer Lipschitz constant upper-bounded by Φ ℓ . Let L : R d → R be an M -Lipschitz loss function. Let X = (G 1 , ..., G n ) and (y 1 , ..., y n ) be a set of (training) graphs and their labels.For each i ∈ [n], let G i = (V i , E i ) and S i ⊂ V i . Suppose that M ℓ∈[L-1] Φ ℓ ≤ c and TMD L w (G i , G i [S i ]) ≤ ε i for all i ∈ [n].Finally, let ĥ ∈ H be a hypothesis with minimum loss over the subsampled training set:ĥ = argmin h∈H i∈[n] L(h(G i [S i ]); y i ).Then ĥ has nearly optimal loss over the original training set as well:1 n i∈[n] L ĥ(G i ); y i ≤ min h∈H 1 n i∈[n] L(h(G i ); y i ) + 2c/n • i∈[n] ε i .Corollary 4.2 illustrates that the loss incurred by training on X ′ as opposed to X is bounded by an error proportional to the average of the distances TMD L w (G i , G i [S i ]). Motivated by this bound, our ideal goal would be to compute the optimal solution to the following optimization problem:EQUATIONWe face two key challenges in computing (approximately) optimal solutions to Equation (1).Challenge 1: exponentially large feasible set. The first challenge is that the number of candidate subsets {S ⊂ V : |S| ≤ k} grows exponentially with k. Therefore, we restrict Equation (1) to an appropriately chosen feasible set S. We discuss a natural heuristic for selecting S in Section 4.3. In this section, we sketch our proof of Theorem 4.4. The proofs in this section are inductive in order to leverage the recursive structure of the TMD (Definition 3.5) and tree distance (Definition 3.4.) However, we emphasize that our final algorithmic result Algorithm 1 is a pure dynamic programming algorithm; i.e., the recursive formulation of TMD is useful for the analysis but is circumvented via dynamic programming in our proposed implementation.Along the way towards proving Theorem 4.4, we prove several properties of the TMD, which may be of independent interest given the broad applications of TMD to out-ofdistribution generalization. Our analysis crucially relies on the following technical lemma.Lemma 4.1. Let w : Z ≥0 → R be a weight function and G = (V, E, f ) a graph. Then, the identity transportation plan I mapping T v (G) to T v (G[S]) for v ∈ S, and T u (G) to ∅ for u / ∈ S is an optimal plan for the OT problemTMD L w (G, G[S]) = OT TDw ρ T L G , T L G[S].(2)Consequently, we can decompose the right-hand-side of (2) into three terms: a sum over the feature norms of deleted nodes v ̸ ∈ S; the cost of transporting the deleted nodes' computation trees (that is, {T v (T L v (G))} v̸ ∈S ) to the empty set; and the cost of transporting the computation trees of the remaining nodes v ∈ S (that is,{T v (T L v (G))} v∈S ) to their computation trees in G[S]. Formally, OT TDw ρ T L G , T L G[S] = v / ∈S ∥f v ∥ + w(L -1) v / ∈S OT TDw ρ T v T L v (G) , ∅ + w(L -1) v∈S OT TDw ρ T v T L v (G) , T v T L v (G[S]) .Proof sketch. We prove this by induction on L. In the base case (L = 1), the statement is equivalent to Definition 3.5. This is because the definition of ρ and the tree distance TD imply that OT TDw (ρ(T 1 G , T 1 G[S])) is equal to the sum of features of the nodes that are inG but not in G[S]. That is, OT TDw (ρ(T 1 G , T 1 G[S] )) = v /∈S ∥f v ∥ = ⟨C, I⟩. This shows the lemma holds in the base case. We prove the inductive step by leveraging recursive formulation of the TD to reduce OT problems of depth L to OT problems of depth L-1, and then applying the inductive hypothesis. ■ Lemma 4.1 implies the following chain rule for the TMD induced by dropping a sequence of nodes from the graph.Lemma 4.2. For T ⊂ S ⊂ V , TMD L w (G, G[S \ T ]) = TMD L w (G, G[S]) + TMD L w (G[S], G[S \ T ]).Proof sketch. To prove the lemma, we apply Lemma 4.1 inductively overL to decompose TMD L w (G, G[S \ T ]) into TMD L w (G, G[S]) and TMD L w (G[S], G[S \T ]).To highlight the intuition, we discuss the base case here and discuss the inductive step in the full proof. If L = 1, then the second and third terms in the expression of Lemma 4.1 are always 0, because T 1 G and T 1 G[S] are multisets of depth 1, soT G (T L v (G)) and T G[S] (T L v (G[S])) are empty sets. Thus,TMD 1 w (G, G[S \ T ]) = v / ∈(S\T ) ∥f v ∥ = v / ∈S ∥f v ∥ + v∈T ∥f v ∥ , Algorithm 1: TreeNorm(G, L, w) Input: Graph G = (V, E, f ) with adjacency matrix A, weights w : {1, ..., L -1} → R + , L ≥ 1 1 Compute x ∈ R |V | to be the vector such that x v = ∥f v ∥ ; 2 Initialize z (0) = x; 3 for ℓ ∈ [L -1] do 4 z (ℓ) ← Az (ℓ-1) ; 5 b ← z (0) + L-1 ℓ=1 ℓ t=1 w(L -t) • z (ℓ) ; return: ∥b∥ 1 (since T ⊂ S, {v ∈ V : v / ∈ (S \T )} = {v ∈ V : v / ∈ S}⊔ {v ∈ T }). Applying Lemma 4.1 to TMD L w (G, G[S \ T ]) and TMD 1 w (G[S], G[S\T ]) implies that the base case holds:EQUATIONTMD 1 w (G[S], G[S \ T ]) = v∈T ∥f v ∥ .The L > 1 case is similar but requires care to handle the recursive OT terms. ■We briefly point out why Lemma 4.2 is surprising. TMD is a pseudo-metric on graphs; so, in general, we should expect the following triangle inequalityTMD L w (G, G[S \ T ]) ≤ TMD L w (G, G[S]) + TMD L w (G[S], G[S \ T ]) to hold. However, Lemma 4.1 says that T ⊂ S ⊂ V is enough to ensure that the inequality is in fact an equality. Intuitively, this means that TMD is additive along any sequence of vertexinduced subgraphs that reduces the full graph G to the empty graph. We use this observation to prove Theorem 4.4.Proof of Theorem 4.4. Let S ∈ S. By Lemma 4.2 with T = V \ S, we have that∥G∥ TN L w = TMD L w (G, G[S]) + ∥G[S]∥ TN L w . Consequently, max S∈S;|S|=k is equivalent to max S∈S;|S|=k ∥G∥ TN L w -TMD L w (G, G[S]), which is in turn equivalent to min S∈S;|S|=k TMD L w (G, G[S]). ■We next leverage Lemma 4.1 to obtain Algorithm 1, a faster algorithm for computing∥G∥ TN L w for G = (V, E, f ).The original implementation would require O(L |V | 4 )time (Chuang & Jegelka, 2022) , whereas ours has runtime O(L|E|). Algorithm 1 is based on the fact that by Lemma 4.1, is that ∥G∥ TN L w is essentially a weighted sum of the number of vertices in all of its L-hop computation trees. For each tree, nodes in level ℓ receive weight w(L -ℓ + 1). Algorithm 1 computes this weighted sum. Its runtime is dominated by the cost of L matrix-vector multiplies with the adjacency matrix, which requires O(|E|)-time. The proof of correctness and runtime (Theorem 4.5) is in Appendix C.We now propose our heuristic for constructing the set of candidate subsets S in our experiments. Most real-world networks are scale-free and small-world, and hence can be well-modeled by random graphs, e.g., preferential attachment, configuration models, and inhomogeneous random graphs (Bollobás & Riordan, 2004; Newman & Watts, 1999) . These graphs converge locally to graph limits (Van Der Hofstad, 2024) . By a "transitive" argument, we can view real-world networks as parts of sequences converging to graph limits in the same sense. Alimohammadi et al. (2023) showed that for such graphs, sampling nodes' local BFS trees (Definition 4.7) asymptotically preserves motifs of the graph limit (see Appendix A). Definition 4.7 (k-BFS subset). Given G = (V, E, f ) and v ∈ V , let ℓ k be the deepest level such that the v-rooted BFS tree of depth ℓ k has at most k nodes (breaking ties in a fixed but arbitrary way). The k-BFS subset of v, denoted S BFS(v;k) , is the set of nodes at distance at most ℓ k from v.We define the set of candidate subsets S = {S BFS(v;k) } v∈V as the nodes' k-BFS subsets. This definition of S is modelagnostic as it is determined solely by the graph structure and is independent of the downstream mode architecture or hyperparameters. In sufficiently sparse graphs S is efficient to compute; however, in denser graphs, there are several previously studied approaches for subsampling the k-BFS subsets for improved efficiency (see, e.g. (Alimohammadi et al., 2023) and references therein) for an overview. 
graph subsampling: 
 We now present an approach that uses TMD to subsample the number of graphs in a dataset while preserving the performance of a downstream GNN trained on the subsampled dataset. Omitted proofs are in Section C. Let X = {G 1 , ..., G n } be a graph dataset. Given a budget k, the goal is to identify a subset I ⊂ [n] of size k such that a GNN obtains similar readouts and loss on the original training set X and on the subsampled set {G i } i∈I . We must ensure the selected set I is representative of the full dataset. We propose selecting I to minimize the medoids objective.Definition 5.1 (Medoids objective). Let X = {G 1 , ..., G n } be a graph dataset and I ⊂ [n] with |I| ≤ k. The medoids objective value of I on X with respect to TMD L w is definedf TMD L w (I; X) = 1/|X| • i∈[n] min j∈I TMD L w (G i , G j ).For j ∈ I, let τ j be the number of graphs in X closest to G j (the j-th medoid's cluster size):τ j := {i : TMD L w (G i , G j ) < TMD L w (G i , G k ), ∀k ̸ = j ∈ I} , breaking ties arbitrarily.The following lemma shows how we can bound the differ-ence in a GNN's loss on the subsampled dataset I ⊂ X and the original dataset X in terms of this medoids objective.Lemma 5.1. Let H be a hypothesis class of (L -1)-layer GNNs h : G → R d with ℓ-th layer Lipschitz constant upperbounded by Φ ℓ . Let L : R d → R be an M -Lipschitz loss function. Let y = (y 1 , ..., y n ) be labels for X and I be a subset of [n] . Then for any GNN h ∈ H and G ∈ G,1 n i∈I τ i L(h(G i ); y i ) - i∈[n] L(h(G i ); y i ) ≤ M • ℓ∈[L-1] Φ ℓ • f TMD L w (I; X).Proof sketch. We use the Lipschitzness and Theorem 4.1 to bound the average deviation in the loss on X and on I.■ Lemma 5.1 implies that if M • ℓ∈[L-1] Φ ℓ • f TMD L w (I; X) is small and we minimize loss over the subsampled training set, we obtain a hypothesis with nearly optimal loss over the entire training set. We summarize this observation below.Corollary 5.2. Suppose M ℓ∈[L-1] Φ ℓ ≤ c and f TMD L w (I; X) ≤ ε.Let ĥ ∈ H be the hypothesis minimizing loss on the subsampled graphs i∈I τ i L(h(G i ); y i ). Then1 n i∈[n] L ĥ(G i ); y i ≤ min h∈H 1 n i∈[n] L(h(G i ); y i ) + 2cε.Corollary 5.2 illustrates that the additional loss incurred by training on the subsampled graphs as opposed to X is bounded by an error proportional to f TMD L w (I; X). This motivates us to formulate the graph subsampling problem as the problem of computing I that minimizes f TMD L w (I; X). Definition 5.3 (Graph subsampling problem). Let X = {G 1 , ..., G n } be a set of graphs and k a graph budget. In the graph subsampling problem, we must select a k-subset I ⊂ X that minimizes the k-medoids objective f TMD L w (I; X). k-medoids is NP-hard (Kazakovtsev & Rozhnov, 2020) , but there are efficient approximations, e.g., in Python (Buitinck et al., 2013) . Consequently, the main computational bottleneck is the computation of TMDs between pairs of graphs G 1 , G 2 ∈ X, which may be large if the graphs have many nodes. However, these computations can be sped up in practice, by combining graph subsampling with our node subsampling approach (see Section 4 and Remark 4.6), which is computationally efficient. See Appendix B for a detailed discussion and experiments. 
experiments: 
 We evaluate our sampling methods on four graph benchmark datasets from the TUDataset graph benchmark library (Morris et al., 2020) . All GNNs were trained for 50 epochs with an 80%-train set and 20%-test split on CPUs. For MUTAG, COX2, and BZX we used a 3-layer Graph Isomorphism Network (Xu et al., 2019) . For IMDB-BINARY we used a 3-layer Graph Convolutional Neural Network (GCN) (Kipf & Welling, 2016) . Hyper-parameter details are in Appendix A.4. In all figures, the line plots and shadows indicate the means and 95% confidence intervals across 256 random trials. Figure 4 compares our TMD-based k-medoids graph subsampling procedure from Section 5 against alternative methods and training on the full dataset.We are unaware of previous work that studies graph subsampling, but we compare against two clustering approaches based on Herding (Welling, 2009) and k-Centers (Farahani & Hekmatfar, 2009) . These approaches are typically used for node subsampling and coreset selection in other domains (see Section 2). We adapt them to graph subsampling by defining distances between graphs using three popular graph kernels: Weisfeiler Lehman kernel (WL) (Shervashidze et al., 2011) , the WL optimal assignment kernel (WL-OT) (Kriege et al., 2016) , graphlet sampling (GS) kernel (Borgwardt & Kriegel, 2005) , and the shortest paths (SP) kernel. Figure 4 shows that our method consistently outperforms uniform random sampling, Herding (with respect to all kernels), and k-Centers (with respect to all kernels) on all four benchmarks. Our approach achieves over 90% of the test accuracy of the full training set for MUTAG, COX2, and BZR using just 1% of the graphs. When graph labeling is very costly, this may be especially powerful.Figure 5 compares our TMD-based node subsampling procedure from Section 4 against several alternative benchmarks and against training on the full dataset. For node subsampling, we only need to compute tree norms, so we use Algorithm 1. We compare against model-agnostic node selection policies Herding (Welling, 2009) and k-Centers (Farahani & Hekmatfar, 2009) , and approaches by Salha et al. (2022) and Razin et al. (2023) . While other approaches exhibit significant variability across datasets and subsampling levels, our approach consistently has the strongest or near-strongest performance across all datasets, sometimes requiring as little as 10% of the original number of nodes to achieve similar test accuracy as training on the full graphs. 
discussion: 
 We proposed new node and graph subsampling approaches for compressing graph datasets while preserving performance of downstream GNNs. Our methods leverage new theoretical analysis of the TMD to ensure the subsampled data maintains structural similarities with the original dataset, and we bound the excess loss incurred by subsampling. Our methods are enabled by an algorithm we developed for selecting subgraphs with minimal TMD from a given graph, providing runtime speedup over prior work. Experiments demonstrate that our method outperforms other comparable model-agnostic methods for node and graph subsampling.A.1. Random graph limits As we discussed in Section 4.3, most real networks are scale-free and small-world, and hence well-modeled by random graph models, such as preferential attachment, configuration models and inhomogenous random graphs. As shown by (Van Der Hofstad, 2024) , such random graph models produce graphs that can be shown to converge locally to a limit (G, o), where G is a graph to which we assign a root node o. By a "transitive" argument, it makes sense to see real networks as parts of sequences converging to graph limits in a similar sense.To be precise, let G * be the set of all possible rooted graphs. A limit graph is defined as a measure over the space G * with respect to the local metricd loc ((G 1 , o 1 ), (G 2 , o 2 )) = 1 1 + inf k {k : B k (G 1 , o 1 ) ̸ ≃ B k (G 2 , o 2 )}where B k (G, v) is the k-hop neighborhood of node v, and ≃ is the graph isomorphism. A sequence of graphs converging to this limit is defined as follows.Definition A.1 (Local convergence (Alimohammadi et al., 2023) ). Let G n = (V n , E n ) denote a finite connected graph. Let (G n , o n ) be the rooted graph obtained by letting o n ∈ V n be chosen uniformly at random. We say that (G n , o n ) converges locally to the connected rooted graph (G, o), which is a (possibly random) element of G * having law µ, when, for every bounded and continuous function h :G * → R, E[h(G n , o n )] → E µ (G, o)where the expectation on the right-hand-side is with respect to (G, o) having law µ, while the expectation on the left-hand-side is with respect to the random vertex o n .A.2. Other graph kernels TMD is one recently proposed pseudometric on graphs, or graph kernel (Chuang & Jegelka, 2022) . It is appealing in our setting because it characterizes the robustness of GNNs (see Theorem 4.1.) Other popular graph kernels include the Weisfeiler Lehman (WL) kernel (Shervashidze et al., 2011) , Weisfeiler Lehman Optimal Transport (WL-OT) (Kriege et al., 2016) , Shortest Paths kernel ( (Borgwardt & Kriegel, 2005) ), and FGW (KAWANO et al., 2024) . To our knowledge, these kernels do not have comparable robustness guarantees to TMD.We use T n 0 to denote n disjoint copies of T 0 . Given two tree multisets T u (T ), T v (T ′ ), we define ρ to be the following augmentation function, which returns two multi-sets of the same size:ρ : (T v (T ′ ), T u (T )) → T v (T ′ ) ∪ T 0 max(|Tu(T )|-|Tv(T ′ )|,0) , T u (T ) ∪ T 0 max(|Tv(T )|-|Tu(T ′ )|,0) .All experiments were run on a 16-core machine with 64 GB of RAM. The GNNs were implemented using Pytorch Geometric (He et al., 2024) . The TMD weight function was set according to Pascal's triangle rule (Chuang & Jegelka, 2022, Theorem 8) and our implementation of TMD was based on (Chuang & Jegelka, 2022) . 3 We used the Graph Kernel Library (Siglidis et al., 2020) for the kernel distances in our experiments.For MUTAG, COX2, and BZX we used a GIN model with the following hyperparameters:• We refrained from using batch normalization in both models. Models were implemented using PytorchGeometric (Fey & Lenssen, 2019) .In this section, we show that the pairwise TMD computations required for our graph subsampling approach can be computed more efficiently using the following two-stage approach. Recall that k-medoids algorithms require querying distances (in our case, TMD) between two graphs G 1 , G 2 . However, if G 1 and G 2 have many nodes, computing TMD L w (G 1 , G 2 ) using the dynamic programming algorithm of (Chuang & Jegelka, 2022 ) may be computationally expensive. Fortunately, we can leverage our efficient methods for node subsampling to efficiently compute approximations of TMD L w (G 1 , G 2 ) as follows. 1. First, we use our node subsampling approach to create new graphs G ′ 1 and G ′ 2 such that TMD L w (G 1 , G ′ 1 ) and TMD L w (G 2 , G ′ 2 ) are minimized while ensuring that G ′ 1 and G ′ 2 have fewer nodes than G 1 and G 2 , respectively. 2. Second, we run a k-medoids clustering algorithm using approximate pairwise distances TMDL w (G ′ 1 , G ′ 2 ) in place of TMD L w (G 1 , G 2 ). Triangle inequality ensures that TMD L w (G ′ 1 , G ′ 2 ) -TMD L w (G 1 , G 2 ) ≤ TMD L w (G 1 , G ′ 1 ) + TMD L w (G 2 , G ′ 2 ).So, whenever TMD L w (G 1 , G ′ 1 ) and TMD L w (G 2 , G ′ 2 ) are not too large (which is the goal of our node-subsampling procedure), we know that TMDL w (G ′ 1 , G ′ 2 ) is a good approximation for TMD L w (G 1 , G 2 ).To illustrate that this method can work well in practice, we conducted an experiment showing the accuracy of our k-medoids graph subsampling approach when using approximate TMD computations on the MUTAG dataset, as described above. For each pair of graphs G 1 , G 2 in the MUTAG dataset, we computed p-approximate TMDs TMD L w (G p 1 , G p 2 ) and ran our graph subsampling approach using TMD L w (G p 1 , G p 2 ) where G p denotes the graph G after using our node subsampling approach to delete a p-fraction of nodes in the graph. We then use TMD L w (G p 1 , G p 2 ) as a proxy for TMD L w (G 1 , G 2 ) in our k-medoids clustering step. Our results indicate that we can afford to delete a large portion (up to 70 percent) of the nodes prior to computing pairwise TMDs with only a small effect on accuracy-and computing TMD L w (G p 1 , G p 2 ) is, on average, eight-fold faster than computing TMD L w (G 1 , G 2 ) when p ≥ .5. See Figure 6 for an illustration of these results.Throughout this section, we use deg G (v) to denote the degree of v ∈ V . i ∈ [n], let G i = (V i , E i ) and S i ⊂ V i . Suppose that M ℓ∈[L-1] Φ ℓ ≤ c and TMD L w (G i , G i [S i ]) ≤ ε i for all i ∈ [n].Finally, let ĥ ∈ H be a hypothesis with minimum loss over the subsampled training set: ĥ= argmin h∈H i∈[n] L(h(G i [S i ]); y i ).Then ĥ has nearly optimal loss over the original training set as well:1 n i∈[n]L ĥ(G i );y i ≤ min h∈H 1 n i∈[n] L(h(G i ); y i ) + 2c/n • i∈[n] ε i .Proof. Consider any h ∈ H. By Theorem 4.1, we have that for each i ∈EQUATIONConsequently, using the fact that L is M -Lipschitz and (3),|L (h(G i ); y i ) -L (h(G i [S i ]))| ≤ M ∥h(G i ; y i ) -h(G i [S i ]; y i )∥ ≤ M   ℓ∈[L-1] Φ ℓ   TMD L w (G i , G i [S i ]) = cTMD L w (G i , G i [S i ]) ≤ cε i .Consequently, by the triangle inequality,1 n i∈[n] L (h(G i ); y i ) - 1 n i∈[n] L (h(G i [S i ])) ≤ 1 n i∈[n] |L (h(G i ); y i ) -L (h(G i [S i ]; y i ))| ≤ c n i∈[n] ε i .So, we know that for any h ∈ HEQUATIONLet h ⋆ be the hypothesis that minimizes average loss across the original dataset X:h ⋆ = argmin h∈H 1 n i∈[n]L(h(G i ); y i ).By definition of ĥ, this means that1 n i∈[n]L( ĥ(G i [S i ];y i ) ≤ 1 n i∈[n] L(h ⋆ (G i [S i ]); y i ) + c n i∈[n] ε i .Applying (4) again to h ⋆ , we have that1 n i∈[n]L( ĥ(G i [S i ]; .y i ) ≤ 1 n i∈[n] L(h ⋆ (G i ); y i ) + c n i∈[n](2)Consequently, we can decompose the right-hand-side of (2) into three terms: a sum over the feature norms of deleted nodes v ̸ ∈ S; the cost of transporting the deleted nodes' computation trees (that is, {T v (T L v (G))} v̸ ∈S ) to the empty set; and the cost of transporting the computation trees of the remaining nodes v ∈ S (that is, {T v ) .In the base case where L = 1, we see that any transportation plan between the nodes of Ḡ[S] and G must transport the excess node feature mass v / ∈S ∥f v ∥ in G to some nodes in G[S]. The transportation plan I has cost exactly equal to v∈V ∥f v -f ′v ∥ = v / ∈S ∥f v ∥, so I must be an optimal transportation plan. Now, if L > 1, consider the OT problem 2, and let C and Γ be the distance matrix and feasible transportation plans for this OT problem. Let γ ∈ Γ be any feasible transportation plan.By the definition of TD, for any i, j ∈ V , transporting γ i,j mass from i to j incurs two costs:(1) γ i,j incurs the cost of transporting f i to f ′j . with ties broken arbitrarily (but consistently with the mapping τ .) First, we can rearrange the sums as follows1 n i∈I τ i L(h(G i ); y i ) - 1 n i∈[n] L(h(G i ); y i ) = 1 n i∈[n]L(h(G κ(i) ); y κ(i) ) -L(h(G i ); y i )≤ 1 n i∈[n]L(h(G κ(i) ; y κ(i) )) -L(h(G i ); y i )≤ M n i∈[n]L(h(G κ(i) ); y κ(i) ) -L(h(G i ); y i ) ,where in the last inequality, we used that L is M -lipschitz. Now, by Theorem 4.1, we haveM n i∈[n] L(h(G κ(i) ); y κ(i) ) -h(G i ; y i ) ≤ M n   ℓ∈[L-1] Φ ℓ   i∈[n] TMD L w (G κ(i) , G i ) = M   ℓ∈[L-1]Φ ℓ   f TMD L w (S; X).Corollary 5.2. Suppose M ℓ∈[L-1] Φ ℓ ≤ c and f TMD L w (I; X) ≤ ε. Let ĥ ∈ H be the hypothesis minimizing loss on the subsampled graphs i∈I τ i L(h(G i ); y i ). Then1 n i∈[n]L ĥ(G i ); y i ≤ min L ĥ(G i ); y i ≤ 1 n i∈S τ i L( ĥ(G i ); y i ) + cε.Let h * be the hypothesis that minimizes average loss over the original dataset X:h * = argmin h∈H 1 n i∈[n]L(h(G i ); y i ).By definition of ĥ, this means that1 n i∈[n]L ĥ(G i ); y i ≤ 1 n i∈S τ i L(h * (G i ); y i ) + cε.Applying Lemma 5.1 again, we have that1 n i∈[n]L ĥ(G i );y i ≤ 1 n i∈[n]L (h * (G i ); y i ) + 2cε.■Stanford University, Stanford, CA United StatesMassachusetts Institute of Technology, Cambridge, MA United StatesJohns Hopkins University, Baltimore, MA United States. Correspondence to: Ishani Karmarkar <ishanik@stanford.edu>.Theorem 8 as stated in(Chuang & Jegelka,2022) is specific to GINs, butChuang & Jegelka (2022) generalize it to other GNNs.2 See, e.g., Section 5 ofChuang & Jegelka (2022) for an empirical analysis of the tightness of these Lipschitz bounds.One consideration with is the runtime required to compute the k-medoids. However, reduced labeling cost, reduced data storage, and privacy are often more significant priorities. 
 
 Appendix:  


Decision: ## Strengths
* The paper is well presented
* Theoretical results provide insights about simplifying computations of TMD between a graph and its subgraph 
* Empirical results are interesting and important, but mostly for graph sampling

## Weaknesses
* Notation-heavy paper, it might be possible to make the presentation simpler
* Experimental results do not contain any discussion about the computing requirement for doing the
  sampling. 
* Experimental results for node sampling (which represents most of the theoretical work) are not so
  convincing. It would be interesting to evaluate the Lipschitz coefficients of the considered
  networks, to assess how strong the theoretical bounds are.

The general sentiment about this paper is mixed. I recommend acceptance as a poster, provided that
authors refine their experimental evaluation.