Interpretable and Convergent Graph Neural Network Layers at Scale

Thanks for the helpful comments. We address each point in turn below.

Comment: [General Weakness and Question 1.] Add a true explainability analysis or remove most of these claims.

Response: Thanks for the suggestion; please see our general response to all reviewers presented above (and we are happy to modify the draft as needed).

Comment: [General Weakness and Question 2.] Add an analysis of the memory required by the model.

Response: Memory requirements were deferred to the supplementary to save space; however, we can summarize here. For the lightweight $\gamma=0$ case, the memory requirement is the same as baselines like GCN with (online) neighbor sampling. Namely, $O(nd)$ of memory is required to store the input features in main memory / disk, and $O(n_sdK)$ of GPU memory is required to store intermediate results (embeddings and gradients) in the device. Here $n$ and $n_s$ are the number of nodes in the full graph and subgraphs, respectively, $d$ is the dimension of input features / hidden embedding, and $K$ is the number of layers. For $\gamma>0$, an extra $O(nd)$ of memory is needed to store the mean vectors $M$, while there is no extra memory needed on the GPU side. As a point of reference, this extra memory requirement is much lower than the $O(ndK)$ required by the GAS-based baselines (Fey et al., 2021) presented in Table 2.

Comment: [General Weakness and Question 3.] Add some ablations on the number of chosen subgraphs.

Response: We note that the number of subgraphs is not a unique issue to MuseGNN, and could potentially affect the performance of any GNN model based on sampling. Still, this is a useful stability metric to check as the reviewer suggests to rule out any appreciable sensitivity. As such, we performed such an ablation with ogbn-products data using 50%, 60%, 70%, 80%, 90%, and 100% of the subgraphs that we originally used in the experiments from Table 2. The accuracy differences were less than 1% across all cases, indicating a degree of robustness to the number of subgraphs. We can add these results to the supplementary in a later revision.

Comment: [General Weakness and Question 4.] In the related works, several papers have proposed subgraph-based GNNs to improve the expressiveness of standard message-passing. Can you compare your approach, either methodologically or experimentally?

Response: This is an interesting point. We surmise from our submission that the reviewer may be referring to the reference (Zeng et al., 2021) or related, which discusses subgraph-based sampling that can influence expressiveness by decoupling depth (number of model layers) from scope (the receptive field size). Similar to Zeng et al., MuseGNN also naturally facilitates the decoupling of depth and scope, where the scope is effectively determined by the underlying sampling method. With regard to depth though, we differ from Zeng et al. in the sense that additional MuseGNN layers push node embeddings closer and closer to a specific target, namely, a minimizer of the objective from Eq.(4). Another more subtle distinction exists w.r.t. oversmoothing: While Zeng et al. relies directly on decoupling to avoid oversmoothing, MuseGNN will be immune to oversmoothing regardless of the particular sampler as long as minimizers of the original energy (4) do not oversmooth.

Thanks for the quick response. Regarding the title, we agree that removing "interpretability" is a good idea. As a more reflective replacement, one option could be something like: "Forming Scalable, Convergent GNN Layers that Minimize a Sampling-Based Energy."

Also, thanks for pointing to the Frasca et al. (2022) reference. Upon quick inspection, we notice that this paper is rich with details that require further time to digest. As one speculative possibility though, the ideas from Frasca et al. could potentially be leveraged to more precisely characterize the expressiveness of MuseGNN's sampling-based energy function.

Thanks for the constructive comments. We address each point in turn below.

Comment: [Weakness 1.] Lack of comparison with representative sampling-based GNNs, such as Cluster-GCN and GraphSAINT.

Response: For the largest-scale benchmarks, neighbor sampling (NS) is typically the method of choice, which is why we adopted it for our baseline comparisons (along with GAS as a more recent popular alternative). Still, as the reviewer suggests, for completeness it can be beneficial to explore other options. While we did not have time to run both during the rebuttal period, we provide new results with Cluster-GCN below.

Note that if the reviewer feels it is important, we can easily add GraphSAINT as well to a later revision, although this approach is not seen on the leaderboards for any of the largest datasets.

Comment: [Weakness 2.] The authors did not conduct experiments on datasets that can illustrate the importance of unfolded GNNs.

Response: This is actually true by design; please see our general response to all reviewers presented above which addresses this issue.

Thanks for the detailed comments. We address each point in turn below.

