GRAPES: Learning to Sample Graphs for Scalable Graph Neural Networks

Thank you for your feedback. Here are our answers to your comments and questions.

1. 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.

2. 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.

3. 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].

4. 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.

5. 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.

6. 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.

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

8. 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.

9. Typo for A_0: Thank you for pointing this out. We fixed this in the paper.

10. 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.

References

[1] Balın, M.F. and Çatalyürek, Ü.V., 2022. yer-neigh (BOR) Sampling: Defusing Neighborhood Explosion in GNNs. arXiv preprint arXiv:2210.13339.

[2] Wenqian Li, Yinchuan Li, Zhigang Li, Jianye Hao, and Yan Pang. Dag matters! Gflownets enhanced explainer for graph neural networks. arXiv preprint arXiv:2303.02448, 2023.

[3] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381–27394, 2021a.

[4] Wenhao Gao, Tianfan Fu, Jimeng Sun, and Connor Coley. Sample efficiency matters: a benchmark for practical molecular optimization. Advances in Neural Information Processing Systems, 35: 21342–21357, 2022.

[5] Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786–9801. PMLR, 2022.

Thank you for your feedback. Here are our answers to your comments and questions.

1. 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.
3. 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.
5. 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.
6. 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.

Thank you for your feedback. Here are our answers to your comments and questions.

1. 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.
2. 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.
3. 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.
4. 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.

References

[1] Geng H, Chen C, He Y, Zeng G, Han Z, Chai H, Yan J. Pyramid Graph Neural Network: A Graph Sampling and Filtering Approach for Multi-scale Disentangled Representations. InProceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining 2023 Aug 6 (pp. 518-530).

[2] Ruiz, L., Chamon, L.F. and Ribeiro, A., 2023. Transferability properties of graph neural networks. IEEE Transactions on Signal Processing.

[3] Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E. J., Everett, K. E., ... & Bengio, Y. (2022, September). GFlowNets and variational inference. In The Eleventh International Conference on Learning Representations.

Thank you for your feedback. Here are our answers to your comments and questions.

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.

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.

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.

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.

References

[1] Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., & Bengio, E. (2023). Gflownet foundations. Journal of Machine Learning Research, 24(210), 1-55.

[2] Ahmed, K., Zeng, Z., Niepert, M., & Van den Broeck, G. (2022, September). SIMPLE: A Gradient Estimator for k-Subset Sampling. In The Eleventh International Conference on Learning Representations.

[3] Abu-El-Haija, S., Perozzi, B., Kapoor, A., Alipourfard, N., Lerman, K., Harutyunyan, H., Ver Steeg, G. and Galstyan, A., 2019, May. Mixhop: Higher-order graph convolutional architectures via sparsified neighborhood mixing. In international conference on machine learning (pp. 21-29). PMLR.

[4] Hamilton, W., Ying, Z. and Leskovec, J., 2017. Inductive representation learning on large graphs. Advances in neural information processing systems, 30.