Paper: GRAPES: Learning to Sample Graphs for Scalable Graph Neural Networks (ICLR-2024)

Reviewer comment (Theory):
Above Eq. 6, it says that $P_B(s_{l} | s_{l+1}) = 1$ -- this implies that the only way to get to $s_{l+1}$ is through $s_{l}$. This strong modeling assumption stems from "s" being the entire path of generated adjacencies $\{A_0, \dots, A_l \}$. This produces a special case of "finite state machines" that specifically have states looking like trees.

Author response:
The tree assumption: Indeed, the graph over which the GFlowNet acts, which can indeed be seen as a finite automata, is a tree. We have highlighted and explained this fact further just above Equation 6. 

   The reason we made this assumption is technical. If we do not add to the states the order when the nodes are added, we must assign some probability to all the previous states that we could come from. However, counting how many such states there are is surprisingly complex, as nodes could be added because they are 1-hop or 2-hops away from the initial target nodes. Ultimately, we do not believe this assumption will impact performance in any way.

Reviewer comment (Experiments):
Only one graph has more than 1 million nodes, whereas one central theme of the paper is about scale.

Author response:
Larger datasets: We agree with the reviewer. However, we followed the recent related work such as GAS, LMC, and LABOR [1] (NeurIPS 2023) and used the datasets they reported.  However, other datasets in our experiments were also relatively large, like Yelp, with more than 700K nodes and almost 7 million edges, and Reddit, with more than 200K nodes and 11 million edges. Note that the full-batch GCN leads to out-of-memory errors with 16GB GPU memory for Reddit, Yelp, and products.

Reviewer comment (Experiments):
The paper should consider applying their method to a problem space beyond graph sampling (or otherwise show compelling use-cases for graph sampling). Specifically, could this method be used to *explain graphs*? For example, in the integrated-gradient sense: determining the presence of which nodes or edges would cause a certain prediction.

Author response:
Application appeal: This paper focuses on improving the scalability of GNNs with sampling. Other applications, such as the explainability of GNNs with GFlowNets, are orthogonal topics to sampling and require further research. These would be interesting for future work. However, adding more applications to our paper beyond sampling would require going less in-depth with the method. It may change the paper into a survey paper on the applications of GFlowNets. For other applications of GFlowNets, as mentioned in the related work, there are several studies, such as graph explainability [2], molecule design, and material science [3, 4, 5].

Reviewer comment (Reproducibility):
While the paper includes information about training, it should also include information on how to do inference. Given a node at a (large) input test graph, are samples taken or the full graph around $n$?

Author response:
The inference specification: Section 5.2 mentions that we use the full graph for the evaluation, which is the same protocol all our baselines have used.

Reviewer comment (Novelty):
The paper is missing references on learnable sampling, such as:  
* DSKReG; CIKM'2021  
* "Performance-adaptive sampling strategy towards fast and accurate GNNs.", KDD'2021  
* Submix; UAI'2023

Author response:
New references: Thank you for the pointers. We now included them in the related work, although the second reference, "Performance-adaptive sampling strategy towards fast and accurate GNNs.", KDD'2021, was already in the paper.

Reviewer comment (Reproducibility):
How is the $GNN_F$ parameterized? Does it train a scalar for every node (i.e., lookup 1D embedding table) or is it a function of features? In my understanding, GNN_F models $P_F$ (correct me if I am wrong), i.e., it should have support on the nodes.

Author response:
Architecture of GNN_F: The architecture of GNN_F is composed of a two-layer GCN with support on the input nodes and its neighbors at different levels, and its outputs are the sampling probabilities for each node. We do not train a scalar for every node as an input to this GCN but instead use the node features present in the datasets we experiment with.

Reviewer comment (Reproducibility):
Is the reward measured only on end states (on the "sum" of the sampled list of adjacency matrices) or on every intermediate state (e.g., sum of adjacencies at that point)?

Author response:
Reward measured on the end state: Yes, in GFlowNet, the reward is given in the end state, not the middle states.

Reviewer comment (Reproducibility):
Can an edge be sampled twice? Does this have any impact on the GCN model?

Author response:
Repeated edges: In each layer, an edge cannot be sampled twice. However, one edge can be sampled in both layers. This should not have a negative effect on the GCN because in the full-batch GCN, after two layers, the target nodes are updated twice by the 1-hop neighbors and once updated by the 2-hop neighbors. Currently, all the layer-wise sampling methods in our baselines (FastGCN, LADIES, AS-GCN) follow the same regime. However, one can design a layer-wise sampling method that samples each edge only once across different layers so that in each layer, only new edges are sampled. Studying that effect is an interesting question that can be pursued in future work.

