SpreadGNN: Serverless Multi-task Federated Learning for Graph Neural Networks

TL;DR

SpreadGNN introduces a serverless, multi-task federated learning framework for GNNs that handles partial labels and decentralized training, outperforming traditional server-dependent methods on molecular property datasets.

Contribution

It presents the first serverless multi-task federated GNN training framework with a novel convergence-guaranteed optimization algorithm, DPA-SGD.

Findings

SpreadGNN outperforms centralized federated GNN models.

Effective in non-I.I.D. distributed graph datasets.

Operates efficiently in constrained topologies.

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