Comment: [From Weakness 1. and Question 1.] One notable weak point of the paper's experimental evaluation is its limited comparison to a rather outdated set of baseline models ... I provide some example papers [1,2,3] that can be used as baselines below. Moreover, in the ogb leaderboard https://ogb.stanford.edu/docs/leader_nodeprop/#ogbn-arxiv there are many new baselines. Could the authors provide more context on their choice of limited and older baselines for comparison? Are there more recent or relevant methods that could have been included for a more comprehensive evaluation?

Response: This is a valid question, and it is definitely true that there now exists a considerable array of different GNN architectures. The ogbn-arxiv leaderboard and references [1,2,3] are nice examples of this diversity as the reviewer mentions. However, many/most of these methods have not been integrated with sampling (or historical caching like GAS) or shown to be competitive on truly large graph datasets as is our focus. Also please kindly note that for the two largest datasets, MAG240M and IGB-full, we are already comparing against the top-ranked homogeneous graph leaderboard model, which is GAT (NS).

As for [1,2,3], these are quite interesting references proposing powerful models (especially for heterophily graphs); however, upon inspection of the empirical evaluations, none of the included benchmarks are especially large or require sampling; in fact, all are amenable to full-graph training such that scalability (our focus) is not a significant factor. As for top-performing models on OGB leaderboards such as ogbn-arxiv, please see the discussion at the beginning of Section 6.

Of course, we in no way intended to imply that [1,2,3] or top ogbn-arxiv models could not eventually be scaled to our large benchmarks. However, it remains an open research question outside of our scope of how to do so, and would likely require considerable engineering bandwidth and resources that we presently do not have. Regardless, our central aim is really to scale unfolded GNNs, so in this regard, the metric of success is more tied to the two criteria mentioned at the beginning of Section 6, which our results thus far have achieved.

Comment: [From Question 2.] Present more results about ablations on $\gamma$, i.e. results on IGB-full, to further demonstrate the effectiveness of $M$.

Response: IGB-full is massive and we do not have the time or computing resources to conduct a careful ablation during the short rebuttal period. However, for now, we can nonetheless use IGB-tiny and IGB-medium as surrogates, as well as ogbn-products as another reference point (note that IGB-medium is actually quite large, with more labeled nodes than even ogbn-papers100M). From the new ablation table below, we observe that it is indeed possible to improve performance with $\gamma > 0$ in all three cases. (Note that for ogbn-products, this accuracy is actually better than reported in our original submission, where we did not carefully explore tuning $\gamma$.)

Comment: [From Question 3.] In Proposition 3.1, what is the definition of $\ell(M)$? What is the purpose of this proposition?

Response: The function $\ell$ is defined by Eq.(2), and $\ell(M)$ merely refers to Eq.(2) with $M$ as the input argument. The paragraph directly above Proposition 3.1, where we describe how the role of $M$ equates to $Y$ when $\gamma = \infty$, provides context for why we introduce $\ell(M)$.

In terms of purpose, Proposition 3.1 serves a key role by quantifying the relationship between ideal full-graph training as instantiated via Eq.(2), and offline sampling using Eq.(4) as used by MuseGNN in our submission. More precisely, when $\gamma = \infty$, Eq.(4) reduces to Eq.(7). From here, Proposition 3.1 then proves that Eq.(7) becomes equivalent to Eq.(2) in expectation (i.e., it is unbiased), provided we conduct sampling independently as described in the proposition statement. We hope this helps to clarify the value of Proposition 3.1, and we are happy to update the draft if necessary to do so. (The end of Section 3.2 very briefly mentions these ideas, but further detail could be helpful.)

Comment: [From Question 4.] According to my understanding, the main novelty of this paper is to incorporate the shared summary embedding matrix $M$ into energy to facilitate controllable linkage between the multiple embeddings that may exist for a given node appearing in different subgraphs (Because the Formula 6 is a commonly used technique for large-scale training).

Response: Our novelty extends beyond incorporating the shared summary embedding matrix $M$ into the stated energy as follows:

• Although offline sampling has been proposed previously, to our knowledge Eq.(6) has not been used in prior work as an energy function for producing GNN layers through minimization (our focus). Please also see Section 3.1 for further details regarding why we chose offline sampling, which differs from prior work.
• Our convergence results, which cover both the $\gamma = 0$ case, meaning Eq.(6), and the more general $\gamma > 0$ case are new. Moreover, by relying on offline sampling these results are agnostic to the underlying sampling method itself.
• Finally, even if we exclude consideration of the empirical efficacy of $\gamma > 0$, Proposition 3.1 quantifies a precise regime whereby our sampling-based energy provides an unbiased estimator of the full-graph energy. This is notable because, generally speaking, it is infeasible to establish notions of unbiasedness in arbitrary GNN models instantiated with sampling (with the exception of unbiasedness narrowly defined within a single model layer).

