序列聚合与重计算：大规模图上图神经网络的分布式全批量训练
机器学习
2022-04-18 v3 人工智能
摘要
我们提出序列聚合与重计算（SAR）方案，用于大规模图上图神经网络（GNNs）的分布式全批量训练。近来 GNN 的大规模训练由基于采样的方法与基于不可学习消息传递的方法主导。而 SAR 是一种分布式技术，可直接在整个大规模图上训练任意类型的 GNN。SAR 的关键创新在于分布式序列重计算方案，其在反向传播过程中顺序重建然后释放 prohibitively 大的 GNN 计算图的片段。这带来了优异的内存扩展特性：即便对稠密连接图，每工作节点的内存消耗也随工作节点数线性下降。利用 SAR，我们报告了迄今最大的全批量 GNN 训练应用，并展示了随工作节点数增加带来的大量内存节省。我们还提出一种基于核融合与注意力矩阵重计算的通用技术，以优化基于注意力模型的运行时间与内存效率。我们表明，与 SAR 结合后，我们优化的注意力核在基于注意力的 GNN 中带来了显著的加速与内存节省。我们已公开 SAR GNN 训练库：\url{https://github.com/IntelLabs/SAR}。
引用
@article{arxiv.2111.06483,
title = {Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs},
author = {Hesham Mostafa},
journal= {arXiv preprint arXiv:2111.06483},
year = {2022}
}