PaperlayerView exact revision on arXiv

Source-linked AI summary

TGL: A General Framework for Temporal GNN Training on Billion-Scale Graphs

Hongkuan Zhou, Da Zheng, Israt Nisa, Vasileios Ioannidis, Xiang Song, George Karypis

arXiv:2203.14883v2cs.LG

TL;DR

Temporal GNN training on large dynamic graphs is challenging, while prior methods were evaluated on small graphs. TGL provides a unified framework with temporal sampling and random chunk scheduling, and reports similar or better accuracy with faster training than open-source baselines.

  • Problem

    Prior TGNNs were evaluated on small dynamic graphs, while training on large graphs must handle chronological dependencies, temporal-neighbor sampling, and varied model strategies.

    Source: §1 · p. 1

  • Method

    TGL is a unified framework for large-scale offline TGNN training, combining temporal sampling, node memory and mailbox components with random chunk scheduling.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    TGL achieves similar or better accuracy than baseline methods, with an average 13× speedup on small datasets and 2.3× speedup using four GPUs on large datasets.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    TGL evaluates five TGNN variants on four small datasets and two billion-edge datasets, providing a framework for single- and multi-GPU training.

    Source: §5 · p. 8

    Sampler performance includes a memory bottleneck when fetching selected edge information and linear workload in generating DGL MFGs as thread count rises.

    Source: §4.2 · p. 7

Abstract

from arXiv · show

Many real world graphs contain time domain information. Temporal Graph Neural Networks capture temporal information as well as structural and contextual information in the generated dynamic node embeddings. Researchers have shown that these embeddings achieve state-of-the-art performance in many different tasks. In this work, we propose TGL, a unified framework for large-scale offline Temporal Graph Neural Network training where users can compose various Temporal Graph Neural Networks with simple configuration files. TGL comprises five main components, a temporal sampler, a mailbox, a node memory module, a memory updater, and a message passing engine. We design a Temporal-CSR data structure and a parallel sampler to efficiently sample temporal neighbors to formtraining mini-batches. We propose a novel random chunk scheduling technique that mitigates the problem of obsolete node memory when training with a large batch size. To address the limitations of current TGNNs only being evaluated on small-scale datasets, we introduce two large-scale real-world datasets with 0.2 and 1.3 billion temporal edges. We evaluate the performance of TGL on four small-scale datasets with a single GPU and the two large datasets with multiple GPUs for both link prediction and node classification tasks. We compare TGL with the open-sourced code of five methods and show that TGL achieves similar or better accuracy with an average of 13x speedup. Our temporal parallel sampler achieves an average of 173x speedup on a multi-core CPU compared with the baselines. On a 4-GPU machine, TGL can train one epoch of more than one billion temporal edges within 1-10 hours. To the best of our knowledge, this is the first work that proposes a general framework for large-scale Temporal Graph Neural Networks training on multiple GPUs.

1 INTRODUCTION

Temporal GNNs model evolving graph relationships, but existing methods and infrastructure had been evaluated or designed mainly for small-scale settings. TGL targets large-scale offline training with a unified framework, temporal sampling and scheduling techniques, and reported accuracy and speed improvements.

1 INTRODUCTION

  • Earlier TGNNs captured temporal, structural, and contextual relationships but were evaluated on graphs with fewer than 10,000 nodes and one million edges.Scaling also requires chronological training, temporal-neighbor sampling, and support for different TGNN memory or snapshot strategies.Source: §1 · p. 1
  • TGL is a unified framework for large-scale offline training on edge-timestamped graphs, supporting TGNN variants that aggregate information from selected temporal neighbors.It targets continuous-time graphs and treats discrete timestamps as granulated continuous time.Source: §1 · p. 2
  • TGL combines a CSR-based temporal-neighbor structure and parallel sampler with random chunk scheduling, and introduces GDELT and MAG datasets with billions of edges.The framework is designed to support diverse TGNN architectures and multi-GPU training.Source: §1 · p. 2
  • TGL reports similar or higher accuracy than baseline open-source methods, with an average 13× speedup on small datasets and 2.3× speedup using four GPUs on large datasets.Figure 1 compares accuracy and per-epoch training time on Wikipedia at batch size 600.Source: §1 · p. 2, Figure 1 · p. 2

2 TEMPORAL GRAPH NEURAL NETWORKS

Temporal GNNs form dynamic node embeddings by processing temporal neighbors with snapshot grouping, time encoding, node memory, or combinations of these strategies. TGL’s background describes node memory, attention aggregation, and timestamp-aware sampling as key components.

2 TEMPORAL GRAPH NEURAL NETWORKS

  • Most TGNN variants can be described through node memory, attention aggregation, and temporal sampling, while snapshot-based methods combine outputs from independently treated snapshots.Some methods combine strategies; CAW-N instead aggregates information from causal anonymous walks.Source: §2 · p. 2

2.1 Node Memory

  • Node memory summarizes a node’s history; new edges trigger an RNN or GRU update, and the resulting memory is combined with node features during message passing.Neighbor input mails are combined and used by the memory updater.Source: §2.1 · p. 3

2.2 Attention Aggregator

  • Attention aggregation gathers and combines information from a node’s hop-1 temporal neighbors using queries, keys, and values.Source: §2.2 · p. 3

2.3 Temporal Sampler

  • Temporal sampling accounts for edge timestamps and may select past neighbors uniformly or choose the most recent ones.The same nodes can appear multiple times as supporting neighbors at different timestamps.Source: §2.3 · p. 3

3 TGL

TGL organizes temporal graph training around reusable sampling, memory, and message-passing components, with methods for efficient temporal-neighbor selection and parallel training. Its design also addresses how snapshot timing and memory dependencies affect training.

