Towards Generalizable Personalized Federated Learning with Adaptive Local Adaptation
Authors: Sijia Chen
Baochun Li
Venue: ICLR 2022 Submitted (2022)
Keywords: Personalized federated learning
Meta-learning
Information theory

Abstract: Personalized federated learning aims to find a shared global model that can be adapted to meet personal needs on each individual device. Starting from such a shared initial model, devices should be able to easily adapt to their local dataset to obtain personalized models. However, we find that existing works cannot generalize well on non-iid scenarios with different heterogeneity degrees of the underlying data distribution among devices. Thus, it is challenging for these methods to train a suitable global model to effectively induce high-quality personalized models without changing learning objectives. In this paper, we point out that this issue can be addressed by balancing information flow from the initial model and training dataset to the local adaptation. We then prove a theorem referred to as the {\em adaptive trade-off theorem}, showing adaptive local adaptation is equivalent to optimizing such information flow based on the information theory. With these theoretical insights, we propose a new framework called {\em adaptive federated meta-learning} (AFML), designed to achieve generalizable personalized federated learning that maintains solid performance under non-IID data scenarios with different degrees of diversity among devices. We test AFML in an extensive set of these non-IID data scenarios, with both CIFAR-100 and Shakespeare datasets. Experimental results demonstrate that AFML can maintain the highest personalized accuracy compared to alternative leading frameworks, yet with a minimal number of communication rounds and local updates needed.

Under review as a conference paper at ICLR 2022
TOWARDS GENERALIZABLE PERSONALIZED FEDER- ATED LEARNING WITH ADAPTIVE LOCAL ADAPTA-
TION
Anonymous authors Paper under double-biee blind review
ABSTRACT
Personalized federated learning aims to find a shared global model that can be adapted to meet personal needs on each individual device. Starting from such a shared initial model, devices should be able to easily adapt to their local dataset to find that existing works cannot generalize obtain personalized models. However, heterogeneity degrees the underlying well on non-iid scenarios with different data distribution among devices. Thus, it is challenging for these methods to train a suitable global model to effectively induce high-quality personalized models without changing learning objectives. In this paper, we point that this issue can be addressed by balancing information flow from the initial model and training dataset the local adaptation. We then prove a theorem referred as the adaptive trade-off theorem, showing adaptive local adaptation is equivalent to optimizing such information flow based on the information theory. With these theoretical insights, we propose a new framework called adaptive federated meta-learning (AFML., designed to achieve generalizable personalized federated learning that maintains solid performance under non-III data scenarios with different degrees of diversity among devices. We test AFML in an extensive set of these non-lID data scenarios, with both CIFAR- 100 and Shakespeare datasets. Experimental results demonstrate that AFML can maintain the highest personalized accuracy compared to alternative leading frameworks, yet with minimal number of communication rounds and local updates needed.
INTRODUCTION
In recent years, research interests in training 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 shared global model in privacy-preserving manner. Recent works Jiang et al. 2011);; Fallah et al. (2020); Mansour et al. (2020) explored personalized federated learning (PFL. Kulkarni et al. (2020), which aims 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, the local adaptation Yu al. (2020). it trains personalized model by performing several local updates with respect to their own data.
On the one hand, Jiang et al. (019); ; Ariazzaaaa et al. (2019) trained global model 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 subset data distribution. On the other hand, the works Fallah et al.) (2020); Khodak et al. (2019); Chen jet al. (2019); Jiang et al. (2019) tended to train a global model with high adaptability that current or new devices expected to produce high-quality personalized models given the heterogeneity of the underlying data distribution for all the users, i.e., independent and identically distributed (non-iid) data.
casily adapt to their datasets in their local adaptation. These approaches are
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

---

