Prototype Transport with Client-Specific Geometry for Non-IID Edge Federated Learning

Abstract

Federated learning has emerged as a dominant paradigm for decentralized machine learning, allowing multiple edge devices to collaboratively train a global model without exchanging raw data. However, deploying federated learning in edge environments presents significant challenges, particularly when dealing with non-independent and identically distributed data across clients. Statistical heterogeneity causes local model updates to diverge, leading to severe performance degradation and convergence instability during global aggregation. To address these issues, prototype-based federated learning has been proposed, where clients share class representations instead of full model parameters. While promising, existing prototype methods often ignore the underlying geometric structures of local feature spaces, leading to suboptimal alignment at the server. This paper introduces a novel framework that leverages optimal transport theory to align prototypes while explicitly accounting for client-specific geometry. By modeling the local feature distributions as metric spaces and computing the Gromov-Wasserstein barycenter at the edge server, our approach effectively harmonizes heterogeneous feature representations without compromising data privacy. Comprehensive theoretical analysis and extensive text-based descriptions of the algorithmic design demonstrate how geometric alignment mitigates weight divergence. Evaluation scenarios across multiple heterogeneous data distributions indicate that the proposed methodology significantly enhances classification accuracy and accelerates convergence while maintaining the stringent communication and privacy constraints inherent to edge networks.