Latency Minimization for Wireless Federated Learning with Heterogeneous Local Updates

In this paper, we study the latency minimization problem for a wireless federated learning (FL) system with heterogeneous computation capability, where different edge devices perform different numbers of local updates in each communication round. We formulate a total latency minimization problem, taking into account both the communication and computation latency in the whole FL procedure. We reveal that decoupling the resource allocation variables from the model convergence is essential to reduce the problem to a single-round latency minimization problem. To solve this simplified problem, we propose an alternating optimization scheme to jointly consider communication and computation resource allocation and mitigate the straggler effect. We prove that the resulting sub-problems, i.e., bandwidth and computation capacity allocation, are both convex and can be optimally solved in closed form, respectively. Simulations show that compared with the baseline scheme that allocates the communication and computation resources equally across edge devices, the proposed scheme can achieve single-round latency reduction.