Task Selection and Resource Optimization in Multi-Task Federated Learning With Model Decomposition

Abstract

In this letter, we investigate the training latency minimization problem for a multi-task federated learning (FL) framework with model decomposition over wireless communication networks. To handle the non-independent and non-identically distributed (non-IID) data, we first transform the multi-class classification task into multiple binary classification tasks. We then introduce sampling equalization to ensure the convergence of FL system. The optimization problem aims to minimize the training latency under energy and FL convergence constraints by optimizing task selection, number of learning iterations, and communication resource allocation. We decompose it into three sub-problems and propose alternating algorithm to address each sub-problem iteratively. Numerical results validate that the proposed algorithm significantly reduces time consumption compared to the conventional algorithms.