Reviewer comment (Writing):
$Z(s_0)$ in the text following Eq. 3 -- it is not clear whether scalar $log Z$ is modeled or if it is a constant and removed.

Author response:
Value of Z: as mentioned in the lines under E.q. 7, Z is a learnable scalar, updated in each backpropagation step. However, the value of Z is kept the same for the different batches. In the paper, we refer to this point as one of the known limitations.

Reviewer comment (Experiments):
It is not clear if the baselines were tuned on a validation set. Why was the batch size fixed to 256 for the main results table?

Author response:
Baselines tuned on the validation set: We used the hyperparameters reported by the baselines. For a fair comparison, we only changed the hidden dimensions so that all the methods have the same GCN hidden dimension, i.e., the same computation capacity. 
2. Why the batch size of 256: We used this batch size because most of the baselines used this number, which required the least amount of change in the baselines and gave more trust in the hyperparameters of the baselines.

Reviewer comment (Evaluation):
There is a concern about the appearance of low F1-score compared to what is reported in other papers. Granted that this is in the transductive setting, it is uncertain if that should cause such a decrease in performance. For instance, according to the GraphSAINT paper, it achieves 96.6% on Reddit in the inductive setting. In this paper, the result is much lower (80.50).

Author response:
The baselines’ performance was lower than reported: Each of the baselines used different GNN architectures, sampling rates, batch sizes, etc., which makes it difficult to judge if the reason for the improvement in the accuracy is the sampling method or these differences. Therefore, for a fair comparison, we kept all the other factors equal for all the methods. For GraphSAINT, significant changes were made in the GNN architecture (from two higher-order aggregators to two GCN layers) and the number of nodes per batch (from thousands to 256+256+256 nodes). Both of these factors, especially the number of nodes, have a high effect on the final accuracy of the GNN with any sampling method. Therefore, the only way to say which method performs the best is to keep these factors the same.

Reviewer comment (Experiments):
The comparison against GraphSAINT may not be fair. Fixing 256 size and 256 samples does not seem appropriate for GraphSAINT, which uses a sample per minibatch instead of such a low number of samples. Also, the GraphSAINT paper shows sample sizes of a few thousands, while this paper uses sample size only up to $2^9$. Additionally, the node sampler used seems to have the lowest performance compared to other samplers of GraphSAINT, so other samplers such as edge, RW, or Multidimensional RW should have been considered.

Author response:
The baselines’ performance was lower than reported: Each of the baselines used different GNN architectures, sampling rates, batch sizes, etc., which makes it difficult to judge if the reason for the improvement in the accuracy is the sampling method or these differences. Therefore, for a fair comparison, we kept all the other factors equal for all the methods. For GraphSAINT, significant changes were made in the GNN architecture (from two higher-order aggregators to two GCN layers) and the number of nodes per batch (from thousands to 256+256+256 nodes). Both of these factors, especially the number of nodes, have a high effect on the final accuracy of the GNN with any sampling method. Therefore, the only way to say which method performs the best is to keep these factors the same. 
4. Why we chose _node-sampler_ in GraphSAINT: One of the challenges we faced was determining how to compare different sampling methods, i.e., node-wise, layer-wise, and subgraph-wise because keeping the number of sampled nodes consistent is difficult across these methods. Our question is, given an equal budget of sampled nodes, which method leads to the best performance. The _node-sampler_ in graphSAINT is the only sampler in this method that allows us to specify how many nodes we want to sample. The _RW-sampler_ starts with $n$ batch nodes and does $k$ walks. Ultimately, it will sample a maximum of $n*(k+1)$ nodes. In Figure 3, we wanted to evaluate different methods when sampling fewer nodes per layer than the batch size (here $n=256$). However, with the _RW-sampler_, the minimum number of nodes sampled in each neighborhood hop is n. We have a paragraph in Appendix 2 (A2) explaining this in the paper.

Reviewer comment (Experiments):
Architectures beyond GCN should be considered. If the sampling approach improves over baselines for multiple architectures such as GAT, GIN, and SAGE, it would create a stronger case for the proposed sampling approach. Currently, the central claim that "GRAPES outperforms state-of-the-art sampling-based methods" does not seem justified. What has been shown is that GRAPES outperforms other methods on a specific GCN architecture under small sample sizes and number of samples.

