Multi-Model Federated Learning Optimization Based on Multi-Agent Reinforcement Learning

Abstract

This paper addresses the problem of Multi-Model Federated Learning (MMFL) in a typical wireless network, where a cellular Base Station (BS) cooperates with multiple clients to simultaneously train several Machine Learning (ML) models. Accordingly, the objective of this paper is to make an efficient joint decision for client association and communication-computation resource allocation to optimize the performance of the MMFL algorithm. In this regard, an optimization problem is formulated to minimize the average global loss of ML models under clients' energy and delay constraints. It is shown that the problem is a mixed-integer optimization whose objective is implicit in terms of the decision variables. To solve the optimization problem, we propose a Multi-Agent Multi-Model Federated Learning (MAMMFL) scheme based on a cooperative multi-agent configuration to intelligently assign models and resources to clients. Specifically, the problem is first converted to a Markov Decision Process (MDP) problem, then it is divided into four sub-MDP problems, where each problem relates to a phase in MMFL. The reinforcement learning algorithm solves each sub-problem, and a team-Q algorithm is adopted to coordinate agents in a cooperative multi-agent setting. Simulation results show that the proposed method can outperform other baselines in terms of average global loss and resource consumption.