3 TGL

  • TGL combines temporal sampling, mailbox and node-memory updates, and message passing to generate temporal embeddings and compute loss.The forward path refreshes memory and the mailbox for the next mini-batch after computing loss.Source: Figure 2 · p. 3, §3 · p. 3
  • Snapshot-based TGNNs use snapshots created from target-node timestamps, rather than static windows, to generate embeddings at arbitrary timestamps.Source: §3 · p. 3
  • TGL’s mailbox stores a fixed number of recent messages; cached messages update involved nodes’ memory before new embeddings are computed.The mailbox is then updated, and the same memory-update scheme is used at inference for consistency.Source: §3 · p. 3

3.1 Parallel Temporal Sampler

  • T-CSR sorts each node’s outgoing edges by timestamp and uses per-node snapshot pointers to locate candidate temporal neighbors.The pointers enable candidate-edge identification in O(1), with O(n|V|) additional storage and O(|E|) maintenance computation per epoch.Source: §3.1 · p. 4, Figure 3 · p. 4
  • The parallel sampler distributes mini-batch root nodes across threads, uses per-node locks for shared targets, and excludes neighbors later than their root timestamps.It generates a DGL message-flow graph for each layer for the trainer’s forward and backward propagation.Source: §3.1 · p. 4

3.2 Parallel Training

  • Larger batches discard more intra-batch dependencies; random chunk scheduling varies chunk placement across epochs to learn more inter-batch dependencies.The method divides training edges into chunks and randomly chooses a starting chunk each epoch.Source: §3.2 · p. 5

4 EXPERIMENTS

Across small- and billion-edge datasets, the experiments evaluate TGL’s sampler, single- and multi-GPU training, and random chunk scheduling. TGL matches or exceeds baseline accuracy while substantially reducing training and sampling time.

4.1 Datasets

  • TGL evaluates four small bipartite datasets alongside GDELT and MAG, which contain 0.2 and 1.3 billion temporal edges, respectively.GDELT spans 2016–2020 and uses actor and event CAMEO features; MAG represents paper citations with abstract embeddings.Source: §4.1 · p. 6, §4.1.1 · p. 6, §4.1.2 · p. 6

4.2 Parallel Temporal Sampler

  • With 32 threads, TGL’s sampler runs 57× faster than TGAT’s and 289× faster than TGN’s, taking under 0.5 seconds per Wikipedia epoch.The reported speedup combines T-CSR, data parallelism, and C++ efficiency over Python.Source: §4.2 · p. 6, Table 4 · p. 6
  • The sampler’s 32-thread speedups are 3.13× for DySAT, 4.20× for TGAT, and 2.42× for TGN, with sub-linear scaling attributed to pointer locks, memory fetching, and MFG generation.Figure 4 breaks runtime down into pointer updates, binary search, neighbor sampling, and DGL MFG generation.Source: §4.2 · pp. 6–7, Figure 4 · p. 6

4.3 Single-GPU Training

  • The single-GPU comparison uses batch size 600, 32 sampler threads, and GPU-resident features, memory, and mailbox; MOOC and LastFM are excluded from baseline comparisons.Those two datasets lack authentic features, and baseline performance evaluation was unavailable.Source: §4.3 · p. 7
  • TGL achieves similar or higher link-prediction AP than five baselines on four small datasets, with an average 13× speedup in per-epoch training time.The result compares accuracy and runtime; TGN has the highest AP on all but LastFM, while JODIE is the fastest-training model among the methods.Source: §4.3 · p. 7, Table 5 · p. 7
  • On Wikipedia and Reddit node classification, TGN and JODIE have the highest AP, and JODIE exceeds other methods by more than 7% AP on Reddit.The authors attribute this result tentatively to noisy Reddit neighbors that hinder high-expressive models from learning general graph-structure patterns.Source: §4.3 · p. 7, Table 6 · p. 7

4.4 Random Chunk Scheduling

  • A batch size of 4800 without random chunk scheduling fails to learn on Reddit and Wikipedia after 5–10 epochs; chunking improves convergence, with Reddit nearing baseline within 30 epochs.On Wikipedia, 16 chunks perform better than no chunks, while 32 chunks reach similar convergence after 30 epochs.Source: §4.4 · p. 8, Figure 6 · p. 8

4.5 Multi-GPU Training

  • On GDELT, four GPUs deliver 1.80×–2.74× speedups across five models; on MAG, APAN runs out of memory.TGN has the highest link-prediction AP and JODIE the fastest training time; at eight GPUs, CPU-memory bandwidth and PCI-E channels are saturated.Source: §4.5 · p. 8, Table 7 · p. 8, Figure 7 · p. 8

5 CONCLUSION

TGL is a unified framework for training TGNN variants at large scale, on single or multiple GPUs using configuration files. Across four small and two billion-edge datasets, it achieves similar or better accuracy with faster training than open-source baselines.

5 CONCLUSION

  • TGL achieves similar or better accuracy on all evaluated datasets with significantly faster training than the open-source baselines.The evaluation covers five TGNN variants, four small-scale datasets, and two large-scale datasets with billions of edges.Source: §5 · p. 8
  • TGL provides a unified framework for training different TGNN variants on single or multiple GPUs through configuration files.The framework includes T-CSR for dynamic graphs, a temporal parallel sampler to reduce sampling overhead, and random chunk scheduling to mitigate dependency loss with large batches.Source: §5 · p. 8
  • TGL evaluates TGNN performance on two large-scale datasets and four small-scale datasets across five TGNN variants.Source: §5 · p. 8

Related papers

Loading 2203.14883v2…