TGLite: A Lightweight Programming Framework for Continuous-Time Temporal Graph Neural Networks

In recent years, Temporal Graph Neural Networks (TGNNs) have achieved great success in learning tasks for graphs that change over time. These dynamic/temporal graphs represent topology changes as either discrete static graph snapshots (called DTDGs), or a continuous stream of timestamped edges (called CTDGs). Because continuous-time graphs have richer time information, it will be crucial to have abstractions for programming CTDG-based models so that practitioners can easily explore new designs and optimizations in this space. A few recent frameworks have been proposed for programming and accelerating TGNN models, but these either do not support continuous-time graphs, lack easy composability, and/or do not facilitate CTDG-specific optimizations.

In this paper, we propose a lightweight framework called TGLite to fill this apparent gap in the status quo. It provides abstractions that serve as composable building blocks for implementing TGNN models for CTDGs. It introduces a novel TBlock representation for capturing message-flow dependencies between nodes, with explicit support for temporal-related attributes, which is well-suited for common TGNN computation patterns. TBlocks serve as a central representation on which many different operators can be defined, such as temporal neighborhood sampling, scatter/segmented computations, as well as optimizations tailored to CTDGs. We use TGLite to implement four existing TGNN models. Compared to the TGL framework, TGLite is able to accelerate runtime performance of training (1.06 -- 3.43×) and inference (1.09 -- 4.65×) of these models on V100 and A100 GPUs across different experimental settings. Notably, when scaling to larger datasets, TGL runs out-of-memory in some cases on the V100 while TGLite is able to run successfully.