Allo: Optimizing Federated Learning via Guided Epoch Allocation

Federated learning has shown great potential for addressing the challenge of data islands while preserving data privacy. Unlike traditional distributed machine learning, federated learning faces the problem of system and statistical heterogeneity, which makes it difficult to deploy at the edge. However, current methods cannot optimize the federated learning by taking the clients' system heterogeneity and statistical heterogeneity together. This paper proposes Allo, a framework based on deep reinforcement learning that can automatically choose the appropriate number of local epochs for each client, which can mitigate the impact of stragglers and unbalanced data distribution. With extensive experiments performed in PyTorch, we show that the performance of federated training can be improved by over 15% on the MNIST dataset and over 12% on CIFAR-10 dataset, as compared to the baseline.