Temporal graph neural networks (TGNNs) achieve superior performance in many dynamic graph learning tasks. These models use node state vectors (i.e., memories) to capture and record graph changes (i.e., events), and then exploit these memories to assist prediction. However, this memory-centric design makes multi-GPU training challenging because predictions for later events often depend on memory updates from earlier events. To maintain model quality, existing multi-GPU TGNN training schemes strictly preserve all event dependencies, which limits parallelism and causes expensive cross-GPU memory synchronization. In this work, we observe that many of these dependencies can be handled more flexibly without harming training quality. Specifically, (a) dependencies in TGNN training have high locality in both time and graph topology, (b) the short-range dependencies missed within a large batch can often be recovered cheaply, and (c) many memory updates are not immediately needed. Based on these observations, we propose DepBloom to accelerate multi-GPU TGNN training with three techniques: (1) temporally bounded spatial partitioning to reduce unnecessary cross-GPU synchronization and improve parallelism; (2) intra-batch memory refinement to recover missing short-range temporal dependencies inside enlarged batches; and (3) critical-path optimization to prioritize urgent memory updates while deferring non-critical work during multi-GPU pipeline execution. Experimental results across representative models and diverse datasets show that DepBloom achieves an average speedup of 2.84× and up to 5.16× over state-of-the-art methods, while maintaining comparable or better model accuracy.
Figures & tables
Figure 1. TGNN computing steps in one batch. In many TGNN training pipelines, the memory update of one batch is deferred and executed as the first step of the next batch.
Figure 2. Comparison of four multi-GPU training paradigms for TGNNs: (a) batch parallelism, (b) epoch parallelism, (c) memory parallelism, and (d) pipeline parallelism. Here, Bji denotes the j -th batch in epoch i , and different colors distinguish batches from different epochs. RW denotes memory read and write operations, where different colors indicate different memory copies. Green arrows indicate inter-GPU data movement.
Figure 3. Latencies of training TGN on WIKI-TALK with different numbers of GPUs. (1) Sample&fetch includes neighbor sampling and input transfer from CPU to GPU; (2) memory_sync denotes the transfer of updated node memories across GPUs; (3) memory_update computes new node-memory states required by subsequent GPUs; and (4) model_train includes forward propagation, backward propagation, and parameter updates.
Figure 4. Trade-off between training efficiency and model quality under different batch sizes on WIKI-TALK .
Figure 5. (a) Cross-chain event ratio versus number of sub-chains, comparing naive sequential partitioning against topology-aware partitioning. (b) Cross-chain event ratio and training loss as a function of temporal window size.
Figure 6. (a) Execution modes under different parallel strategies. (b) Speedup versus training loss increase under larger batch sizes. Arrows show the percentage loss increase from BS=1200 to BS=4800.
Figure 7. Urgency mismatch in batch-produced memory updates. (a) Distribution of updated nodes immediately required by the next batch ( Hot Node ) and those that are not ( Cold Node ). (b) Breakdown of the critical path latency.
Figure 8. Design overview of DepBloom.
Figure 9. Comparison between the original pipeline parallelism and Topology-Aware Hierarchical Multi-Chain Parallelism. Using eight consecutive batches (batch size=2) as an example, the original design executes batches in a single pipeline and incurs frequent unnecessary synchronization. Our design partitions events into multiple topology-aware sub-chains and maps each sub-chain to a GPU group for parallel execution across sub-chains with minimal dependency loss.
Figure 10. Comparison of three multi-GPU pipeline execution strategies for TGNN training. (a) Pipelined Small Batches : each GPU executes one small batch per stage in strict sequential order. (b) Pipelined Large Batches (Local Parallelism) : micro-batches within each large batch are scheduled in parallel, each consuming the same stale memory from the previous batch. (c) Pipelined Large Batches (Local Refinement) : same scheduling as (b), but a refinement step (green) is inserted between micro-batches to propagate short-range dependencies locally before the next micro-batch executes. : actual scheduling dependencies that enforce serial execution order; : data usage relationships among micro-batches, which allow parallel execution ; : sample & fetch; : memory update; : memory refinement; : training & backward.
Figure 11. Comparisons between unified node update and priority-aware decoupled node update.
Figure 12. Overall performance comparison on NVIDIA A100 and RTX 4090 across four datasets and five TGNN models. The first two rows show end-to-end speedup over DistTGL 1 , and the last two rows show average precision (AP).
Sample Message
Memory Update
Node Embedding
JODIE
most recent (n=1)
RNN
Identity
TGN
most recent (n=1)
GRU
GAT
APAN
most recent (n=10)
Transformer
Identity
TGAT
uniform (n=10)
Identity
2-layer GAT
DySAT
uniform (n=10)
GAT
RNN
Table 2. Details of TGNN models, with dout=100
Dataset
∣V∣
∣E∣
∣dv∣
∣de∣
WIKI
9,227
157,474
172
100*
REDDIT
11,000
672,447
172
100*
WIKI-TALK
1,140,149
7,833,140
32*
100*
GDELT
16,682
191,290,882
186
413
Table 3. Dataset statistics. ∣dv∣ and ∣de∣ show the dimensions of node features and edge features, respectively.
Figure 13. Scalability on WIKI-TALK and GDELT with APAN , under single- and multi-node GPU configurations. For multi-node, inter-node memory parallelism is combined with DepBloom’s intra-node dependency-aware execution.
Figure 14. Effect of temporally bounded spatial partitioning. (a) Average precision (AP). (b) End-to-end speedup.
Figure 15. Effect of intra-batch memory refinement. SB uses batch size 600, LB uses batch size 2400, and LB+IMR(4) splits each large batch into four chunks of size 600. (a) Normalized validation loss. (b) End-to-end speedup.
Figure 16. Effect of critical-path optimization on multi-GPU pipeline execution. Comparison of pipeline execution with ( W ) and without ( W/O ) the optimization.
Figure 17. Sensitivity of intra-batch refinement and temporal bound. (a) Sensitivity to refinement granularity under different batch sizes on WIKI and REDDIT . Top: Normalized training loss. Bottom: Speedup. (b) Sensitivity to the temporal bound on REDDIT and WIKI-TALK . Normalized training loss.
Figure 18. Latency breakdown of DepBloom across datasets and TGNN models, including preprocessing for temporally bounded spatial partitioning and extra computation from intra-batch memory refinement.
Temporal graph neural networks model how node relationships, graph structure, and features change over time to support prediction on evolving networks.
Temporal Graph Neural Networks (TGNNs) are widely used for learning from dynamic graphs in applications such as recommendation, social network analysis, and traffic forecasting. However, scaling TGNN training to large dynamic graphs remains challenging due to three intertwined bottlenecks: memory I/O, irregular computation, and temporal neighbor sampling. Existing systems often optimize these stages in isolation, leaving substantial performance headroom on the table. We present FAST, a holistic framework that accelerates end-to-end TGNN training by jointly optimizing sampling, memory I/O, and computation. FAST introduces SlimCache, which exploits within-batch compression and cross-batch caching to reduce host-device data movement under limited GPU memory budgets. It further designs thread-efficient graph operators tailored to sparse temporal subgraphs, improving GPU cache locality and reducing the latency of aggregation and edge softmax. In addition, FAST employs a topology-aware sampling strategy that improves CPU cache locality and accelerates temporal neighbor sampling. Extensive experiments on real-world large dynamic graphs show that FAST achieves an average of 2.1x (up to 4.7x) speedup over state-of-the-art systems without sacrificing model accuracy.
Yushu Cai, Qingrui Zhu, Lei Liu +3
Xidian University Guangzhou, Guangdong, China · Hunan University Changsha, Hunan, China
Memory-based temporal graph neural networks (TGN) rely on recurrent node memories to capture evolving interaction histories. However, temporal mini-batching introduces a trade-off between computational efficiency and temporal fidelity: when a node appears multiple times within the same batch, later interactions are evaluated using stale pre-batch memory because the intermediate memory updates have not yet been committed. This problem becomes more severe as temporal batch size increases. To overcome this challenge, we propose AnchorRM, an anchor-based residual memory mechanism that reconstructs event-specific intra-batch node states without executing sequential recurrent updates. AnchorRM treats the last committed node memory at the beginning of each batch as a pre-batch fixed anchor and uses causally preceding same-batch interactions to infer the missing state evolution relative to that anchor. The learned residual then calibrates the anchor to reconstruct an event-specific node memory for each interaction. Because each event-specific residual is predicted directly from the fixed anchor and its causal history, reconstructed states do not depend on previously reconstructed states and can therefore be computed in parallel within the batch while preserving temporal causality. Experiments across four temporal graph datasets and three memory-based backbones demonstrate that AnchorRM achieves an average 6.16x training speedup (up to 10.05x) and 28.67x batch scaling (up to 40x). These results demonstrate that AnchorRM substantially improves the scalability and training efficiency of memory-based temporal graph learning while maintaining competitive predictive performance.
Haroon Gharwi, Kai Shu, Yue Dai
Department of Computer Science, Illinois Institute of Technology, Chicago, IL, USA · Department of Computer Science, Emory University, Atlanta, GA, USA
Temporal graph neural networks (TGNNs) have shown strong performance on dynamic graph learning tasks by continuously updating node memories from evolving interactions (i.e., events) and leveraging these memories for downstream predictions. As graphs grow rapidly in modern applications, training requires processing an increasing number of events; meanwhile, each event can update node memories used by later events, creating dependencies that constrain parallel execution and limit the benefits of additional computing resources. These costs motivate an upstream question: Are all events, and the dependencies they introduce, necessary for learning? We investigate this question through two roles of temporal events: prediction targets and historical events sampled as inputs. We observe that targets involving stable node memories tend to exhibit lower prediction loss, while filtering historical events associated with stable memories can retain similar training-loss trajectories in several evaluated settings. Together, these observations suggest an opportunity to identify input redundancy through changes in node memories. Building on this opportunity, we propose RAFT, a redundancy-aware input-filtering framework for TGNN training. RAFT tracks memory changes during training and uses them to guide two filters: a root-event filter removes selected prediction targets before training, and a support-event filter excludes selected historical context during sampling. Filtering inputs before their downstream preparation and processing reduces the work they trigger. Across three TGNN models and five datasets, RAFT reduces training inputs by 35.23% on average and up to 97.11%, while reducing recorded training time by up to 52.56% with negligible changes to model precisions.
Yue Dai, Sarah Zhou, Youtao Zhang
Department of Computer Science, Illinois Institute of Technology, Chicago, USA · BASIS Independent Silicon Valley, San Jose, CA, USA · Department of Computer Science, University of Pittsburgh, Pittsburgh, USA