Resource Allocation for Multi-Task Federated Learning Algorithm over Wireless Communication Networks

The multi-task federated learning (FL) problem in the wireless communication system is investigated in this paper. The base station (BS) and wireless users cooperatively perform a two-task FL algorithm in the established model. Users use their local datasets to train two local models of two different tasks. The trained local model of only one task is transmitted to the BS at each time and the BS aggregates the obtained models to calculate a global model, which will be sent back to all users. Since the resources for wireless transmission, such as transmit power and number of subcarriers are limited, the BS have to allocate resources reasonably to minimize the time consumption of the FL procedure while meeting the required learning performance. On the other hand, users are dynamically arranged to participate in different tasks in each iteration. This resource allocation and users arrangement problem is formulated as an optimization problem which aims to minimize time consumption of the two-task FL procedure. To address this nonconvex problem, we first decompose it into two convex sub-problems. Then we propose an iterative algorithm to solve this problem via iteratively obtaining the optimal solution of the joint power control and communication round optimization subproblem, and user arrangement subproblem. Simulation results of this multi-task FL system show that the proposed algorithm can reduce 7.02% and 9.67% completion time compared to the uniform and random user selection schemes respectively.