arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2409.05477cs.LG
收藏

用Transformer改造时序图神经网络

Retrofitting Temporal Graph Neural Networks with Transformer

School of Computer Science, Wuhan University(武汉大学计算机学院);Centre for Perceptual and Interactive Intelligence (CPII)(感知与交互智能中心);ETH Zürich(苏黎世联邦理工学院);eBay(易贝);Peking University(北京大学)

Qiang Huang, Xiao Yan, Xin Wang, Susie Xi Rao, Zhichao Han, Fangcheng Fu, Wentao Zhang, Jiawei Jiang

AI总结:

本文提出TF-TGN,用Transformer解码器作为时序图神经网络骨干,通过序列化消息聚合、并行CSR转换与图采样,实现超2.2倍训练加速且精度不降。

摘要:

时序图神经网络(TGNNs)通过将时间信息纳入基于图的操作中,优于常规GNN。然而,TGNNs采用专门模型(如TGN、TGAT和APAN),并需要定制训练框架(如TGL和ETC)。在本文中,我们提出TF-TGN,使用Transformer解码器作为TGNN的骨干模型,以利用Transformer的代码库进行高效训练。特别是,Transformer在语言建模中取得了巨大成功,因此社区开发了高性能内核(如flash-attention和memory-efficient attention)以及高效分布式训练方案(如PyTorch FSDP、DeepSpeed和Megatron-LM)。我们观察到TGNN类似于语言建模,即TGNN中按时间顺序发生的节点与其时序邻居之间的消息聚合操作可以结构化为序列建模。除了这种相似性,我们还纳入了一系列算法设计,包括后缀填充、带自环的时序图注意力以及因果掩蔽自注意力,以使TF-TGN正常工作。在训练过程中,现有系统在转换图拓扑和进行图采样方面速度较慢。因此,我们提出了并行化CSR格式转换和图采样的方法。我们还调整了Transformer代码库,以使用多个GPU高效训练TF-TGN。我们在9个图上进行实验,并与2个最先进的TGNN训练框架进行比较。结果表明,TF-TGN可以将训练加速超过2.20倍,同时提供与现有SOTA TGNN相当甚至更优的准确性。TF-TGN可在https://github.com/qianghuangwhu/TF-TGN获取。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