PyTorch Geometric Temporal Index (PGT-I)

• PGT-I is an indexing extension that dynamically reconstructs spatiotemporal snapshots from compact index arrays, reducing memory overhead in ST-GNN training.
• It implements index-batching and distributed-index-batching using NumPy views and DDP, significantly improving runtime and scalability over traditional preprocessing.
• Experimental evaluations show dramatic speedups and memory reductions, enabling training on large datasets like full PeMS with models such as DCRNN.

Searching arXiv for the specified papers and related terminology to ground the article in the cited literature. PyTorch Geometric Temporal Index (PGT-I) is an extension to PyTorch Geometric Temporal (PGT) for making spatiotemporal graph neural network (ST-GNN) training scalable on HPC systems. It integrates distributed data parallel training and two novel strategies—index-batching and distributed-index-batching—and replaces explicit materialization of overlapping spatiotemporal snapshots with dynamic runtime construction from compact index arrays (Ockerman et al., 15 Jul 2025). In historical terms, PGT-I extends the framework introduced in "PyTorch Geometric Temporal: Spatiotemporal Signal Processing with Neural Machine Learning Models" (Rozemberczki et al., 2021), which provided temporal snapshot iterators, graph snapshot objects, benchmark dataset loaders, and recurrent or attention-based spatiotemporal neural layers, but did not define a formal indexing scheme under the name "PyTorch Geometric Temporal Index."

1. Terminological scope and historical position

The 2021 PGT framework is a Python deep learning framework for spatiotemporal graph learning built on top of the PyTorch/PyTorch Geometric ecosystem. Its stated goal is to make temporal geometric deep learning available in a unified, easy-to-use framework, with support for temporal graph data handling, spatiotemporal neural network layers, benchmark datasets, GPU acceleration, and simple training and evaluation pipelines (Rozemberczki et al., 2021). The framework was presented as the first open-source GPU-accelerated deep learning library for parametric spatiotemporal graph learning.

Within that earlier framework, the closest precursor to an indexing concept was the snapshot-based temporal organization mechanism: datasets are represented as sequences of snapshots, iterators return snapshots in temporal order, and direct indexing allows retrieval of a specific time step. The paper explicitly states, however, that PGT is not itself an “index” in the database sense and does not define a separately named taxonomy called PGT-I. In that context, “PGT-I” can only be interpreted informally as the library’s internal organization of temporal graph data and models.

The 2025 paper formalizes the term. "PGT-I: Scaling Spatiotemporal GNNs with Memory-Efficient Distributed Training" (Ockerman et al., 15 Jul 2025) introduces PyTorch Geometric Temporal Index as a named extension of PGT. Its central motivation is that conventional ST-GNN pipelines are constrained primarily by preprocessing-induced memory blow-up rather than by model definition alone. The extension therefore reframes indexing from a convenience feature into the core systems abstraction for scalable ST-GNN training.

2. Spatiotemporal data model and the indexing principle

PGT defines three spatiotemporal input types through a discrete temporal snapshot view. The first is a dynamic graph with temporal signal,

$$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}_1),\dots,(\mathcal{G}_T,\mathbf{X}_T)\right\},$$

with constant vertex set and $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$. The second is a dynamic graph with static signal,

$$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$

again with constant vertex set. The third is a static graph with temporal signal,

$$\mathcal{D}=\left\{ (\mathcal{G},\mathbf{X}_1),\dots,(\mathcal{G},\mathbf{X}_T)\right\},$$

with $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$ for all $t$ (Rozemberczki et al., 2021). These definitions are the backbone of the PGT data model.

PGT operationalizes these regimes through three types of Spatiotemporal Signal Iterators. These iterators avoid redundant storage, return a graph snapshot at each time step, can be indexed directly for access to a specific snapshot, and support advanced temporal batching. A graph snapshot contains node features, labels, edge indices, and edge weights. Internally, these snapshots are stored as NumPy arrays and returned as PyTorch Geometric Data objects, enabling easy GPU transfer and integration with PyTorch workflows (Rozemberczki et al., 2021).

