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: The paper frames the personalized federated learning (PFL) problem as a meta-learning task. In order to achieve the appropriate trade-off between global optimization and local adaptability, the authors utilize the recent work [1] to add a meta-learning regularization term for better personalization effect. The paper then demonstrates the benefits of their proposed method over the other alternatives by implementing extensive experiments.  [1] Yin M, Tucker G, Zhou M, et al. Meta-learning without memorization[C]//International Conference on Learning Representations. 2020.