Reinforcement Learning-based Layer-wise Aggregation for Personalized Federated Learning
Primary Area: general machine learning (i.e., none of the above)
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Keywords: Personalized Federated Learning, Reinforcement Learning, Layer-wise Aggregation
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
Abstract: A key challenge in Federated Learning (FL) is statistical heterogeneity, which may result in slow convergence and accuracy reduction. To tackle this problem, personalized federated learning (PFL) aims to adapt the global model to the individual data distribution of each client. One approach for this is personalized aggregation, which automatically determines how much each client can benefit from other clients' models. This paper proposes a new PFL method based on two principles: a) shared knowledge and personalized knowledge are reflected in different layers of the network and b) clients with more data should contribute more to shared knowledge, while knowledge transfer from similar clients can boost personalization. Based on these, we propose a Reinforcement Learning-based Layer-wise Aggregation method (pFedRLLA) that applies different mechanisms for different neural network layers. For layers representing shared knowledge, aggregation is carried out based on the size of the local data samples of the client. For layers representing personalized knowledge, a deep reinforcement learning (DRL) agent is used to generate personalized aggregation weights. To ascertain efficiency and scalability, we train a single DRL agent (for all users) that operates on the server-side and takes as input a subset of user models. To further reduce its state-space, we design a multi-head auto-encoder to obtain low-dimensional embeddings of user models. Extensive experiments on benchmark datasets for variable data heterogeneity levels reveal that the proposed algorithm consistently outperforms baselines in terms of both higher accuracy (up to +3.1\%) and faster convergence (a reduction of global rounds by up to 20.5\%).
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 4696
Loading