# Learning Large Graph Property Prediction via Graph Segment Training

arXiv: https://arxiv.org/abs/2305.12322 · HTML: https://mlsanity.com/paper/2305.12322/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 21 May 2023.

## Builds on (most cited first, 29 of 29)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Recipe for a General, Powerful, Scalable Graph Transformer](https://mlsanity.com/paper/2205.12454/) (2022)
- [Benchmarking Graph Neural Networks](https://mlsanity.com/paper/2003.00982/) (2020)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling](https://mlsanity.com/paper/1801.10247/) (2018)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [Do Transformers Really Perform Bad for Graph Representation?](https://mlsanity.com/paper/2106.05234/) (2021)
- [Design Space for Graph Neural Networks](https://mlsanity.com/paper/2011.08843/) (2020)
- [Adaptive Sampling Towards Fast Graph Representation Learning](https://mlsanity.com/paper/1809.05343/) (2018)
- [Layer-Dependent Importance Sampling for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1911.07323/) (2019)
- [Scaling Graph Neural Networks with Approximate PageRank](https://mlsanity.com/paper/2007.01570/) (2020)
- [Machine Learning on Graphs: A Model and Comprehensive Taxonomy](https://mlsanity.com/paper/2005.03675/) (2020)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [The Adverse Effects of Code Duplication in Machine Learning Models of Code](https://mlsanity.com/paper/1812.06469/) (2018)
- [GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings](https://mlsanity.com/paper/2106.05609/) (2021)
- [The Implicit and Explicit Regularization Effects of Dropout](https://mlsanity.com/paper/2002.12915/) (2020)
- [Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets](https://mlsanity.com/paper/2203.04810/) (2022)
- [A Large-Scale Database for Graph Representation Learning](https://mlsanity.com/paper/2011.07682/) (2020)
- [PipeGCN: Efficient Full-Graph Training of Graph Convolutional Networks with Pipelined Feature Communication](https://mlsanity.com/paper/2203.10428/) (2022)
- [BNS-GCN: Efficient Full-Graph Training of Graph Convolutional Networks with Partition-Parallelism and Random Boundary Node Sampling](https://mlsanity.com/paper/2203.10983/) (2022)
- [GraphFM: Improving Large-Scale GNN Training via Feature Momentum](https://mlsanity.com/paper/2206.07161/) (2022)
- [Graph Classification via Deep Learning with Virtual Nodes](https://mlsanity.com/paper/1708.04357/) (2017)
- [Understanding GNN Computational Graph: A Coordinated Computation, IO, and Memory Perspective](https://mlsanity.com/paper/2110.09524/) (2021)
- [Graph Traversal with Tensor Functionals: A Meta-Algorithm for Scalable Learning](https://mlsanity.com/paper/2102.04350/) (2021)

## Built on by (most cited first, 9 of 9)

- [On the Scalability of GNNs for Molecular Graphs](https://mlsanity.com/paper/2404.11568/) (2024)
- [Spatio-Spectral Graph Neural Networks](https://mlsanity.com/paper/2405.19121/) (2024)
- [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://mlsanity.com/paper/2308.13490/) (2023)
- [LayerDAG: A Layerwise Autoregressive Diffusion Model for Directed Acyclic Graph Generation](https://mlsanity.com/paper/2411.02322/) (2024)
- [Nonparametric Teaching for Graph Property Learners](https://mlsanity.com/paper/2505.14170/) (2025)
- [Graph neural networks with configuration cross-attention for tensor compilers](https://mlsanity.com/paper/2405.16623/) (2024)
- [ReHub: Linear Complexity Graph Transformers with Adaptive Hub-Spoke Reassignment](https://mlsanity.com/paper/2412.01519/) (2024)
- [Fractal Graph Contrastive Learning](https://mlsanity.com/paper/2505.11356/) (2025)
- [Better Call Graphs: A New Dataset of Function Call Graphs for Malware Classification](https://mlsanity.com/paper/2512.20872/) (2025)

## Abstract (arXiv)

Learning to predict properties of large graphs is challenging because each prediction requires the knowledge of an entire graph, while the amount of memory available during training is bounded. Here we propose Graph Segment Training (GST), a general framework that utilizes a divide-and-conquer approach to allow learning large graph property prediction with a constant memory footprint. GST first divides a large graph into segments and then backpropagates through only a few segments sampled per training iteration. We refine the GST paradigm by introducing a historical embedding table to efficiently obtain embeddings for segments not sampled for backpropagation. To mitigate the staleness of historical embeddings, we design two novel techniques. First, we finetune the prediction head to fix the input distribution shift. Second, we introduce Stale Embedding Dropout to drop some stale embeddings during training to reduce bias. We evaluate our complete method GST-EFD (with all the techniques together) on two large graph property prediction benchmarks: MalNet and TpuGraphs. Our experiments show that GST-EFD is both memory-efficient and fast, while offering a slight boost on test accuracy over a typical full graph training regime.