Author response:
Using different GNN architectures: We have tested GRAPES on several different datasets with different structures and connectivities, from single-label to multi-label nodes. GRAPES has shown superior performance to the baselines in these settings. While experiments with a different GNN architecture are interesting, including them would require re-executing all our main experiments. We consider that the fact that GRAPES improves upon the baselines under different datasets and tasks provides strong evidence of performance generalization. Nonetheless, we will include experiments with a different GNN architecture in the camera-ready. However, testing on different architectures is impossible in a week because we need to tune the parameters for GRAPES, as well as run the baselines. We will include these results in the camera-ready version. About the small sample size point, we argue that having a small sample size is more challenging than a large sample size because less data is available to learn from. Increasing the batch size will also increase memory usage, which is the main problem that graph sampling tries to overcome. Moreover, the smaller batch size allows a larger embedding size or input features. This can be especially useful in graphs with multimodal nodes that need large feature encoders.

Reviewer comment (Evaluation):
The downside of being adaptive is that there is extra computation involved per batch. However, no discussion of training time has been presented. This would have helped with understanding execution time - F1 tradeoff.

Author response:
Runtime report: A full training and testing experiment of GRAPES on our largest dataset, ogbn-products, with 100 epochs takes around 100 minutes. While other methods, such as GAS, have optimized implementations with low-level code that lead to lower runtimes, we found the runtime of GRAPES reasonably low to allow for extensive experimentation. It’s important to note that our main focus was GPU memory optimization. We achieved this goal while being able to run GRAPES on large graphs within a feasible runtime. In the future, we plan to explore optimizations such as using more efficient parallelization when sampling in GRAPES.

Reviewer comment (Novelty):
Some related work is missing, and perhaps also a comparison with other graph sampling baselines from the graph signal processing literature. Check, e.g., "Efficient Sampling Set Selection for Bandlimited Graph Signals Using Graph Spectral Proxies", by Anis and others, and papers therein (specifically, the works of Kovacevic and Moura; Chamon and Ribeiro; Segarra, Marques and Ribeiro; etc.). These papers are part of a subfield of graph signal processing—graph signal sampling—which studies how to sample graphs so as to maximize the preservation of their spectra. Since graph spectral information is typically very correlated with performance in graph machine learning tasks, I believe these are important references/comparisons to include.

Author response:
Graph signal processing literature: Thank you for the pointers; indeed, relevant works. The mentioned paper looks into sampling nodes in a graph for a unique and stable graph signal reconstruction. This work has been used in [1] as a node sampling technique for graph neural networks. We included these works in the related work of our revised paper. Also, the study on “transferability properties of GNNs” [2] uses the weights of a GNN trained on a mid-sized graph and transfers them to larger graphs given the graphon similarity between the graphs. Although this work doesn’t focus on sampling, it’s an interesting approach to scaling GNNs to larger graphs. We included this paper in the related work of our revised paper.

Reviewer comment (Writing):
The explanation of why the method is trained off-policy is not very clear for readers not familiar with reinforcement learning. There is a result which is only mentioned in passing—"Importantly, GFlowNets [...] can learn from off-policy distributions without adjusting the objective"—which is important in justifying the choice of off-policy training, and hence should be described in further detail (perhaps a short subsection) in the camera-ready. It would also be interesting to see empirical comparisons between training off-policy and using gradient estimation methods.

Author response:
On expanding on the off-policy setup: This is a fair point. We changed the description of why Trajectory Balance has favorable properties according to the literature. The sentence we changed is now: “[3] showed that the trajectory balance loss is stable when learning from off-policy distributions without adjusting the objective with importance weights”. We also clarified Section 4.3 and added a larger Appendix section (Appendix A) going into more detail for our off-policy setup, discussing the challenges with computing importance weights for RL being intractable.

Reviewer comment (Experiments):
The numerical experiments only consider node classification tasks.

Author response:
Only node classification: Our method is general enough to support many tasks. We focused on node classification, which already required running a significant number of experiments. However, exploring GRAPES in other tasks is a good direction for future work.

Reviewer comment (Novelty):
Other relevant line of related work is that on the "transferability properties of GNNs". See e.g. the work of Ruiz et al.

Author response:
Graph signal processing literature: Thank you for the pointers; indeed, relevant works. The mentioned paper looks into sampling nodes in a graph for a unique and stable graph signal reconstruction. This work has been used in [1] as a node sampling technique for graph neural networks. We included these works in the related work of our revised paper. Also, the study on “transferability properties of GNNs” [2] uses the weights of a GNN trained on a mid-sized graph and transfers them to larger graphs given the graphon similarity between the graphs. Although this work doesn’t focus on sampling, it’s an interesting approach to scaling GNNs to larger graphs. We included this paper in the related work of our revised paper.