PGT-I preserves the snapshot view but changes how snapshots are represented in memory. Typical ST-GNN workflows use sliding window analysis to turn a long spatiotemporal sequence into many overlapping input and target pairs. The paper argues that existing open-source implementations often store all of these snapshots explicitly, causing the same raw temporal data to be copied repeatedly across overlapping windows and inflating both the $x$ input array and the $y$ label array (Ockerman et al., 15 Jul 2025). The core indexing principle of PGT-I is therefore to stop materializing large numbers of redundant spatiotemporal snapshots during preprocessing and instead construct them on demand at runtime from compact index arrays.

Under index-batching, the system stores one copy of the original time-series graph data and a list or array of indices indicating where each snapshot begins. If the horizon is $horizon$, then a snapshot starting at index start is reconstructed as feature input data[start : start + horizon] and label target data[start + horizon : start + (2 * horizon)]. The paper highlights that NumPy views are used, so slices are references rather than copied arrays (Ockerman et al., 15 Jul 2025). This makes the indexing abstraction the central mechanism by which PGT-I reduces memory overhead.

3. Architecture, workflow, and distributed execution

The original PGT library is organized around four major technical building blocks: spatiotemporal signal iterators, neural network layer definitions, integrated benchmark dataset loaders, and train-test temporal splitting (Rozemberczki et al., 2021). Neural network layer definitions follow design principles of non-proliferation of classes, hyperparameter inspection and type hinting, limited public methods, and auxiliary standalone layers. Train-test temporal splitting separates the last fraction of snapshots for testing, preserves iterator type, and does not interfere with semi-supervised strategies like node masking.

PGT-I retains the general PGT workflow but replaces the expensive “materialize everything first” preprocessing strategy with index-based snapshot reconstruction. Its high-level workflow is: load the raw spatiotemporal dataset once; standardize data using only a single stored copy; store the original data together with an array of starting indices or graph IDs instead of every overlapping $x$ and $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$0 snapshot; reconstruct each snapshot dynamically using slicing or views during training; and optionally run this workflow on GPU only or distribute it across workers with DDP (Ockerman et al., 15 Jul 2025). The paper explicitly states that PGT-I is integrated into the PGT codebase with minimal changes to the existing workflow.

PGT-I uses Dask together with PyTorch Distributed Data Parallel through dask-pytorch-ddp. For distributed-index-batching, the dataset is shuffled globally at the start of each epoch, partitioned across workers, and processed such that each worker independently performs forward and backward passes while gradients are synchronized via AllReduce. Parameters are updated locally on each worker after gradient averaging (Ockerman et al., 15 Jul 2025). The paper notes that each worker maintains its own in-memory copy of the dataset; because the dataset is compact under index-batching, this becomes feasible. The system thereby avoids inter-worker communication for sample loading during the epoch, leaving DDP gradient synchronization as the primary communication path.

Component	Role
Spatiotemporal Signal Iterators	Return temporal graph snapshots in order
Index-batching	Store original data plus index array
Distributed-index-batching	Combine index-batching with DDP across workers

This architecture is significant because it shifts the systems bottleneck from redundant preprocessing storage to dynamic reconstruction at batch time. A plausible implication is that the practical scalability of ST-GNN training depends not only on model parallelism or graph partitioning strategies, but on how temporal examples are indexed and reconstructed.

4. Mathematical and algorithmic formulation

The PGT-I paper describes an ST-GNN forecasting setting with a static graph and dynamic signals,

$\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$1

where $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$2 denotes nodes, $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$3 denotes edges, and $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$4 is a weighted adjacency matrix (Ockerman et al., 15 Jul 2025). Given past observations $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$5, the model predicts future observations according to

$\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$6

