# Retrofitting Temporal Graph Neural Networks with Transformer

arXiv: https://arxiv.org/abs/2409.05477 · HTML: https://mlsanity.com/paper/2409.05477/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 38 references that are in the map.
- First posted on arXiv 9 Sep 2024.

## Builds on (most cited first, 17 of 17)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [Temporal Graph Networks for Deep Learning on Dynamic Graphs](https://mlsanity.com/paper/2006.10637/) (2020)
- [Inductive Representation Learning on Temporal Graphs](https://mlsanity.com/paper/2002.07962/) (2020)
- [Predicting Dynamic Embedding Trajectory in Temporal Interaction Networks](https://mlsanity.com/paper/1908.01207/) (2019)
- [OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs](https://mlsanity.com/paper/2103.09430/) (2021)
- [Efficient Training of Language Models to Fill in the Middle](https://mlsanity.com/paper/2207.14255/) (2022)
- [Inductive Representation Learning in Temporal Networks via Causal Anonymous Walks](https://mlsanity.com/paper/2101.05974/) (2021)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [Towards Better Evaluation for Dynamic Link Prediction](https://mlsanity.com/paper/2207.10128/) (2022)
- [Motifs in Temporal Networks](https://mlsanity.com/paper/1612.09259/) (2016)
- [APAN: Asynchronous Propagation Attention Network for Real-time Temporal Graph Embedding](https://mlsanity.com/paper/2011.11545/) (2020)
- [Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect](https://mlsanity.com/paper/1903.04611/) (2019)

## Built on by (most cited first, 5 of 5)

- [Guiding LLM-based Smart Contract Generation with Finite State Machine](https://mlsanity.com/paper/2505.08542/) (2025)
- [When Speed meets Accuracy: an Efficient and Effective Graph Model for Temporal Link Prediction](https://mlsanity.com/paper/2507.13825/) (2025)
- [BiTA: Bidirectional Gated Recurrent Unit-Transformer Aggregator in a Temporal Graph Network Framework for Alert Prediction in Computer Networks](https://mlsanity.com/paper/2604.22781/) (2026)
- [Text-attributed Graph Condensation via Text Selection and Attribute Matching](https://mlsanity.com/paper/2606.03839/) (2026)
- [GTIN: A Unified Framework for Joint Event and Time Prediction in Temporal Graphs](https://mlsanity.com/paper/2607.23556/) (2026)

## Abstract (arXiv)

Temporal graph neural networks (TGNNs) outperform regular GNNs by incorporating time information into graph-based operations. However, TGNNs adopt specialized models (e.g., TGN, TGAT, and APAN ) and require tailored training frameworks (e.g., TGL and ETC). In this paper, we propose TF-TGN, which uses Transformer decoder as the backbone model for TGNN to enjoy Transformer's codebase for efficient training. In particular, Transformer achieves tremendous success for language modeling, and thus the community developed high-performance kernels (e.g., flash-attention and memory-efficient attention) and efficient distributed training schemes (e.g., PyTorch FSDP, DeepSpeed, and Megatron-LM). We observe that TGNN resembles language modeling, i.e., the message aggregation operation between chronologically occurring nodes and their temporal neighbors in TGNNs can be structured as sequence modeling. Beside this similarity, we also incorporate a series of algorithm designs including suffix infilling, temporal graph attention with self-loop, and causal masking self-attention to make TF-TGN work. During training, existing systems are slow in transforming the graph topology and conducting graph sampling. As such, we propose methods to parallelize the CSR format conversion and graph sampling. We also adapt Transformer codebase to train TF-TGN efficiently with multiple GPUs. We experiment with 9 graphs and compare with 2 state-of-the-art TGNN training frameworks. The results show that TF-TGN can accelerate training by over 2.20 while providing comparable or even superior accuracy to existing SOTA TGNNs. TF-TGN is available at https://github.com/qianghuangwhu/TF-TGN.