LEARNING SHAREABLE BASES FOR PERSONALIZED FEDERATED IMAGE CLASSIFICATION

Abstract

Personalized federated learning (PFL) aims to leverage the collective wisdom of clients' data while constructing customized models that are tailored to individual client's data distributions. The existing work of PFL mostly aims to personalize for participating clients. In this paper, we focus on a less studied but practically important scenario-generating a personalized model for a new client efficiently. Different from most previous approaches that learn a whole or partial network for each client, we explicitly model the clients' overall meta distribution and embed each client into a low dimension space. We propose FEDBASIS, a novel PFL algorithm that learns a set of few, shareable basis models, upon which each client only needs to learn the coefficients for combining them into a personalized network. FEDBASIS is parameter-efficient, robust, and more accurate compared to other competitive PFL baselines, especially in a low data regime, without increasing the inference cost. To demonstrate its applicability, we further present a PFL evaluation protocol for image classification, featuring larger data discrepancies across clients in both the image and label spaces as well as more faithful training and test splits.

1. INTRODUCTION

Recent years have witnessed a gradual shift in computer vision and machine learning from simply building a stronger model (e.g., image classifier) to taking more users' aspects into account. For instance, more attention has been paid to data privacy and ownership in collecting data for model training (Jordan & Mitchell, 2015; Papernot et al., 2016) . Building models that are tailored to users' data, preferences, and characteristics have been shown to greatly improve user experience (Rudovic et al., 2018) . Personalized federated learning (PFL) is a relatively new machine learning paradigm that can potentially fulfill the demands of both worlds (Kulkarni et al., 2020) . On the one hand, it follows the setup of federated learning (FL): training models with decentralized data held by users (i.e., clients) (Kairouz et al., 2019) . On the other hand, it aims to construct customized models for individual clients that would perform well for their respective data distributions. While appealing, existing work of PFL has mainly focused on how to train the personalized models, e.g., via federated multi-task learning (Li et al., 2020a; Smith et al., 2017) , model interpolation (Mansour et al., 2020) , fine-tuning (Chen & Chao, 2022; Yu et al., 2020) , etc. Specifically, existing algorithms mostly require saving for each client a whole or partial model (e.g., a ConvNet classifier or feature extractor). This implies a linear parameter complexity with respect to the number of clients, which is parameter-inefficient and unfavorable for personalized cloud service -the overall system needs a linear space of storage, not to mention the efforts for profiling, versioning, and provenance, for every client. Less attention has been paid to how to deploy and maintain the personalized system. A practical challenge of previous work is how to fulfill new clients' queries, who did not involve in the training phase. Beyond training personalized models for the participated clients only, we focus on preparing to serve new clients with fast, data-efficient personalization. A promising solution is Model Agnostic Meta-Learning (MAML) (Finn et al., 2017) that aims to learn a good initialization such that it can be adapted to a new task fast, e.g., in a few SGD steps. The model-based idea has been inserted into PFL as well, by learning a model ready to be fine-tuned on each client's local data (Fallah et al., 2020) . However, it still learns the parameters of a whole or partial model for each client. Several recent studies (Pillutla et al., 2022; Wu et al., 2022; Fallah et al., 2020) show that when individual clients' data are scarce, fine-tuning may suffer from overfitting and being sensitive