Reviewer comment (Evaluation):
Have you analyzed the specific subgraphs that are sampled by GRAPES in different tasks? What are their characteristics (are they connected? do the sampled nodes have high centrality? etc.). GRAPES sounds like a nice tool for understanding which characteristics of a graph are most important in a given task.

Author response:
Analysis of the sampled nodes: This is a very good question. We plan to focus on this in future work. We believe that evaluating GRAPES on knowledge graphs, which have different relation and node types and multimodal information, can give us more insight about the type of information that is influential in the graph for a GNN.

Reviewer comment (Theory):
The description of the algorithm is problematic. As far as I know, GFlowNet is a method proposed to sample for an energy-based model: it addresses a distribution approximation problem. It is a special case of an RL algorithm. In my view, the paper is a pure RL problem especially since the reward function is clearly defined. I think an RL formulation is straightforward from that. The formulation with GFlowNet is very misleading -- I spent hours before realizing that this is not a distribution approximation problem. Actually, the reward scaling in 4.2 could be avoided within an RL formulation.

Author response:
1. Why we use GFlowNets instead of RL: This is a great question with a nuanced answer. The main difference is in the objective the two approaches aim to optimize. Roughly, RL is about finding a policy that minimizes the loss and converges to a deterministic solution. Instead, GFlowNet aims to sample in proportion to low loss and usually converges to a stochastic solution. The reviewer is correct that GFlowNets are a special case of RL in that recent theoretical results show GFlowNets optimizes a maximum-entropy RL objective (see e.g. [1], Section 7.2).

   The second reason why GFlowNets are useful in our setting is more technical: Its off-policy properties. RL requires a policy distribution. In setup, we need a distribution for sampling exactly $k$ out of $n$ nodes. This is a tricky distribution that scales polynomially in computation in both $k$ and $n$ [1] (we tried to use the codebase of [2], but this failed already in medium-sized graphs). Instead, we used a distribution that is a product of independent Bernoulli’s and used the Gumbel-max trick to sample exactly $k$ positive values from this distribution. This is off-policy: We sample from a different policy than the distribution we optimize. In GFlowNet loss functions, we can use any off-policy distribution without changing the loss, while in RL loss functions, off-policy sampling requires importance weighting, which would again require computing the distribution presented in [2].

   We modified the paper to highlight why we chose GFlowNets instead of (regular) RL. (Section 2.2), and added extra explanations on the off-policy sampling in Section 4.3 and Appendix A.

Reviewer comment (Reproducibility):
The method is much more complex than previous methods because it has this extra learnable component. I don't know whether it is easy for others to apply such a model to a different application. To me, the simplicity of model tuning is more appealing than minor performance improvement: one may not see the improvement if the model cannot be well-tuned.

Author response:
2. Using a more complex method: We agree with the reviewer that our model is more complex than the heuristics-based methods, i.e., FastGCN, LADIES, and GraphSAINT. However, the accuracy improvement of GRAPES compared to these baselines is not minor. We argue that to have an adaptive method, one needs a learnable component in sampling, which necessarily adds complexity. The same happens with AS-GCN, which needs to train an attention layer for sampling and uses significantly more memory than GRAPES. About the reviewer's concern about the other applications of GRAPES, the design of our model is not dependent on any specific architectures or loss functions; therefore, it can be easily applied to other tasks like link prediction.

Reviewer comment (Experiments):
The performance values of baseline methods reported in Table 1 are much lower than those reported in their original papers. I don't know how much I can trust the comparison. For example, Graph-SAINT has f1 scores, 0.511±0.001, 0.966±0.001, and 0.653±0.003 on the Flickr, Reddit, and Yelp datasets. These numbers are much higher than the reported numbers in the submission.

Author response:
3. The baselines’ performance was lower than reported in the original papers: Each baseline used different GNN architectures, sampling rates, batch sizes, etc., which makes it difficult to judge if the reason for their good performance is the sampling method or the mentioned different factors. Therefore, for a fair comparison, we kept all the other factors the same for all the methods. For GraphSAINT, the major changes were the GNN architecture (from two higher order aggregator layers [3, 4] to two GCN layers), the number of nodes per batch (from thousands to $BatchSize+2*SampleSize=256+256+256$ nodes), and the depth of sampling (for some datasets from 4 to 2 layers of depth). Using this uniform setup, where only the sampling method varies, we can clearly show that GRAPES as a sampling method outperforms the baselines.

Reviewer comment (Experiments):
Can you put data statistics in the experiment section? These numbers are important to the understanding of experiment results.

Author response:
4. Dataset statistics: Because of the lack of space in the main body of the paper and the nine-page limit, we put this in the appendix.