[Paper Review] 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 Networks68 references29 citations

TL;DR

SpreadGNN enables serverless, multitask federated training of Graph Neural Networks with partial labels, using Decentralized Periodic Averaging SGD and task relationship regularization to outperform centralized FL baselines in non-IID molecular graph datasets.

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

Motivation & Objective

• Address privacy/regulatory constraints hindering central data sharing for molecular graphs.
• Develop a serverless Federated Multi-Task Learning framework for GNNs that handles partial labels across clients.
• Provide a convergence-guaranteed optimization method suitable for decentralized networks.
• Demonstrate that serverless SpreadGNN can outperform central-server FedAvg in non-IID, partially labeled settings.

Proposed method

• Extend Federated Multi-Task Learning to graph neural networks with partial labels via a centralized formulation (FedGMTL) and then remove the central server to obtain SpreadGNN.
• Introduce Decentralized Periodic Averaging SGD (DPA-SGD) where clients perform local SGD and synchronize with neighbors every τ iterations.
• Use per-client task covariance matrices Ωk and a global task-relationship regularizer Tr(Φ_task Ω−1 Φ_task^T) to share task information across clients.
• Propose an alternating optimization over network weights Wk and Ωk, with communication-based alignment of Ωk across neighbors (Algorithm 1 in Appendix).
• Provide a convergence analysis for DPA-SGD showing conditions under which the averaged model converges in non-convex settings (Theorem 1).
• Evaluate on MoleculeNet datasets (SIDER, Tox21, MUV, QM8) with non-IID, partially labeled partitions and two GNN backbones (GraphSAGE, GAT).

Experimental results

Research questions

• RQ1 Can a serverless, decentralized federated multitask learning framework train GNNs effectively with partial labels?
• RQ2 Does DPA-SGD with task-relationship regularization improve performance over centralized FedAvg in non-IID molecular graph settings?
• RQ3 Is SpreadGNN agnostic to GNN architecture and robust to varying client connectivity/topology?
• RQ4 What is the impact of topology, neighborhood size, and communication period on learning performance and convergence?
• RQ5 Can clients learn to predict tasks they do not have in their local labels through cross-client task relationships?

Key findings

• SpreadGNN outperforms FedAvg on molecular property prediction tasks when all clients can communicate, despite the absence of a central server.
• SpreadGNN achieves competitive or better accuracy than FedGMTL (server-assisted) and can approach central-server performance under realistic serverless topologies.
• The framework is agnostic to the choice of GNN model (GraphSAGE or GAT) and remains effective under constrained client connectivity.
• The task-relationship regularizer (Ω) and decentralized averaging enable learning across partially labeled tasks, improving performance in non-IID settings.
• Convergence analysis (Theorem 1) provides conditions under which DPA-SGD converges in the non-convex, decentralized FL setting.