最近,研究人员开发了时序图神经网络(TGNNs),利用动态图中的时间信息,建立在图神经网络(GNNs)在学习静态图表示方面的成功基础上。TGNNs在各种动态图上的诸如时序链接预测和动态节点分类等下游任务上表现出了优越的准确性,显著超过了静态GNNs和其他传统方法。
在动态图上,随着时间的推移,每个节点上有更多的相关事件。当这个数字很高时,TGNNs无法使用基于时间注意力的聚合或历史邻居采样技术来完全捕捉历史信息。研究人员创建了基于记忆的时序图神经网络(M-TGNNs),它们存储节点级别的记忆向量来总结独立的节点历史,以弥补丢失的历史。
尽管M-TGNNs取得了成功,但它们的可扩展性差使得它们难以在大规模的生产系统中实施。由于辅助节点记忆产生的时间依赖性,训练小批量必须简短并按照时间顺序安排。在M-TGNN训练中利用数据并行性特别困难,有两个方面:
- 仅仅提高批量大小会导致信息丢失和事件之间时间依赖性信息的丢失。
- 所有训练器必须访问和维护节点记忆的统一版本,这会在分布式系统中产生大量的远程流量。
南加州大学和AWS的新研究提供了DistTGL,一种在分布式GPU集群上进行M-TGNN训练的可扩展和有效的方法。DistTGL从三个方面改进了当前的M-TGNN训练系统:
- 模型:通过引入更多静态节点记忆,提高了M-TGNNs节点记忆的准确性和收敛速度。
- 算法:为了解决分散设置中的准确性损失和通信开销问题,团队提供了一种新颖的训练算法。
- 系统:为了减少小批量生成相关的开销,他们开发了一个使用预取和流水线技术优化的系统。
DistTGL在收敛速度和训练吞吐量方面显著优于先前的方法。DistTGL是第一个将M-TGNN训练扩展到分布式GPU集群的工作。Github上公开提供了DistTGL。
他们基于M-TGNN训练的独特特性,提出了两种创新的并行训练方法——纪元并行和记忆并行,它们使得M-TGNNs能够在多个GPU上捕捉与单个GPU上相同数量的依赖图事件。根据数据集和硬件特征,他们提供了选择最佳训练设置的启发式建议。
研究人员对节点记忆上的记忆操作进行了序列化,并通过一个单独的守护进程有效地执行它们,消除了复杂和昂贵的同步,以重叠小批量创建和GPU训练。在试验中,当扩展到多个GPU时,DistTGL在收敛速度方面比最先进的单机方法提高了10倍以上。”