Source-linked AI summary
SpreadGNN: Serverless Multi-task Federated Learning for Graph Neural Networks
Chaoyang He, Emir Ceyani, Keshav Balasubramanian, Murali Annavaram, Salman Avestimehr
TL;DR
Centralized GNN training is difficult when molecular graphs are private, distributed, non-I.I.D., and partially labeled. SpreadGNN combines serverless multi-task federated learning with DPA-SGD, and experiments show it can outperform FedAvg even under constrained client communication.
Problem
Molecular graph data cannot always be centralized because of privacy, regulatory, and competitive constraints, while federated GNN training must handle non-I.I.D. data, partial labels, and missing central coordination.
Method
SpreadGNN formulates partial-label federated multi-task learning for GNNs and solves the serverless optimization problem with Decentralized Periodic Averaging SGD, supported by convergence analysis.
Results
SpreadGNN outperforms FedAvg on graph-level molecular property prediction and regression tasks both with all-to-all communication and when clients communicate with only subsets of peers.
Takeaways & Limitations
The framework supports molecular property prediction with GNNs without requiring a centralized topology, including under non-I.I.D. dataset size and label distributions.
Takeaways & Limitations
Federated learning does not guarantee that data or models cannot be leaked, so production deployment requires additional security and privacy components such as Differential Privacy and Secure Aggregation.
Abstract
from arXiv · showhide
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
1 Introduction
Molecular graph learning involves heterogeneous, privately held data with multiple and often incomplete labels, making centralized GNN training impractical. SpreadGNN addresses this setting with serverless multi-task federated learning and reports better performance than FedAvg across unconstrained and constrained client topologies.
- Motivation: Molecular graphs are structurally heterogeneous, continuously generated, and labeled along multiple orthogonal dimensions, motivating a multi-task learning formulation.Their small individual graph size is offset by the large and varied population of molecules.
- Motivation: Private molecular graphs and partially labeled molecules cannot always be shared, although collective training could accelerate drug discovery.Proprietary data, privacy, regulatory, and competitive concerns constrain collaboration.
- Research gap: Federated GNNs face non-I.I.D. data, incomplete labels, and the practical absence of a central server in molecular settings.Different owners may hold different molecules or different subsets of labels, while competing entities may not accept centralized coordination.
- Approach: SpreadGNN combines a partial-label multi-task formulation with decentralized periodic averaging stochastic gradient descent and a convergence guarantee.The framework is designed specifically for serverless federated GNN training.
- Evaluation: Across graph-level molecular property prediction and regression tasks, SpreadGNN outperforms FedAvg both with all-to-all communication and when clients communicate with only subsets of peers.The evaluation uses synthetic non-I.I.D. and partially labeled datasets derived from MoleculeNet [69].
2 SpreadGNN Framework
SpreadGNN formulates graph-level federated multi-task learning for decentralized clients with heterogeneous, partially labeled molecular graph data. It replaces server-based aggregation with DPA-SGD, neighbor-based task-correlation exchange, and alternating optimization, with convergence analysis for the decentralized algorithm.
- Graph classifier: Each client uses an L-layer message-passing GNN, readout pooling, and a task classifier to produce graph-level predictions.Message passing aggregates neighboring node information, while pooling creates a graph embedding that the task classifier maps to labels.
- Problem setting: SpreadGNN addresses federated GNN learning when clients have private graph datasets, incomplete task labels, and no central server for parameter aggregation.The setting targets graph-level learning across decentralized datasets whose data cannot be centralized because of privacy, regulation, or intellectual-property constraints.
- Federated multi-task formulation: FedGMTL regularizes task weights through a covariance matrix so clients can relate locally observed tasks to tasks absent from their own labels.The task-relationship regularizer constrains task weights using Ω, whose entries reflect task relatedness; this formulation still relies on a central server.
- Serverless optimization: DPA-SGD performs local SGD and periodically synchronizes parameters with neighboring clients, while each client maintains and exchanges a local task-correlation matrix Ωk.Synchronization occurs every τ iterations, and the exchange mechanism aligns differently sized local covariance matrices across connected clients.
- Optimization: SpreadGNN alternates updates of client model weights and local correlation matrices to optimize a non-convex decentralized multi-task objective covering the union of client task sets.The local task matrix includes the client and its neighbors, allowing collective learning of tasks that individual clients may not observe.
- Convergence properties: The convergence analysis characterizes DPA-SGD as trading off convergence rate against communication efficiency and convergence speed under network-connectivity assumptions.The analysis uses a symmetric node-connection matrix whose spectral properties describe decentralized communication, with identity and fully connected matrices representing limiting cases.
3 Experiments
Experiments evaluate SpreadGNN on partially labeled, non-I.I.D. molecular property datasets using multiple GNN architectures and communication topologies. SpreadGNN outperforms centralized FedAvg in complete and constrained settings, while performance varies with regularizer, topology, and communication period.
- 3.1 Setup: The evaluation uses MoleculeNet datasets with quantity skew and mutually exclusive client label masks, testing GraphSAGE and GAT under complete and constrained topologies.The setup includes Tox21, MUV, QM8, and SIDER experiments with client-specific topology and optimizer configurations.
- 3.2 Results: SpreadGNN outperforms centralized FedAvg when all clients communicate, while GraphSAGE and GAT both benefit from the framework.The comparison uses a complete topology and shows that the framework remains effective across the two tested GNN architectures.
- 3.3 Sensitivity Analysis: Classification requires more careful tuning of the task-relationship regularizer than regression, which is more robust to variation in its coefficient.The reported behavior suggests greater independence among regression properties than among classification properties.
- 3.2 Results: In constrained communication, SpreadGNN still outperforms FedAvg, with ring topology outperforming random topology when each client has two neighbors.A ring provides an indirect path between every pair of clients, whereas random connections do not guarantee this property.
- 3.3 Sensitivity Analysis: Lower communication periods generally perform better, although the experiments report exceptions to this trend.The communication period τ is treated as an important hyperparameter and is examined through ablation studies.
4 Related Works
Prior work spans molecular representation learning, federated learning, federated GNNs, and distributed SGD. These studies provide related modeling and optimization foundations, but address different combinations of graph learning, decentralization, privacy, and non-convex federated training.
- Molecular Representation Learning: Molecular representation methods use atom-neighbor encodings, CNN fingerprints, graph convolutions, and message passing to represent molecules for property prediction.
- FL: Federated learning research includes multi-task frameworks and optimization algorithms, but prior strong-duality guarantees apply only when the objective is convex.
- Federated Graph Neural Networks: Federated GNN studies exchange graph statistics, use cryptography, combine federated and meta-learning, or partition models across edge and cloud systems.
- Stochastic Gradient Descent Optimization: Distributed optimization work studies synchronized mini-batch SGD, local SGD, FedAvg, and decentralized SGD to trade communication frequency against computation and convergence.
5 Conclusion
SpreadGNN enables serverless federated GNN training for molecular property prediction with partial labels and non-I.I.D. client data, while DPA-SGD provides convergence analysis.
- SpreadGNN can outperform a central server-dependent baseline even when clients communicate only with a few neighbors.The framework addresses non-I.I.D. dataset size and label distributions across clients.
Broader impact
SpreadGNN supports serverless GNN training for molecular property prediction, but federated learning does not guarantee that data or models will not leak.
- Federated learning protects privacy only to a certain extent and cannot guarantee that data or models are not leaked.Deployment as a public service requires comprehensive security and privacy-risk analysis.
- Differential Privacy and Secure Aggregation are recommended as additional security and privacy components for production deployment.
Appendix
The appendix provides supplementary proof, algorithm, dataset, hyperparameter, and ablation details for SpreadGNN.
- The appendix includes the proof of Theorem 1 and an algorithm sketch for SpreadGNN using DPA-SGD.
- It also provides detailed dataset descriptions, experiment hyperparameters, and ablation studies of communication period and network topology.
6 Algorithm Sketch
SpreadGNN performs local stochastic-gradient updates at each node and periodically aggregates and aligns optimization variables with neighboring nodes.
- 6 Algorithm Sketch: Each node iterates through global rounds, client epochs, and mini-batches while computing gradients and updating local optimization variables.
- 6 Algorithm Sketch: The algorithm runs these node-level updates and periodic neighbor communications until the maximum number of global iterations is reached.
- 6 Algorithm Sketch: Every communication period τ, each node performs aggregation and alignment over its neighbors.
7 Dataset Details
The experiments use molecular-property datasets from MoleculeNet and represent molecules as graphs built from RDKit-derived atom features and NetworkX graph construction.
- Dataset Details: The benchmark includes SIDER, Tox21, MUV, and QM8 molecular-property datasets from MoleculeNet.SIDER covers adverse drug reactions, Tox21 records compound toxicity, MUV contains 17 virtual-screening tasks, and QM8 models quantum-mechanical properties.
- Feature Extraction Procedure for Molecules: Molecules are converted into graph data through two stages: atom-level feature extraction with RDKit and graph construction with NetworkX.
- Feature Extraction Procedure for Molecules: The atom-feature set matches the features used in.
- Feature Extraction Procedure for Molecules: Table 3 lists the atom features used to represent molecular nodes.
8 Model Hyperparameters
SpreadGNN uses GAT or GraphSAGE followed by a readout that aggregates node embeddings and transforms them into task outputs, with hyperparameters selected by grid search.
- Model Architecture: The model combines a GAT or GraphSAGE GNN with a readout function for graph-level prediction.The GNN maps atom features to node embeddings, which the readout uses to produce graph-level outputs.
- Readout: The readout concatenates node embeddings with input features before mean pooling and task-specific transformation.Including input features in the readout improves generalization according to the passage.
- Training Hyperparameters: Grid search selects task-specific hyperparameters, while experiments use batch size 1 and two GNN layers.The two-layer choice avoids the over-smoothing reported for deeper GNNs [36].
9 Detailed Ablation Studies
The ablations examine communication period and network topology, showing that both affect performance and that their impact depends on the communication budget and network size.
- 9.1 Effect of Communication Period τ: Increasing communication period τ generally decreases performance, although τ = 5 can outperform averaging and exchanging every round.Thus, τ controls a trade-off between model performance and running time.
- Experimental Configuration: The experiments use communication period τ = 1 and a complete client topology for SpreadGNN’s reported configurations.
- 9.2 Effect of Serverless Network Topology: In a four-client GraphSAGE network, two-neighbor random and ring topologies show minimal performance differences.The total number of neighbors matters more than topology when the network has fewer clients.
- 9.2 Effect of Serverless Network Topology: With more neighbors, each client’s local model becomes more robust, while topology matters more in larger networks.In an eight-client network with two neighbors, random connections perform worse than a ring topology.
- Theoretical Analysis: DPA-SGD trades convergence rate against communication efficiency and convergence speed under stated stochastic-gradient and objective-function assumptions.The analysis assumes unbiased stochastic gradients with bounded variance, and that the objective is Lipschitz and lower bounded.