Under review as a conference paper at ICLR 2022
results, we clearly point out that this is caused by targeting unsuitable trade-off between the optimality and adaptability of the global model, which is referred to as the inappropriate optimality-
adaptability trade-off.
The most recent works Hanzely & Richtárik 2200); Mansour et al. (0200;; : Deng et al. (2020)
potentially alleviated this issue by mixing the global model and personalized local models one
objective function for joint optimization. However, these works lack a clear insight of generalized PFL
and -ddataabblyy 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, rethink the local adaptation primarily from an information theory perspective Yin et al. 22200). We argue that the information required to learn a personalized model ø in local adaptation is derived from both the initial model 0 and the local dataset D. Thus, high ( 0) and T(o; 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 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 AFMI. 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 local dataset with personal preferences and objectives. Limited by the privacy constraint, we cannot access the whole data distribution of clients to predetermined the global model targeted for optimality or adaptability to facilitate the local adaptation each client. Thus, the major problem 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 clients with the whole data distribution (2) in the decentralized system. For an individual device with index j C, we assume the device-specific task objective 1' is sampled from a task objective distribution pTT), which follows the definition of the meta-learning Triantafillou
al. 2019). The local dataset M of this client contains the train dataset D { (xx y)² and the
test dataset D' {(x', "3k,")'))) that are disjoint. Both of them are sampled from the client-speciffc data distribution DJ_ D'j (X, y C (X,Y) and have N and N samples, respectively. Besides, PFL aims to train a shared global model A that is regarded as the initialization for producing high- quality personalized models o¹ 6²
Non-iid rid data with diversity quantity Ç. The basic non-lID scenario utilized in our work is the distribution- based label non-III data setting. The distribution of classes within one client follows the
2

---

Under review as a conference paper at ICLR 2022
Dirichlet distribution with parameter 0.5. We define 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-exchuee
devices all C' devices. Thus, the diversity of the non-IID data can be quantified C. (= means that local datasets of participating clients subsets of the whole data distribution ()) which defines a general non-IID data scenario in the FL.
Then, the 0 trained to get good personalized models by using ming L(Ø) := E L' (o)) oi =
0 - VV LV (0) where vn is the update rate and L' R R is the loss function of M in device j. Thus, local adaptation, the quality of &ي depends on and M2, On the one hand, when ( = the
learning of Φ benefits from training 0 to get optimality the large global dataset { j=1 On the other hand, 0 with high adaptability can be easily adapted to the local data in each client to obtain better &² when Ç 0.
The object of our generalizable PFL is to obtain high-quality personalized models efficiently under non-Ill) data scenarios with any diversity Ç [0. 1j. 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, 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 ISSUR FROM EXPERIMENTS
200
(a) Accuracy curve for two federated (b) Personalized test accuracy under (c) Accuracy of two federated mod- models under different
different personalization steps and Ç.els local, rraindd-room-scril models under different Ç.
Illustration of the issue in the FedAvg-FT and FedAvg-Meta by performing experiments under non-lID 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.
Figure 1:
Our conducted experiments of the CIFAR100 dataset shown Fig. present the unstable perfor- mance of optimality-oriente with fine-tuning method FedAvg-FT Jiang et al. (2019) and adaptability- oriented FedAvg-Metu Chen et al. (2019) method in terms of the averaged personalized accuracy of clients and the learning efficiency. The local models were trained locally clients' own data for 200 epochs with a learning rate of 0.005.
Both two methods fail when there is high diversity among clients, i.c. Ç = 0.1. The upper left subfigure of Eq. [ (a). (b) presents that they cannot converge while Eq. (c) further shows that they obtain bad personalized accuracy compared with the model trained locally.
3

---

Under review as a conference paper at ICLR 2022
Toward optimality. FedAvg-FT FT obtains low accuracy and training speed in the medium diversity conditions, as shown by Fig. [ = 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. (Tb) 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-F1 outperforms the adaptability-oriente method with a large margin in all metrics. Especially in Fig Fig. [b), personalization of the FedAvg- FT can start from a powerful initial model to reach the accuracy of 0.6 in only seven personalization epoches.
Toward adaptability. FedAvg- Mettvv-Meta can quickly converge to the best accuracy when the diversity is relatively high. Fig. (bb)) shows that the initial model with high adaptability can be trained with the local dataset to have 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. [c) further shows that both two methods are less accurate than the local models for the majority of participants when the diversity is high the FedAvg-Mete 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.
0.1, 4, 0.6. But consistent with the above analysis,
Both methods is obviously not conducive to generalizable PFL. When the diversity among clients is high, targeting the global model 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 weak model with respect to a small-scale local dataset produces 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-lll) data scenarios with different diversities. Therefore, this motivates us to propose the idea of adaptive local adaptation in which the global model to train to obtain a suitable trade-off' between optimality and adaptability based on the non-IID data condition.
3.2 INSIGHTSS THROOGH THE INFORMATION THEORY
In the local adaptation of PFL., the initial model updated based on the local data. Thus, from the information theory perspective, the information for learning the personalized model derives from the initial model and the local dataset. This makes us give insights into the discussed issue through the information flow in the local adaptation.
When < - 1, FedAvg-Mee with the meta-learning framework, tends to produce a complex global model with parameters 6, which will only overfit datasets in part of clients containing similar task objective P. This phenomenon is defined as task overfitting in the meta-learning research Yin et al. (2020). Thus, in the local adaptation, 0 effectively performs the test dataset without utilizing local training samples. This can be denoted as q(T'l [x' ()q(() ( 0) = q(ī") |x', ()q(o(0). in which the local model Φ is independent of the dataset 1) in the client, such that q( |x' D) = q(V' x' 0) where y' is the predicted label. Then, the participanting clients is unable to adjust the initial model based on personal local dataset.
In contrast, when there is high diversity Ç + O among clients, solely targeting the optimality in FeT makes weights of the global model provide sparse prior information for personalized learning. The main reason is that the local data distribution deviates a lot from the global distribution. For instance, in the extreme case, the shared global model can be regarded as randomly initialized weights for each client to train on their own small-size local dataset. Thus, the training for personalized model heavily relies on the information of the local dataset. This can be denoted as 405' |x,,, o ()) D, D, q(v'l |æ', Φ)((o) D), in which weights ò of the local model only dependent on the local dataset D. Personalized learning is equivalent to training random weights based on D, which damages the training speed and the performance.
An unsuitable trade-off hinders the generalizable PFL by causing the over- dependeedd on the global model or the local dataset in local adaptation. Therefore, argue that the generalizable PFL. can be
4

---

Under review as a conference paper at ICLR 2022
achieved by implementing the adaptive local adaptation that is formulated to balance the contributions of these two information sources in personalization.
3.3 ADAPTIVE LOCAL ADAPTATION BALANCING INFORMATION FLOWSS
The Bayesian inference shown by the graphical model in Fig. 2 TRAIN part presents the pre- diction process of the personalized local model 6). The initial model / updated with the local dataset D to generate the σ², and then the x'' is classified by the g² Thus, predicted label y is obtained based
Obsaaeed
Sutver
information derived from training data D (i.e., D - oJ - yy) and 0 (i.e., + D'). As discussed in Section 3.2, target- ing the optimality and adaptability of the global model corresponds to over-depending on the former and latter information path in the local adaptation, respectively. These, we can achieve the adaptive local adaptation by balancing these two information flows.
Ddge deviree
TRAM
P(L
Figure 2: The graphical model that expresses the learning process the personalized FL and our proposed method.
Motivated by the work in Yin et al. (2020) that utilizes the mutual information to express the relationship between variables, naturally for- mulate the dependencies on two information flows as the mutual information [(; 6) and ((o:D), Therefore, the optimal information flow balance can be obtained by solving a constrained optimization problem, shown in Eq. [ This problem aims prevent the learning process from solely relying on one single information flow.
max I(V': ; 10, D) S.t.I D,0, x') S Le
where I, is the constant information constraint and is a intermediate variable that is the hidden representation of the input x'.
Given 0 and D, the information flow path 0 > y is motivated by maximizing mutual information between the target y' and the intermediate variable z that is generated by 0. This encourages train a powerful global model that performs well on the global distribution. However, given the input x', we limit the dependence of local learning 0 by restricting the mutual information between the prediction and f in the reasoning process. This mechanism encourages the initial model not to excessively store information of device-specifee data its parameter 8, thereby increasing the importance of the local dataset in learning. Such constrain induces the global model with high adaptability.
3.4 AN ADAPTIVE TRADE-OFF THEOREM
Here, we solve the optimization problem in Eq by introducing Lagrange multiplier 3. Then, the problem transformed into maximizing the formula, A T(y'; z, 0, D) BT D; (')) Then, by modeling the & as the stochastic variable following Gaussian distribution, the variational inference is utilized to get the solution. After the mathematical derivation shown in Appendix A, we can determine a lower bound of this formula as:
2 Ear o ( [log q( | 00,c))
E [ [[K") ()]lr(=))] 3E [
(2)
+ [(<':710) I
where r(0) (0, is variational approximation to the target distribution 6. In the first term of the equation, we present the detailed operation of the reparameterization trick while we ignore this detail in the following terms.
5

---

Under review as a conference paper at ICLR 2022
If the y is O, the model predictions do not depend the given x', leading low accuracy. Thus, maximizing [(x: y (0) term is equivalent to minimize training losses for high accuracy. Similarly. the high accuracy across clients leads to model predictions that are less dependent on local training data sets, making smaller T(y'; D)0). Therefore, we conclude that maximizing the lower bound in Eq can be replaced by maximizing training accuracy and three regularizers shown by the theorem below.
Theorem 1 (Adaptive trade-off theorem). In non-lID data scenarios with different diversity, suitable trade-off for learning high-quality personalization models is obtained by targeting the following equation during the training process.
max E E,-,(o.)) aw 02,00,)) (1)r(() | [0)
8
-3E (KLp(( x' (')) ((m.))pr())
(3)
E( v'r (D'1) [L] (
where 0 represents the purameters of the global model truined on local data of decentralized clients 4(0) {M²), and o) ( 0).
The above theorem demonstrates that the adaptive local adaptation can be implemented by maximizing a novel objective function training. Then, the learning system is able find suitable trade-off to produce high-quality personalized models in non-iid data scenarios with any degree of diversity.
3.5 ADAPTIVE
ALGORITHM
As pointed by federate meta-lcarning methods Chen ct al. 2019); Khodak et al. (2019); Fallah et al. (2020), clients in the FL paradigm can be regarded as tasks in the meta-learning framework. have the support set D, and the query set D'. Based these definitions, the last two terms in Eq 3 of Theorem are naturally general training objectives of meta-learning. The initial model FF is updated using the support to minimize the loss in the query set. Therefore, based on Theorem 1. we propose the adaptive meta-learde peeaaleeriingg framework (AFML) to achieve generalizable PFL with adaptive local adaptation. The training schema of AFML naturally follows the meta meta-learning framework Finn et al. (2017).
Then,
Theorem also introduces three additional terms, which can be expressed as three trade-off regular- ization terms. Thus, the regularizer LroPF is formulated as follows.
Ltorr me' Y'ED"
a N
Σ
Ex-p(e) [- log q(' = v'10 xcc] LENG
Σ KL (p(Z)22)
LCONS
(4)
N
+ 8KL(p(0),
LCMPX
where the LENG encourages encoding that is maximally informative about the target. LCONS gives constraints to the latent representation of ž, which is similar to the work Alemi et al. 2016). LCMPX motivates an initial model with low complexity.
For the training process based on all participating client datasets, the third item LeMPx can be further written as E [K (p(( M) (r())). Therefore, when we 9 N(T) with T (9a, 0a), the LCMPX can be put in the server and computed as L = (L)p( 9 (()))
Finally, to achieve the Adaptive trade-off theorem, design split the global model into two modules, including the encoding network f and the classification network h. The encoding network trained to leam the hidden representation of the observation æ., while h is used to make the prediction based on ć. One intuition of such design is to alleviate the inconsistent between different local datasets by learning compact shared space. Then, add the regularization weights of the encoding network and leave the weights of h unrestricted.
The corresponding training algorithm is shown by Algorithm I.

---

Under review as a conference paper at ICLR 2022
Algorithm 1: Learning algorithm of AFML
input =OuterLoop/Server learning rate 77 725 InnerLoop/Client learning rate ry: Regularization coefficient 93 #C participating clients
Initialization: Initialize weights distribution ~ (T)) with r = (0w, 0,) for the encoding
network /: Initialize weights 0 for classification network h. Initialize local models $6) }:
Φ b= { :...,
Φi, ...,"
for communication round t = 1,2,... do
Randomly select K clients C, from the total clients;
Sample 0g from ((T) with reparameterization;
Distribute O, and 0 selected clients; for each client j Ct in parallel do 0g AdientUpdaae (j. O.):
end
i 0
111 K Ejec, ge:
Aggregating the updates for the global classification network
10
M2 K
Ejec, o/y
Aggregating the updates for the global encoding network
11
end
Function ClientUpdate (j,0,0)
12
M - (Di = (x) D = (x"5, y));;
13
Encode observation to the hidden feature z 20
Update the initial model, o² < θ nV n.. L = 1 % Ez W'D: (2(:2)) 33)
15
Obtain the loss on the query set, L ET 1 ED① (h (z') 6¹) y²))
Compute the gradients for the classification network, gf = 7 V Di:i
17
Compute the gradients for the global model, gh = (IS- + alhOFF) where
LTOFF L'EEN((D', + L'oNs(D', + (L)M
return g6, 06-9/;:
19
4 RELATED WORK
Federated learning (FL) Zhao et al. (2018); Kairouz et al. (2019): ; Li et al. (2020) machine learning paradigm in which devices collaboratively train global model while keeping the training data decentralized. With conventional federated learning (FL) McMahan et al. (2017). the global model is only trained develop a common output for all devices. However, the inconsistent learning objectives and data heterogeneity among devices in non-iid scenarios introduce a new personalization requirement Chen et al.] (2019). Thus, we have witnessed significant progress works Tiang et al. (001); Fallah et al. (2020: Mansour et al. (2020); Yu et al. (2020); Khodak et al. (2019) of personalized federated learning (PFL).
Due to inherent diversity among local data shards and objectives, many works Chen et al. (019);
Jiang et al. (2019); Fallah et al. al. (2020) designed target functions that maximum the adaptability of the global model. Especially, Chen et al. (2019) firstly noticed that such target function could be achieved by introducing the meta-learning framework to PFL.
However, it is still challenging for existing methods maintain performance when there non-iid scenarios with different degrees of diversity among devices. AArivzhhagan et al, (2019) discussed that sole pursuing adaptability damages the personalization when devices require powerful initial model in local adaptation. They solved this problem by training a base layer to optimal and changing only the last layer personalize in individual devices.
7

---

Under review as a conference paper at ICLR 2022
The latest works proposed to maintain personalization performance in non-iid scenarios with various diversities were Deng et al. (2020); Hanzely& & Richtárik (2020); Khodak et al. (2019); Mansour et al. (2020). Khodak al. (2019) specifically pointed out that existing PFL methods performance is unstable when diversity among different devices' data changes. These works alleviated the issue by seeking joint optimization of the global model and personalized models. The proposed methods, such as Deng et al. (2020) heavily relied on trainable mixing weights that are updated during the learning process.
5 EXPERIMENTS
(a)
(b)
(c)
(d)
Figure 3: Test accuracy of personalized models (a) (b) and Communication rounds (c) (d) required to reach target aceuracy under Ç E [0, 1]. Our AFML is compared with FedAvg-FT FedAvg-Mett
and FedMAML based on the CIFAR100 (a) (c) and Shakespeare (b) (d) datasets.
In this section, we evaluate the performance the proposed AFML on CIFAR- and Shakespeare Caldas al. (208)) datasets under the basic non-III setting described in the Section 2 with the diversity Ç E 00 1J. Then, following the three metrics described in the Section 3 of the Appendix, AFML compared with three state-of-the-art methods, including FedAvg-FT Jiang et al. (2019). FedPer-Meta Fallah et al. (2020), and FedMAML Deng et al.] (2020) where FedPer-Meta is an enhancement work of FedAvg-Mets Mettg-M-ttaa Chen al. (2019). Limited by the space, more settings and experimental results are described in Section 3 of the appendix.
5.1 RESULTS
Accuracy performance. Fig presents that AFML achieves highest accuracy compared to three leading frameworks under non-Ill) scenarios with any degrees of diversity, yet with minimum communication rounds used.
In the non-lID scenario with highly diversity Ç S 0.3, the average personalization accuracy of AFML around 56% and 39% CIFAR-100 and Shakespeare, respectively. However, other frameworks cannot converge. In the non-lID scenario with low diversity Ç 2 0.5, AFML has a significant performance improvement and gets the highest accuracy, which is 5.25% higher than FedMAML, 9.79% higher than FedPer-Meta and 18.69% higher than FedAvg-Fi FT on CIFAR100, AFML also performs well on the Shakespeare dataset. Moreover, compared with other methods, after performing five personalization epochs (Pers.5 Acc), the average personalization accuracy of AFMI. is 5.2% and 7.4% higher than its Pers.1 Ace on CIFAR- and Shakespeare, respectively.
The main reason is that AFML not only provides a strong point for personalization but also improves adaptability the global model. The insight observation from our theorem is that AFML trains an optimal global model with low generalization error while also targets adaptability. And it can achieve a suitable trade-off under non-iid scenarios with different diversity. This also contributes to the high convergence speed.
Efficiency performance. As shown in Fig. compared with other algorithms, the computation cost of AFML achieves significant reduction. In particularly, the Bytes of AFML is respectively 22.5%, 55.7%, and 67.3% less than FedMAML, FedPer-Mett Meta and FedAvg-FI average on both datasets.
Fairness Comparison. The results presented in Fig. 4Tb and Fig. demonstrates that AFML has the ability to leam high-quality personalized models for majority clients. For both two datasets, AFML leads to more clients with higher accuracy compared to the model trained locally. Moreover,
8

---

Under review as a conference paper at ICLR 2022
proposed AFML does not encourage the sacrifice of faimess to obtain higher average accuracy. reflecting the effectiveness of trade-off regularization terms. The main reason is that each device can optimize its information flow to obtain suitable trade-off for better personalization.
250
CPer302
Shalempia's
Fralay P. Fries
Spertm DAVeHAA
C-12
(-06
20
Sytbem CUNTAENE
250
10
C-08
200
(a) System overhead for achieving a target accuracy (b) Accuracy of four models vs. local. trained-from- scratch models under different
in different methods
Figure 4: Comparison of our AFML with three leading PFL methods on CIFARIOO and Shakespeare datasets, For System overhead in (a), we set the target accuracy as its test accuracy when convergence. (The symbol "x" represents that the algorithm cannot converge.)
C=0.1
030
(=06
C=9
< 0.0.1
0.4
0.6
{ 0.8
Figure 5: Accuracy distribution of devices in AFML compared to FedAvg-FF LadPer-etaa and FedMAML under different non-lID settings. The first row presents the results of the CIFARI00, while the second row presents the Shakespeare.
6 CONCLUSION
We have pointed out the connection between the generalizable PFL and learning objectives (i.e.. optimality and adaptability) trade-off of the global model. An in-depth study through experiments and mathematical analysis shows that a suitable trade-off can be obtained by implementing adaptive local adaptation. A proven theorem referred to as the adaptive trade-off theorem further shows balancing information flow in local adaptation contributes to the desired suitable trade-off. These theoretical insights guide our design of novel adaptive federated meta-learning(ATI (FFM)) framework that achieves generalizable PFL with adaptive local adaptation. Through the empirical experiments on non-iid data scenarios with extensive diversity among devices, we have demonstrated that the proposed AFML is superior to alternative learning frameworks in terms of personalized learning quality and convergence speed.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410, 2016.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Feder- ated learning with personalization layers. arXiv preprint 19Xxvv/912.0000) 2019.
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 1812:01812.01097, 2018.
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated meta-learning with fast convergence and efficient communication. arXiv preprint arXiv: 1807777777 2019.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 2020.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta- learning approach. arXiv preprint arXiv:2002.07948, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. arXiv preprint arXivv 2017.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv preprint wwXiv:2002.05516, 2020.
Yihan Jiang. Jakub Konečny, Keith Rush, Ruiyu Li, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv: 2019.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint 11Xxv:1112..4977 2019.
Mikhail Khodak, Maria-Florina F Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta- learning methods. In Advances in Neural Information Processing Systems, pp. 5917-5928, 2019.
Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant. Survey of personalization techniques for federated learning. arXiv preprint arXiv:2003.088: 08673,, 2020.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-60, 2020.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002.10619 2020.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-efficii learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp.
2017.
Eleni Triantafillou, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Utku Evei, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, et al. Meta-dataset: A dataset of datasets for learning learn from few examples. arXiv preprint arXiv:1903.03094 2019.
Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Meta-learning without memorization. arXiv preprint arXiv: 192.0219,, 2020.
Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adapta- tion. CoRR, abs/2002.044588, 2020. /ttp://///bpppuuuttttrrcdooooooo abs- 200--588-
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.00588 2018.
10


Reviews:
[{"review_id": "sWaYOSptltI", "reviewer": "Reviewer_uetq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed loss with regularizer, to combat memorization (also maximize the information flow) is an interesting and promising method. Apart from the positive parts, I have a few comments:\n\n(1) The dataset with different zeta is good for synthetic dataset (which AFML adapts well with differently labeled dataset). However, FL typically works on real world dataset. Applying AFML on more real world dataset (e.g. EMNIST-62 without randomize the labels) could further strengthen the position of this paper.\n\n(2) The exact differentiate factors between this paper and Yin et al could be made, to make this paper's contribution more easy to digest. I need to cross-read this paper and Yin et al multiple times to understand the differentiating factor.\n\n(3) The hyperparameter alpha and beta in AFML is pretty critical for final performance. I see on appendix that it is fixed, a badly configured alpha and beta will not work. Some discussion on the loss hyperparameter could be helpful (e.g when will the performance on personalization hurts global model performance, there should be some trade-off, which is missing in discussion). The potential side effects of combating memorization (e.g hurt performance under some cases, since memorization could be very useful in many cases) is not shown in this paper.", "summary_of_the_paper": "This paper is a Federated Learning (FL) extension version of \"Meta Learning without Memorization\" (Yin et al 2020). The author uses FL personalization loss, together with the three regularizers (information loss(L_ENC), activation regularization(L_CONS) and weight regularization(L_CMPX)). The activation and weight regularization are very close to Yin et al, while the information loss L_ENC is the difference, since Meta Learning aims only at fast adaptation, while FL personalization also aims to learn a good global model without adaptation: L_ENC could be pretty useful to learn a feature extraction network. \n\nThe proposed loss function and regularizer, show some improvement of personalization on non-IID dataset with different level of zeta, which indicates that panelizing memorization could improve personalization performance (test acc on average and individual improvement, and communication rounds). It is an interesting work interplaying between meta learning and FL.", "summary_of_the_review": "Combating memorization is very important to improve FL personalization performance, the proposed method is promising. \nHowever, some experiments on real-world experiments are missing to support the claim. \nMoreover, the linkage between meta learning work (Yin et al) and the proposed work could be better addressed.\nFinally, the negative effects of combating memorization is missing, which could further strengthen the paper. \n\nOverall, I give score 6, and happy to conduct discussion with the authors and other reviewers. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eMXRg9YQoYc", "reviewer": "Reviewer_qKrd", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Contributions\n\n1. Using mutual information to deal with the global learning vs personalization trade-off in federated learning is a novel approach and may yield promising results with further development.\n2. The empirical results presented are somewhat promising, showing accuracy and efficiency improvements over FedAvg+FT and PerFedAvg with very limited details provided.\n\nMajor issues\n\n1. The writing is very poor and makes the technique difficult to understand. The problem formulation is not clear, for example, it is not defined what the global dataset is or what the local loss functions are. The motivation for equation (1) is very vague, and intuition for the derivation of (2) and Theorem 1 are not clear. It is not clear what the algorithm is actually doing (e.g. what is the encoding network?). There are many grammatical mistakes.\n\n2. Theorem 1 cannot be considered a theorem. \"A suitable trade-off for learning personalized models\" is not defined so the statement is meaningless.\n\n3. Very few experimental details are provided. Granted not all experimental details need to be provided in the main body, but for a paper whose main contributions are empirical, there are far too few details given. It is not explained how accuracies are computed, how many clients are in the dataset, how many samples are in each client's datasets, how fine-tuning is implemented for FedAvg+FT, what the difference between FedAvg-Meta (known as PerFedAvg in the original paper) and FedMAML is, and how the hyperparameters for these methods are chosen. It is not explained why FedAvg+FT does not converge in some cases, or what not converging means (and the lack of convergence is strange -- I suspect this is a failure of initialization). Without these details it is difficult to glean the meaning of the results. \n\n4. Given the breadth of personalized federated learning methods, it is insufficient to compare against only three baselines (especially since two of the baselines are very similar).\n\nMinor issues\n\n1. The writing in the plots is too small.\n\n2. The claim of improved fairness due to the proposed method is not substantiated with a definition of fairness and a corresponding metric showing the improvement of the proposed method.\n\nComparison with related work\n\nThe authors discuss works that combine federated learning and meta-learning at a high level (Fallah et al. 2020, Chen et al. 2019, Jiang et al. 2019) and mention other personalized federated learning techniques without explanation. Also, they ignore works on personalized FL from 2021. The use of mutual information in federated learning is novel to the best of my knowledge, but the proposed approach draws inspiration from Yin et al. 2020. The authors mention this work but do not clearly explain how their technique differs.\n\nYin et al., META-LEARNING WITHOUT MEMORIZATION, ICLR 2020 ", "summary_of_the_paper": "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.", "summary_of_the_review": "Considering the above mentioned points, I believe that this paper is below the bar. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "MK6UxYLx8pL", "reviewer": "Reviewer_Mvwu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of using information theory to motivate appropriate regularization in meta-learning seems interesting. However, this contribution is mainly made in [1]. The paper borrows the techniques from [1] to solve personalized federated learning (PFL) tasks. Unfortunately, although the authors have put much effort in explaining their algorithms, I still have little insight about the motivation of algorithm design and derivation of the regularization term. At least to me, the main techniques used in this paper seem very similar to [1]. I suggest the authors to explain what is the key difference between your paper and [1], and what are challenges of applying [1] to PFL problem and how did you solve them.\n\nAnother main concern is about tuning the parameters in your proposed algorithms. I don’t see any suggestions on how to choose tuning parameters such as $\\alpha$ and $\\beta$. Thus, I have doubts that the superiority of experimental performance may come from flexible tuning parameter space. I would suggest the authors to carefully design experiments to test the robustness of parameter choice or at least explain how you choose hyperparameters of your algorithms and competitors.\n\nIn addition, the writing is not in good shape. I observed a lot of typos and grammar mistakes through the paper. I pointed two places below, but there are much more errors. Please check them carefully. And the figures shown in the paper are too small to read. Besides, although the authors explained the information theory motivation, I still have little insight. I think one important ingredient is using generative model to learn some latent representation. However, the authors did not carefully explain that. I would suggest the authors to explain your algorithm design step by step so that people who are not familiar with relevant works can also follow. The high-level intuition can come after or accompany the detailed algorithm description.\n\nMinor issue:\n1. Grammar mistakes and typos. For example, in page 1, “Then, in the local adaptation Yu et al. (2020), it trains …” should be “Then, in the local adaptation, Yu et al. (2020) trains ….”. Citation should be treated as a subject. Please check similar problems in the paper. And in page 4, “Both methods is obviously not conducive to generalizable PFL” should be “Both methods are obviously not conducive to generalizable PFL”. There are many grammar mistakes. Please have a check on writing.\n2. Since the whole paper talks about diversity level, I think it will be better to sue $\\zeta$ instead of $1-\\zeta$ to refer to it, as I feel it more natural to think about a larger $\\zeta$ indicates a larger diversity level rather than the opposite.\n\n[1] Yin M, Tucker G, Zhou M, et al. Meta-learning without memorization[C]//International Conference on Learning Representations. 2020.", "summary_of_the_paper": "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.\n\n[1] Yin M, Tucker G, Zhou M, et al. Meta-learning without memorization[C]//International Conference on Learning Representations. 2020.", "summary_of_the_review": "Given the main techniques is borrow from existing literature, I have concerns about original contribution. Besides, the writing is not in good shape and needs major revision. I would suggest the paper for further improvement rather than accepting it at current state.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper approaches personalized federated learning from the perspective of meta-learning and use the mutual information framework developed in a recent work to regularize local model training. All the reviewers consider the writing very poor and hard to understand, and the contributions not sufficient for acceptance.
Decision: Reject