Comment: [From Question 4.] But existing ablation on $\gamma$ are not sufficient to demonstrate the general effectiveness of this additional constraint term.

Response: Regarding additional ablations, please see our response to Question 2 above.

Thanks for the detailed comments. We address each point in turn below.

Comment: [From Weakness 1.] The paper emphasizes Interpretable in the title and illustrates the explainability of energy minimizers. The key concern is the energy minimization is not aligned with the common concept of GNNs interpretation. For example, some of existing works of interpretable GNNs highlight neighborhood structure leading to the node label classification. The superficial conclusion of node embedding information stemming from the node itself or neighbors makes no sense to the real interpretable applications.

Response: We apologize for any confusion the title may have caused relative to existing work on interpretable GNNs. Please see our general comments to all reviewers above which directly address this issue. In any event, we are happy to modify the title and framing to increase clarity.

Comment: [From Weakness 2.] In addition, the absence of case study to validate the explanations provided by the energy function raises questions about one of the paper's main claims.

Response: Please see the general comments above to all reviewers which address this issue in depth. But in brief, our starting assumption is that prior work has already demonstrated the value of unfolded GNNs formed from interpretable energy functions; we make no claims of demonstrating this ourselves. Instead, our focus herein is entirely on scaling this class of GNN models to the largest possible graphs, with provable convergence guarantees.

Comment: [From Weakness 3.] The incorporation of the energy minimization into GNNs is not novel.

Response: Absolutely, we 100% agree. These models were originally invented, motivated, and validated in prior work, e.g., the many references we cite in Sections 2.1 and 2.2. That being said, our contribution is not the invention or motivation of these models, rather, it is scaling these models to huge graphs with convergence guarantees that account for the required sampling involved.

Comment: *[From Weakness 3.] The Dirichilet energy of node embeddings, i.e., the second term in Eq.(2), has been extensively studied in graph domains and GNNs. For example, it has been used to analyze over-smoothing issue in [1]. Following this work, several novel GNNs to optimize this energy function have been proposed.

$1$ Cai, Chen, and Yusu Wang. "A note on over-smoothing for graph neural networks." arXiv preprint arXiv:2006.13318 (2020).*

Response: Exactly, the Dirichlet energy, often in combination with other terms as in Eq.(2) of our submission, is the starting point for inspiring a wide variety of existing unfolded GNN architectures, e.g., see references in Section 2.1. And indeed as the reviewer correctly suggests (and the reviewer provided reference [1] shows), this Dirichlet energy can be used to analyze oversmoothing. In fact, this is a good example of the very type of interpretability that these energy-based models possess. In this case, the Dirichlet energy alone (akin to the second term in Eq.(2)) leads to oversmoothing since constant node embeddings (i.e., oversmoothing) will achieve the minimum. In contrast, it has also been observed in prior work that when additional terms involving a base model are included as in Eq.(2), oversmoothing no longer happens precisely because global minimizers transparently no longer oversmooth. Thanks for pointing out reference [1], this is a great example to add to our related work.

Comment: [From Weakness 3.] The constraint of closeness to base model embedding, i.e., the first term of Eq. (2), has been implicitly included in models like SIGN [2], SAGN [3], DAGNN[4].

Response: While these architectures may incorporate some form of base model or precomputed features, to our knowledge the forward pass in each case does not actually minimize an explicit energy function. Hence we cannot directly interpret their output embeddings as energy function minimizers as we can with the GNN architectures formed by minimizing Eq.(2). Even so, they can be cited for added context if the reviewer suggests.

Thanks for the detailed comments. We address each point in turn below.

Comment: [From 1st, 2nd Weakness and 1st Question] Novelty: What are the limitations of neighborhood sampling that require the development of the proposed offline sampling scheme? From a practical point of view, what are the actual disadvantages of NS that motivate the development of MuseGNN?

Response: Just to clarify, we are not proposing a new sampling scheme per se. Rather, we are integrating an existing offline sampling approach within a new role, namely, as the basis for the new energy function from Eq.(4), which then leads to a new scalable unfolded GNN architecture and MuseGNN. In this sense, the motivation of MuseGNN is unrelated to any disadvantages of NS. On this last point though, some additional context may be useful.

There are two completely separate aspects of sampling considered in our paper that are relevant to understanding MuseGNN:

