OptimES: Optimizing Federated Learning Using Remote Embeddings for Graph Neural Networks

2025-09-26

AI Summary

To address the high communication overhead and slow convergence in federated graph neural network (GNN) training caused by remote embedding sharing, this paper proposes OptimES, an optimization framework. Its core contributions are: (i) dynamic pruning of remote neighborhoods based on boundary nodes to eliminate redundant embedding transmissions; (ii) asynchronous embedding pushing overlapped with local training computation to improve hardware utilization; and (iii) an adaptive, demand-driven embedding pulling strategy that selectively updates critical embeddings. All mechanisms preserve data privacy under federated learning constraints. Experiments demonstrate that OptimES significantly reduces communication volume and training latency. On large-scale graphs, it achieves up to 3.5× speedup and 16% higher accuracy over EmbC; on sparse graphs, it reaches target accuracy 11× faster. These results substantially alleviate the efficiency bottlenecks inherent in existing federated GNN systems.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Distributed Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Graph Neural Networks (GNNs) have experienced rapid advancements in recent years due to their ability to learn meaningful representations from graph data structures. However, in most real-world settings, such as financial transaction networks and healthcare networks, this data is localized to different data owners and cannot be aggregated due to privacy concerns. Federated Learning (FL) has emerged as a viable machine learning approach for training a shared model that iteratively aggregates local models trained on decentralized data. This addresses privacy concerns while leveraging parallelism. State-of-the-art methods enhance the privacy-respecting convergence accuracy of federated GNN training by sharing remote embeddings of boundary vertices through a server (EmbC). However, they are limited by diminished performance due to large communication costs. In this article, we propose OptimES, an optimized federated GNN training framework that employs remote neighbourhood pruning, overlapping the push of embeddings to the server with local training, and dynamic pulling of embeddings to reduce network costs and training time. We perform a rigorous evaluation of these strategies for four common graph datasets with up to $111M$ vertices and $1.8B$ edges. We see that a modest drop in per-round accuracy due to the preemptive push of embeddings is out-stripped by the reduction in per-round training time for large and dense graphs like Reddit and Products, converging up to $approx 3.5 imes$ faster than EmbC and giving up to $approx16%$ better accuracy than the default federated GNN learning. While accuracy improvements over default federated GNNs are modest for sparser graphs like Arxiv and Papers, they achieve the target accuracy about $approx11 imes$ faster than EmbC.

Problem

• Optimizing federated GNN training with remote embeddings
• Reducing communication costs in privacy-preserving graph learning
• Accelerating convergence while maintaining model accuracy

Innovation

• Uses remote neighborhood pruning to reduce network costs
• Overlaps embedding push with local training for efficiency
• Employs dynamic pulling of embeddings to cut training time