Jointly Optimizing Dataset Size and Local Updates in Heterogeneous Mobile Edge Learning

TL;DR

This paper introduces a heterogeneity-aware optimization framework for federated learning on resource-constrained wireless edge devices, jointly tuning dataset size and local updates to improve model accuracy.

Contribution

It develops a convex optimization approach to jointly optimize local updates and dataset sizes considering device heterogeneity, enhancing federated learning performance.

Findings

Heterogeneity-aware approach outperforms heterogeneity-unaware methods.

Optimal local updates depend on device capabilities and data distribution.

The proposed method effectively balances communication and computation constraints.

Abstract

This paper proposes to maximize the accuracy of a distributed machine learning (ML) model trained on learners connected via the resource-constrained wireless edge. We jointly optimize the number of local/global updates and the task size allocation to minimize the loss while taking into account heterogeneous communication and computation capabilities of each learner. By leveraging existing bounds on the difference between the training loss at any given iteration and the theoretically optimal loss, we derive an expression for the objective function in terms of the number of local updates. The resulting convex program is solved to obtain the optimal number of local updates which is used to obtain the total updates and batch sizes for each learner. The merits of the proposed solution, which is heterogeneity aware (HA), are exhibited by comparing its performance to the heterogeneity unaware (HU) approach.