Year: 2021

Abstract:
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 (FED-ROD), the learned model can simultaneously achieve state-ofthe-art generic and personalized performance, essentially bridging the two tasks.

text: “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 (FED-ROD), the learned model can simultaneously achieve state-ofthe-art generic and personalized performance, essentially bridging the two tasks.

heading: 1 INTRODUCTION
text: 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 FEDAVG (McMahan et al., 2017) take a multi-round approach shown in Figure 1. Within each round, the server first 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 FEDAVG.
Concretely, algorithms designed for generic FL (G-FL) often discard the local models {wm} 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 {wm} 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.
At first 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 optimize1. 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. Specifically for classification, the nonIID clients can result from non-identical class distributions or non-identical class-conditional data distributions. One way to mitigate their influences 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.
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 FEDAVG 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 Ro-
bust Decoupling (FED-ROD), 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. Specifically 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 fine-tune the models given new clients’ data.
We validate FED-ROD on multiple datasets under various non-IID settings. FED-ROD consistently outperforms existing generic and personalized FL algorithms in both setups. Moreover, FED-ROD is compatible with and can further improve advanced generic FL algorithms like FEDDYN (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 FED-ROD to excel on both at the same time. We validate FED-ROD with extensive experiments.
1However, the literature does not evaluate G-FL algorithms in the P-FL setup by their local models.
• 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. • FED-ROD enables zero-shot adaptation and much effective fine-tuning for new clients.

heading: 2 RELATED WORK (A DETAILED VERSION IS IN APPENDIX A)
text: Generic federated learning. FEDAVG (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, FEDPROX (Li et al., 2020a) and FEDDYN (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. (2020); Hanzely et al. (2020); 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, FED-ROD shares a single feature extractor among clients, inspired by Caruana (1997); Zhang et al. (2014). 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, FED-ROD simultaneously excels in both FL setups.
Instead of designing specific algorithms for personalized FL, Cheng et al. (2021); Wang et al. (2019); Yu et al. (2020) showed that performing post-processing (e.g., fine-tuning) to a generic FL model (e.g., w̄ in FEDAVG) leads to promising personalized accuracy. We further showed that, the local models {wm} learned in FEDAVG 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 sacrificing either of them.
A recent work PFEDHN (Shamsian et al., 2021) also applies hypernetworks (Ha et al., 2017) but in a very different way from FED-ROD. PFEDHN 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.

heading: 3 PERSONALIZED MODELS EMERGE FROM GENERIC FEDERATED LEARNING
text: 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.

heading: 3.1 BACKGROUND
text: Generic federated learning. In a generic FL setting with M clients, where each client has a data set Dm = {(xi, yi)}|Dm|i=1 , the optimization problem to solve can be formulated as
min w L(w) = M∑ m=1 |Dm| |D| Lm(w), where Lm(w) = 1 |Dm| ∑ i `(xi, yi;w). (1)
Here, w is the model parameter; D = ∪mDm is the aggregated data set from all clients; Lm(w) is the empirical risk computed from client m’s data; ` is a loss function applied to each data instance.
Federated averaging (FEDAVG). As clients’ data are separate, Equation 1 cannot be solved directly. A standard way to relax it is FEDAVG (McMahan et al., 2017), which iterates between two steps, local training and global aggregation, for multiple rounds of communication
Local: wm = arg minw Lm(w), initialized with w̄; Global: w̄ ← M∑ m=1 |Dm| |D| wm. (2)
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 wm. 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, wm 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 wm, 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) define the optimization problems similar to the following
min {Ω,w1,··· ,wM} M∑ m=1 |Dm| |D| Lm(wm) +R(Ω,w1, · · · ,wM ), (3)
whereR is a regularizer; Ω is introduced to relate clients. The regularizer is imposed to prevent wm from over-fitting 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 modelwm.
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 define Ω 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 L2 regularizer between Ω and each wm. The corresponding local training step thus could generally be formulated as
Local: w(t+1)m = arg minw Lm(w) + λ
2 ‖w −Ω‖22, initialized with w(t)m , (4)
wherew(t)m denotes the local model after the t-th round; λ is the regularization coefficient. 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 specifically used for each client’s test data.

heading: 3.2 LOCAL MODELS OF GENERIC FL ALGORITHMS ARE STRONG PERSONALIZED MODELS
text: 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 FEDAVG (i.e., Equation 2) notably outperforms using its global model in the P-FL setup. At first glance, this may not be surprising, as local training in FEDAVG is driven by clients’ empirical risks. What really surprises us, as will be seen in section 5, is that FEDAVG’s local models outperform most of the existing personalized FL algorithms, even if no explicit regularization is imposed in Equation 2.

heading: 3.3 INITIALIZATION WITH WEIGHT AVERAGE IS A STRONG REGULARIZER
text: To gain a further understanding, we plot FEDAVG local models’ accuracy on clients’ training and test data. We do so also for a state-of-the-art personalized FL algorithm DITTO (Li et al., 2021a), whose local training step for producing personalized models is similar to Equation 4. As shown in Figure 3, FEDAVG has a lower training but higher test accuracy, implying that FEDAVG’s local training is more regularized than Equation 4.
We attribute this effect to the initialization in Equation 2. Specifically, by initializing w with w̄, we essentially impose an L2 regularizer λ 2 ‖w − w̄‖ 2 2 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 ‖w−w̄‖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 FEDDYN (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.

heading: 4 FEDERATED ROBUST DECOUPLING (FED-ROD)
text: 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.

heading: 4.1 IMPROVING GENERIC FL WITH BALANCED RISK MINIMIZATION (BRM)
text: We first analyze what factors may lead to non-IID conditions. Suppose the data instance (x, y) of clientm is sampled from a client-specific joint distribution Pm(x, y) = Pm(x|y)Pm(y), the non-IID distributions among clients can result from non-identical class distributions Pm(x|y), non-identical class-conditional data distributions Pm(y), or both. All these cases can make Lm(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 influence of non-IID data is to make Lm(w) align with each other. This can be challenging to achieve if clients have different Pm(x|y): without knowing clients’ data2, it is hard to design such an aligned Lm(w). However, when clients have different Pm(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 classifier 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-field. Re-weighting and re-sampling (Buda et al., 2018) are the most fundamental techniques. Denote by Nm,c the number of training instances of class c for client m, these techniques adjust Lm(w) in Equation 1 into
LBRm (w) ∝ ∑ i qyi`(xi, yi;w), where qyi is usually set as 1 Nm,yi or 1√ Nm,yi . (5)
Namely, they mitigate the influence of Pm(y) by turning the empirical risk Lm into a balanced risk LBRm , 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 ` (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 define LBRm using these losses, e.g., the balanced softmax (BSM) loss (Ren et al., 2020)
LBRm (w) ∝ ∑ i `BSM(xi, yi;w), where `BSM(x, y;w) = − log Nγm,y exp(gy(x;w))∑ c∈CN γ m,c exp(gc(x;w)) . (6)
Here, gc(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 gy(x;w) in training to overcome feature deviation (Ye et al., 2020) in testing. We list other class-balanced losses in the appendix.
2Clients having different Pm(x|y) is related to domain adaptation (Gong et al., 2012) and generalization (Muandet et al., 2013), which require knowing the distributions of all/some clients for algorithm design.
3This is indeed the main cause of non-IID data distributions in the literature of FL (Hsu et al., 2019; 2020).
We take advantage of these existing efforts by replacing the empirical risk Lm in Equation 2 with a balanced risk LBRm , 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 LBRm can easily be incorporated into advanced FL algorithms like FEDDYN (Acar et al., 2021), because these algorithms are agnostic to the local objectives being used.

heading: 4.2 LOCAL TRAINING AND LOCAL MODEL DECOUPLING WITH ERM AND BRM
text: The use of balanced risk LBRm 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 wm’s personalized performance, since it is no longer optimized towards client’s empirical risk Lm. To address these contrasting pursuits of generic and personalized FL, we propose a unifying FL framework named Federated Robust Decoupling (FED-ROD), 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 FED-ROD. The overall training process of FED-ROD follows FEDAVG, 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 FEDAVG. 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 hG(z;ψ) and hP (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 = hG(z;ψ). For personalized prediction, we perform f(x;θ), followed by ŷP = hG(z;ψ) + hP (z;φm). That is, hP is an add-on to hG, providing personalized information that is not captured by the generic head. The overall objective. FED-ROD learns the generic model with the balanced risk LBRm and the personalized predictor with the empirical risk Lm. That is, different from Equation 1, FED-ROD aims to solve the following two optimization problems simultaneously
min θ,ψ L({θ,ψ}) = M∑ m=1 |Dm| |D| LBRm ({θ,ψ}) and min φm Lm({θ,ψ,φm}),∀m ∈ [M ]. (7)
We note that, Lm is only used to learn the personalized head parameterized by φm.
Learning. Equation 7 cannot be solved directly in federated learning, so FED-ROD follows FEDAVG to learn iteratively between the local training and global aggregation steps
Local: θ?m,ψ ? m = arg minθ,ψ LBRm ({θ,ψ}), initialized with θ̄, ψ̄, (8)
φ?m = arg minφm Lm({θ,ψ,φm}), initialized with φ ′ m, (9)
Global: θ̄ ← M∑ m=1 |Dm| |D| θ?m, ψ̄ ← M∑ m=1 |Dm| |D| ψ?m, (10)
where φ′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 Lm({θ,ψ,φm}). The θ and ψ in Equation 9 thus come dynamically from the SGD updates of Equation 8. In other words, Equation 9 is not merely fine-tuning on top of the generic model. In the end of federated learning, we will obtain θ̄ and ψ̄ (Equation 10) for generic predictions and {θ?m,ψ?m,φ?m}Mm=1 (Equation 8 and Equation 9) for personalized predictions, respectively. Please be referred to the appendix for the pseudocode.

heading: 4.3 ADAPTIVE PERSONALIZED PREDICTORS VIA HYPERNETWORKS
text: 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, FED-ROD 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 byHP (am;ν) the meta-model parameterized by ν, whose output isφm. Here, am ∈ R|C| is the |C|-dimensional vector that records the class distribution of client m; i.e., the c-th dimension am[c] =
Nm,c∑ c′ Nm,c′ . Accordingly, the local training step of φm in Equation 9 is replaced by
Local: ν?m = arg minν Lm({θ,ψ,ν}), initialized with ν̄; Global: ν̄ ← M∑ m=1 |Dm| |D| ν?m. (11)
We implement HP by a lightweight hypernetwork (Ha et al., 2017) with two fully-connected layers. With the learned ν̄, the meta-model HP can locally generate φm based on am, making it adaptive to new clients simply by class distributions. The parameter φm can be further updated using clients’ data. We name this version FED-ROD (hyper); the previous one, FED-ROD (linear). Please see Figure 4 (c-d) for an illustration. We include more details in the appendix.

heading: 5 EXPERIMENT (MORE DETAILS AND RESULTS IN THE APPENDIX)
text: 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 qc is drawn from Dir(α) for class c, and we assign data of class c to client m proportionally to qc[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. (2020), 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 five 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 1: Results in G-FL accuracy and P-FL accuracy (%). ?: methods with no G-FL models and we combine their P-FL models. §: official implementation. Blue/bold fonts highlight the best baseline/our approach.
Dataset EMNIST FMNIST CIFAR-10 CIFAR-100
Non-IID Writers Dir(0.1) Dir(0.3) Dir(0.1) Dir(0.3) Dir(0.1) Dir(0.3)
Test Set G-FL P-FL G-FL P-FL G-FL P-FL G-FL P-FL G-FL P-FL G-FL P-FL G-FL P-FL
Method / Model GM GM PM GM GM PM GM GM PM GM GM PM GM GM PM GM GM PM GM GM PM
FEDAVG 97.0 96.9 97.2 81.1 81.0 91.5 83.4 83.2 90.5 57.6 57.1 90.5 68.6 69.4 85.1 41.8 41.6 70.2 46.4 46.2 61.7 FEDPROX 97.0 97.0 97.0 82.2 82.3 91.4 84.5 84.5 89.7 58.7 58.9 89.7 69.9 69.8 84.7 41.7 41.6 70.4 46.5 46.4 61.5 SCAFFOLD 97.1 97.0 97.1 83.1 83.0 89.0 85.1 85.0 90.4 61.2 60.8 90.1 71.1 71.5 84.8 42.3 42.1 70.4 46.5 46.5 61.7 FEDDYN § 97.3 97.3 97.3 83.2 83.2 90.7 86.1 86.1 91.5 63.4 63.9 92.4 72.5 73.2 85.4 43.0 43.0 72.0 47.5 47.4 62.5 MTL ? 75.4 75.0 85.6 36.1 36.0 87.3 53.1 53.4 78.3 12.1 12.7 90.6 13.5 13.7 80.2 9.5 9.3 60.7 10.8 10.7 49.9 LG-FEDAVG ?§ 80.1 80.0 95.6 54.8 54.5 89.5 66.8 66.8 84.4 29.5 28.8 90.8 46.7 46.2 82.4 23.5 23.4 66.7 34.5 33.9 55.4 FEDPER ? 93.3 93.1 97.2 74.5 74.4 91.3 79.9 79.9 90.4 50.4 50.2 89.9 64.4 64.5 84.9 37.6 37.6 71.0 40.3 40.1 62.5 PER-FEDAVG 95.1 - 97.0 80.5 - 82.8 84.1 - 86.7 60.7 - 82.7 70.5 - 80.7 39.0 - 66.6 44.5 - 58.9 PFEDME § 96.3 96.0 97.1 76.7 76.7 83.4 79.0 79.0 83.4 50.6 50.7 76.6 62.1 61.7 70.5 38.6 38.5 63.0 41.4 41.1 53.4 DITTO 97.0 97.0 97.4 81.5 81.5 89.4 83.3 83.2 90.1 58.1 58.3 86.8 69.7 69.8 81.5 41.7 41.8 68.5 46.4 46.4 58.8 FEDFOMO ? 80.5 80.4 95.9 34.5 34.3 90.0 70.1 69.9 89.6 30.5 31.2 90.5 45.3 45.1 83.4 35.4 35.3 68.9 39.6 39.3 58.4 FEDREP ?§ 95.0 95.1 97.5 79.5 80.1 91.8 80.6 80.5 90.5 56.6 56.2 91.0 67.7 67.5 85.2 40.7 40.7 71.5 46.0 46.0 62.1 Local only - - 64.2 - - 85.9 - - 85.0 - - 87.4 - - 75.7 - - 40.0 - - 32.5
FED-ROD (linear) 97.3 97.3 97.5 83.9 83.9 92.7 86.3 86.3 94.5 68.5 68.5 92.7 76.9 76.8 86.4 45.9 45.8 72.2 48.5 48.5 62.3 FED-ROD (hyper) 97.3 97.3 97.5 83.9 83.9 92.9 86.3 86.3 94.8 68.5 68.5 92.5 76.9 76.8 86.8 45.9 45.8 72.3 48.5 48.5 62.5
+ FEDDYN 97.4 97.4 97.5 85.9 85.7 95.3 87.5 87.5 94.6 68.2 68.2 92.7 74.6 74.6 85.6 46.2 46.2 72.5 48.4 48.4 62.5
Pm(y) and average the weighted accuracy across M clients as 1M ∑ m ∑ i Pm(yi)1(yi=ŷi)∑ i Pm(yi) . Here, i is the instance index. This evaluation is more robust (essentially as the expectation) than assigning each client a specific 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 LBRm and report the FED-ROD (hyper) version (cf. subsection 4.3). Table 2 provides the ablation study.
Baselines. For G-FL methods including FEDAVG (McMahan et al., 2017), FEDPROX (Li et al., 2020a), SCAFFOLD (Karimireddy et al., 2020b), and FEDDYN (Acar et al., 2021), we use their global models w̄ for G-FL evaluation; their local models (i.e., wm in Figure 1) for P-FL evaluation. For P-FL methods, to evaluate their G-FL performance, we use the available global models in PFEDME (Dinh et al., 2020) and DITTO (Li et al., 2021a) or average the final personalized models for MTL (Smith et al., 2017), FEDPER (Arivazhagan et al., 2019), LG-FEDAVG (Liang et al., 2020), FEDFOMO (Zhang et al., 2021), and FEDREP (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 FEDAVG has been applied to P-FL in literature.

heading: 5.1 RESULTS
text: FED-ROD 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, FEDPROX, and FEDDYN) outperforms FEDAVG and personalized methods, and our FED-ROD can have further gains by using balanced risk minimization (BRM). We also investigate combining FED-ROD and FEDDYN (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 FEDAVG: 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 benefits of federated collaboration4. For generic FL methods, using PMs (i.e., local models {wm}) 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 FED-ROD 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 benefit from implicit regularization; (b) the personalized head learned with clients’ empirical risks.
4PMs of “local only” could outperform the GM of FEDAVG on the P-FL accuracy, especially when the non-IID condition becomes severe (e.g., Dir (0.1)): it is hard to train a single GM to perform well in P-FL.
Table 4: P-FL accuracy on future non-IID clients (Dir(0.3) for all datasets). Each cell is before/after local fine-tuning.
Method FMNIST CIFAR-10 CIFAR-100
FEDAVG 80.3/87.2 56.2/76.0 38.9/54.3 FEDDYN 82.3/87.6 61.7/76.2 40.1/57.2 PER-FEDAVG 82.1/89.6 60.0/79.8 37.6/55.6
FED-ROD (linear) 83.5/91.3 62.4/80.2 40.0/58.2 FED-ROD (hyper) 88.9/91.4 75.7/81.5 40.7/59.0
+FEDDYN 89.2/91.3 77.1/83.5 41.4/59.5
Figure 5: Upper: G-FL test accuracy along the training rounds before/after averaging the local models. Lower: variances ofwm − w̄ across clients.
BRM effectively reduces the variance of G-FL accuracy and local gradients. To understand why FED-ROD improves G-FL, we visualize the global model w̄’s and each local model wm’s G-FL accuracy on CIFAR-10 (Dir(0.3)), in Figure 5 (upper). FED-ROD 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 wm deviates from w̄ after local training in Figure 5 (lower). FED-ROD 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. FED-ROD benefits from decoupling. We compare several variants of FED-ROD (cf. Figure 4), with one head (reduced to FEDAVG) or different networks (linear/hyper). We evaluate on CIFAR-10 (Dir(0.3)). As shown in Table 2, FEDAVG with BRM significantly improves G-FL but degrades in P-FL. FED-ROD remedies this by training a decoupled personalized head. We note that, FED-ROD does not merely fine-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).
FED-ROD (hyper) benefits 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 fine-tuning it with clients’ data for several steps. Table 4 and Figure 6 shows the averaged accuracy on new clients. Without fine-tuning, FED-ROD (hyper) can already generate personalized models, and outperforms others methods stably with fine-tuning.
More results and analyses in the Appendix. The Appendix includes studies with class-imbalanced global distributions and we show that FED-ROD 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 Equation 4 to 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.

heading: 6 CONCLUSION
text: 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 FED-ROD. 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. FED-ROD seamlessly incorporates these two observations to excel in both FL settings, and further enables fast adaptation to new clients via an adaptive module.

heading: Appendix
text: We provide details omitted in the main paper.
• Appendix A: additional comparison to related work (cf. section 2 and section 3 of the main paper).
• Appendix B: additional details of FED-ROD (cf. section 3 and section 4 of the main paper). • Appendix C: details of experimental setups (cf. section 5 of the main paper). • Appendix D: additional experimental results and analysis (cf. section 3 and section 5 of the
main paper).

heading: A COMPARISON TO RELATED WORK

heading: A.1 FOLLOW-UP WORKS OF FEDAVG
text: 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 wm 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, FEDPROX (Li et al., 2020a) introduced a regularizer with respect to w̄,
min w Lm(w) +
λ 2 ‖w − w̄‖2.
FEDDYN (Acar et al., 2021) further added a dynamic term based on the local model of the previous round w′m,
min w Lm(w) + 〈∇Lm(w′m),w〉+
λ 2 ‖w − w̄‖2.
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 Lm(w) still plays an important role in driving the local model update. Since Lm(w) directly reflects the (non-IID) client data distribution, the learned local models are indeed strong candidates for personalized models.

heading: A.2 GENERIC FEDERATED LEARNING
text: FEDAVG (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 (Konečnỳ 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 FEDAVG. 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; FEDPROX (Li et al., 2020a), FEDDANE (Li et al., 2019), and FEDDYN (Acar et al., 2021) employed regularization toward the global model; SCAFFOLD (Karimireddy et al., 2020a) MIME (Karimireddy et al., 2020b) leveraged control varieties and/or server statistics to correct local gradients; (Wang et al., 2020b; Yao et al., 2019) modified 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 reflect 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 FEDDYN (Acar et al., 2021) and SCAFFOLD (Karimireddy et al., 2020a) to boost the generic performance.

heading: A.3 PERSONALIZED FEDERATED LEARNING
text: 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), FED-ROD 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 FEDPER (Arivazhagan et al., 2019) and FEDREP (Collins et al., 2021) which also learned a shared feature extractor, FED-ROD 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 specific algorithms for personalized FL, (Wang et al., 2019; Yu et al., 2020) showed that performing post-processing (e.g., fine-tuning) to a generic FL model (e.g., w̄ learned by FEDAVG) already leads to promising personalized accuracy. In this work, we further showed that, the local modelswm learned in FEDAVG 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 sacrificing either of them.
PFEDHN (Shamsian et al., 2021) also applies hypernetworks (Ha et al., 2017) but for a very different purpose from FED-ROD. Specifically, PFEDHN 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.

heading: A.4 SYSTEMATIC OVERHEAD
text: FED-ROD has similar computation cost, communication size, and number of parameters as FEDAVG. We discuss the difference between FED-ROD and existing generic FL methods from a system view. FEDPROX (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. FEDDYN (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.

heading: A.5 CLASS-IMBALANCED LEARNING
text: 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 reflect class imbalance so as to train a model that is less biased toward major classes.

heading: A.6 ZERO-SHOT LEARNING
text: 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 classifier 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 efficient parameterization.

heading: B ADDITIONAL DETAILS OF FED-ROD

heading: B.1 ADDITIONAL BACKGROUND (CF. SUBSECTION 3.1 OF THE MAIN PAPER)
text: 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 classification problem whose label space is C, a commonly used loss is the cross entropy,
`(x, y;w) = − log exp(gy(x;w))∑ c∈C exp(gc(x;w)) , (12)
where gc(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, FEDAVG (McMahan et al., 2017) aims for generic FL but it creates both the global and local models.

heading: B.2 OVERVIEW OF FED-ROD
text: For generic predictions, FED-ROD performs feature extraction z = f(x;θ), followed by hG(z;ψ). For personalized predictions, FED-ROD performs z = f(x;θ), followed by hG(z;ψ) +hP (z;φm). The element-wise addition is performed at the logit level. That is, gc(x;w) in Equation 12 can be re-written as
gc(x; {θ,ψ,φm}) = { hGc (z;ψ) Generic model, hGc (z;ψ) + h P c (z;φm) Personalized model,
(13)
where z = f(x;θ) is the extracted feature.
The overall training process of FED-ROD iterates between the local training and global aggregation steps. In local training, FED-ROD aims to minimize the following objective
LBRm ({θ,ψ}) + Lm({θ,ψ,φm}). (14) The empirical risk Lm(wm = {θ,ψ,φm}) is defined as 1|Dm| ∑ i `(xi, yi;wm), where Dm = {(xi, yi)}|Dm|i=1 is the training data of client m. We will introduce more options of the balanced risk LBRm 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 FEDAVG. For Lm({θ,ψ,φm}), we do not derive gradients w.r.t. θ and ψ. We emphasize that, according to subsection 4.2 of the main paper, the finally learned parameters of FED-ROD (linear) are θ̄, ψ̄, and {φ?m}Mm=1. We then plug them into Equation 13 for predictions. In algorithm 1 and algorithm 2, we provide pseudocode of our FED-ROD algorithm.
Algorithm 1: FED-ROD (linear) (Federated Robust Decoupling) Server input : initial global model parameter θ̄ and ψ̄; Client m’s input : initial local model parameter φ?m, local step size η, local labeled data Dm; for r ← 1 to R do
Sample clients S ⊆ {1, · · · , N}; Communicate θ̄ and ψ̄ to all clients m ∈ S; for each client m ∈ S in parallel do
Initialize θ ← θ̄, ψ ← ψ̄, and φm ← φ?m; {θ?m,ψ?m,φ?m} ← Client local training({θ,ψ,φm},Dm, η); [Equation 8 and Equation 9] Communicate θ?m and ψ?m to the server;
end Construct θ̄ = ∑ m∈S |Dm|∑ m′∈S |Dm′ | θ?m;
Construct ψ̄ = ∑ m∈S |Dm|∑ m′∈S |Dm′ | ψ?m;
end Server output :θ̄ and ψ̄; Client m’s output :{θ?m,ψ?m,φ?m}.
Algorithm 2: FED-ROD (hyper) (Federated Robust Decoupling) Server input : initial global model parameter θ̄, ψ̄, and ν̄; Client m’s input : local step size η, local labeled data Dm; for r ← 1 to R do
Sample clients S ⊆ {1, · · · , N}; Communicate θ̄, ψ̄, and ν̄ to all clients m ∈ S; for each client m ∈ S in parallel do
Initialize θ ← θ̄, ψ ← ψ̄, and ν ← ν̄; {θ?m,ψ?m,ν?m} ← Client local training({θ,ψ,ν},Dm, η); [Equation 11] Communicate θ?m, ψ?m, and ν?m to the server;
end Construct θ̄ = ∑ m∈S |Dm|∑ m′∈S |Dm′ | θ?m;
Construct ψ̄ = ∑ m∈S |Dm|∑ m′∈S |Dm′ | ψ?m;
Construct ν̄ = ∑ m∈S |Dm|∑ m′∈S |Dm′ | ν?m;
end Server output :θ̄, ψ̄, and ν̄ (for personalized model generation). Client m’s output :{θ?m,ψ?m,ν?m}.

heading: B.3 BALANCED RISK MINIMIZATION (BRM)
text: To learn a generic model, standard federated learning (e.g., FEDAVG (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[Lm(w)] for client m, if client m’s data Dm are IID partitioned from D. Here, the expectation is over different Dm partitioned from D. In reality, Lm could diverge from L due to non-IID partitions of the aggregated data D into clients’ data. That is,
E[Lm(w)] 6= E[Lm′(w)] 6= L(w). We mitigate the non-IID situation by directly adjusting Lm such that E[Lm(w)] ≈ E[Lm′(w)] ≈ L(w). Essentially, Lm(w) is the client’s empirical risk, which could be different among clients if their class distribution Pm(y) are different. We, therefore, propose to turn the empirical risk Lm into a class-balanced risk LBRm by replacing ` 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.
One may wonder what if the global distribution is class-imbalanced? Will BRM still be beneficial to FL? In subsection D.5, we perform experiments to show that FED-ROD with BRM can still improve on FEDAVG 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.

heading: B.4 ON FEDERATED LEARNING FOR THE PERSONALIZED HEAD WITH HYPERNETWORKS
text: One drawback of existing personalized methods is that the personalized models are only available for clients involved in training or with sufficient 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 specific prediction head φm for each client m, we propose to learn a meta-model HP (am;ν) with a shared meta-parameter ν. The input to HP is a vector am ∈ R|C|, which records the proportion of class c ∈ C in client m’s data. The output of HP is φm for hP . In other words, HP can adaptively output personalized prediction heads for clients given their local class distributions am.
We implement the meta-model HP by a hypernetwork (Ha et al., 2017), which can be seen as a lightweight classifier generator given am. 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 am 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 FED-ROD. Hypernetworks add only a small overhead to the original model.

heading: B.5 EXTENSION WITH META-LEARNING FOR THE IMPROVED BSM LOSS
text: FED-ROD 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 ∑ i qm,yi`
BSM(xi, yi;w), where qm,yi is meta-learned with a small balanced meta dataset Dmeta provided by the server. (See Table 5 for a comparison.) The Dmeta 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 `BSMγ = − log Nγm,y exp(gy(x;w))∑ c∈CN γ m,c exp(gc(x;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-specific γm with meta-learning for `BSMγm . More specifically, given a meta-learning rate η, the meta-learning process involves the following iterative steps:
1. Compute the Meta-BSM loss with a mini-batch B ∼ Dm; i.e., ∀(x, y) ∈ B, compute `BSMγm (x, y;wm).
2. Predict the example weights with qm,y = MWNet(`BSMγm (x, y;wm); ζm), ∀(x, y) ∈ B. 3. Re-weight the Meta-BSM loss: LBRm,B(wm) = ∑ (x,y)∈B qm,yi` BSM γm (x, y;wm), and per-
form one step of gradient descent to create a duplicated model w̃m = wm − η∇wmLBRm,B . 4. Computes the loss on the meta dataset Dmeta using the duplicated model: LBRm,Dmeta(w̃m) =∑ (x,y)∈Dmeta qm,y` BSM γm (x, y, w̃m), followed by updating γm ← γm − η∇γmL BR m,Dmeta and
ζm ← ζm − η∇ζmLBRm,Dmeta .
5. Update the model: wm ← wm − η∇wmLBRm,B(wm).
Throughout the federated learning process, γm and qm,y are dynamically learned with meta-learning for different clients and rounds.
Results of FED-ROD 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, FED-ROD 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.

heading: B.6 CONNECTION TO INVARIANT RISK MINIMIZATION GAMES (IRMG)
text: FED-ROD 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 classifier 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 Equation 7 in the main paper
min θ,ψ LIRM(θ,ψ) = M∑ m=1 Lm(θ,ψ), (15) s.t ψ ∈ arg minψ′ Lm(θ,ψ′),∀m ∈ [M ]. (16)
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 classifier ψ as an ensemble of environment-specific classifiers (by averaging over model weights) φ̄ = 1M ∑ m φm:
min θ,φ̄ LIRMG(θ, φ̄) = M∑ m=1 Lm(θ, φ̄), (17)
s.t φm ∈ arg minφm′=m Lm(θ, {φm′} M m′=1),∀m ∈ [M ]. (18)
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 classifiers. 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-specific classifiers φm on the data of each environment Dm, and (3) updating the main classifier through weight averaging φ̄ = 1M ∑ m φm.
We highlight the similarity between IRMG and FED-ROD: both are training a strong generic feature extractor and a set of personalized classifiers. For predictions on data of client (environment) m in Equation 18, IRMG uses ŷ = 1M (φ > mz + ∑ m′ 6=m φ > m′z); FED-ROD’s personalized model is ŷ = hG(z;ψ) + hP (z;φm). We can connect IRMG to FED-ROD by re-writing its prediction as hG(z; φ̄) := φ̄>z = 1M ∑ m φ ′> mz and h P (z;φm) := 1 M (φ > mz − φ′m > z), where φ′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 IMPLEMENTATION DETAILS
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, and 64 channels, 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 FED-ROD, 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 flipped 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 five times of experiments with different random seeds.
For FEDPROX (Li et al., 2020a), the strength of regularization λ is selected from [1e−2, 1e−3, 1e−4]. For FEDDYN (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 PER-FEDAVG (Fallah et al., 2020), the meta-learning rate β̂ is selected from [1e−2, 1e−3, 1e−4]. For PFEDME (Dinh et al., 2020), the strength of regularization λ is selected from [15, 20, 30]. FED-ROD introduces no extra hyperparameters on top of FEDAVG.
For the generic and personalized heads of FED-ROD, 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 FED-ROD, 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 Dtest for evaluation on generic FL (G-FL):
G-FL accuracy : 1 |Dtest| ∑ i 1(yi = ŷi,G), (19)
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 Dtest but weight instances w.r.t. each client’s class distribution:
P-FL accuracy : 1
M ∑ m ∑ i Pm(yi)1(yi = ŷi,P )∑ i Pm(yi) . (20)
We do so instead of separating Dtest 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 verified 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 Table 1 in Table 7, Table 8, and Table 9 here.

heading: D ADDITIONAL EXPERIMENTS AND ANALYSES
text: 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 FED-ROD 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 FED-ROD’s generic and personalized heads (cf. subsection 4.2 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 FED-ROD with other G-FL algorithms (cf. subsection 5.1
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 FED-ROD (cf. subsection 5.1 in the
main paper).

heading: D.1 PERSONALIZED MODELS EMERGE FROM LOCAL TRAINING OF GENERIC FEDERATED LEARNING
text: 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 FEDAVG do not impose extra regularization but re-initialize the local models at every round. Here in Figure 7, we monitor the two loss terms, ∑ m |Dm| |D| Lm(wm) and ∑ m |Dm| |D| ‖wm − w̄‖ 2 2 (cf. Equation 3 and Equation 4 of the main paper), for FEDAVG 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 L2 regularizer in training the personalized models.) Ditto achieves a lower empirical risk (i.e., the first term), likely due to the fact that it does not perform re-initialization. Surprisingly, FEDAVG 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 coefficient λ as infinity. We note that, the reason that the regularization term of Ditto increases along the communication rounds is because ever time the global model w̄

heading: FMNIST CIFAR-10
text: is updated, it moves sharply away from the local model wm. 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.

heading: D.2 BALANCED RISK MINIMIZATION (BRM) IMPROVES GENERIC-FL PERFORMANCE
text: To understand why FED-ROD outperforms other generic methods in the G-FL accuracy, we visualize each local modelwm’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 ∆wm = wm − w̄ across clients are smaller for FED-ROD, which result from a more consistent local training objective.
In Figure 10, we further compare the G-FL accuracy among FEDAVG, FED-ROD with the original BSM loss, and FED-ROD with the Meta-BSM loss introduced in subsection B.5 along the training process (i.e., training curve). The local models of FEDAVG 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 FED-ROD (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, FED-ROD (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.

heading: D.3 THE ROLES OF FED-ROD’S GENERIC AND PERSONALIZED HEADS
text: To demonstrate that FED-ROD’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.

heading: D.4 PERSONALIZATION WITH HYPERNETWORKS
text: FED-ROD (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 FED-ROD (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. FED-ROD (hyper) notably outperforms other methods before or after local training.

heading: D.5 CLASS-IMBALANCED GLOBAL DISTRIBUTIONS
text: 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
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 FED-ROD remains robust on both generic accuracy and client accuracy consistently. We see that FEDDYN also performs well, especially on FMNIST of which the setup has more clients (100) but a lower participation rate (20%). By combining FEDDYN with FED-ROD, we achieve further improvements.
Essentially, the generic FL methods (except for FED-ROD) 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. FED-ROD still outperforms other methods, demonstrating that FED-ROD learns a robust, generic, and strong model.

heading: D.6 COMPATIBILITY OF FED-ROD WITH OTHER G-FL ALGORITHMS
text: As mentioned in the main paper, other G-FL algorithms like FEDDYN (Acar et al., 2021) can be incorporated into FED-ROD to optimize the generic model (using the balanced risk). We show the results in Table 12, following Table 1 of the main paper. Combining FED-ROD with SCAFFOLD (Karimireddy et al., 2020b), FEDDYN (Acar et al., 2021), and FEDPROX (Li et al., 2020a) can lead to higher accuracy than each individual algorithm along in several cases.

heading: D.7 COMPARISON TO PERSONALIZED FL ALGORITHMS
text: From Table 1 in the main paper and Table 10, the personalized FL algorithms are usually outperformed by local models of generic FL algorithms in terms of the P-FL accuracy (i.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.g., MTL, LG-FEDAVG, and PFEDME). Some P-FL methods can not even outperform local training alone. A similar observation is also reported in FEDFOMO (Zhang et al., 2021). These observations justify the benefits of FL that similar clients can improve each other by aggregating a global model and updating it locally, while the benefits 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 DITTO (Li et al., 2021a), a state-ofthe-art personalized FL algorithm. We found that DITTO learns two local models. One of them is used to build the global model exactly like FEDAVG. 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 FEDAVG (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
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.

heading: D.8 ADDITIONAL STUDIES AND DISCUSSIONS
text: Different network architectures. FED-ROD can easily be applied to other modern neural network architectures. In Table 15, we show that FED-ROD can be used with deeper networks. FED-ROD is not merely fine-tuning. FED-ROD is not merely pre-training the model with BSM and then fine-tuning it with ERM for two reasons. First, for FED-ROD (linear), the P-head is learned dynamically with the updating feature extractor across multiple rounds. Second, for FED-ROD (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 fine-tuning baseline. On CIFAR-10 Dir(0.3), it has 84.5% for P-FL (PM), lower than 86.4% and 86.8% by FED-ROD (linear) and FED-ROD (hyper). Note that, hypernetworks allow fast adaptation for new clients. Comparison to the reported results in other personalized FL papers. Existing works usually report FEDAVG’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 FEDAVG’s local model wm (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 “FEDAVG + update”, which fine-tunes FEDAVG’s global model w̄ with only few mini-batches for each client. The resulting personalized models thus capture less personalized information than wm in FEDAVG. For a fair comparison, we also strengthen PER-FEDAVG (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 confirm our observation when clients have insufficient samples, we further conduct the following experiments. First, we enlarge the number of clients for CIFAR-10 and FMNIST experiments by five times. That is, each client’s data size becomes one-fifth 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 FEDAVG still perform on a par with personalized FL algorithms like FEDPER, and FED-ROD can still achieve the best P-FL
accuracy. We attribute the superior personalized performance by FEDAVG and FED-ROD to the implicit regularization discussed in subsection 3.3.
We also want to point out that, even if each client has insufficient data, the P-FL performance of local only may still have higher accuracy than the GM of FEDAVG on the personalized accuracy, especially when the non-IID condition becomes severe (e.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.

Review summaries:
FL trains models both on the server and clients. Such model training becomes problematic when the data distributions on each of the clients diverge from each other. The paper answers the question of prioritizing the generic model at the server or the personalized model on each client. The paper proposes to decouple a model’s dual duties with two prediction tasks. 

This paper focuses on a challenging problem in federated learning setting where data across different local workers are imbalanced, and most of the current work in the community are optimizing either generic performance or personalized performance. This paper proposes a new framework, Fed-RoD, which combines both global objective and local objective, to achieve that the trained model is suitable for both generic and personalized task. The main contributions are: 1. Focuses on an interesting problem which aims to increase both local and global performance in imbalanced federated learning setting; 2. Designed a two-loss, two-predictor framework to achieve the goal; 3. Based on the extensive experiments, the performance achieves state-of-the-art on both sides.

The paper proposes Federated Robust Decoupling (FED-ROD) to excel on both generic FL and personalized FL. With extensive results, it illustrates that strong personalized models emerge from the local training step of generic FL algorithms, and show that class-balanced objectives are effective for improving the generic FL performance when clients have different class distributions. Lastly, the paper demonstrates that FED-ROD enables zero-shot adaptation and much effective fine-tuning for new clients.