Asynchronous Federated Learning with non-convex client objective functions and heterogeneous dataset

TL;DR

This paper advances asynchronous federated learning by addressing non-convex objectives and data heterogeneity, introducing new aggregation and learning rate strategies, with theoretical convergence guarantees and practical validation.

Contribution

It extends AFL to non-convex and heterogeneous data settings, proposing staleness-aware aggregation and adaptive learning rates, with comprehensive convergence analysis and real-world experiments.

Findings

Improved convergence stability with staleness-aware aggregation.

Enhanced scalability and performance in heterogeneous environments.

Validated effectiveness through experiments in PyTorch.

Abstract

Federated Learning (FL) enables collaborative model training across decentralized devices while preserving data privacy. However, traditional FL suffers from communication overhead, system heterogeneity, and straggler effects. Asynchronous Federated Learning (AFL) addresses these by allowing clients to update independently, improving scalability and reducing synchronization delays. This paper extends AFL to handle non-convex objective functions and heterogeneous datasets, common in modern deep learning. We present a rigorous convergence analysis, deriving bounds on the expected gradient norm and studying the effects of staleness, variance, and heterogeneity. To mitigate stale updates, we introduce a staleness aware aggregation that prioritizes fresher updates and a dynamic learning rate schedule that adapts to client staleness and heterogeneity, improving stability and convergence. Our framework accommodates variations in computational power, data distribution, and communication delays, making it practical for real world applications. We also analyze the impact of client selection strategies-sampling with or without replacement-on variance and convergence. Implemented in PyTorch with Python's asyncio, our approach is validated through experiments demonstrating improved performance and scalability for asynchronous, heterogeneous, and non-convex FL scenarios.