mlsanityData to 14 Oct 2026

PGT-I: Scaling Spatiotemporal GNNs with Memory-Efficient Distributed Training

arXiv 2507.11683 ↗First posted 15 Jul 2025cs.DC cs.AI cs.LGcited by 3

Not on the map. Why?

Abstract (arXiv)

Spatiotemporal graph neural networks (ST-GNNs) are powerful tools for modeling spatial and temporal data dependencies. However, their applications have been limited primarily to small-scale datasets because of memory constraints. While distributed training offers a solution, current frameworks lack support for spatiotemporal models and overlook the properties of spatiotemporal data. Informed by a scaling study on a large-scale workload, we present PyTorch Geometric Temporal Index (PGT-I), an extension to PyTorch Geometric Temporal that integrates distributed data parallel training and two novel strategies: index-batching and distributed-index-batching. Our index techniques exploit spatiotemporal structure to construct snapshots dynamically at runtime, significantly reducing memory overhead, while distributed-index-batching extends this approach by enabling scalable processing across multiple GPUs. Our techniques enable the first-ever training of an ST-GNN on the entire PeMS dataset without graph partitioning, reducing peak memory usage by up to 89% and achieving up to a 11.78x speedup over standard DDP with 128 GPUs.

Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Builds on 34 of its 78 references that are in the map. First posted on arXiv 15 Jul 2025.

Builds on

34 of 78 references.

References · 34

  1. 2020Scaling Laws for Neural Language Modelscited by 6,586
  2. 2017Inductive Representation Learning on Large Graphscited by 6,320
  3. 2019Fast Graph Representation Learning with PyTorch Geometriccited by 2,434
  4. 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,492
  5. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
  6. 2017Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecastingcited by 1,136
  7. 2019Graph WaveNet for Deep Spatial-Temporal Graph Modelingcited by 829
  8. 2017Large Batch Training of Convolutional Networkscited by 535
  9. 2020Adaptive Graph Convolutional Recurrent Network for Traffic Forecastingcited by 530
  10. 2018T-GCN: A Temporal Graph ConvolutionalNetwork for Traffic Predictioncited by 401
  11. 2014One weird trick for parallelizing convolutional neural networkscited by 264
  12. 2021Dynamic Graph Convolutional Recurrent Network for Traffic Prediction: Benchmark and Solutioncited by 157
22 more
  1. 2020DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphscited by 139
  2. 2021PyTorch Geometric Temporal: Spatiotemporal Signal Processing with Neural Machine Learning Modelscited by 118
  3. 2020PyTorch Distributed: Experiences on Accelerating Data Parallel Trainingcited by 90
  4. 2020Normalization Techniques in Training DNNs: Methodology, Analysis and Applicationcited by 90
  5. 2024Spatial-Temporal Large Language Model for Traffic Predictioncited by 83
  6. 2020A3T-GCN: Attention Temporal Graph Convolutional Network for Traffic Forecastingcited by 54
  7. 2022Learning the Evolutionary and Multi-scale Graph Structure for Multivariate Time Series Forecastingcited by 36
  8. 2022Distributed Graph Neural Network Training: A Surveycited by 32
  9. 2019GADMM: Fast and Communication Efficient Framework for Distributed Machine Learningcited by 30
  10. 2021Conditional Local Convolution for Spatio-temporal Meteorological Forecastingcited by 30
  11. 2022Parallel and Distributed Graph Neural Networks: An In-Depth Concurrency Analysiscited by 29
  12. 2020AST-GCN: Attribute-Augmented Spatiotemporal Graph Convolutional Network for Traffic Forecastingcited by 27
  13. 2019Learning Internal Representations (COLT 1995)cited by 25
  14. 2017Convergence Analysis of Distributed Stochastic Gradient Descent with Shufflingcited by 23
  15. 2021Chickenpox Cases in Hungary: a Benchmark Dataset for Spatiotemporal Signal Processing with Graph Neural Networkscited by 19
  16. 2019Graph-Partitioning-Based Diffusion Convolutional Recurrent Neural Network for Large-Scale Traffic Forecastingcited by 16
  17. 2022STGC-GNNs: A GNN-based traffic prediction framework with a spatial-temporal Granger causality graphcited by 12
  18. 2023DistTGL: Distributed Memory-Based Temporal Graph Neural Network Trainingcited by 12
  19. 2019Incrementally Improving Graph WaveNet Performance on Traffic Predictioncited by 10
  20. 2022Graph Convolutional Networks for Traffic Forecasting with Missing Valuescited by 7
  21. 2024Beyond Gradient Averaging in Parallel Optimization: Improved Robustness through Gradient Agreement Filteringcited by 3
  22. 2024Accelerating Biological Spatial Cluster Analysis with the Parallel Integral Image Techniquecited by 1

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 1; 2026: 2.

3 papers.

  1. 2025A Distributed Framework for Causal Modeling of Performance Variability in GPU Tracescited by 0
  2. 2026Learning from Complexity: Exploring Dynamic Sample Pruning of Spatio-Temporal Trainingcited by 0
  3. 2026Two-level domain-decomposition AdaGrad method for scalable training of graph neural networkscited by 0