O N B RIDGING G ENERIC AND P ERSONALIZED F EDERATED L EARNING FOR I MAGE C LASSIFICATION
A BSTRACT
Federated learning is promising for its capability to collaboratively train models with multiple clients without accessing their data, but vulnerable when clients’ data distributions diverge from each other.
This divergence further leads to a dilemma: “Should we prioritize the learned model’s generic performance (for future use at the server) or its personalized performance (for each client)?”
These two, seemingly competing goals have divided the community to focus on one or the other, yet in this paper we show that it is possible to approach both at the same time.
Concretely, we propose a novel federated learning framework that explicitly decouples a model’s dual duties with two prediction tasks .
On the one hand, we introduce a family of losses that are robust to non-identical class distributions, enabling clients to train a generic predictor with a consistent objective across them.
On the other hand, we formulate the personalized predictor as a lightweight adaptive module that is learned to minimize each client’s empirical risk on top of the generic predictor.
With this two-loss, two-predictor framework which we name Federated Robust
Decoupling (F ED -R O D) , the learned model can simultaneously achieve state-ofthe-art generic and personalized performance, essentially bridging the two tasks.
I NTRODUCTION
Large-scale data are the driving forces for modern machine learning but come with the risk of data privacy.
In applications like health care, data are required to be kept separate to enforce ownership and protection, hindering the collective wisdom (of data) for training strong models.
Federated learning (FL), which aims to train a model with multiple data sources ( i e. , clients) while keeping their data decentralized, has emerged as a popular paradigm to resolve these concerns (Kairouz et al., 2019).
The standard setup of FL seeks to train a single “global” model that can perform well on generic data distributions (Kairouz et al., 2019), e.g. , the union of clients’ data.
As clients’ data are kept separate, mainstream algorithms like F ED A VG (McMahan et al., 2017) take a multi-round approach shown in
Figure 1.
Within each round, the server ﬁrst broadcasts the “global” model to the clients, who then independently update the model locally using their own (often limited) data.
The server then aggregates the “local” models back into the “global” model and proceeds to the next round.
This pipeline is shown promising if clients’ data are IID ( i.e. , with similar data and label distributions) (Stich, 2019; Zhou & Cong, 2017), which is, however, hard to meet in reality and thus results in a drastic performance drop (Li et al., 2020b; Zhao et al., 2018).
Instead of sticking to a single “global” model that features the generic performance, another setup of FL seeks to construct a “personalized” model for each client to acknowledge the heterogeneity among clients (Dinh et al., 2020; Hanzely et al., 2020; Smith et al., 2017).
This latter setup (usually called personalized FL ) is shown to outperform the former (which we name generic FL ) regarding the test accuracy of each client alone.
So far, these two seemingly contrasting FL setups are developed independently.
In this paper, we however found that they can be approached simultaneously by generic FL algorithms like F ED A VG .
Concretely, algorithms designed for generic FL (G-FL) often discard the local models { w m } after training (see Figure 1).
As a result, when they are evaluated in a personalized setting (P-FL), it is the global model ¯ w being tested (Arivazhagan et al., 2019;
Dinh et al., 2020; Fallah et al., 2020; Li et al., 2021a; Liang et al., 2020; Smith et al., 2017;
Zhang et al., 2021).
Here, we found that if we instead keep { w m } and evaluate them in P-FL, they outperform nearly all the existing P-FL algorithms.
In other words, personalized models seem to come for free from the local training step of generic FL .
[Figure]
At ﬁrst glance this may not be totally surprising: local training in G-FL algorithms is driven by the client’s empirical risk, which is what a personalized model strives to optimize 1 .
What really surprises us is that even without an explicit regularization term imposed by most P-FL algorithms (Dinh et al., 2020;
Smith et al., 2017), the local models of G-FL algorithms can achieve better generalization performance.
We conduct a detailed analysis and argue that global aggregation — taking average over model weights — indeed acts like a regularizer for local models.
Moreover, applying advanced G-FL algorithms (Acar et al., 2021) to improve the G-FL accuracy seems to not hurt the “local” models’ P-FL accuracy.
Building upon these observations, we dig deeper into generic FL.
Speciﬁcally for classiﬁcation, the nonIID clients can result from non-identical class distributions or non-identical class-conditional data distributions .
One way to mitigate their inﬂuences is to make the local training objectives more aligned among clients.
While this can hardly be achieved for the latter case without knowing clients’ data, we can do so for the former case by setting a consistent goal among clients — the learned local models should classify every class well, even if clients’ data have different class distributions.
We realize this by viewing each client’s local training as an independent class-imbalanced problem (Cui et al., 2019;
He & Garcia, 2009) and applying objective functions dedicated to it (Cao et al., 2019;
Ren et al., 2020).
As will be shown in section 5, these class-balanced objectives lead to much consistent local training among clients, making the resulting “global” model more robust to non-IID conditions.
The use of class-balanced objectives, nevertheless, degrades the local models’ P-FL performance.
This is because the local models are no longer learned to optimize clients’ empirical risks.
[Figure]
To address this issue, we propose a unifying framework for G-FL and P-FL which explicitly decouples a local model’s dual duties : serving as the personalized model and the ingredient of the global model.
Concretely, we follow the F ED A VG pipeline and train the local model with the class-balanced loss , but on top of the feature extractor, we introduce a lightweight personalized predictor and train it with client’s empirical risk (see Figure 2).
With this two-loss, twopredictor framework which we name Federated Robust Decoupling (F ED -R O D) , the resulting global model can be more robust to non-identical class distributions; the personalized predictor can lead to decent P-FL accuracy due to the implicit regularization and the empirical loss.
Speciﬁcally for the personalized predictor, we propose to explicitly parameterize it with clients’ class distributions via a hypernetwork (Ha et al., 2017).
That is, we learn a shared meta-model that outputs personalized predictors for clients given their class distributions.
This not only enables zero-shot model adaptation to new clients (without their data but class distributions), but also provides a better initialization to ﬁne-tune the models given new clients’ data.
We validate F ED -R O D on multiple datasets under various non-IID settings.
F ED -R O D consistently outperforms existing generic and personalized FL algorithms in both setups.
Moreover, F ED -R O D is compatible with and can further improve advanced generic FL algorithms like F ED D YN
(Acar et al., 2021) whenever non-identical class distributions occur.
Our contributions are three-folded:
• Unlike most of the previous works that focus on either generic FL or personalized FL we propose F ED -R O D to excel on both at the same time.
We validate F ED -R O D with extensive experiments.
• We show that strong personalized models emerge from the local training step of generic FL algorithms, due to implicit regularization.
We further show that class-balanced objectives are effective for improving the generic FL performance when clients have different class distributions.
• F ED -R O D enables zero-shot adaptation and much effective ﬁne-tuning for new clients.
R ELATED W ORK (A DETAILED VERSION IS IN A PPENDIX A)
Generic federated learning.
F ED A VG (McMahan et al., 2017) is the standard algorithm, and many works are proposed to improve it, either in the global aggregation step (Chen & Chao, 2021;
Hsu et al., 2019; Lin et al., 2020; Reddi et al., 2021;
Wang et al., 2020a; Yurochkin et al., 2019) or local training step (Malinovskiy et al., 2020; Wang et al., 2020b; Yuan & Ma, 2020;
Zhao et al., 2018).
For example, to reduce local models’ drifts from the global model, F ED P ROX (Li et al., 2020a) and F ED D YN (Acar et al., 2021) employed regularization toward the global model;
SCAFFOLD (Karimireddy et al., 2020a) leveraged control variates to correct local gradients.
We also aim to reduce local models’ drifts but via a different way.
We apply objective functions in class-imbalanced learning (He & Garcia, 2009), which are designed to be robust to class distribution changes.
The closest to ours is (Hsu et al., 2020), which used a traditional class-imbalanced treatment named re-weighting.
We show that more advanced techniques can be applied to further improve the performance, especially under extreme non-IID conditions where re-weighting is ineffective.
Personalized federated learning.
Many approaches for personalized FL
(Kulkarni et al., 2020) are based on multi-task learning (MTL) (Ruder, 2017; Zhang & Yang, 2017).
For instance, Smith et al.
(2017) encouraged related clients to learn similar models; Dinh et al.
Hanzely et al.
Li et al.
(2021a) regularized local models with a learnable global model.
Our approach is inspired by MTL as well but has notable differences.
First, we found that global aggregation in generic FL already serves as a strong regularizer.
Second, instead of learning for each client a feature extractor (Bui et al., 2019; Liang et al., 2020) or an entire model, F ED -R O D shares a single feature extractor among clients, inspired by Caruana (1997); Zhang et al.
This reduces the total parameters to be learned and improves generalization.
Compared to (Arivazhagan et al., 2019; Collins et al., 2021) which also learned a shared feature extractor, F ED -R O D simultaneously excels in both FL setups.
Instead of designing speciﬁc algorithms for personalized FL, Cheng et al.
; Wang et al.
Yu et al.
(2020) showed that performing post-processing ( e.g. , ﬁne-tuning) to a generic FL model ( e.g. , ¯ w in F ED A VG ) leads to promising personalized accuracy.
We further showed that, the local models { w m } learned in F ED A VG and other generic FL algorithms are strong personalized models.
We note that, while many personalized FL algorithms also produce a global model, it is mainly used to regularize or construct personalized models but not for evaluation in the generic setup.
In contrast, we learn models to excel in both setups via a single framework without sacriﬁcing either of them.
A recent work P F ED HN (Shamsian et al., 2021) also applies hypernetworks (Ha et al., 2017) but in a very different way from F ED -R O D. P F ED HN learns a hypernetwork at the server to aggregate clients’ updates and produce entire models for them for the next round.
In contrast, we learn the hypernetwork locally to construct the personalized predictors, not the entire models, for fast adaptation to clients.
P ERSONALIZED M ODELS E MERGE FROM G ENERIC F EDERATED L EARNING
In this section, we show that personalized FL (P-FL) models emerge from the training process of generic FL (G-FL) algorithms.
To begin with, we review representative G-FL and P-FL algorithms.
3. B ACKGROUND
Generic federated learning.
In a generic FL setting with M clients, where each client has a data set
D m = { ( x i , y i ) } |D m | i =1 , the optimization problem to solve can be formulated as min w L ( w ) =
M (cid:88) m =
|D m | |D| L m ( w ) , where L m ( w ) =
|D m | (cid:88) i (cid:96) ( x i , y i ; w ) .
Here, w is the model parameter; D = ∪ m D m is the aggregated data set from all clients; L m ( w ) is the empirical risk computed from client m ’s data; (cid:96) is a loss function applied to each data instance.
Federated averaging (F ED A VG ).
As clients’ data are separate, Equation 1 cannot be solved directly.
A standard way to relax it is F ED A VG (McMahan et al., 2017), which iterates between two steps, local training and global aggregation, for multiple rounds of communication
Local: w m = arg min w L m ( w ) , initialized with ¯ w ;
[Equation]
The local training is performed at all (or part of) the clients in parallel, usually with multiple epochs of SGD to produce the local model w m .
The global aggregation is by taking element-wise average over model weights.
Since local training is driven by clients’ empirical risks, when clients’ data are non-IID, w m would drift away from each other, making ¯ w deviate from the solution of Equation 1.
Personalized federated learning.
Personalized FL learns for each client m a model w m , whose goal is to perform well on client m ’s data.
While there is no agreed objective function so far, many existing works (Dinh et al., 2020; Hanzely & Richtárik, 2020; Hanzely et al., 2020; Li & Wang, 2019;
Li et al., 2021a; Smith et al., 2017) deﬁne the optimization problems similar to the following
[Equation]
where R is a regularizer; Ω is introduced to relate clients.
The regularizer is imposed to prevent w m from over-ﬁtting client m ’s limited data.
Unlike Equation 1, Equation 3 directly seeks to minimize each client’s empirical risk (plus a regularization term) by the corresponding personalized model w m .
In practice, personalized FL algorithms often run iteratively between the local and global steps as well, so as to update Ω according to clients’ models.
One example is to deﬁne Ω as a global model (Dinh et al., 2020; Hanzely & Richtárik, 2020; Hanzely et al., 2020; Li et al., 2021a), e.g. , by taking average over clients’ models, and apply an L 2 regularizer between Ω and each w m .
The corresponding local training step thus could generally be formulated as
[Equation]
where w ( t ) m denotes the local model after the t -th round; λ is the regularization coefﬁcient.
It is worth noting that unlike Equation 2, w in Equation 4 is initialized by w ( t ) m , not by Ω (or ¯ w ).
Terminology.
Let us clarify the concepts of “global” vs. “local” models, and “generic” vs. “personalized” models.
The former corresponds to the training phase: local models are the ones after every round of local training, which are then aggregated into the global model at the server (Equation 2).
The latter corresponds to the testing phase: the generic model is used at the server for generic future test data, while personalized models are speciﬁcally used for each client’s test data.
3. L OCAL MODELS OF GENERIC FL ALGORITHMS ARE STRONG PERSONALIZED MODELS
Building upon the aforementioned concepts, we investigate the literature and found that when generic
FL algorithms are evaluated in the P-FL setup, it is their global models being tested.
In contrast, when personalized FL algorithms are applied, it is their local models ( e.g. , Equation 4) being tested.
This discrepancy motivates us to instead evaluate generic FL algorithms using their local models.
Figure 1 summarizes the results (see section 5 for details).
Using local models of F ED A VG ( i.e. ,
Equation 2) notably outperforms using its global model in the P-FL setup.
At ﬁrst glance, this may not be surprising, as local training in F ED A VG is driven by clients’ empirical risks.
What really surprises us, as will be seen in section 5, is that F ED A VG ’s local models outperform most of the existing personalized FL algorithms, even if no explicit regularization is imposed in Equation 2.
3. I NITIALIZATION WITH WEIGHT AVERAGE IS A STRONG REGULARIZER
To gain a further understanding, we plot F ED A VG local models’ accuracy on clients’ training and test data.
We do so also for a state-of-the-art personalized FL algorithm D ITTO (Li et al., 2021a), whose local training step for producing personalized models is similar to Equation 4.
As shown in Figure 3,
F ED A VG has a lower training but higher test accuracy, implying that F ED A VG ’s local training is more regularized than Equation 4.
We attribute this effect to the initialization in Equation 2.
Speciﬁcally, by initializing w with ¯ w , we essentially impose an L 2 regularizer λ 2 (cid:107) w − ¯ w (cid:107) 22 with λ → ∞ at the beginning of each round of local training, followed by resetting λ to be 0 .
We found that this implicit regularization leads to a smaller value of (cid:107) w − ¯ w (cid:107) 22 at the end of each local training round, compared to Equation 4.
Due to the page limit, we leave additional analyses in the appendix.
We note that, advanced generic FL algorithms like SCAFFOLD (Karimireddy et al., 2020b) and
F ED D YN (Acar et al., 2021) still apply this initialization and learn with the empirical risk during local training.
Thus, their local models are strong personalized models as well.
F EDERATED R OBUST D ECOUPLING (F ED -R O D)
The fact that personalized models emerge from generic FL algorithms motivate us to focus more on how to improve the latter, especially when clients have non-IID data distributions.
4. I MPROVING GENERIC FL WITH B ALANCED R ISK M INIMIZATION (BRM)
We ﬁrst analyze what factors may lead to non-IID conditions.
Suppose the data instance ( x , y ) of client m is sampled from a client-speciﬁc joint distribution
P m ( x , y )
=
P m ( x | y ) P m ( y ) , the non-IID distributions among clients can result from non-identical class distributions P m ( x | y ) , non-identical class-conditional data distributions P m ( y ) , or both.
All these cases can make L m ( w ) deviate from
L ( w ) in Equation 1, which is the main cause of degradation in generic FL (Li et al., 2020a;b).
One way to mitigate the inﬂuence of non-IID data is to make L m ( w ) align with each other.
This can be challenging to achieve if clients have different P m ( x | y ) : without knowing clients’ data 2 , it is hard to design such an aligned L m ( w ) .
However, when clients have different P m ( y ) 3 , i.e. , different and hence imbalanced class distributions, we can indeed design a consistent local training objective by setting a shared goal for the clients — the learned local models should classify all the classes well.
It is worth noting that setting such a goal does not require every client to know others’ data.
Learning a classiﬁer to perform well on all classes irrespective of the training class distribution is the main focus of class-imbalanced learning (He & Garcia, 2009; Japkowicz, 2000; Johnson &
Khoshgoftaar, 2019).
We therefore propose to treat each client’s local training as a class-imbalanced learning problem and leverage techniques developed in this sub-ﬁeld.
Re-weighting and re-sampling (Buda et al., 2018) are the most fundamental techniques.
Denote by N m,c the number of training instances of class c for client m , these techniques adjust L m ( w ) in Equation 1 into
[Equation]
where q y i is usually set as
[Equation]
Namely, they mitigate the inﬂuence of P m ( y ) by turning the empirical risk L m into a balanced risk
L BRm , such that every client solves a more consistent objective that is robust to the class distributions.
Recently, many class-imbalanced works proposed to replace the instance loss (cid:96) ( e.g. , cross entropy) with a class-balanced loss (Cao et al., 2019;
Kang et al., 2020; Khan et al., 2017; Ren et al., 2020; Ye et al., 2020), showing more promising results than re-weighting or re-sampling.
We can also deﬁne
L BRm using these losses, e.g. , the balanced softmax (BSM) loss (Ren et al., 2020)
L BRm ( w ) ∝ (cid:88)
i
[Equation]
BSM ( x i , y i ; w ) , where (cid:96) BSM ( x , y ; w ) =
[Equation]
log
[Equation]
Here, g c ( x ; w ) is the logit for class c , C is the label space, and γ is a hyper-parameter.
The BSM loss is an unbiased extension of softmax to accommodate the class distribution shift between training and testing.
It encourages a minor-class instance to claim a larger logit g y ( x ; w ) in training to overcome feature deviation (Ye et al., 2020) in testing.
We list other class-balanced losses in the appendix.
[Figure]
We take advantage of these existing efforts by replacing the empirical risk L m in Equation 2 with a balanced risk L BRm , which either takes the form of Equation 5 or applies a class-balanced loss ( e.g. , Equation 6), or both.
We note that, a variant of Equation 5 has been used in (Hsu et al., 2020).
However, our experiments show that it is less effective than class-balanced losses in extreme non-IID cases.
Interestingly, we found that L BRm can easily be incorporated into advanced FL algorithms like
F ED D YN (Acar et al., 2021), because these algorithms are agnostic to the local objectives being used.
4. L OCAL TRAINING AND LOCAL MODEL DECOUPLING WITH ERM AND BRM
The use of balanced risk L BRm in local training notably improves the resulting global model ¯ w ’s generic performance, as will be seen in section 5.
Nevertheless, it inevitably hurts the local model w m ’s personalized performance, since it is no longer optimized towards client’s empirical risk L m .
To address these contrasting pursuits of generic and personalized FL, we propose a unifying FL framework named Federated Robust Decoupling (F ED -R O D) , which decouples the dual duties of local models by learning two predictors on top of a shared feature extractor:
one trained with empirical risk minimization (ERM) for personalized
FL (P-FL) and the other with balanced risk minimization (BRM) for generic FL (G-FL).
Figure 4 (c-d) illustrates the model and local training objective of
F ED -R O D.
The overall training process of F ED -R O D follows F ED A VG , iterating between local training and global aggregation.
As mentioned in subsection 4.1, other generic FL algorithms (Acar et al., 2021; Karimireddy et al., 2020b; Li et al., 2020a) can easily be applied to the BRM branch to further improve the generic performance.
Without loss of generality, we focus on the basic version built upon F ED A VG .
We start with the model in Figure 4 (c).
Notations.
We denote by f ( x ; θ ) the shared feature extractor parameterized by θ , whose output is z .
We denote by h G ( z ; ψ ) and h P ( z ; φ m ) the generic and personalized prediction heads parameterized by ψ and φ m , respectively; both are fully-connected (FC) layers.
In short, our generic model is parameterized by { θ , ψ } ; our personalized model for client m is parameterized by { θ , ψ , φ m } .
Predictions.
For generic prediction, we perform z = f
( x ; θ ) , followed by ˆ y G = h G ( z ; ψ ) .
For personalized prediction, we perform f ( x ; θ ) , followed by ˆ y P = h G ( z ; ψ ) + h P ( z ; φ m ) .
That is, h P is an add-on to h G , providing personalized information that is not captured by the generic head.
The overall objective.
F ED -R O D learns the generic model with the balanced risk L BRm and the personalized predictor with the empirical risk L m .
That is, different from Equation 1, F ED -R O D aims to solve the following two optimization problems simultaneously
[Equation]
We note that, L m is only used to learn the personalized head parameterized by φ m .
Learning.
Equation 7 cannot be solved directly in federated learning, so F ED -R O D follows F ED A VG to learn iteratively between the local training and global aggregation steps
[Equation]
where φ (cid:48) m is learned from the previous round, similar to w ( t ) m in Equation 4.
That is, the personalized head will not be averaged globally but kept locally.
In our implementation, Equation 8 and
Equation 9 are solved simultaneously via SGD, and we do not derive gradients w.r.t.
θ and ψ from
L m ( { θ , ψ , φ m } ) .
The θ and ψ in Equation 9 thus come dynamically from the SGD updates of
Equation 8.
In other words, Equation 9 is not merely ﬁne-tuning on top of the generic model.
In the end of federated learning, we will obtain ¯ θ and ¯ ψ (Equation 10) for generic predictions and { θ (cid:63)m , ψ (cid:63)m , φ (cid:63)m } Mm =1 (Equation 8 and Equation 9) for personalized predictions, respectively.
Please be referred to the appendix for the pseudocode.
4. A DAPTIVE PERSONALIZED PREDICTORS VIA HYPERNETWORKS
In subsection 4.2, the parameter φ m of the personalized predictor is learned independently for each client and never shared across clients.
In other words, for a new client not involved in the training phase, F ED -R O D can only offer the global model for generic prediction.
In this subsection, we investigate learning a shared personalized predictor that can adapt to new clients.
Concretely, we propose to learn a meta-model which can generate φ m for a client given the client’s class distribution.
We denote by H P ( a m ; ν ) the meta-model parameterized by ν , whose output is φ m .
Here, a m ∈ R | C | is the | C | -dimensional vector that records the class distribution of client m ; i.e. , the c -th dimension a m [ c ] =
N m,c (cid:80) c (cid:48) N m,c (cid:48) .
Accordingly, the local training step of φ m in Equation 9 is replaced by
Local: ν (cid:63)m = arg min ν
L m ( { θ , ψ , ν } ) , initialized with ¯ ν ; Global: ¯ ν ←
[Equation]
We implement H P by a lightweight hypernetwork (Ha et al., 2017) with two fully-connected layers.
With the learned ¯ ν , the meta-model H P can locally generate φ m based on a m , making it adaptive to new clients simply by class distributions.
The parameter φ m can be further updated using clients’ data.
We name this version F ED -R O D (hyper); the previous one, F ED -R O D (linear).
Please see
Figure 4 (c-d) for an illustration.
We include more details in the appendix.
E XPERIMENT (M ORE D ETAILS AND R ESULTS IN T HE A PPENDIX )
Datasets, models, and settings.
We use CIFAR-10/100 (Krizhevsky et al., 2009) and FashionMNIST (FMNIST)
(Xiao et al., 2017).
We also include a realistic EMNIST (Cohen et al., 2017) dataset, which collects hand-written letters of thousands of writers.
To simulate the non-IID data distributions on CIFAR and FMNIST, we follow Hsu et al.
(2019) to create a heterogeneous partition for M clients: an M -dimensional vector q c is drawn from Dir ( α ) for class c , and we assign data of class c to client m proportionally to q c [ m ] .
The resulting clients have different numbers of total images and different class distributions.
With α < 1 , most of the training examples of one class are likely assigned to a small portion of clients.
Similar to Lin et al.
, we use M = 100 clients for
FMNIST and M = 20 for CIFAR-10/100, and sample 20% / 40% clients at every round, respectively.
For EMNIST, we use the digit images, follow Caldas et al.
(2018) to construct 2 , 185 clients (each is a writer), and sample 5% clients at every round.
We use a ConvNet (LeCun et al., 1998) similar to (Acar et al., 2021; McMahan et al., 2017).
It contains 3 convolutional layers and 2 fully-connected layers.
We train every FL algorithm for 100 rounds, with 5 local epochs in each round.
We report the mean accuracy of ﬁve times of experiments with different random seeds.
We evaluate the generic performance (G-FL) using the generic model (GM) on the standard generic test set.
For
FMNIST and CIFAR-10/100, we evaluate the personalized performance (P-FL) using personalized models (PM) on the same set, but re-weight the accuracy according to clients’ class distributions
[Table]
[Table]
[Table]
the instance index.
This evaluation is more robust (essentially as the expectation) than assigning each client a speciﬁc test set.
For EMNIST, each client has its own test set with the same writing style.
Our variants.
We mainly use Equation 6 with γ = 1 as the L BRm and report the F ED -R O D (hyper) version (cf. subsection 4.3).
Table 2 provides the ablation study.
Baselines.
For G-FL methods including F ED A VG (McMahan et al., 2017), F ED P ROX (Li et al., 2020a), SCAFFOLD (Karimireddy et al., 2020b), and F ED D YN
(Acar et al., 2021), we use their global models ¯ w for G-FL evaluation; their local models ( i.e. , w m in Figure 1) for P-FL evaluation.
For P-FL methods, to evaluate their G-FL performance, we use the available global models in
P F ED M E (Dinh et al., 2020) and D ITTO (Li et al., 2021a) or average the ﬁnal personalized models for MTL (Smith et al., 2017), F ED P ER (Arivazhagan et al., 2019), LG-F ED A VG (Liang et al., 2020),
F ED FOMO (Zhang et al., 2021), and F ED R EP (Collins et al., 2021).
To illustrate the difference between applying GMs and PMs in a P-FL setting, we also evaluate the
P-FL performance using GMs, which is how F ED A VG has been applied to P-FL in literature.
5. R ESULTS
F ED -R O D bridges G-FL and P-FL and consistently outperforms all generic and personalized
FL methods.
Table 1 summarizes the results.
In terms of G-FL accuracy , advanced local training ( i.e. , SCAFFOLD, F ED P ROX , and F ED D YN ) outperforms F ED A VG and personalized methods, and our F ED -R
O D can have further gains by using balanced risk minimization (BRM).
We also investigate combining F ED -R O D and F ED D YN (Acar et al., 2021), using the latter to optimize the generic model with BRM, which outperforms either ingredient in many cases.
We report the G-FL accuracy of personalized FL algorithms mainly to investigate if they have similar properties like
F ED A VG : an algorithm designed for one setup can also construct models for the other setup.
In terms of P-FL accuracy , by using PMs most methods outperform the baseline of local training with individual client’s data without communication ( i.e. , local only), justifying the beneﬁts of federated collaboration 4 .
For generic FL methods, using PMs ( i.e. , local models { w m } ) clearly outperforms using GMs ( i.e. , ¯ w ), which supports our claims and observations in Figure 1 and subsection 3.2.
It is worth noting that the local models from generic FL methods are highly competitive to or even outperform personalized models produced by personalized FL methods.
This provides generic FL methods with an add-on functionality to output personalized models by keeping the checkpoints on clients after local training.
Our F ED -R O D achieves the highest P-FL accuracy and we attribute this to (a) the shared feature extractor learned with the balanced risk and re-initialized every round to beneﬁt from implicit regularization; (b) the personalized head learned with clients’ empirical risks.
[Table] [Table]
[Table] [Table]
BRM effectively reduces the variance of G-FL accuracy and local gradients.
To understand why F ED -R
O D improves
G-FL, we visualize the global model ¯ w ’s and each local model w m ’s G-FL accuracy on CIFAR-10 (Dir(0.3)), in Figure(upper).
F ED -R O D not only learns a better global model for GFL, but also has a smaller variance of accuracy across the local models’ generic heads (as their objectives are more aligned).
We also show how w m deviates from ¯ w after local training in Figure 5 (lower).
F ED -R O D has a smaller variance.
This coincides with the study in (Kong et al., 2021): lower variances of the local gradients could imply better generic performance.
F ED -R O D beneﬁts from decoupling.
We compare several variants of F ED -R O D (cf.
Figure 4), with one head (reduced to F ED A VG ) or different networks (linear/hyper).
We evaluate on CIFAR-10 (Dir(0.3)).
As shown in Table 2, F ED A VG with BRM signiﬁcantly improves G-FL but degrades in P-FL.
F ED -R O D remedies this by training a decoupled personalized head.
We note that, F ED -R O D does not merely ﬁne-tune the global model with clients’ data (cf. subsection 4.2).
We also compare different balanced losses in Table 3: advanced losses outperforms importance re-weighting (Hsu et al., 2020).
F ED -R O D (hyper) beneﬁts future clients.
To validate the generalizability to new clients, we build on the Dir(0.3)
nonIID setting for FMNIST and CIFAR-10/100, but split the training data into 100 clients ( 50 are in training; 50 are new).
We train on the 50 training clients for 100 rounds (sampling 20 of them every round).
We then evaluate on the 50 new clients individually, either using the global model directly or ﬁne-tuning it with clients’ data for several steps.
Table 4 and Figure 6 shows the averaged accuracy on new clients.
Without ﬁne-tuning,
F ED -R O D (hyper) can already generate personalized models, and outperforms others methods stably with ﬁne-tuning.
More results and analyses in the Appendix.
The Appendix includes studies with class-imbalanced global distributions and we show that F ED -R O D still performs well.
BRM can be further improved with meta-learned hyper-parameters.
We validate that re-initializing the local models by the global model at every round ( i.e. , Equation 2) does lead to a much smaller regularization loss than Equationto support our claim in subsection 3.3.
More comprehensive results regarding more clients, deeper backbones, compatibility with other methods, robustness against adversaries, etc, are also provided.
C ONCLUSION
Most of the existing work in federated learning (FL) has been dedicated to either learning a better generic model or personalized models.
We show that these two contrasting goals can be achieved simultaneously via a novel two-loss, two-predictor FL framework F ED -R O D. Concretely, we show that strong personalized models emerge from the local training of generic FL algorithms, due to implicit regularization; imposing class-balanced objectives further improves the generic FL accuracy when clients have non-IID distributions.
F ED -R O D seamlessly incorporates these two observations to excel in both FL settings, and further enables fast adaptation to new clients via an adaptive module.
A CKNOWLEDGMENTS
This research is partially supported by NSF IIS-2107077, NSF OAC-2118240, NSF OAC-2112606, and the OSU GI Development funds.
We are thankful for the generous support of the computational resources by the Ohio Supercomputer Center and AWS Cloud Credits for Research.
R EPRODUCIBILITY S TATEMENT
We report the results with the average over 5 runs of different random seeds.
We exhaustively provide the information about the hyperparameters, datasets, evaluation, and other details in sectionand Appendix C, which should be comprehensive for reproducibility.
We also provide our code in https://github.com/hongyouc/Fed-RoD
A C OMPARISON TO R ELATED W ORK A. F OLLOW - UP WORKS OF F ED A VG
Several recent works (Karimireddy et al., 2020a;b; Zhao et al., 2018) have shown that, with multiple steps of local SGD updates, the local model w m would drift away from each other, leading to a degenerated global model ¯ w that deviates from the solution of Equation 1 of the main text.
One way to mitigate this is to modify the local training objective (cf.
Equation 2 of the main text).
For instance, F ED P ROX (Li et al., 2020a) introduced a regularizer with respect to ¯ w ,
[Equation]
F ED D YN (Acar et al., 2021) further added a dynamic term based on the local model of the previous round w (cid:48) m ,
[Equation]
These regularizers aim to stabilize local training and align the objectives among clients.
Some other works did not change the objectives but introduced control variates or momentum to correct the local gradient (Karimireddy et al., 2020a;b), or designed a new optimizer more suitable for decentralized learning (Malinovskiy et al., 2020; Pathak & Wainwright, 2020; Yuan & Ma, 2020).
It is worth mentioning, in most of these works, the empirical risk L m ( w ) still plays an important role in driving the local model update.
Since L m ( w ) directly reﬂects the (non-IID) client data distribution, the learned local models are indeed strong candidates for personalized models.
A. G ENERIC FEDERATED LEARNING
F ED A VG (McMahan et al., 2017) is the standard algorithm, which involves multiple rounds of local training and global aggregation.
Many works have studied its convergence (Haddadpour & Mahdavi, 2019; Khaled et al., 2020), robustness (Bonawitz et al., 2019), communication (Koneˇcn`y et al., 2016;
Reisizadeh et al., 2019), especially for non-IID clients (Li et al., 2020a;b; Zhao et al., 2018).
Many other works proposed to improve F ED A VG .
In terms of global aggregation , (Wang et al., 2020a;
Yurochkin et al., 2019) matched local model weights before averaging.
(Chen & Chao, 2021; He et al., 2020; Lin et al., 2020;
Zhou et al., 2020) replaced weight average by model ensemble and distillation.
(Hsu et al., 2019; Reddi et al., 2021) applied server momentum and adaptive optimization to improve the global model update.
In terms of local training , (Liang et al., 2019; Malinovskiy et al., 2020; Pathak & Wainwright, 2020; Yuan & Ma, 2020) improved the optimizer.
To reduce local models’ drifts from the global model, (Zhao et al., 2018) mixed client and server data in local training; F ED P ROX (Li et al., 2020a), F ED DANE (Li et al., 2019), and
F ED D YN (Acar et al., 2021) employed regularization toward the global model; SCAFFOLD (Karimireddy et al., 2020a)
M IME (Karimireddy et al., 2020b) leveraged control varieties and/or server statistics to correct local gradients; (Wang et al., 2020b; Yao et al., 2019) modiﬁed the local model update rules.
For most of them, the empirical risks on clients’ data are the major forces to drive local training.
We also aim to reduce local models’ drifts but via a different way.
We directly bypass the empirical risks that reﬂect clients’ data distributions.
Instead, we apply objective functions in class-imbalanced learning (He & Garcia, 2009), which are designed to be robust to the change of class distributions.
Our approach is different from (Duan et al., 2020; Wang et al., 2020c; Yang et al., 2020), which monitored and resolved class imbalance from the server while we tackled it at the clients.
Our approach is also different from agnostic FL
(Deng et al., 2020b;
Mohri et al., 2019), whose local training is still built on empirical risk minimization.
The closest to ours is (Hsu et al., 2020), which used a traditional class-imbalanced treatment, re-weighting, to mitigate non-identical class distributions.
We show that more advanced techniques can be applied to further improve the performance, especially under extreme non-IID conditions where re-weighting is less effective.
Moreover, our method is compatible with existing efforts like F ED D YN (Acar et al., 2021) and SCAFFOLD (Karimireddy et al., 2020a) to boost the generic performance.
A. P ERSONALIZED FEDERATED LEARNING
Personalized FL (Kulkarni et al., 2020) learns a customized model for each client.
Many approaches are based on multi-task learning (MTL) (Evgeniou & Pontil, 2007; 2004; Jacob et al., 2009; Ruder, 2017; Zhang & Yang, 2017; Zhang & Yeung, 2010) — leveraging the clients’ task relatedness to improve model generalizability.
For instance, (Smith et al., 2017) encouraged related clients to learn similar models; (Corinzia & Buhmann, 2019;
Dinh et al., 2020; Hanzely & Richtárik, 2020; Hanzely et al., 2020; Li & Wang, 2019;
Li et al., 2021a) regularized local models with a learnable global model, prior, or set of data logits.
(Arivazhagan et al., 2019; Bui et al., 2019; Li et al., 2021b; Liang et al., 2020) designed the model architecture to have both personalized (usually the feature extractor) and shareable components.
(Huang et al., 2021; Zhang et al., 2021) constructed for each client an initialized model or regularizer based on learnable bases.
Our approach is inspired by MTL as well but has several notable differences from existing works.
First, we found that the global aggregation step in generic FL already serves as a strong regularizer.
Second, instead of learning for each client a personalized feature extractor (Bui et al., 2019; Liang et al., 2020) or an entire independent model that can operate alone (Dinh et al., 2020; Hanzely et al., 2020; Smith et al., 2017), F ED -R O D shares a single feature extractor among all clients, inspired by invariant risk minimization (Ahuja et al., 2020; Arjovsky et al., 2019) and domain generalization (Ghifary et al., 2015; Muandet et al., 2013).
This reduces the total parameters to be learned and improves model’s generalizability.
Compared to
F ED P ER (Arivazhagan et al., 2019) and F ED R EP (Collins et al., 2021) which also learned a shared feature extractor, F ED -R O D simultaneously outputs a single, strong global model to excel in the generic FL setup.
Some other approaches are based on mixture models.
(Agarwal et al., 2020;
Deng et al., 2020a; Mansour et al., 2020; Peterson et al., 2019; Zec et al., 2020) (separately) learned global and personalized models and performed a mixture of them for prediction.
(Reisser et al., 2021) learned a sets of expert models and used them to construct personalized models.
Meta-learning is also applied to learn a good initialized model that can be adapted to each client with a few steps of local training (Chen et al., 2018;
Fallah et al., 2020; Jiang et al., 2019; Khodak et al., 2019).
Instead of designing speciﬁc algorithms for personalized FL, (Wang et al., 2019; Yu et al., 2020) showed that performing post-processing ( e.g. , ﬁne-tuning) to a generic FL model ( e.g. , ¯ w learned by
F ED A VG ) already leads to promising personalized accuracy.
In this work, we further showed that, the local models w m learned in F ED A VG and other generic FL algorithms are indeed strong personalized models.
We note that, while many personalized FL algorithms also produce a global model, it is mainly used to regularize or construct personalized models but not for evaluation in the generic setup.
In contrast, we learn models to excel in both the setups via a single framework without sacriﬁcing either of them.
P F ED HN (Shamsian et al., 2021) also applies hypernetworks (Ha et al., 2017) but for a very different purpose from F ED -R O D. Speciﬁcally, P F ED HN learns a hypernetwork at the server to aggregate clients’ model updates and produce their entire models for the next round.
In contrast, we learn the hypernetwork locally to construct the personalized predictors, not the entire models, for fast adaptation to clients.
A. A VERAGING MODEL WEIGHTS AS A REGULARIZER
In subsection 3.3, we demonstrate that taking the average over model weights indeed acts as a regularizer for local models to improve their individual personalized performance. In more traditional machine learning, the regularization effects of averaging multiple independentlytrained models have been observed in some techniques like bagging (Poggio et al., 2002; Skurichina & Duin, 1998). Indeed, in several recent works of FL (Chen & Chao, 2021; He et al., 2020; Lin et al., 2020; Zhou et al., 2020), the authors replaced weight average by bagging/model ensemble to improve the generic performance on the global test set.
That is, they found that performing the model ensemble over clients’ models can yield more robust predictions on the global test set than the global model, which is generated by averaging the client models’ weights.
Here, we however study a different regularization effect, in personalized FL on local test sets.
As reviewed in subsection 3.1, personalized FL algorithms often impose a regularizer on the local/personalized models to overcome the fact that clients usually have limited data (please see Equation 3 and Equation 4 and the surrounding text).
What we claim is that even without such an explicit regularizer, the model weight average before local training (Equation 2) already serves as an implicit regularizer to the local models for their individual personalized performance, as we discussed in subsection 3.3 (Figure 3) and empirically veriﬁed in subsection D.2 and Figure 7.
A. S YSTEMATIC OVERHEAD
F ED -R O D has similar computation cost, communication size, and number of parameters as F ED A VG .
We discuss the difference between F ED -R O D and existing generic FL methods from a system view.
F ED P ROX (Li et al., 2020a) proposes a proximal term to prevent client from diverging from the server model, which is more robust to the heterogeneous system.
SCAFFOLD (Karimireddy et al., 2020b) imposes a gradient correction during client training.
Maintaining such a correction term, however, doubles the size of communication.
F ED D YN (Acar et al., 2021) resolves the communication cost issue by introducing a novel dynamic regularization.
However, it requires all users to maintain their previous models locally throughout the FL process, which is not desired when users have memory and synchronization constraints.
A. C LASS - IMBALANCED LEARNING
Class-imbalanced learning attracts increasing attention for two reasons.
First, models trained under this scenario using empirical risk minimization perform poorly on minor classes of scarce training data.
Second, many real-world data sets are class-imbalanced by nature (Gupta et al., 2019; Van Horn & Perona, 2017; Van Horn et al., 2018).
In this paper, we employ a mainstream approach, costsensitive learning (Cao et al., 2019;
Li et al., 2020c;
Ren et al., 2020; Tan et al., 2020; Ye et al., 2020), which adjusts the training objective to reﬂect class imbalance so as to train a model that is less biased toward major classes.
A. Z ERO - SHOT LEARNING
Our design choice of parameterizing the personalized prediction head with clients’ class distributions is reminiscent of zero-shot learning (Changpinyo et al., 2016; 2017; 2020; Lampert et al., 2013; Xian et al., 2018), whose goal is to build an object classiﬁer based on its semantic representation.
The key difference is that we build an entire fully-connected layer for FL, not just a single class vector.
We employ hypernetworks (Ha et al., 2017) for efﬁcient parameterization.
B A DDITIONAL D ETAILS OF F ED -R O D B. A DDITIONAL BACKGROUND ( CF . SUBSECTION 3.1 OF THE MAIN PAPER )
In the generic federated learning (FL) setting, the goal is to construct a single “global” model that can perform well for test data from all the clients.
Let w denote the parameters of the model, for a classiﬁcation problem whose label space is C , a commonly used loss is the cross entropy,
[Equation]
where g c ( x ; w ) is the model’s output logit for class c .
We note that, the concepts of global vs. local models and generic vs. personalized models should not be confused.
No matter which task (generic or personalized) an FL algorithm focuses on, as long as it has the local training step, it generates local models; as long as it has the global aggregation step (of the entire model), it generates a global model.
For instance, F ED A VG (McMahan et al., 2017) aims for generic FL
but it creates both the global and local models.
B. O VERVIEW OF F ED -R O D
For generic predictions, F ED -R O D performs feature extraction z
= f ( x ; θ ) , followed by h G ( z ; ψ ) .
For personalized predictions, F ED -R O D performs z
= f ( x ; θ ) , followed by h G ( z ; ψ ) + h P
( z ; φ m ) .
The element-wise addition is performed at the logit level.
That is, g c ( x ; w ) in Equation 12 can be re-written as
[Equation]
where z = f ( x ; θ ) is the extracted feature.
The overall training process of F ED -R O D iterates between the local training and global aggregation steps.
In local training, F ED -R O D aims to minimize the following objective
[Equation]
[Equation]
The empirical risk L m ( w m = { θ , ψ , φ m } ) is deﬁned as 1 |D m | (cid:80) i (cid:96) ( x i , y i ; w m ) , where D m = ( x i , y i ) } |D m | i =1 is the training data of client m .
We will introduce more options of the balanced risk
L BRm in subsection B.3.
We optimize Equation 14 via stochastic gradient descent (SGD).
We updates θ , ψ , and φ m in a single forward-backward pass, which consumes almost the same computation cost as F ED A VG .
For L m ( { θ , ψ , φ m } ) , we do not derive gradients w.r.t.
θ and ψ .
We emphasize that, according to subsection 4.2 of the main paper, the ﬁnally learned parameters of
F ED -R O D (linear) are ¯ θ , ¯ ψ , and { φ (cid:63)m } M m =1 .
We then plug them into Equation 13 for predictions.
In algorithm 1 and algorithm 2, we provide pseudocode of our F ED -R O D algorithm.
|D m | (cid:80) m (cid:48)∈S |D m (cid:48) |
|D m
(cid:80) m
D
B. B ALANCED RISK MINIMIZATION (BRM)
To learn a generic model, standard federated learning ( e.g. , F ED A VG (McMahan et al., 2017)) aims to optimize L ( w ) in Equation 1 of the main paper.
In theory, the overall objective L ( w ) is equal to the expected objective E [ L m ( w )] for client m , if client m ’s data D m are IID partitioned from D .
Here, the expectation is over different D m partitioned from D .
In reality, L m could diverge from L due to non-IID partitions of the aggregated data D into clients’ data.
That is, E [ L m ( w )] (cid:54) =
E [ L m (cid:48) ( w )] (cid:54) =
L ( w ) .
We mitigate the non-IID situation by directly adjusting L m such that E [ L m ( w
[Equation] ] ≈ E [ L m
[Equation] ( w [Equation] ] ≈ L ( w ) .
Essentially, L m ( w ) is the client’s empirical risk, which could be different among clients if their class distribution P m ( y ) are different.
We, therefore, propose to turn the empirical risk L m into a class-balanced risk L BR m by replacing (cid:96) in Equation 12 with a class-balanced loss (Cao et al., 2019; Cui et al., 2019;
Kang et al., 2020; Khan et al., 2017; Ren et al., 2020; Ye et al., 2020).
The class-balanced loss attempts to make the learned model robust to different training class distributions, such that the learned model can perform well for all the test classes.
In other words, the class-balanced loss is designed with an implicit assumption that the test data will be class-balanced, even though the training data may not be.
Table 5 summarizes some popular class-balanced losses.
We also include some extensions with meta-learning.
See subsection B.5.
[Table]
[Table]
One may wonder what if the global distribution is class-imbalanced?
Will BRM still be beneﬁcial to
FL?
In subsection D. [Table] we perform experiments to show that F ED -R O D with BRM can still improve on F ED A VG since BRM seeks to learn every class well.
Even though the global distribution might be skewed, BRM provides a novel alternative to mitigate the non-IID problem by making every client optimize a more consistent objective as we discussed above.
Designing better losses for BRM in FL will be interesting future work
[Table]
B. O N FEDERATED LEARNING FOR THE PERSONALIZED HEAD WITH HYPERNETWORKS
One drawback of existing personalized methods is that the personalized models are only available for clients involved in training or with sufﬁcient training data.
When new clients arrive in testing, is it possible for the federated system to provide corresponding personalized models?
To this end, instead of learning a speciﬁc prediction head φ m for each client m , we propose to learn a meta-model H P ( a m ; ν ) with a shared meta-parameter ν .
The input to H P is a vector a m ∈ R | C | , which records the proportion of class c ∈ C in client m ’s data.
The output of H P is φ m for h P .
In other words, H P can adaptively output personalized prediction heads for clients given their local class distributions a m .
We implement the meta-model H P by a hypernetwork (Ha et al., 2017), which can be seen as a lightweight classiﬁer generator given a m .
This lightweight hypernetwork not only enables clients to collaboratively learn a module that can generate customized models, but also allows any (future) clients to immediately generate their own personalized predictors given their local class distribution a m as input, even without training.
We construct the hypernetwork by two fully-connected (FC) layers (with a ReLU nonlinear layer in between).
Table 6 summarizes the number of parameters of each part in F ED -R O D. Hypernetworks add only a small overhead to the original model.
B. E XTENSION WITH META - LEARNING FOR THE IMPROVED BSM LOSS
F ED -R O D incorporates a balanced loss to learn the generic model.
Here we study a more advanced way to derive such balanced loss with meta-learning.
Inspired by (Ren et al., 2018; Shu et al., 2019) and the FL scenario proposed by (Zhao et al., 2018), we seek to combine the BSM loss and re-weighting as (cid:80) i q m,y i (cid:96) BSM ( x i , y i ; w ) , where q m,y i is meta-learned with a small balanced meta dataset D meta provided by the server.
(See Table 5 for a comparison.)
The D meta should have a similar distribution to the future test data.
We implement this idea with the Meta-Weight Net (MWNet) (Shu et al., 2019) with learnable parameter ζ .
In addition, we notice that the original BSM loss (cid:96) BSM γ
[Equation]
x [Equation] w )) has a hyperparameter γ which is set to be 1 via validation (Ren et al., 2020).
However, in federated learning it can be hard to tune such a hyperparameter due to the large number of non-IID clients.
Therefore, we propose to learn a client-speciﬁc γ m with meta-learning for (cid:96) BSM γ m .
More speciﬁcally, given a meta-learning rate η , the meta-learning process involves the following iterative steps:
Compute the Meta-BSM loss with a mini-batch B ∼ D m ; i.e. , ∀ ( x , y ) ∈ B , compute (cid:96) BSM γ m ( x , y ; w m ) .
Predict the example weights with q m,y =
MWNet ( (cid:96) BSM γ m ( x , y ; w m ); ζ m ) , ∀ ( x , y ) ∈ B . 3.
Re-weight the Meta-BSM loss: L BR m,B ( w m ) =
(cid:80) ( x
,y ) ∈ B q m,y i (cid:96) BSM
γ m ( x , y ; w m ) , and perform one step of gradient descent to create a duplicated model ˜ w m = w m − η ∇ w m L BR m,B . 4.
Computes the loss on the meta dataset D meta using the duplicated model: L BR m, D meta ( ˜ w m ) =
(cid:80) ( x ,y ) ∈D meta q m,y (cid:96) BSM γ m ( x , y, ˜ w m ) , followed by updating γ m ← γ m − η ∇ γ m L BR m, D meta and ζ m ← ζ m − η ∇ ζ m L BR m, D meta . 5.
Update the model: w m ← w m − η ∇ w m L BR m,B ( w m ) .
Throughout the federated learning process, γ m and q m,y are dynamically learned with meta-learning for different clients and rounds.
Results of F ED -R O D with Meta-BSM We sample 10 images for each class (only 0 .
2% of the overall training set) from the training set as the meta set.
We compare to (Zhao et al., 2018) that concatenates the meta set to clients’ local data.
The results in Table 10 and Table 12 are encouraging.
With a very small meta set, F ED -R O D outperforms (Zhao et al., 2018) by 1% to 14% on accuracy across different settings, validating the importance of balanced losses and how to set them up dynamically via meta-learning.
B. C ONNECTION TO I NVARIANT R ISK M INIMIZATION G AMES (IRMG)
F ED -R O D is inspired by a recently proposed machine learning framework Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) and its extension Invariant Risk Minimization Games (IRMG) (Ahuja et al., 2020).
Suppose that the whole dataset is collected from many environments, where data from each environment is associated with its characteristic, IRM introduces the concept of learning an invariant predictor.
(Note that, in IRM the learner can access data from all the environments; thus, it is not for an FL setting.)
Given the training data partition, IRM aims to learn an invariant feature extractor z = f
( x ; θ ) and a classiﬁer h ( z ; ψ ) that achieves the minimum risk for all the environments.
The concept of environments can be connected to clients’ private local data in FL which are often non-IID.
That is, given M environments, we can re-write IRM in a similar expression to Equationin the main paper
[Equation]
Unfortunately, IRM is intractable to solve in practice given the constraint that every environment relies on the same parameters (Ahuja et al., 2020).
IRMG relaxes it by reformulating the classiﬁer ψ as an ensemble of environment-speciﬁc classiﬁers (by averaging over model weights) [Equation] (cid:80)
[Equation]
IRMG is proved to optimize the same invariant predictor of IRM when it converges to the equilibrium in game theory, and it holds for a large class of non-linear classiﬁers.
IRMG is solved through iterative optimization: (1) training the feature extractor θ with centralized data ( i.e. , aggregated data from all environments), (2) training the environment-speciﬁc classiﬁers φ m on the data of each environment
D m , and (3) updating the main classiﬁer through weight averaging ¯ φ = 1 M (cid:80) m φ m .
We highlight the similarity between IRMG and F ED -R O D: both are training a strong generic feature extractor and a set of personalized classiﬁers.
For predictions on data of client (environment) m in Equation 18, IRMG uses ˆ y = 1 M ( φ (cid:62) m z
+ (cid:80) m (cid:48) (cid:54)
= m φ (cid:62) m (cid:48) z ) ; F ED -R O D’s personalized model is ˆ y =
h G
z ; ψ ) + h P
( z ; φ m ) .
We can connect IRMG to F ED -R O D by re-writing its prediction as h G ( z ; ¯ φ ) :
=
¯ φ (cid:62) z = 1 M (cid:80) m φ (cid:48)(cid:62) m z and h P ( z ; φ m ) :
= 1 M ( φ (cid:62) m z − φ (cid:48) m (cid:62) z ) , where φ (cid:48) m is the client m ’s model in the previous round/iteration of learning.
IRMG can not be applied directly to federated learning for the following reasons.
First, centralized training of the feature extractor is intractable since clients’ data are not allowed to be aggregated to the server.
Second, to perform the iterative optimization of IRMG, the clients are required to communicate every step, which is not feasible in FL due to communication constraints.
C I MPLEMENTATION D ETAILS
Implementation.
We adopt ConvNet (LeCun et al., 1998) following the existing works (Acar et al., 2021; McMahan et al., 2017; TensorFlow team, 2016).
For EMNIST/FMNIST, it contains 2 Conv layers and 2 FC layers.
The Conv layers have 32 and 64 channels, respectively.
The FC layers are with 50 neurons as the hidden size and 10 neurons for 10 classes as outputs, respectively.
For
CIFAR-10/100, it contains 3 Conv layers and 2 FC layers.
The Conv layers have 32 , 64 , andchannels, respectively.
The FC layers are with 64 neurons as the hidden size and 10 / 100 neurons for 10 / 100 classes as outputs, respectively.
To implement hypernetworks in F ED -R O D, we use a simple 2-FC ReLU network with hidden size 16 for EMNIST/FMNIST/CIFAR-100 and 32 for CIFAR-10.
We use standard pre-processing, where EMNIST/FMNIST and CIFAR-10/100 images are normalized.
EMNIST/FMNIST is trained without augmentation.
The 32 × 32 CIFAR-10/100 images are padded 2 pixels each side, randomly ﬂipped horizontally, and then randomly cropped back to 32 × 32 .
We train every method for 100 rounds.
We initialize the model weights from normal distributions.
As mentioned in (Li et al., 2020b), the local learning rate must decay along the communication rounds.
We initialize it with 0 . 01 and decay it by 0 .
99 every round, similar to (Acar et al., 2021).
Throughout the experiments, we use the SGD optimizer with weight decay 1e − 5 and a 0 . 9 momentum.
The mini-batch size is 40 ( 16 for EMNIST).
In each round, clients perform local training for 5 epochs.
We report the mean over ﬁve times of experiments with different random seeds.
For F ED P ROX (Li et al., 2020a), the strength of regularization λ is selected from [1e − 2 , 1e − 3 , 1e − 4] .
For F ED D YN (Acar et al., 2021), the strength of regularization λ is selected from [1e − 1 , 1e − 2 , 1e − 3] as suggested in (Acar et al., 2021).
For P ER -F ED A VG (Fallah et al., 2020), the meta-learning rate ˆ β [Table] is selected from [1e − 2 , 1e − 3 , 1e − 4] .
For P F ED M E (Dinh et al., 2020), the strength of regularization λ is selected from [15 , 20 , 30] .
F ED -R O D introduces no extra hyperparameters on top of F ED A VG .
For the generic and personalized heads of F ED -R O D, we study using 1 ∼ 4 FC layers but do not see a notable gain by using more layers.
We attribute this to the well-learned generic features.
Thus, for all our experiments on F ED -R O D, we use a single FC layer for each head.
We run our experiments on four GeForce RTX 2080 Ti GPUs with Intel i9-9960X CPUs.
Evaluation.
Both datasets and the non-IID Dirichlet simulation are widely studied and used in literature (Acar et al., 2021;
Hsu et al., 2019; Lin et al., 2020).
We use the standard balanced test set
D test for evaluation on generic FL (G-FL):
[Equation]
where ˆ y i,G here is the predicted label ( i.e. , arg max over the logits).
For evaluation on personalized
FL (P-FL), we still apply D test but weight instances w.r.t.
each client’s class distribution:
[Equation]
accuracy
[Equation] We do so instead of separating D test into separate clients’ test sets in order to avoid the variance caused by how we split test data (except the EMNIST dataset that each client has its own test set with the writer’s styles).
What we compute is essentially the expectation over the splits.
We have veriﬁed that the difference of the two evaluation methods is negligible.
In Table 1 of the main paper and some other tables in the appendix, we evaluate G-FL by an
FL algorithm’s generic (usually the global) model, denoted as GM .
We evaluate P-FL by an FL algorithm’s personalized models (or local models of a G-FL algorithm), denoted as PM .
For P-FL, we also report the generic model’s accuracy following the literature to demonstrate the difference.
Due to the space limit of the main paper, we provide the standard deviations of the results of Tablein Table 7, Table 8, and Table 9 here.
[Table]
[Table]
[Figure]
D A DDITIONAL E XPERIMENTS AND A NALYSES
Here we provide additional experiments and analyses omitted in the main paper.
We validate our claims in the main paper and the designs of our proposed F ED -R O D via the following experiments:
• subsection D.1: personalized models emerge from local training of generic federated learning (cf. subsection 3.2, subsection 3.3, and subsection 5.1 in the main paper).
• subsection D.2: balanced risk minimization (BRM) improves generic-FL performance (cf. subsection 5.1 in the main paper).
• subsection D.3: the roles of F ED -R O D’s generic and personalized heads (cf. subsection 4.
in the main paper).
• subsection D.4: personalization with hypernetworks (cf. subsection 5.1 in the main paper).
• subsection D.5: robustness to class-imbalanced global data.
• subsection D.6: compatibility of F ED -R O D with other G-FL algorithms (cf. subsection 5.
in the main paper).
• subsection D.7: comparison to personalized FL algorithms (cf. subsection 5.1 in the main paper).
• subsection D.8: ablation studies and discussions on F ED -R O D (cf. subsection 5.1 in the main paper).
D. P ERSONALIZED MODELS EMERGE FROM LOCAL TRAINING OF GENERIC FEDERATED LEARNING
As mentioned in section 3 in the main paper, personalized FL algorithms usually impose an extra regularizer (cf.
Equation 3 and Equation 4 of the main paper) during local training, but do not re-initialize the local models by the global models at every round.
In contrast, generic FL algorithms like F ED A VG do not impose extra regularization but re-initialize the local models at every round.
Here in Figure 7, we monitor the two loss terms, (cid:80) m |D m | |D| L m ( w m ) and (cid:80) m |D m | |D| (cid:107) w m − ¯ w (cid:107) 22 (cf.
Equation 3 and Equation 4 of the main paper), for F ED A VG and a state-of-the-art personalized FL algorithm Ditto (Li et al., 2021a) at the end of each local training round.
(Ditto does include the
L 2 regularizer in training the personalized models.)
Ditto achieves a lower empirical risk ( i.e. , the ﬁrst term), likely due to the fact that it does not perform re-initialization.
Surprisingly, F ED A VG achieves a much smaller regularization term ( i.e. , the second term) than Ditto, even if it does not impose such a regularizer in training.
We attribute this to the strong effect of regularization by re-initialization: as mentioned in subsection 3.3 of the main paper, re-initialization is equivalent to setting the regularization coefﬁcient λ as inﬁnity.
We note that, the reason that the regularization term of Ditto increases along the communication rounds is because ever time the global model ¯ w
[Figure] [Figure]
is updated, it moves sharply away from the local model w m .
Thus, even if the regularization term is added into local training, it cannot be effectively optimized.
This analysis suggests that the local models of generic FL algorithms are more regularized than the personalized models of personalized
FL algorithms.
The local models of generic FL algorithms are thus strong candidates to be evaluated in the personalized FL setting.
D. B ALANCED RISK MINIMIZATION (BRM) IMPROVES GENERIC -FL PERFORMANCE
To understand why F ED -R O D outperforms other generic methods in the G-FL accuracy, we visualize each local model w m ’s G-FL accuracy after local training in Figure 8 (both datasets with Dir(0.3)).
Methods rely on ERM suffer as their local models tend to diverge.
Figure 9 further shows that the variances of local weight update ∆ w m = w m − ¯ w across clients are smaller for F ED -R O D, which result from a more consistent local training objective.
In Figure 10, we further compare the G-FL accuracy among F ED A VG , F ED -R O D with the original
BSM loss, and F ED -R O D with the Meta-BSM loss introduced in subsection B.5 along the training process ( i.e. , training curve).
The local models of F ED A VG tend to diverge from each other due to the non-IID issue, resulting in high variances and low accuracy of G-FL.
The global aggregation does improve the G-FL accuracy, validating its importance in federated learning.
The local training in
F ED -R O D (BSM) not only leads to a better global model, but also has smaller variances and higher accuracy for the local models (as their objectives are more aligned).
With the help of meta dataset and meta-learning, F ED -R O D (Meta-BSM) yields even better G-FL performance for both global models and local models, and has much smaller variances among local models’ performance, demonstrating the superiority of using meta-learning to learn a balanced objective.
D. T HE ROLES OF F ED -R O D’ S GENERIC AND PERSONALIZED HEADS
To demonstrate that F ED -R O D’s two heads learn something different, we plot in Figure 11 every local model’s generic prediction and personalized prediction on its and other clients’ data ( i.e. , P-FL accuracy).
The generic head performs well in general for every client’s test data.
The personalized head could further improve for its own data (diagonal), but degrade for others’ data.
D. P ERSONALIZATION WITH HYPERNETWORKS
F ED -R O D (hyper) learns the personalized head with hypernetworks as introduced in subsection B.4.
The goal is to learn a hypernetwork such that it can directly generate a personalization prediction head given client’s class distribution, without further local training.
Figure 12 shows the training (convergence) curves on CIFAR-10 Dir(0.3).
The hypernetwork (globally aggregated, before further local training) can converge to be on par with that after local training.
In the main paper (cf.
Figure 6), we also show that it servers as a strong starting point for future clients — it can generate personalized models simply with future clients’ class distributions.
That is, the clients may not have labeled data, but provide the hypernetwork with their preference/prior knowledge.
It can also be used as the warm-start model for further local training when labeled data are available at the clients.
Table 4 provides the P-FL results for the new 50 clients studied in subsection 5.1 and Figure 6 in the main paper.
Except for F ED -R O D (hyper), the accuracy before local training is obtained by the global model.
The best personalized model after local training is selected for each client using a validation set.
F ED -R O D (hyper) notably outperforms other methods before or after local training
. D. C LASS - IMBALANCED GLOBAL DISTRIBUTIONS
In the real world, data frequency naturally follows a long-tailed distribution, rather than a classbalanced one.
Since the server has no knowledge and control about the whole collection of the clients’ data, the clients data may collectively be class-imbalanced.
This adds an additional challenge for the server to learn a fair and class-balanced model.
We follow the setup in (Cao et al., 2019) to transform
FMNIST and CIFAR-10 training sets into class-imbalanced versions, in which the sample sizes per class follow an exponential decay.
The imbalanced ratio (IM) is controlled as the ratio between
[Table]
sample sizes of the most frequent and least frequent classes.
Here we consider IM = 10 and IM = 100 .
The generic test set remains class-balanced.
Table 10 shows that
F ED -R
O D remains robust on both generic accuracy and client accuracy consistently.
We see that F ED D YN also performs well, especially on FMNIST of which the setup has more clients ( 100 ) but a lower participation rate ( 20% ).
By combining F ED D YN with F ED -R O D, we achieve further improvements.
Essentially, the generic FL methods (except for F ED -R O D) are optimizing toward the overall classimbalanced distribution rather than the class-balanced distribution.
In Table 11, we further examine the G-FL accuracy on a class-imbalanced test set whose class distribution is the same as the global training set.
F ED -R O D still outperforms other methods, demonstrating that F ED -R O D learns a robust, generic, and strong model.
D. C OMPATIBILITY OF F ED -R O D WITH OTHER G-FL ALGORITHMS
As mentioned in the main paper, other G-FL algorithms like F ED D YN (Acar et al., 2021) can be incorporated into F ED -R O D to optimize the generic model (using the balanced risk).
We show the results in Table 12, following Table 1 of the main paper.
Combining F ED -R O D with SCAFFOLD (Karimireddy et al., 2020b),
F ED D YN (Acar et al., 2021), and F ED P ROX (Li et al., 2020a) can lead to higher accuracy than each individual algorithm along in several cases.
[Table] [Table]
[Table]
D. C OMPARISON TO PERSONALIZED FL ALGORITHMS
From Table 1 in the main paper and Table
[Table]
the personalized FL algorithms are usually outperformed by local models of generic FL algorithms in terms of the P-FL accuracy ( i
[Table]
e. , the PM column).
The gap is larger when client data are more IID, especially for P-FL methods whose personalized models do not explicitly rely on weight averaging of other clients’ models ( e
[Table]
g. , MTL, LG-F ED A VG
[Table]
and
P F ED M E ).
Some P-FL methods can not even outperform local training alone.
A similar observation is also reported in F ED FOMO (Zhang et al., 2021).
These observations justify the beneﬁts of FL that similar clients can improve each other by aggregating a global model and updating it locally [Table] while the beneﬁts might decay for very dissimilar clients.
To further demonstrate the effect of building a global model and re-initialing the local/personalized models using it (cf. section 3 in the main paper), we investigate D ITTO (Li et al., 2021a), a state-ofthe-art personalized FL algorithm.
We found that D ITTO learns two local models.
One of them is used to build the global model exactly like F ED A VG .
The global model is then used to regularize the other local model (cf.
Equation 4 in the main paper), which is used for personalized prediction.
To differentiate these two local models, we call the former the local model (LM), and the latter the personalized model (PM).
We note that, the PM model is kept locally and is never re-initialized by the global model.
In Table 13, we show the P-FL accuracy using the LM and PM models.
The LM model trained in the same way as F ED A VG (with re-initialization) surprisingly outperforms the PM model.
We further replicate the experiments in (Li et al., 2021a) on robustness against adversary attacks in Table 14.
Besides comparing LM and PM, we also evaluate the global model GM for P-FL accuracy.
With out adversarial attacks, the LM model outperforms the PM model.
However, with
[Table] [Table]
[Table]
adversarial attacks, the PM model notably outperforms the other two models.
We surmise that, when there are adversarial clients, the resulting generic model will carry the adversarial information; re-initializing the local models with it thus would lead to degraded performance.
D. A DDITIONAL STUDIES AND DISCUSSIONS
Different network architectures.
F ED -R O D can easily be applied to other modern neural network architectures.
In Table 15, we show that F ED -R O D can be used with deeper networks.
F ED -R O D is not merely ﬁne-tuning.
F ED -R O D is not merely pre-training the model with BSM and then ﬁne-tuning it with ERM for two reasons.
First, for F ED -R O D (linear), the P-head is learned dynamically with the updating feature extractor across multiple rounds.
Second, for F ED -R O D (hyper), the hypernetwork has to undergo the local training and global aggregation iterations over multiple rounds.
In Table 2 of the main paper, we report the ﬁne-tuning baseline.
On CIFARDir(0.3), it has 84 .
5% for P-FL (PM), lower than 86 .
4% and 86 .
8% by F ED -R O D (linear) and
F ED -R O D (hyper).
Note that, hypernetworks allow fast adaptation for new clients.
Comparison to the reported results in other personalized FL papers.
Existing works usually report F ED A VG ’s personalized performance by evaluating its global model ( i.e. , the GM column in Table 1 of the main paper).
In this paper, we evaluate F ED A VG ’s local model w m ( i.e. , the PM column in Table 1 of the main paper), which is locally trained for epochs.
We see a huge performance gap between these two models.
In (Fallah et al., 2020), the authors investigated a baseline “F ED A VG + update”, which ﬁne-tunes F ED A VG ’s global model ¯ w with only few mini-batches for each client.
The resulting personalized models thus capture less personalized information than w m in F ED A VG .
For a fair comparison, we also strengthen P ER -F ED A VG (Fallah et al., 2020) by updating with more epochs.
Effects of local sample size to P-FL performance In subsection 3.2 and Table 1, we show that local models of generic FL algorithms are strong personalized models.
Indeed, the local sample size is an important factor in the P-FL performance.
If a client has enough training samples, training its own model (the local only baseline) can already be strong without any federated learning.
On the other hand, when each client does not have enough samples to train a good model on its own.
It will be crucial to have a generic model learned from federated learning as the starting point of personalization.
To conﬁrm our observation when clients have insufﬁcient samples, we further conduct the following experiments.
First, we enlarge the number of clients for CIFAR-10 and FMNIST experiments by ﬁve times.
That is, each client’s data size becomes one-ﬁfth on average.
Second, we point out that the experiments on CIFAR-100 in Table 1 are with 20 clients.
CIFAR-100 has the same total number of training images as CIFAR-10 but with 10 times more classes.
In other words, the number of images per class is one-tenth.
Table 16 shows the results: all the experiments are based on Dir(0.3).
Even when the local only models perform worse in P-FL, the local models of F ED A VG still perform on a par with personalized FL algorithms like F ED P ER , and F ED -R
O D can still achieve the best P-FL
[Table]
accuracy.
We attribute the superior personalized performance by F ED A VG and F ED -R O D to the implicit regularization discussed in subsection 3.3.
We also want to point out that, even if each client has insufﬁcient data, the P-FL performance of local only may still have higher accuracy than the GM of F ED A VG on the personalized accuracy, especially when the non-IID condition becomes severe (e
[Table]
g., Dir (0.1)).
When the non-IID condition is severe, it is harder to train a single GM model to perform well in the personalized setting
[Table]