Quality-Aware Distributed Computation and Communication Scheduling for Fast Convergent Wireless Federated Learning

In wireless federated learning (WFL), machine learning (ML) models are trained distributively on wireless edge devices without the need of collecting data from the devices. In such a setting, the quality of a local model update heavily depends on the variance of the local stochastic gradient, determined by the mini-batch data size used to compute the update. In this paper, we explore quality-aware distributed computation for WFL where user devices share limited communication resources, using mini-batch size as a "knob" to control the quality of users’ local updates. In particular, we study joint mini-batch size design and communication scheduling, with the goal of minimizing the training loss as well as the training time of the FL algorithm. For the case of IID data, we first characterize the optimal communication scheduling and the optimal minibatch sizes. Then we develop a greedy algorithm that finds the optimal set of participating users with an approximation ratio. For the case of non-IID data, we first characterize the optimal communication structure and the optimal mini-batch sizes. Then we develop algorithms that find the optimal communication order for some special cases. Our findings provide useful insights for the computation-communication co-design for WFL. We evaluate the proposed mini-batch size design and communication scheduling using simulations, which corroborate improved learning accuracy and learning time.