FedSA: A staleness-aware asynchronous Federated Learning algorithm with non-IID data

Abstract

This paper presents new asynchronous methods to the Federated Learning (FL), one of the next-generation paradigms for Artificial Intelligence (AI) systems. We consider the two-fold challenges lay ahead. First, non-IID (non-Independent and Identically Distributed) data across devices cause unstable performance. Second, unreliable and slow environments not only slow the convergence but also cause staleness issues. To address these challenges, this study uses a bottom-up approach for analysis and algorithm design. We first reformulate FL by unifying both synchronous and asynchronous updating schemes with an asynchrony-related parameter. We theoretically analyze this new form and find practical strategies for optimization. The key findings include: 1) a two-stage training strategy to accelerate training and reduce communication overhead; 2) strategies of choosing key hyperparameters optimally for these stages to maintain efficiency and robustness. With these theoretical guarantees, we propose FedSA (Federated Staleness-Aware), a novel asynchronous federated learning algorithm. We validate FedSA on different tasks with non-IID/IID and staleness settings. Our results indicate that, given a large proportion of stale devices, the proposed algorithm presents state-of-the-art performance by outperforming existing methods on both non-IID and IID cases.