The paper also gives a standard preprocessing algorithm: load file, extract node data, create $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$7 and $\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$8 with sliding windows, stack them, and normalize using training statistics,

$\mathbf{X}_t\in\mathbb{R}^{|V|\times d}$9

The size of the preprocessed representation is derived as

$$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$0

This expression captures the duplication caused by overlapping windows together with the factor of $$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$1 from separately storing $$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$2 and $$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$3 (Ockerman et al., 15 Jul 2025).

Index-batching reduces the representation to

$$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$4

Instead of multiplying storage by horizon and duplicating labels, it stores one original dataset copy and one index array. The runtime strategy is correspondingly to reconstruct windows at batch time from indices, use NumPy views or tensor slicing to avoid copies, and, on GPU, move the compact dataset to GPU first and keep all operations there (Ockerman et al., 15 Jul 2025).

Within the broader PGT framework, training for node-level regression or forecasting is described in terms of mean squared error over snapshots,

$$\mathcal{D}=\left\{ (\mathcal{G}_1,\mathbf{X}),\dots,(\mathcal{G}_T,\mathbf{X})\right\},$$5

and two optimization regimes are presented: incremental backpropagation, in which loss is computed, backpropagated, and weights are updated after each snapshot; and cumulative backpropagation, in which losses are accumulated across snapshots in an epoch and backpropagation occurs once at epoch end (Rozemberczki et al., 2021). This suggests a methodological continuity: PGT-I changes data representation and distributed execution while remaining compatible with established spatiotemporal training pipelines.

5. Experimental evaluation and scaling characteristics

All PGT-I experiments were run on ALCF Polaris, with one compute node comprising an AMD EPYC Milan 7543P 32-core CPU at 2.8 GHz, 512 GB DDR4 RAM, and 4 NVIDIA A100 GPUs, connected through HPE Slingshot 11 in a Dragonfly topology with adaptive routing. Distributed experiments used up to 128 GPUs and 32 compute nodes (Ockerman et al., 15 Jul 2025). The datasets used were Chickenpox-Hungary, Windmill-Large, METR-LA, PeMS-BAY, PeMS-All-LA, and the full PeMS dataset. The paper reports dramatic preprocessing expansion, including PeMS-All-LA from 2.12 GB to 102.08 GB and PeMS from 8.71 GB to 419.46 GB.

The main model in the PGT-I study is DCRNN. Compared implementations include baseline PyTorch DCRNN, PGT-DCRNN, PGT-DCRNN with index-batching, PGT-DCRNN with GPU-index-batching, distributed-index-batching, baseline DDP, and generalized-distributed-index-batching for larger-than-memory discussion (Ockerman et al., 15 Jul 2025). Additional validation uses A3T-GCN on METR-LA and ST-LLM on PeMS-BAY. The common training setup uses a standard split of 70% train, 10% validation, and 20% test, Adam optimization, hyperparameters mostly from prior work, and results averaged over 10 runs in single-GPU benchmark experiments.

A central case study compares DCRNN with PGT-DCRNN on PeMS-All-LA for a single epoch. DCRNN records runtime of 68.48 min, max system memory of 371.25 GB, and max GPU memory of 24.84 GB, whereas PGT-DCRNN records runtime of 4.48 min, max system memory of 259.84 GB, and max GPU memory of 1.58 GB (Ockerman et al., 15 Jul 2025). This establishes that the baseline PGT framework already improves runtime and memory relative to a baseline PyTorch implementation, before the introduction of PGT-I indexing methods.

For single-GPU index-batching, memory reductions are substantial while runtime and accuracy remain nearly unchanged. On Chickenpox-Hungary, memory changes from 1093 MB to 1089 MB; on Windmill-Large, from 2455 MB to 1304 MB, a reduction of about 46.88%; and on PeMS-BAY, from 4497 MB to 1335 MB, a reduction of about 70.31% (Ockerman et al., 15 Jul 2025). Runtime changes by less than 1% overall, and MAE remains essentially the same.

