Towards Generalizable Personalized Federated Learning with Adaptive Local Adaptation

1 INTRODUCTION . In recent years , research interests in training a machine learning model on edge devices motivated the paradigm of federated learning ( FL ) Li et al . ( 2020 ) , which makes it feasible for multiple devices to collaboratively train a shared global model in a privacy-preserving manner . Recent works Jiang et al . ( 2019 ) ; Fallah et al . ( 2020 ) ; Mansour et al . ( 2020 ) explored personalized federated learning ( PFL ) Kulkarni et al . ( 2020 ) , which aims to produce high-quality personalized models based on device-specific datasets and objectives . Specifically , each device uses the global model shared by the central server for its initial local model . Then , in the local adaptation Yu et al . ( 2020 ) , it trains a personalized model by performing several local updates with respect to their own data . On the one hand , Jiang et al . ( 2019 ) ; Arivazhagan et al . ( 2019 ) trained a global model to obtain the best accuracy on the whole data distribution of all participants . Thus , such global model with high optimality can be used as a strong start initialization point in personalization for each client with a subset data distribution . On the other hand , the works Fallah et al . ( 2020 ) ; Khodak et al . ( 2019 ) ; Chen et al . ( 2019 ) ; Jiang et al . ( 2019 ) tended to train a global model with high adaptability that current or new devices can easily adapt to their datasets in their local adaptation . These approaches are expected to produce high-quality personalized models given the heterogeneity of the underlying data distribution for all the users , i.e. , non independent and identically distributed ( non-iid ) data . However , we observe that existing frameworks lack the capability to adapt or generalize well between various non-iid data scenarios with different degrees of data and objectives diversity among devices . The learning efficiency and personalized performance of existing approaches are not perfectly performed to the changes in such diversity of non-iid data scenarios . With quantitative experimental results , we clearly point out that this is caused by targeting an unsuitable trade-off between the optimality and adaptability of the global model , which is referred to as the inappropriate optimalityadaptability trade-off . The most recent works Hanzely & Richtárik ( 2020 ) ; Mansour et al . ( 2020 ) ; Deng et al . ( 2020 ) potentially alleviated this issue by mixing the global model and personalized local models in one objective function for joint optimization . However , these works lack a clear insight of generalized PFL and optimality-adaptability trade-off . In addition , the proposed methods relied upon the combination formula with mixing weights that are hard to update in different non-IID data scenarios . In this paper , we rethink the local adaptation primarily from an information theory perspective Yin et al . ( 2020 ) . We argue that the information required to learn a personalized model φ in local adaptation is derived from both the initial model θ and the local dataset D. Thus , high I ( φ ; θ ) and I ( φ ; D ) correspond to the global model trained to be high optimality and adaptability on the whole data distribution , respectively . This leads us to identify suitable trade-offs in different on-IID data scenarios by balancing two information flow sources . Therefore , with the idea of avoiding the local update from solely relying on one information source , we formalize the learning process with an optimization problem by adding mutual information constraints . The key insight of this approach is that it can maintain performance by adaptively adjusting two information flows in the local adaptation based on the non-IID data condition . Our mathematical analysis of this optimization problem further gives a lower bound that consists of the meta-learning training loss Finn et al . ( 2017 ) and three regularizers . Our proposed adaptive trade-off theorem merges these principles to an embodiment . It theoretically ensures that our major contribution in this paper is to achieve generalizable PFL with adaptive local adaptation . These theoretical insights lead to our design of a novel adaptive federated meta-learning ( AFML ) framework . The basic learning architecture of AFML naturally inherits the meta-learning PFL framework that has been tested by fed-MAML and proven to converge Fallah et al . ( 2020 ) . Then , the corresponding training algorithm is proposed based on deduced regularization terms of the theorem . We tested AFML on CIFAR-100 and Shakespeare datasets . Our experimental results on non-iid scenarios with extensive degrees of diversity show that AFML maintains the highest personalized accuracy . Moreover , AFML achieves a reduction in required communication rounds by 25.5 % and 11.6 % , and a reduction in communication cost by 1.6 and 3.9 times compared with alternative leading PFL frameworks . 2 PROBLEM STATEMENT . We focus on the standard personalized federated learning ( PFL ) problem , which aims to train a global model utilized as the initialization to produce the personalized model based on the local data of each client . In practical applications , each participating client can contain a local dataset with personal preferences and objectives . Limited by the privacy constraint , we can not access the whole data distribution of clients to predetermined the global model targeted for optimality or adaptability to facilitate the local adaptation in each client . Thus , the major problem is to adaptively achieve the suitable trade-off between these two objectives of the global model , thus maintaining the solid personalized local models under non-IID data scenarios with different degrees of diversity among clients . Basically , there are C clients with the whole data distribution ( X , Y ) in the decentralized system . For an individual device with index j ∈ C , we assume the device-specific task objective Γj is sampled from a task objective distribution p ( Γ ) , which follows the definition of the meta-learning Triantafillou et al . ( 2019 ) . The local dataset M j of this client contains the train dataset Dj : { ( x , y ) j } and the test dataset D′j : { ( x′ , y′ ) j } that are disjoint . Both of them are sampled from the client-specific data distribution Dj , D′j ∼ ( X , Y ) j ⊂ ( X , Y ) and have N j and N ′j samples , respectively . Besides , PFL aims to train a shared global model θ that is regarded as the initialization for producing high-quality personalized models φ1 , ... , φC . Non-iid data with diversity quantity ζ . The basic non-IID scenario utilized in our work is the distribution-based label non-IID data setting . The distribution of classes within one client follows the Dirichlet distribution with parameter 0.5 . We define a mutually exclusive set in which each client augments the samples and randomly relabels the assigned classes . Thus , any two clients in this set potentially have different learning objectives and data distribution . Then , let ζ denote the proportion of non-mutually-exclusive devices in all C devices . Thus , the diversity of the non-IID data can be quantified as 1− ζ. ζ = 1 means that local datasets of participating clients are subsets of the whole data distribution ( X , Y ) , which defines a general non-IID data scenario in the FL . Then , the θ is trained to get good personalized models by using minθ L ( θ ) : = 1C ∑ j L j ( φj ) , φj : = θ − ηlOLj ( θ ) where ηl is the update rate and Lj : Rd → R is the loss function of M j in device j . Thus , in local adaptation , the quality of φj depends on θ and M j . On the one hand , when ζ = 1 , the learning of φj benefits from training θ to get optimality on the large global dataset { M j } C j=1 . On the other hand , a θ with high adaptability can be easily adapted to the local data in each client to obtain better φj when ζ = 0 . The object of our generalizable PFL is to obtain high-quality personalized models efficiently under non-IID data scenarios with any diversity ζ ∈ [ 0 , 1 ] . We argue that adaptive local adaptation induces a suitable trade-off between the optimality and adaptability of the global model , which contributes to generalizable PFL . Thus , our goal is to achieve generalizable PFL by implementing adaptive local adaptation . 3 METHODOLOGY . In this section , we first present the issue of solely targeting the adaptability or optimality of the global model in leading works . Then , we provide insights into this issue through experiments and mathematical analysis . Based on the findings , we prove a theorem referred to as the adaptive trade-off theorem , which motivates the design of the adaptive federated meta-learning ( AFML ) framework . 3.1 ANALYSIS THE ISSUE FROM EXPERIMENTS 500 1000 15000.10 0.15 0.20 0.25 0.30 = 0.1 1000 1500 20000.1 0.2 0.3 0.4 0.5 = 0.4 1500 1750 2000 2250 25000.2 0.3 0.4 0.5 0.6 = 0.6 1500 20000.2 0.3 0.4 0.5 0.6 0.7 = 0.9 Cifar100 learning curve of personalization FedAvg-FT , K=3 FedAvg-FT , K=10 FedAvg-Meta , 80 % support FedAvg-Meta , 20 % support ( a ) Accuracy curve for two federated models under different ζ . 0 1 3 5 7 9 11 0.05 0.10 0.15 0.20 0.25 0.30 Te st A cc ur ac y = 0.1 0 1 3 5 7 9 11 0.1 0.2 0.3 0.4 0.5 = 0.4 0 1 3 5 7 9 11 personalization epochs 0.2 0.3 0.4 0.5 0.6 Te st A cc ur ac y = 0.6 0 1 3 5 7 9 11 personalization epochs 0.3 0.4 0.5 0.6 0.7 = 0.9 FedAvg-FT , K=3 FedAvg-FT , K=10 FedAvg-Meta , 80 % support FedAvg-Meta , 20 % support ( b ) Personalized test accuracy under different personalization steps and ζ . 30 35 40 45 50 55 60 65 70 5 10 15 20 25 30 35 40 45 Ac c_ lo ca l Ac c_ pe r = 0.1 30 35 40 45 50 55 60 65 70 10 5 0 5 10 15 20 25 30 35 40 = 0.4 30 35 40 45 50 55 60 65 70 Acc_local 25 20 15 10 5 0 5 10 15 Ac c_ lo ca l Ac c_ pe r = 0.6 30 35 40 45 50 55 60 65 70 Acc_local 35 30 25 20 15 10 5 0 5 = 0.9 FedAvg-FT , Pers=5 FedAvg-Meta , Pers=5 ( c ) Accuracy of two federated models vs. local , trained-from-scratch models under different ζ . Figure 1 : Illustration of the issue in the FedAvg-FT and FedAvg-Meta by performing experiments under non-IID data scenarios with different degrees of diversity ζ = 0.1 , 0.4 , 0.6 , 0.9 among clients . The E and Pers denote the local update steps and personalization steps , respectively . s % support in FedAvg-Meta is the proportion of support set used for training . Our conducted experiments of the CIFAR100 dataset shown in Fig . 1 present the unstable performance of optimality-oriented with fine-tuning method FedAvg-FT Jiang et al . ( 2019 ) and adaptabilityoriented FedAvg-Meta Chen et al . ( 2019 ) method in terms of the averaged personalized accuracy of clients and the learning efficiency . The local models were trained locally on clients ’ own data for 200 epochs with a learning rate of 0.005 . Both two methods fail when there is a high diversity among clients , i.e . ζ = 0.1 . The upper left subfigure of Eq . 1 ( a ) , ( b ) presents that they can not converge while Eq . 1 ( c ) further shows that they obtain a bad personalized accuracy as compared with the model trained locally . Toward optimality . FedAvg-FT obtains low accuracy and training speed in the medium diversity conditions , as shown by Fig . 1 ζ = 0.4 , 0.6 . This illustrates that the initial model trained to be optimal in the whole data distribution is hard to adapt to the local dataset . For example , Fig . 1 ( b ) shows that the corresponding accuracy of the personalized model improves less than 0.1 than the initial one . But when there is low diversity ζ = 0.9 , FedAvg-FT outperforms the adaptability-oriented method with a large margin in all metrics . Especially in Fig Fig . 1 ( b ) , personalization of the FedAvgFT can start from a powerful initial model to reach the accuracy of 0.6 in only seven personalization epoches . Toward adaptability . FedAvg-Meta can quickly converge to the best accuracy when the diversity is relatively high . Fig . 1 ( b ) shows that the initial model with high adaptability can be trained with the local dataset to have a 0.3 improvement in accuracy within 7 personalization steps . Otherwise , the personalization of FedAvg-Meta presents an opposite trend when the diversity is low ζ = 0.9 . It is trained by 1500 communication rounds to obtain an unstable convergence while requiring 11 epochs to update the initial model with low performance to reach a competitive accuracy . Fig . 1 ( c ) further shows that both two methods are less accurate than the local models for the majority of participants when the diversity is high ζ = 0.1 , 0.4 , 0.6 . But consistent with the above analysis , the FedAvg-Meta is significantly better than the FedAvg-FT . This means that most clients can update the initial model with high adaptability to obtain a high-quality personal model . However , starting from a powerful initial model is preferable when ζ = 0.9 because it is hard to enhance a weak initial model based on the small-size local dataset . Both methods is obviously not conducive to generalizable PFL . When the diversity among clients is high , targeting the global model to perform well on the whole data distribution ( i.e. , optimality ) makes it hard to be adapted to the local dataset . On the contrary , even though the global model can be highly adaptable , updating a weak model with respect to a small-scale local dataset produces a poor personalized model and damages the efficiency . As pointed by experimental results , solely targeting optimality ( FedAvg-FT ) or adaptability ( FedAvg-Meta ) of the global model leads to an unstable performance in non-IID data scenarios with different diversities . Therefore , this motivates us to propose the idea of adaptive local adaptation in which the global model is to train to obtain a suitable trade-off between optimality and adaptability based on the non-IID data condition .

Summary: This paper poses an "optimality vs adaptability" tradeoff for federated learning, where optimality corresponds to strong performance on a global dataset and adaptability corresponds to a model's ability to adapt to perform well on local datasets. To deal with this tradeoff, this paper seems to draw motivation from the modified MAML approach of Yin et al. 2020 that is motivated by preventing the meta-parameters from memorizing the meta-training data. In the federated learning setting considered here, the current paper tries to balance the mutual information between the user-updated parameters and the global data ("optimality") and the user-updated parameters and the meta-parameters ("adaptability"). The proposed approach is similar to MAML in that the meta-parameters are a globally shared initialization for gradient descent, and the user-updated parameters are the result of gradient descent from the shared initialization (i.e. the meta-parameters). Some theoretical analysis and preliminary empirical results are provided.