CMoS: Customizing Model Structures for Personalized Federated Learning

Abstract

Personalized Federated Learning (PFL) has recently received significant interest due to its capability of generating customized models for data-heterogeneous clients. Previous work focuses on weight adaptation to fit the data distribution while ignoring the correlation between various data distributions and their corresponding model structures. In this paper, we introduce a new paradigm called data-aware structure personalization, aiming to personalize the model structure according to the data distribution of each client during FL. Specifically, we develop CMoS, a framework to Customize Model Structures for heterogeneous client data in a learning-based manner. The key idea is to gradually learn a data-aware sub-structure from the original model at each federated round, which is achieved by introducing a newly designed binary mask and a sparsity loss during local training. Extensive experiments on multiple scenarios demonstrate the effectiveness of CMoS in generating personalized models with superior performance.