[论文解读] SpreadGNN: Serverless Multi-task Federated Learning for Graph Neural Networks

Chaoyang He, Emir Ceyani+3 more|arXiv (Cornell University)|Jun 4, 2021

Advanced Graph Neural Networks参考文献 68被引用 29

一句话总结

SpreadGNN 使无服务器、多任务联邦学习的图神经网络在部分标签下进行训练，使用去中心化的周期性平均 SGD 和 任务关系正则化，在非 IID 分子图数据集上优于集中式 FL 基线。

ABSTRACT

Graph Neural Networks (GNNs) are the first choice methods for graph machine learning problems thanks to their ability to learn state-of-the-art level representations from graph-structured data. However, centralizing a massive amount of real-world graph data for GNN training is prohibitive due to user-side privacy concerns, regulation restrictions, and commercial competition. Federated Learning is the de-facto standard for collaborative training of machine learning models over many distributed edge devices without the need for centralization. Nevertheless, training graph neural networks in a federated setting is vaguely defined and brings statistical and systems challenges. This work proposes SpreadGNN, a novel multi-task federated training framework capable of operating in the presence of partial labels and absence of a central server for the first time in the literature. SpreadGNN extends federated multi-task learning to realistic serverless settings for GNNs, and utilizes a novel optimization algorithm with a convergence guarantee, Decentralized Periodic Averaging SGD (DPA-SGD), to solve decentralized multi-task learning problems. We empirically demonstrate the efficacy of our framework on a variety of non-I.I.D. distributed graph-level molecular property prediction datasets with partial labels. Our results show that SpreadGNN outperforms GNN models trained over a central server-dependent federated learning system, even in constrained topologies. The source code is publicly available at https://github.com/FedML-AI/SpreadGNN

研究动机与目标

• 解决阻碍分子图数据集中共享的隐私/法规约束。
• 开发一个无服务器的 GNN 联邦多任务学习框架，能够处理来自各客户端的部分标签。
• 提供适用于去中心化网络的收敛性有保证的优化方法。
• 证明无服务器的 SpreadGNN 在非 IID、部分标签设置下能够超过中心服务器的 FedAvg。

提出的方法

• 将带部分标签的联邦多任务学习扩展到图神经网络，采用集中式公式（FedGMTL），然后移除中心服务器以得到 SpreadGNN。
• 引入去中心化的周期性平均 SGD（DPA-SGD），其中客户端执行本地 SGD，并在每 τ 次迭代与邻居同步。
• 使用每个客户端的任务协方差矩阵 Ωk 与全局任务关系正则项 Tr(Φ_task Ω−1 Φ_task^T) 在客户端之间共享任务信息。
• 提出对网络权重 Wk 与 Ωk 的交替优化，在邻居之间进行基于通信的 Ωk 对齐（Appendix 的 Algorithm 1）。
• 提供 DPA-SGD 的收敛性分析，给出平均模型在非凸情形下收敛的条件（定理 1）。
• 在 MoleculeNet 数据集（SIDER、Tox21、MUV、QM8）上进行评估，使用非 IID、部分标签分区和两个 GNN 骨干网络（GraphSAGE、GAT）。

实验结果

研究问题

• RQ1 一个无服务器、去中心化的联邦多任务学习框架是否能够在部分标签下有效训练 GNN？
• RQ2 在非 IID 的分子图设置中，带任务关系正则化的 DPA-SGD 是否能优于集中式 FedAvg？
• RQ3 SpreadGNN 是否对 GNN 架构无关且对客户端连接性/拓扑的变化具有鲁棒性？
• RQ4 拓扑、邻居规模和通信周期对学习性能与收敛有什么影响？
• RQ5 客户端是否能够通过跨客户端的任务关系来学习预测本地标签中没有的任务？

主要发现

• 在所有客户端均可通信时，SpreadGNN 在分子性质预测任务上优于 FedAvg，尽管不存在中心服务器。
• SpreadGNN 的准确率与 FedGMTL（有服务器的）相当甚至更好，在现实的无服务器拓扑下可接近中心服务器的性能。
• 该框架对 GNN 模型的选择（GraphSAGE 或 GAT）具有无关性，在受限的客户端连接性下仍然有效。
• 任务关系正则项（Ω）和去中心化平均实现跨部分标签任务的学习，在非 IID 设置下提高了性能。
• 收敛性分析（定理 1）给出在非凸、去中心化联邦学习设置下 DPA-SGD 收敛的条件。