[Paper Review] Convergence Time Optimization for Federated Learning over Wireless Networks

Mingzhe Chen, H. Vincent Poor+2 more|arXiv (Cornell University)|Jan 22, 2020

Privacy-Preserving Technologies in Data31 references33 citations

TL;DR

The paper proposes a joint design of user selection, wireless resource allocation, and ANN-based prediction to accelerate federated learning convergence over wireless networks while reducing training loss.

ABSTRACT

In this paper, the convergence time of federated learning (FL), when deployed over a realistic wireless network, is studied. In particular, a wireless network is considered in which wireless users transmit their local FL models (trained using their locally collected data) to a base station (BS). The BS, acting as a central controller, generates a global FL model using the received local FL models and broadcasts it back to all users. Due to the limited number of resource blocks (RBs) in a wireless network, only a subset of users can be selected to transmit their local FL model parameters to the BS at each learning step. Moreover, since each user has unique training data samples, the BS prefers to include all local user FL models to generate a converged global FL model. Hence, the FL performance and convergence time will be significantly affected by the user selection scheme. Therefore, it is necessary to design an appropriate user selection scheme that enables users of higher importance to be selected more frequently. This joint learning, wireless resource allocation, and user selection problem is formulated as an optimization problem whose goal is to minimize the FL convergence time while optimizing the FL performance. To solve this problem, a probabilistic user selection scheme is proposed such that the BS is connected to the users whose local FL models have significant effects on its global FL model with high probabilities. Given the user selection policy, the uplink RB allocation can be determined. To further reduce the FL convergence time, artificial neural networks (ANNs) are used to estimate the local FL models of the users that are not allocated any RBs for local FL model transmission at each given learning step, which enables the BS to enhance its global FL model and improve the FL convergence speed and performance.

Motivation & Objective

• Motivate and model federated learning over realistic wireless networks with limited resource blocks (RBs).
• Minimize FL convergence time and training loss through joint user selection and RB allocation.
• Enable prediction of non-transmitting users' local models via neural networks to improve the global model.
• Analyze how scheduling, resources, and prediction affect convergence speed and learning performance.

Proposed method

• Model FL with a BS collecting local models from a subset of users per iteration due to RB limits.
• Formulate a joint optimization to minimize iteration time and training loss, given user selection and RB allocation.
• Introduce a probabilistic user selection scheme that prioritizes users whose local updates strongly affect the global model.
• Use an interior-point approach to solve the RB allocation problem at each iteration after fixing user selection.
• Employ an artificial neural network (MLP) to predict non-assistive users’ local models to enhance the global model.
• Update the global model using a combination of transmitting and predicted local models with a prediction error check.

Experimental results

Research questions

• RQ1 How can user selection and RB allocation be designed to minimize FL convergence time under wireless constraints?
• RQ2 What is the impact of predicting non-transmitting users’ local models on FL convergence and accuracy?
• RQ3 What role does the learning method (full vs stochastic gradient) play in convergence under the proposed wireless framework?
• RQ4 How effective is the proposed probabilistic user selection in ensuring all users contribute to the global model?

Key findings

• Convergence time of FL over wireless networks can be reduced by up to 56% with the proposed design.
• ANN-based prediction of non-transmitting users’ models can improve global model quality and convergence speed.
• Probabilistic user selection increases the chance of including impactful local models while ensuring every user has non-zero access probability.
• RB allocation via linearized constraints enables efficient optimization and faster per-iteration updates.
• Including predicted local models in the global update can reduce training loss and enhance accuracy for downstream tasks (e.g., handwritten digit recognition) by up to 3% over standard FL.
• The analysis shows that full/SGD, RB allocation, user selection, and prediction accuracy significantly influence convergence behavior.