The full PeMS dataset is the decisive scale test. Standard PGT-DCRNN crashes during preprocessing on the full PeMS dataset, whereas index-batching enables training with maximum memory of 45.75 GB (Ockerman et al., 15 Jul 2025). The paper further states that PGT-I enables the first-ever training of an ST-GNN on the entire PeMS dataset without graph partitioning. Under GPU-index-batching, runtime improves from 332.79 min to 290.54 min, a speedup of 12.70%, CPU memory decreases from 45.75 GB to 18.18 GB, a reduction of 60.25%, and GPU memory increases from 5.43 GB to 18.60 GB.

Distributed-index-batching shows strong scaling. Compared to single-GPU index-batching on PeMS, it achieves up to 88.33x speedup including preprocessing and up to 99.81x speedup for training time alone with 128 GPUs. Compared with baseline DDP, the speedup ranges from 2.18x on 4 GPUs to 13.10x on 128 GPUs (Ockerman et al., 15 Jul 2025). For larger-than-memory scenarios, generalized-distributed-index-batching improves epoch time by up to 2.28x over baseline DDP and reduces memory usage from 479.66 GB to 53.28 GB with four workers. At the same time, as the number of GPUs increases, optimal training MAE rises from 1.84 on 1 GPU to 2.46 on 128 GPUs; the paper attributes this to large-batch training effects and notes that learning-rate scaling reduces the degradation.

6. Model ecology, applications, and limitations

The PGT model ecology is broad. The 2021 paper categorizes implemented spatiotemporal models by temporal layer, GNN layer, proximity order, and multi-type edge support, with examples including DCRNN, GConvGRU, GConvLSTM, GC-LSTM, DyGrAE, LRGCN, EGCN-H, EGCN-O, T-GCN, A3T-GCN, AGCRN, MPNN LSTM, STGCN, ASTGCN, MSTGCN, GMAN, MTGNN, and AAGCN (Rozemberczki et al., 2021). The case study architecture pattern is a recurrent or spatiotemporal graph layer followed by nonlinearity, dropout, and a final feedforward prediction layer. This model taxonomy is the organizational backdrop against which PGT-I introduces an indexing and training systems contribution rather than a new spatiotemporal layer family.

The application scope of the PGT ecosystem is equally broad. Empirical evaluation in the 2021 paper includes epidemiological forecasting, ride-hail or delivery demand prediction, web-traffic management, and social media interaction prediction, with benchmark datasets such as Chickenpox Hungary, Windmill Large/Medium/Small, Pedal Me Deliveries, Wikipedia Math, Twitter Tennis RG, Twitter Tennis UO, Covid19 England, Montevideo Buses, and MTM-1 Hand Motions (Rozemberczki et al., 2021). The 2025 PGT-I study concentrates on scaling workloads such as Chickenpox-Hungary, Windmill-Large, METR-LA, PeMS-BAY, PeMS-All-LA, and full PeMS, thereby emphasizing large-scale traffic and sensor-network forecasting (Ockerman et al., 15 Jul 2025).

Two limitations are explicit in the literature. First, the 2021 experiments report that most recurrent GNNs have similar predictive performance across node regression tasks; there is no universal “silver bullet” model, and on Wikipedia Math the cumulative backpropagation strategy can hurt performance relative to incremental updates (Rozemberczki et al., 2021). Second, the 2025 scaling study observes that accuracy worsens somewhat as GPU count increases because of large-batch optimization effects, and that baseline DDP can exhibit lower memory use at very large worker counts because its per-worker data partitions shrink (Ockerman et al., 15 Jul 2025). These observations caution against interpreting PGT-I as a universal improvement along every axis. Its principal contribution is the conversion of ST-GNN training from a memory-bound preprocessing problem into an index-based runtime problem, with substantial gains in feasibility, system memory, and throughput on large spatiotemporal workloads.