Reinforcement-Learning-Based Layer-Wise Aggregation for Personalized Federated Learning

Abstract

A key challenge in classical federated learning (FL) is statistical heterogeneity, which may lead to slow convergence and low accuracy. To tackle this, personalized FL (PFL) accounts for the individual data distribution of each client. This article proposes a new PFL method that relies on two principles: 1) shared knowledge and personalized knowledge can be reflected in different layers of the network and 2) clients with more data can contribute more to shared knowledge, while knowledge transfer from similar clients can boost personalization. We propose a novel method that applies aggregation based on the local data sizes for the shared knowledge layers and uses a deep reinforcement learning (DRL) agent for aggregating the layers pertaining to personalized knowledge. To ascertain efficiency and scalability, we train a single DRL agent (for all users) that operates on the server side, taking as input the subset of models corresponding to participants in the previous round. To reduce the dimensionality of its state space, we design a multihead autoencoder (MHAE). Extensive experiments on benchmark datasets for variable data heterogeneity levels reveal benefits over leading baselines in terms of both higher accuracy (up to +3.71%) and faster convergence (a reduction of global rounds by up to 30.6%). Our code is accessible at: https://github.com/fdksd/pFedRLLA.