1. There is the distinction between online vs offline sampling as first mentioned in Section 2.3, where the online version involves dynamically obtaining a new set of random samples during each training epoch, while the offline version is predicated on a fixed set of samples obtained prior to training. We chose the latter for MuseGNN, and the reasons are detailed in Section 3.1.
2. There is the actual design of the sampler itself, which can be applied in either an online or offline fashion. Neighbor sampling (NS) is arguably the most common choice in the literature, but there are many others such as those used by Cluster-GCN or Shadow-GNN etc. Regardless, as pointed out in our paper, MuseGNN can adopt any such sampler, and the convergence guarantees introduced in Section 5 will still hold.

From this distinction then, we hope the nature of MuseGNN relative to prior uses of sampling is more clear, but we are happy to provide further details if needed.

Comment: [From 4th Weakness and 1st Question (last part)] Please add significance intervals for GAT with neighbourhood sampling for GAT (NS) and SAGE (NS) and GCN (NS).

Response: We apologize that the original submission inadvertently did not provide sufficient details regarding these numbers. In fact, for the two largest datasets, we relied on numbers from the official IGB/OGB leaderboards for GAT (NS), SAGE (NS), and GCN (NS); we did not compute them all from scratch because of the significant time and cost involved (a common practice in the GNN literature). For example, just a single training run for these models requires a very large GPU instance (such as a p4dn.24xlarge AWS instance) with extra large RAM (over 1TB). So instead we just focused on computing MuseGNN results for these datasets without bandwidth for multiple runs across all baseline models as would be needed for significance intervals. We will include these details in the revision. Thanks for pointing out our oversight.

In contrast, for speed results (i.e., epoch timing) alone we do not require training until convergence. Instead, we only need to average over a few epochs to obtain the results in Table 3, which is much more affordable.

Comment: [2nd Question] What are the number of training epochs for GCN (NS)?

Response: The number of epochs changes from dataset-to-dataset, and can even vary a bit from trial-to-trial due to the use of early stoppage (based on the validation set) as commonly adopted during training. As reference though, for ogbn-arxiv it is roughly 30 epochs, for IGB-tiny about 20 epochs, and for ogbn-papers100M it is about 40 epochs. Incidentally, MuseGNN is similar (e.g., roughly 40 epochs for ogbn-papers100M).

Comment: [3rd Question] What are the actual training and inference times for the baselines and MuseGNN? This is also a relevant question because the sampling schemes themselves could take different amounts of time.

Response: The average training time for each epoch is reported in Table 3, where MuseGNN is shown to be roughly comparable to the commonly-used NS baselines (the total numbers of epochs are similar; see above). And indeed, while different sampling schemes may require different amounts of time, NS is arguably the fastest (at least conditioned on those in common use), e.g., in our tests, it is faster than ShadowKHop on baseline GNNs. Hence there are unlikely to be any significantly faster variants than what we have already shown in Table 3. Please also see Sections B.2 and B.3 for additional ablations related to sampling.

As for inference, although we did not report specific results in the paper, they follow the same trend as training. Basically, both MuseGNN and the baseline GNNs all skip the backward pass and remain similar in run-time for the forward inference pass. However, if the reviewer feels these numbers are important, we can easily add them to a revision.

Thanks for the follow-up comments. Per the reviewer's suggestion, we have quickly uploaded a new version of the paper with interpretability de-emphasized. This mainly involved editing Section 2.2 and modifying a few other phrases. There is now no overlap/conflation with existing GNN work on interpretability, with supporting claims for unfolded GNNs deferred to prior work as needed. The intended message of our paper is surely much more clear now, good suggestion.

In terms of conceptual novelty, we emphasize that, out of all the possible options for scaling GNNs (and unfolded GNNs in particular), we uniquely offer the following:

1. Introduce a novel offline-sampling-based energy formulation, including a penalty term for favoring consistency among the different node representations with the same node index across sampled subgraphs. To our knowledge, this form of regularization has never been used before.
2. Derive an unbiased estimate of the full graph energy for certain choices of sampling (Proposition 3.1).
3. Accommodate new convergence guarantees that are independent of the specific choice of sampling method itself (Section 5), and account for sampling and the regularization mentioned above.
4. Demonstrate SOTA performance on the two largest publicly-available graph benchmarks, and competitive performance on others using just a single, simple/transparent architecture.

To reflect the de-emphasis of "interpretability", we have now uploaded a new draft of the paper per Reviewer PP1g's request (any other reviewer suggested changes can easily be incorporated later, as time is now limited and the discussion window will soon close). The required updates mainly involved editing Section 2.2 and modifying a few other phrases. There is now no overlap/conflation with existing GNN work on interpretability, with supporting claims for unfolded GNNs deferred to prior work as needed. Thanks to reviewers for constructively pointing out this source of confusion, and the current openreview draft is notably improved because of it.