Diurnal or Nocturnal? Federated Learning of Multi-branch Networks from Periodically Shifting Distributions
Authors: Chen Zhu
Zheng Xu
Mingqing Chen
Jakub Konečný
Andrew Hard
Tom Goldstein
Venue: ICLR 2022 Poster (2022)
Keywords: Federated Learning
Peroredical Distribution Shift

Abstract: Federated learning has been deployed to train machine learning models from decentralized client data on mobile devices in practice. The clients available for training are observed to have periodically shifting distributions changing with the time of day, which can cause instability in training and degrade the model performance. In this paper, instead of modeling the distribution shift with a block-cyclic pattern as previous works, we model it with a mixture of distributions that gradually shifts between daytime and nighttime modes, and find this intuitive model to better match the observations in practical federated learning systems. 
Furthermore, we propose to jointly train a clustering model and a multi-branch network to allocate lightweight specialized branches to clients from different modes. A temporal prior is used to significantly boost the training performance.
Experiments for image classification on EMNIST and CIFAR datasets, and next word prediction on the Stack Overflow dataset show that the proposed algorithm can counter the effects of the distribution shift and significantly improve the final model performance. 

Published as a conference paper at ICLR 2022
DIURNAL OR NOCTURNAL? FEDERATED LEARNING OF MULTI-BRANCH NETWORKS FROM PERIODICALLY SHIFTING DISTRIBUTIONS
Chen Zhu'; Zheng Xu², Mingqing Chen', Jakub Konecny", Andrew Hard², Tom Goldstein University of Maryland, College Park, Google
ABSTRACT
Federated learning has been deployed to train machine learning models from de- centralized client data on mobile devices in practice. The clients available for training are observed to have periodically shifting distributions changing with the time of day, which can cause instability in training and degrade the model perfor- mance. In this paper, instead of modeling the distribution shift with block-cyclic pattern as previous works, we model it with a mixture of distributions that grad- ually shifts between daytime and nighttime modes, and find this intuitive model to better match the observations in practical federated learning systems. Further- more, we propose to jointly train a clustering model and multi-branch network allocate lightweight specialized branches to clients from different modes. A tem- poral prior is used significantly boost the training performance. Experiments for image classification on EMNIST and CIFAR datasets, and next word predic- tion the Stack Overflow dataset show that the proposed algorithm can counter the effects of the distribution shift and significantly improve the final model per- formance.
INTRODUCTION
In Federated Learning (FL), many clients collaboratively train machine learning model with de- centralized data under the orchestration of a central server (Kairouz et al., 2019). FL is designed for privacy protection: the private data of local clients will never be directly transferred the server or shared with other clients, which follows the principle of data minimization and keeps the attack surface of the system small (Wang et al.] 2021). Initially introduced for decentralized training on mobile devices (McMahan et al. 2017), FL has been widely applied for various different appli- cations including finance, health, digital assistance and personalized recommendations (see a few recent surveys (Yang et al. 2219:; Kairouz et al. 2019 Li et al. 2020a Lim et al. 2020 Wang et al., 202T). Specifically, cross-device FL has been used practice improve utilityy and privacy of applications such as next word prediction HHard et al. 2018), emoji suggestion (Ramaswamy et al.) 2019), query suggestion (Yang et al.) 2018), out of vocabulary word discovery (Chen et al. 2019).
and keyword trigger models (Grangvist et al., 2020 Hard et al. 2020).
A typical communication round FL starts with a server broadcasting a global model to clients. Clients then perform local computation private data and only send back aggregated model up- dates. Finally, the server aggregates the client updates and apply them to the global model before beginning the next round. In practical FL systems (Bonawitz et al. 2019:; can only participate when the local criteria is met, such as when mobile devices are charging, idle, and connected participate training at different times of the day are usually from different time zones that can have significant differencee, which can cause periodically shifting data distribution that may degrade the training stability and final model performance (Yang et al. 2018 Eichner et al. 2019). For cen- tralized systems where client data can be collected, such problem may be mitigated by caching and uniformly sampling from cached data. However, due to the privacy and system constraints (Wang
al.] [021]. clients
unmetered network. For the server, clients that satisfy their local criteria and
Work done as intern Google. Correspondence chen chenzhu@und. edu kuzheng@google comm

---

Published as a conference paper at ICLR 2022
et al. 2021p. the orchestrator (server) in FL systems is not allowed collect the raw user data, and must deal with such non-lID, heterogeneous data distribution.
To our knowledge, there only few previous works (Eichner et al. 2019 Ding al. 2020) dis- cussing periodical distribution shift of client population in federated learning. These works assume a block-cyclic structure where daytime clients and nighttime clients alternately participate in training. Eichner et al. (2019) proposed the semi-cyclic SGD approach, where clients participating training at different time slots will contribute and only use the corresponding model of the group. By learn- ing separate models for different blocks, they can obtain the same guarantee as the i.i.d. non-cyclic setting under their assumptions. However, there are several caveats of semi-cyclic SGD that makes it difficult to apply in practice: (1) It assigns models to clients based on their participation time, but not all clients will participate in federated learning, hence is hard to decide the correct group for these clients. (2) It maintains version of the full model for each clients group, which potentially increases the communication cost or privacy risk. (3) The assumption of the abrupt switch from the daytime group to the nighttime group at specific time of a day is unintuitive in practice. (Ding) ctt al. 2020 is a variant of semi-cyclic SGD that inherits these issues. We provide discussions of more related works such as heterogeneity, clustering, and multi-brance networks in Appendix A.7
In this paper, we study periodical distribution shift of clients, and make the following contributions:
We revisit the periodical distribution shift. Instead of adopting the block-cyclic structure (Fichner et al. 2019), we assume a smooth transition between the daytime mode and nighttime mode, and empirically verify through simulation that its impact on training better matches the observation in practical FL systems.
2. We propose to jointly train multi-brane network and clustering model to select the branch that better fits the client's distribution based on the feature representations. The lightweight branches for the day and night modes only slightly increase the communication cost, but significantly improve the model performance. Unlike (Mansour et al.] 2020: Ghosh et al.] [2020: Marfoq et al., )02T), the masure--based clustering model does not rely labelled data, and can be easily applied for inference on new clients.
3. We propose to use the temporal prior of the client distribution to enhance the clustering models. We assume participating clients per communication round is a mixture of daytime and nighttime clients, and the number of participating clients from the daytime group will gradually increase as time goes from nighttime daytime, and vice versa. By exploiting this prior, we can train models that are even more accurate than models trained with uniformly sampled clients in conventional federated simulation.
We provide simulations the distribution shift on three benchmark datasets (EMNIST, CIFAR and Stack Overflow) to evaluate the empirical performance of FL algorithms under the periodic distribution shift with smooth transition. We perform extensive experiments, where the multi- branch networks trained by our method outperform the distribution-oblivionn baselines by large margin: 3-5% on EMNIST, 2- 1414 on CIFAR, and 0.4-1.35% on the challenging Stack Overflow dataset under various degrees of distribution shifts. By leveraging the temporal priors, the pro- posed method can take advantage of the periodic distribution shift and beat the strong baselines of training with uniformly sampled clients by 4%, 4% and 0.45%, respectively.
MODELING PERIODICAL DISTRIBUTION SHIFT
FL setting, We consider federated learning algorithms for a set of clients I, where the i-th client has data D, sampled IID (independent and identically distributed) from its own distribution, but the distribution of different clients can be heterogeneous. We minimize the expected loss on all clients,
minimize L(w)= w
> (w)) where L((w) IEI
1
Σ ((w,E) and Σn =1 = (1) EED:
[D,]
TET
where p: is the weight (probability) of client i, and & (ax, y) is training sample pair of data and label on a client. By setting p; = [D.1/ Σjer [D,]. the federated training loss recovers the empirical risk minimization (ERM) objective on all client samples. For simplicity, we abuse notation and use Dr to denote training sample (without label) for client i.
2

---

Published as a conference paper at ICLR 2022
Training Accuracy in Real System
Training accuracy (smooth)
Training accuracy on Iblocketyclcc
0.200
0.30
0.25
0.20
0.20
0.125
0:15
0.15
0.300
Trameed only during rightlime Tranne onl during stytime Traiiing all the time
Training onl during daylime
The ning only during daytine
0.05
Truining only during rightcima
0.05
The rina only duttne rignttime The nini al the time
0.025
0.00
Truining all the time
0.00
10 12
14
500 750 1200 1250 1500 3750 2000 Communizatio Rourds
250
750 1000 1250 1500 1750 2002 Communicatinn hounds
Number days since tacks statted
Figure I: Training accuracy language models. (Left): on-device training a practical FL system. (Middle): simulation of the smooth distribution shift with (t.)(t) on Stack Overflow (7' = 256). In the beginning of each period, the probability of clients coming from nighttime mode is 1, which linearly decreases to 0 at the middle of each period so that all clients are from daytime mode, corresponding to the peaks and valleys of the curves respectively. (Right): simulation of the block-cyelic shift Stack Overflow (T 256).
Periodic distribution shift. A subset of clients (t) C I(t) available for training in com- munication round I. I() periodically changes with most of the clients from a daytime (nighttime) client group in midday (midnight). Figure (left) (left) shows that the training loss in a cross-device FL system has daily oscillation. Such oscillations was also observed by Yang et al. (2018), who con- jecture it is due the domain differences between clients from different time zones. [Eichnrr et al.) (019) study the block-cyclic structure, where the model is trained for I" rounds each day, and the clients are from the day mode and night mode alternately, each last continuously for T/2 rounds.
We plot the training curves of block-cyclic structure in Figure [ (right), and observe it is different from the (left) curves from a practical FL systems.
Smooth transition. We also assume the distribution changes periodically with a period of T. Unlike (Fichner et al., 2019), we assume clients round are a mixture of daytime clients and nighttime clients, denoted as In (t) and (t), respectively. Intuitively, since the available population is usually large around the clock, the population distribution available clients should shift gradually, rather than abruptly from one mode another in the block-cyclic structure. To better approximate the behavior in practice, we assume that in each period, clients come from the day mode I, (t) with probability qt) that varies smoothly between 0 and 1. Specifically, with two disjoint sets of clients with different data distributions, and define R + [0. 1] to be a periodic function with period of T. At each round 4, we sample the clients from the following distribution
simulate I (t) and I2(t) (t)
P(i E (t) = q(t), P(i Iz(t)) = 1 1-g(t)). g(t).
(2)
We consider a periodic Linear function (L) and a smooth Cosine function (C) for q()), each further
parameterized by an exponent factor p > 0 to control smoothness of the transition,
(1.p(1)= 2- mod T T
1 1 (cos (2xl/T) + 11
(1).p(1)=
We visualize the transition probability g(t) in Figure 7 the appendix. When p < 1, more daytime clients are available in T rounds, and when p >> 1, more nighttime clients are available. This can simulate the difference in the number of completed training rounds during daytime and nighttime observed in practice (Yang et al. 2018).
Observation and insight. Figure 1 (middle) simulates the training curve with q (1) to control the probability for sampling from I (t), which more accurately approximates the curves from a practical FL system. The three curves show that training a single model around the clock achieves the best performance on both modes despite the domain differences, which motivates us to train a model with large amount of shared weights. Semi-cyclic SGD (Eichner et al., 2019) provides the worst-case guarantees when the domains of each block are unrelated, and argues that training single model on both modes is not optimal, which is different from our observations. Another important observation is that the training accuracy reaches its minima (maxima) when the client population is most biased towards the day mode night mode, e.g., round 1024 and 1152 in the middle figure, from which we can infer the peak moment of daytime clients and nighttime clients in practice and use it to define strong prior to improve the learning process.
3

---

Published as a conference paper at ICLR 2022
3 LEARNING FROM PERIODICAL DISTRIBUTION SHIFT
We consider applications in which clients collaboratively train a model with a common input and output space, and the distribution of clients changes with time. We propose to jointly train a multi- branch network and a clustering model to assign specialized branches for prediction on clients from different modes, guided by the temporal prior of the client distributions. We consider two methods for enforcing the temporal prior: FEDTEM (Section 3.T, which is based on a Gaussian Mixture Model (GMM), and FEDTKM Section
3.2 which is based on K-means clustering.
Multi-branch network. Eichner et al. (0019) showed the merits of training separate models for different data distributions when the distributions can be identified during evaluation. However, training single model with shared feature extraction layers on all available data can usually im- prove the data efficiency of representation learning. For example, for vision tasks, it helps to learn to extract common low-level features, while for language tasks it helps to learn shared embeddings and grammars from the context. To handle the distribution shift while learning shared feature repre- sentations and alleviating the communication overhead, we adopt the weight-sharing strategy from multi-task learning (Caruana 1997) to train a multi multi-branch network with shared feature extraction layers f f x) followed by one of the specialized prediction branches gl (Wk. f(w x)) for clients from each cluster k (1 < & < K). We set each prediction branch Oa to be a single linear layer, which is more communication efficient and data efficient than having K versions of the same model.
Temporal prior. The temporal prior i(t) is an estimate for the ratio of clients coming from the daytime cluster, (t). From the observations in Section 2 we can locate the time when g(t) is most likely to be 0 or by observing when the minima and maxima occur from the training curve. In between these minima and maxima, we consider three types of q(t) in our current experiments: 1) Linear: (() = 1L.1 (t); 2) Cosine: ğ(t)= = qC,, (t); 3) Soft: see Appendix A.3
Training Objective. Our model assumes clients are from either the daytime or the nighttime cluster (K' = 2), so it has two branches, and each branch is used for one group of clients. Let ki be the branch index of client i chosen by the clustering model, and K be the index of the other branch. Each client trains the network with the following objective
minimize [:(w) =
1 [Di
Σ EED.
(celon (w, x) + deCE(gr; (w,x), s(, s(c, v)),
(4)
the label smoothing function for one-hot vector y e E [0, 1] deter- mines the amount of label smoothing, and À > is the regularization strength. The label smoothing regularization updates the other branch jointly with the feature extractor prevent staleness, while encouraging the two branches to specialize in different feature subspaces: branch k is trained to become less certain on features of samples from cluster K
where s(E, y) =
3.1 LEARNING A GAUSSIAN MIXTURE MODEL WITH PER-CLIENT TEMPORAL STATISTICS
We propose Federated Expectation- Maximmizaiiiniiiinnnnn with Temporal prior (FEDTEM) to learn a Gaus- sian Mixture Model (GMM) to infer the cluster client is from, and select the corresponding branch in the multi-branch network. We define discrete latent variables z and Ç to represent which cluster sample and a client is from, respectively. FEDTEM assumes samples on the same client are from the same cluster, so the GMM prior P(c k) = P(= = k) for any k. We define P(x|2 = k) as a Gaussian distribution (f 2=)|4,, σ in the feature space. For efficiency, assume and learn a diagonal covariance dy for the Gaussian models. Algorithm ! summarizes the training process,
and the details of each step are provided in the following sections.
3.1.1 MODELING THE CLIENT DISTRIBUTIONS AND SELECTING BRANCHES FOR TRAINING
The prior P(() of the client distribution are constantly changing due to the periodic distribution shift. To stay up-to-date, the i-th client estimates its probability of coming from the k-th cluster based its data D, before the local update steps during training. Specifically, given the GMM parameters, since P(() = P(z), the local maximum likelihood estimation (MLE) of P(G) on each new client is equal to the MLE of p(z) on client i. Let π. be the MILE of P(z) on client where
4

---

Published as a conference paper at ICLR 2022
the k-th dimension of n,; denoted as The is the MLE of P(= kc) client i. Then
TWN 2)) (5)
P(G: = k) mi
P(==
DDj
D
x)),
BED:
#ED
For completeness, give the derivation of Eq. in Appendix A.1 where also compare with the empirical results of using the posterior P(((D;) instead of the MLE मद for branch selection.
We select the branch based the MLE of P())) and train the network by optimizing Eq. 4 We can either greedily select the branch kg arg max; ;. or sample from the discrete distribution Tς. We show the greedy approach achieves better empirical results in Figure and use it by default.
Algorithm FEDTEM: Federated EM with Temporal Prior (Training)
1: Input: A stream of clients I()) with periodical distribution shift; Number of communication rounds N;
Number of rounds per day T.
2: Output: Network parameters w (wff" wf, w GMM parameters @ (e*,, oN, >T").
3: for
N do
A set m clients (t) C [t)) participates training;
5: Server broadcasts parameters of the network w" and the GMM 0³ to (');;
6: for clients i E I' (f) in parallel do
7:
Estimate
D., and choose the branch KC arg max Tiki
Eq.
8:
Given Ar run local updates for the network by optimizing Eq. and get w On D;, compute the MLE #े and all feasible Gaussian parameters (p ((1))
9:
see Eq
10:
Server aggregates the model update Σ E:((( -w')
11:
Update network parameters to u with the prescribed server optimizer.
12:
Server collects GMM params [T; P
++1 a(+1 (t)};
13:
Update the GMM parameters (p g!+1 n (+1) with the temporal prior on म.:
b see Eq. 10
3.1.2
UPDATING THE MIXTURE MODEL PARAMETERS
We introduce a federated Expectation-Maxiin Maximizatiinn (EM) (Dempster et al. 1977) enhanced by the temporal prior to update the parameters of the GMM. The temporal prior is enforced via bottom- up approach: we start by running EM on each client and send all possible locally optimal GMM parameters the server, and then select and aggregate the GMM parameters with the temporal prior on the server. We find the optimal GMM parameters after the local update steps as described in Section 3.1.1 that the GMM is updated in the feature space of the updated network. We give details of the process the following.
E step. For each client i I(t), evaluate the posterior (x=( for each sample infer the probability of æ coming from cluster k based on the locally updated network with parameters w(+1
T$ (()((( , o%) "
(x)) = P(= = k k/ar)
1+1 1oru(=",))ll x5)
M step on clients. To learn GMM that better distinguishes the daytime and nighttime clus- ters, we incorporate the temporal prior q(t) into the M step by considering the posterior P(ç = DD, 0r, q())), where 0 (n', o', π²) is the collection of all GMM parameters. The M step optimizes the following objective
K
maximize Σ Σ P(s k|D, (e.t)"""()))))))))))))) P(D,, kl0. q(t).
kE
Since the clients cannot share their data with the server, we first find locally optimal parameters 0 argmaxg ΣK P(s = k)D,, 0', q()) log P = k00, (()) on each client i, and then average on the server. However, it is hard to evaluate the posterior P(ç k[D,, 0' q(()) locally without additional communication rounds, since q(t) is an estimate of the ratio of clients from one cluster, and the posterior for one client depends
the estimates of other clients. A viable approach is to assume the posterior to be one-hot, ie.. for a certain cluster k., P(s = kDD,, 0', (())) k k., and P(C k)D;, 0' [()) if l * k.t Under this modeling assumption, we can learn
5

---

Published as a conference paper at ICLR 2022
more distinctive features for the clusters at communication overhead. To see this, when k* is given, we only need to solve
0 -
P k" [Di, 0', (())) log < - k)),
(8)
where its optimal mean and variance, and the MLE of the GMM prior for cluster hs on client are
G++1
Σ f(w x), 1 zED;
Σ (f(w) ++1 -))-i())) A [D; a>D;
Σ mED:
Yike (x).
Dal
DDd
(9)
Note Mik- and a are identical for different k"; see Appendix A.2 for derivations. The clients send the optimal mean and variance to the server. To enforce temporal priors on the server, clients also send the MLE of the prior on Dj for each cluster k; i.e., Mee- to the server In this way, the total number of parameters sent to the server can be less than a single GMM model.
M step on server with temporal prior. The temporal prior is enforced via a bottom-up approach: each client sends all possible solutions and MLEs { [+1
2++1 TT E I't), k E {1, 2} to the server; the server then estimates the one-hot posterior P()iDi, 0¹, (f(.(t)) based -.p- A set Zi (t) of Lū(t)- [T) + ] clients with highest Th will have P(S 1 0 T())) 1, while the remaining clients, constituting set (4), will have P(C = 2/D;, 0' i())) = Then, similar as FedAvg (McMahan et al.) (2017), the server updates the GMM parameters for each mode k as the weighted average of the optimal GMM parameters of each client i E T (t):
ik
ik
D.l t+1
D
M
Σ
M Pi
1
Σ
M
4-11 ih
J
m 1+1
Σラ
M
(10)
(E)
(EIL(0
where Mt (iEI(()) [D:] is the total number of samples from clients assigned to mode Note
o+ is not necessarily an unbiased estimate, but find this estimate gives good results in practice.
We also use a running average of the GMM prior (-1 during training and set it to be a uniform distribution during inference; see Appendix A.4 for details.
3.2 LEARNING A CLUSTERING MODEL WITH AGGREGATED TEMPORAL STATISTICS
FEDTEM collects per-client म, to the temporal prior and update the GMM parameters, which may not satisfy the strong aggregation-only data minimization principle (Bonawitz et al. 2021 Wang et al. 2021). We propose alternative Federated K-Means algorithm augmented by the Temporal prior (FEDTKM)), where the temporal prior enforced based on only aggregated results, described in Algorithm - FEDTKM clusters the ith client and selects the branches based on the averaged distance of the features to the k cluster centers,
du((w")
Σ axllf(w', -ccll -call RED,
(11)
D.l
where the superscript denotes the communication round, and og distance scalar factors
trol the "prior" of the clusters and enforce the temporal priors. During training, client i computes select the branch with minimal de, while test time, we use a minibatch
da on its training to estimate d, In the case of two centers, we fix d1 1, and lear scalar 0z > O to rescale the distances the second cluster C2: when 0r > makes the FEDTKM algorithm more likely to assign the client to the first cluster, and vice versa. > 1 is updated by quantile-base (private) estimator inspired by (Andrew et al. 2021).
The clients train the network with their selected branches by Eq. 11 through the objective in Eq. After the local training steps, each client estimates its cluster assignment again with the locally up- dated parameters wf. and computes an indicator variable (t) 1 [da (ι))) diz (wx() for quantile estimation, and the feature means f(rwf) P. ΣeED, f(wff, x) for k-means centers. Then (x())
and f(w); can be sent to a trusted, secure aggregator to get the ratio of clients assigned to cluster 1, q(1), and the averaged k-means centers c++ ++1 as
((t)
1
Σ
(t),,
ct
Σ
Pak Dj1
f(w)). k or 2,
(12)
It(t)
Exer ri D
(E
(t)
TThe MLE is defined in Eq.

---

Published as a conference paper at ICLR 2022
where r.1 (t), = (t))..
The distance factor 02 is updated by tracking the difference between the empirical quantile q(t) and the oracle quantile q(t) estimated by the temporal prior,
a2 s+1 = exp {n' [ī(+) - {(t)]} a²
(13)
where '/' 2 0 is the step size of the geometric update step 4. In all experiments, we use the periodical linear function ((t) = (t) as the temporal prior.
Algorithm FEDTKM: Federated k-means with Temporal Prior
1: Input: A stream of clients Z(()) from a periodically shifting distribution; Number of communication rounds
N: Number of rounds per day T'.
Output: Network parameters w (wf", wi,, wh"). ke-maans cluster centers ct and e² NN
2).
distance scalar 02
3: for to do
4: A set of m clients (t) C I() participate training;
5: Server broadcasts parameters of the network w and the m-means parameters (ci, a') I')
6: for elients i E I' (f) in parallel do
7:
Compute average distance du (w=" = Po EeED, ax ofll .)) ckll;
8:
Choose the branch with minimal distance Kj arg min; dia dix (wn)
9:
Given ki, run local updates
network by optimizing Eq. 1 and get wf:
10:
Compute q. (f) = 1 [dd < {c=((x))] to be aggregated for at);
11:
Compute f(w))) [Pal ΣMED, f(wif, .z) be aggregated for c or c2 t
12:
Server receives aggregated network and k-means parameters {w', c o a(t) from the clients;
13:
Update the distance scalar 4.1 with the temporal prior q(t): al 2411 = exp (ví (a(t) - (())] ai;
14:
Update the k-means centers into c+1 C2 using Eq. 12
15:
Update network parameters to w 1+1 using wi w" with the prescribed server optimizer.
3.3
COMPARISONS AND DISCUSSIONS
Privacy. In our simulation experiments (Section[4) FEDTEM often performs better than FEDTKM. However, the server has to observe the GMM parameters for participating client in FEDTEM, which may require new techniques to satisfy the strong data minimization and data anonymization princi- ples (Wang et al., 2021 Bonawitz et al., 2021). On the other hand, FEDTKM only uses aggregated results, which is compatible with the strong data minimization principle, and casier to be further pro- tected by other privacy techniques such as secure aggregation (Bonawitz al. 2019) and differential privacy (MMMahan et al., 2018 Kairouz et al.) 2021).
Differences from previous methods. Compared with previous methods that jointly train a cluster- ing model with multiple networks, our method has four key differences that better suits cross-device FL. under temporal distribution shift. First, our method introduces the temporal prior to regularize the clustering model, which is missing in previous works. Second, our b-branch network with shared feature extractor is more communication and data efficient than using networks. Third, our clus- tering model selects the branches based on the feature, which eliminates the need for labeled data for clients not participate in training. Clustering models of (Ghosh et al. 2020 Marloq et al. 202T) are based on the loss, which is impractical for clients without labeled data. Fourth, our method does not maintain states on clients. (Marfog et al.) 202T) assumes stateful clients, which is impractical for cross-device setting where each client only participate limited times (Wang et al. 2021). We give more detailed comparison with previous methods in Appendix A.7
4 EXPERIMENTS
Dataset. We consider two image classification datasets, EMNIST and CIFAR, and one next word prediction task Stack Overflow (SO). The split of the day mode (I,) and night mode (Iz), and other statistics, are shown in Table the Appendix. On SO, we use a vocabulary size of 10K and report the test accuracy without special tokens. We truncate each sentence have more than 20 tokens. For the distribution shift, we use T = 256 for the majority of our results, and we compare results under both linear and cosine distribution shifts with p (0.1, 0.25, 0.5, 1, 2, 10).
Architecture of the multi-branch networks. On EMNIST, we train LeNet with 2 Conv layers and 2 FC layers. On CIFAR, we train ResNet-18 with the Batch Norm replaced by Group Norm (Wul
7

---

Published as a conference paper at ICLR 2022
Accuraty EMMST (Uiear SNI
Accurscy OFFRR (Linear Shift)
Accuracy on Stack Overloow (Lineur Shift)
92
Tishes Cretar
Win Lnds DHCA
K-naar
ScsGp
Hin coe deCAo
Varilla
TRCEN
TaEEN
FISTEM (iver)
FRETEM (o07
FedTEM (ooine)
FECTICN
24.25
No D61 5011
0.23
02
26.05
0.25
6.5
2.
10.0
Accuracy ENNIST iCoste SNIf
Accuracy on CIFAR (Cosine
Accuracy on Stack Overllow (Cosine Shint)
40.0
5.50
9.25
SC5.0
19
0.00
Tishit Kemeara
26:35
Vin Lans J0C0
FestH
24.50
Vih los IFCA
Fedi4
POSTES (hraer)
75.0
TeSTION
72.5
No CiRe 57
No Dai. Stil.
1: Did, Stet
6.1
0.25
25
1.0
24
44
0.2
2.2
0.3
10
20
45
15
23 0.25
100 Lo
24
40
Figure 2: Comparing the results of FEDTEM and FEDTKM with baseline methods. On EMNIST and CIFAR, FEDTKM can achieve similar performances as FEDTEM and outperforms most of the other methods. On Stack Overllow, FEDTKM is not as good as FEDTEM, but it still outperforms other methods, and can be even better than the model trained without distribution shift when the training population is less biased (p is closer to I).
& He 2018) for stability in federated learning. We use the last FC layer of these convolutional networks the multi-branch part and share all the remaining layers. We use the output from the shared feature extractor for clustering, which are 128 and 512 dimensional for EMNIST and CIFAR respectively. On SO, train single-layer LSTM (Hochreiter & Schmidhuber 1997) with hid- den size of 670 and embedding size of 96. To alleviate the communication overhead and prevent overfitting, define the branches to be the last projection layer before the final prediction layer, which is 15x smaller than the final prediction layer. The communication overheads of the extra branch are 0.7%, 0.5% and 1.6% respectively for models on EMNIST, CIFAR and SO.
Hyperparameters. We use FedAdam (Reddi et al., 2021 as the optimizer. Table 2 shows the training hyperparameters, with implementation details and practical notes given in Section For Min Loss, FEDTEM and FEDTKM, we also do a grid search on the label smoothing parameters e and A. For each hyperparameter setting, we run experiments with different random seeds and report their mean and standard error.
Ablation study and temporal prior. In Appendix A.5 and Figure 3] we provide ablation studies on the temporal prior to quantify its efficacy, and see which temporal prior (Linear, Cosine, or Soft) works better in practice. From the results we can see: 1) all three types of priors improve the baseline in most cases; 2) linear prior works better on image datasets, while cosine prior works better on SO; 3) with weaker assumptions, Soft prior is able to fit the distribution shift and improves the results. With these observations, for FEDTEM, we choose the linear prior (t) on EMNIST and CIFAR and the cosine prior (t) on SO. For FEDTKM, we only consider the linear prior (t).
Baselines. For fair comparisons, unless otherwise stated, adapt all methods into settings to maintain two important properties: 1) training the same multi-branch network to keep model capacities the same; 2) labels are not available for test clients. For the baseline without any branch selection technique, we train the same multi-branch network by taking the averaged output from both branches for prediction. Such networks trained under shifting data distribution are denoted as "Vanilla"" while those trained without distribution shift are denoted as "No Dist. Shift" (NDS).In Appendix A we also provide results of training single-branch models in these two settings. which obtained almost identical results. The baseline methods are: 1) "SCSGD", where similar as Semi- Cyelic SGD (Fichner et al., 2019), we train one model during first half of each period, and the other model on the other half. During test time, since there is no indicator of which mode the clients come from in practice, we select the models according to the certainty of their predictions, measured by evaluating the KLD between the prediction and uniform distribution over all labels. 2) T-shot K-means, which is an enhanced variant of one-shot K-means (Dennis et al.] 2021). collecting the cluster centers
data from all participating clients during a whole period of T rounds before training. For both training and test, it selects the branches for each sample according to the nearest cluster center. 3) "Min Loss (IFCA)", where same IFCA (Ghosh et al. 2020), choose the
8

---

Published as a conference paper at ICLR 2022
branch with minimum loss on the local training set of each client during training, but different from IFCA, we choose branches with highest certainty on the unlabeled test set clients, using the same criterion 1). We applied the label smoothing regularization (Eq. 4 to IFCA, which improves the results. 4) FedEM (Marfoq et al. (02T), where we use the same algorithm on participating clients during training. FedEM is implemented with privilege to access sample labels for test clients, to enable its EM steps which uses the loss.
Main Results. The results of all methods are shown in Figure 2 By comparing results of Vanilla and NDS, we find the temporal distribution shift indeed degrades vanilla training, causing the worst- case decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and SO. While with FEDTEM and FEDTKM, the accuracy can be even higher than models trained in the NDS setting in most cases, demonstrating the efficacy of the temporal prior for learning more distinctive feature representations for the daytime and nighttime modes with specialized prediction branches. The improvement on the strong NDS baseline can be as high as 4%, 4% and 0.45% with FEDTEM on the three datasets. FEDTEM is not better than NDS when p deviates too much from as the data distribution is extremely skewed in such settings, while our methods are still better than most other methods in most cases. For baseline methods, SCSGD performs surprisingly well on EMNIST when p is large, but the performance quickly drops on the more complicated CIFAR and SO. T-shot K-means achieves improvements for extreme p's on EMNIST, but it is not significantly better than "Vanilla" on CIFAR, due to the difficulty in reliable clustering in more complicated image spaces. We did not implement T-shot K-means for SO due to the difficulty of clustering the raw data of the language task. Min Loss (IFCA) alleviates the drop and sometimes competes with NDS, but cannot achieve improvements over NDS. By contrast, FedEM often achieves inferior results even with the privilege access to the test labels for EM estimation. This is probably due to the staleness of the priors in the challenging periodic distribution shift setting, as we simulate cross-device FL on a large population where clients only participate training for few rounds on EMNIST and CIFAR, and at most one round on SO.
Effect of label smoothing regularization. Shown in Figure Since EMNIST is relatively easy and the network on it is small, the label smoothing does not show significant improvements. However, it is critical to the success CIFAR and SO. Without the regularization, one branch will not be updated simultaneously with the feature extractor, resulting in instability and low test accuracy.
MLE for Evaluation. For FEDTEM, during evaluation, we find it beneficial for image classification to use the MLE in Eq. [5 on the minibatches (typically of size no larger than 64) to select the best branch, compared with the sample-level inference. We show its benefit in Figure 4 in the appendix. In practice, similar approaches can be realized through caching. We find Min Loss does not show much difference when similar approach is applied, where branches chosen by the averaged certainty on minibatches. The baseline model even gets worse performance with this approach, as shown in Figure 4 indicating our model is much better distinguishing two modes in expectation. However, this approach does not show improvements for FEDTEM on SO.
Effect of T. As shown in Figure g for FEDTEM with linear prior on EMNIST, smaller T tends to decrease the performance when p deviates too much from 1. In such scenarios, the distribution changes frequently and abruptly. However, the performance of FEDTEM is maintained when the distribution changes in frequent but more balanced way.
5 CONCLUSIONS
In this paper, we showed the influence of a smoothly shifting distribution between daytime and nighttime modes better matches the phenomenon in practical FL systems, and developed algorithms to train multi-branch networks to tackle the distribution shift. Our methods incorporates priors of the temporal distribution shifts to learn mixture clustering model to guide the network and select corresponding branches for clients from different modes during both training and inference. The clustering models are defined in the feature space and does not require labelled data for inference, hence is ready to be deployed on clients without labeled data. The branches are lightweight with little communication overhead, and the model demonstrates significant improvements in test accuracy on the benchmark datasets. Our methods also satisties many other real-world constraints like single round communication and stateless clients (Bonawitz et al. 2019 Paulik al. 2021 Wang et al. 2021), which makes it possible to be deployed in practical FL, systems.

---

Published as a conference paper at ICLR 2022
6 REPRODUCIBILITY STATEMENT
Our code is available on GitHub For reproducibility, we have reported the mean and standard error for 3 experiments with different random seeds for settings. The dataset we used for simulations are all available in Tensorflow Federated, and our implementation is largely based on the code from Adaptive Federated Optimization (Reddi et al. 2021 3
We have listed our hyperparameters in Table 2 We have given the pseudo code in Algorithm | and the detailed formulas and derivations of our algorithm Section 3.1
7 ACKNOWLEDGEMENTS
The authors would like to thank Galen Andrew and Brendan McMahan for helpful discussions. Goldstein was supported by the National Science Foundation #1912866, the DARPA GARD pro- gram, and the Sloan Foundation.
REFERENCES
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Representations (ICLR), 2021.
Mathieu Andreux, Jean Ogier du Terrail, Constance Beguier, and Eric W Tramel. Siloed federated learning for multi-centric histopathology datasets. In Domain Adaptation and Reprexentation Transfer. and Distributed and Collaborative Learning, pp. 129-139. Springer, 2020.
Galen Andrew, Om Thakkar, Swaroop Ramaswamy, and Hugh Brendan McMahan. Differentially private learning with adaptive clipping. In Thirty-Fifth Conference on Neural Information Pro- cessing Systems, 2021.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary, Fed- erated learning with personalization layers. arXiv preprint arXiv:1912.00818, 2019.
Kallista Bonawitz, Peter Kairouz, Brendan McMahan, and Daniel Ramage. Federated learning and privacy: Building privacy-rreeervig systems for machine learning and data science on decentral- ized data. Quene, 19(5):87-114, 2021.
Keith Bonawitz, Hubert Fichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, H Brendan McMahan, et al. Towards federated learning at scale: System design. SysML, 2019.
Christopher Briggs, Zhong Fan, and Peter Andras. Federated learning with hierarchical clustering of local updates improve training on non-iid data. In 2020 International Joint Conference on Neural Networks (IJCNN, pp. 1-9. IEEE, 2020.
Rich Caruana. Multitask learning. Machine learning, 28(1):41-75, 1997.
Mingqing Chen, Rajiv Mathews, Tom Ouyang, and Françoise Beaufays. Federated learning of out- of-vocabulary words. arXiv preprint 1tt.:/9903.10333 2019.
Arthur P Dempster, Nan M Laird, and Donald Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society: Series (Methodological), 39(1): 1-22, 1977.
Don Kurian Dennis, Tian Li, and Virginia Smith. Heterogeneity for the win: One-shot federated clustering. ICML, 2021.
10

---

Published as a conference paper at ICLR 2022
Yucheng Ding, Chaoyue Niu, Yikai Yan, Zhenzhe Zheng, Fan Wu, Guihai Chen, Shaojie Tang, and Rongfei Jia. Distributed optimization over block-cyclic data. arXiv:2002.07454, 2020.
Moming Duan, Duo Liu, Xinyuan Ji, Yu Wu, Liang Liang, Xianzhang Chen, and Yujuan Tan. Flexible clustered federated learning for client-level data distribution shift. arXiv preprint arXiv:2108. 2021.
Hubert Eichner, Tomer Koren, Brendan McMahan, Nathan Srebro, and Kunal Talwar. Semi-cyclic stochastic gradient descent. ICML, 2019.
Yanan Fu, Xuefeng Liu, Shaojie Tang. Jianwei Niu, and Zhangmin Huang. Cic-fl: Enabling class imbalance-aware clustered federated learning over shifted distributions. In International Confer- ence on Database Systems for Advanced Applications, pp. 37-52. Springer, 2021.
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. NeurlPS, 33, 2020.
Filip Granqvist, Matt Seigel, Rogier van Dalen, Áine Cahill, Stephen Shum, and Matthias Paulik. Improving on-device speaker verification using federated learning with privacy. arXiv preprint wrXiv:2008.02651 2020.
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: 18/ 1.1.1.00 1.03604, 2018
Andrew Hard, Kurt Partridge, Cameron Nguyen, Niranjan Subrahmanya, Aishanee Shah, Pai Zhu, Ignacio Lopez Moreno, and Rajiv Mathews. Training keyword spotting models on non-iid data with federated learning. arXiv preprint arXiv:2005.10406, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Kevin Hsich, Amar Phanishayce, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of decentralized machine learning. In International Conference Machine Learning, pp. 4387- 4398. PMLR, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real-world data distribution. In ECCV, pp. 76-92. Springer, 2020.
Peter Kairouz, Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems federated learning. arXiv preprint arXiv: 192.04991 2019.
Peter Kairouz, Brendan McMahan, Shuang Song, Om Thakkar, Abhradeep Thakurta, and Zheng
Xu. Practical and private (deep) learning without sampling or shuffling. International Conference on Machine Learning (ICML), 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastie controlled averaging for on-device federated learning. International Conference on Machine Learning (ICML), 2020.
Mikhail Khodak, Maria-Florina Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta- learning methods. In Advances in Neural Information Processing Systems, 2019.
Tian Li. Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-66 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Sys- tems, 2020b.
Paul Pu Liang. Terrance Liu, Liu Ziyin, Nicholas B. Allen, Randy P. Auerbach. David Brent, Ruslan Salakhutdinov, and Louis Philipp Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523, 2020.

---

Published as a conference paper at ICLR 2022
Wei Yang Bryan Lim, Nguyen Cong Luong, Dinh Thai Hoang, Yutao Jiao, Qiang Yang, Dusit Niyato, and Chunyan Miao. Federated learning in mobile edge networks: A comprehensive survey. IEEE Communications Surveys & Tutorials, 2020.
Liang,
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications federated learning. 1xXi:2202211019, 2020.
Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. arXiv:2108. 10252, 2021.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-eficcie learning of deep networks from decentralized data. In AISTATS, pp- 1273-1282. PMLR, 2017.
Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. In International Conference on Learning Representations (ICLR), 2018. URL /att:a::/openreeveee...tttt
Matthias Paulik, Matt Seigel, Henry Mason, Dominic Telaar, Joris Kluivers, Rogier Dalen, Chi Wai Lau, Luke Carlson, Filip Granqvist, Chris Vandevelde, et al. Federated evalua- tion and tuning for on-device personalization: System design & applications. arXiv preprint arXiv:2102.088 085003 2021.
Swaroop Ramaswamy, Rajiv Mathews, Kanishka Rao, and Françoise Beaufays. Federated learning for emoji prediction in mobile keyboard. arXiv preprint arXiv: 04329, 2019.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. ICLR, 2021.
Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaie. Robust federated learning: The case of affine distribution shifts. In NeurIPS, 2020.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. IEEE transactions on ral networks and learning systems, 2020.
Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, Keith Rush, and Sushant Prakash. Federated reconstruction: Partially local federated learning. arXiv preprint arXiv:2102.03448
2021.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Advances in Neural Information Processing Systems, 2017.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Blaise Aguera Arcas, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, et al. field guide to federated optimization. arXiv:2107.06911 2021.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-, 2018.
Ming Xie, Guodong Long, Tao Shen, Tianyi Zhou, Xianzhi Wang. Jing Jiang, and Chengqi Zhang. Multi-center federated learning. arXiv preprint 2 186.08647 2021.
Qiang Yang. Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Timothy Yang. Galen Andrew, Hubert Eichner, Haicheng Sun, Wei Li, Nicholas Kong, Daniel Ra- mage, and Françoise Beaufays. Applied federated learning: Improving google keyboard query suggestions. arXiv preprint arXiv: htff::///229900 2018
Jachong Yoon, Wonyong Jeong, Giwoong Lee, Eunho Yang, and Sung Ju Hwang. Federated contin- ual learning with weighted inter-client transfer. In International Conference on Machine Learn- ing, pp. 22073-12066 PMLR, 2021.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou, Data-free knowledge distillation for heterogeneous federated learning. ICMI., 2021.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A. FURTHER DISCUSSIONS OF BRANCH SELECTION WITH GMM
Derivation of the MLE. Recall the parameters of the GMM model. We consider maximizing the expectation of the complete-data log likelihood log P(D., Z under the posterior distribution
P(Z)D;, 0), where the collection latent variables associated with each of the sample æ E D;. Formally, is solving the following constrained optimization problem
maximize Ez~r(Z (P, [log P(D:, Z
K
(14)
subject to
m = m, 2 0, for 1 S K. Σ
We introduce an indicator function Jama] which is 1 if 2 = k or 0 otherwise. In this way, the log likelihood can be represented
[D.
log P(D., 0 log
II
II
P( 2n 0))
n=1 A=
(15)
P.l K
1 (an M)
[log nx + log P(æ,, 0)]
Meanwhile,
= P(2n k|/,, 0) = T P( 2n k 0)
Ez-PiZID.0) [1 (2. m]]
(16)
E P( (2 2n = =j 0)
Plug Eq. 15 and Eq. 16 back into Eq.
14 and ignore terms that do not depend on π,, we find we are
solving the following problem for
DDl K
maximize
P(zn = k|zn, 0) log Th,
R=1k=1
(17)
subject Σ m. = 1, fo all 13j K.
The Lagrangian multiplier of this problem is
K
h(TT, A) Σ
Σ
P(zn
k|
1og TT + X(1 Σ Ty).
(18)
==1k=1
Let the derivative w.r.t. ne be 0,
D
ah
1 A= 0. P(zn k| k| 0) T&k
(19)
OTk
Multiply both sides of Eq.
19 with
The
k K, and apply the constraint that
ΣK K
TTk 1, we have
D.l
P(z. k|w., 0) = [D.].
(20)
1=1
Plug this back into Eq.
19 get the equation desired.
P
k( 0)
(21)
Dr
D=
13

---

Published as a conference paper at ICLR 2022
An alternative for branch selection. Instead of computing the MLE, with the assumption that sam- ples in Dj are IID, we can compute P(G]DD; = P(z(D;) using Bayes' theorem and use P(G]/D,) for branch selection. Specifically,
ILEDD P(al= = k)
K P(s: P=== 2 T. LLED, P(æ|= - j)
=
(22)
As shown in Figure we find MIE achieves better empirical results than this alternative.
A.2 DERIVATION OF THE LOCALLY OPTIMAL GMMS
We give the derivation of the optimal mean and variance of the GMM on each client, under the one-hot assumption that P(C k|D,, 0', .())) = if k = k", and P(s = kjD,, 0", .(()) = 0 if o k*, for some given k.. Under this assumption, we need solve
0+ ak = argmax log P(D,, = k0 a(t))
argmax log P(C k- 6= q(t)))))))))))))))))))))
Σ log P(=|5 k" k²,))((t)
mED:
(23)
(lf(w
1,)2 + log lo- 2
Σ BED,
= argmax log TK=
2 [
where d is the dimension of the features, [px-] [ox- and [f(w]+1 æ)); denote the j-th dimension of PR, Th [w x)) respectively. From Eq. 23 we can see the optimal solution for client is to set the mean and variance for the k*-th mode
p
1
Σ ØED:
f(w5+1 ,.), 12
1
Σ (f(w 2x) 12
(24)
[Dg
[Di
A.3 THE SOFT PRIOR
The soft prior (t) gives strong signals only at time steps where we believe q(t) = 0 or (()) 1. In between, it only requires the ratio to be non-increasing non-decreasing, where the ratio is estimated using the posterior from Eq. as
(al E D,, (a())) (2())))
ds(())
(25)
DDj
i.e., the ratio of samples whose posterior has higher probability mode 1. In this way, the model can still learn to distinguish the modes but we do not enforce a strong prior at every time step. Specifically, it is defined as
1,
if mod = 1 if t mod T = T' if < t mod T < T if mod
0. min (t, (t),,,
(t)
(26)
max(@s (t. (t)).
A.4 IMPLEMENTATION DETAILS AND PRACTICAL NOTES
Setting the G.MM prior. During test, to compare the model performance with established baselines and ensure fairness, we consider static test set where the number of clients and samples from both modes are roughly the same. Therefore, we use uniform distribution for the priors on the test set. This does not make the GMM worse if it estimates P(x|=) accurately the test set. During training, we estimate the prior r' in every step to match the shifting distribution. However, we find beneficial to use a running average of и ff+1 3 + (1 - M / Mt) during training. which achieves better results than setting / = or using uniform prior in our preliminary experiments. We use B 0.99 for the moving average in all experiments.
14

---

Published as a conference paper at ICLR 2022
Accuracy on EMNIST (Linear Shift)
Accuracy CIFAR (Linear
Accuracy on Stack Overflow (Linear Shift)
25.8
25.6
80
25.4
Vantlla
Vanile
EM Onty
SM Only
25.0
Only
FosTEMM
PoSTEM
Olecar
24.8
linear
24.6
0.25
0.5
2.0
0.1
0.25
0.5
2.0
2.0
0.1
0.25
2.5
4.2
10.0
Accuracy EMNIST [Cosine Shift)
Accuracy (Cosine Snift)
Accuracy on Stack Overtoo (Cosine Shifti
02
Vanila
12
36
23.25
Accurecy
23.00
%
Verlii
24.75
Variile
80
EM Only
24.50
bx Only
FedliEM EM (incart
24.25
[ineat)
FedTEM Joosieet
[cssine]
78
No Dis. Shin
24.00
0.25
0.5
2.0
0.1
0.25
0.5
1.0
4.0
0.25
0.5
2.0
Figure 3: The effect of the temporal prior. Note we always use the same periodical linear or cosine priors, (t) or ac, (t), under all types and p's for the distribution shift.
Setting Te for FEDTKM. Although we can identify the modes through the local minima and max- ima of the training curve, the distribution shift can happen in arbitrary ways in between. As a result, the certainty of the temporal prior should also be lower in between the modes. To incorporate this heuristic, instead of using a constant 1y', we let
ry = 2 (f where Mmax > is a constant.
(27)
We find this give better results when p * i.e., when the underlying distribution shift is biased and different from the linear temporal prior.
A.5 ABLATION STUDIES: THE EFFECT OF TEMPORAL PRIOR
To quantify the effect of temporal prior on training, we compare four versions of FRDTEM: 1) only apply the EM part of our algorithm without the temporal priors, where the only difference from TED is that the server, in the M step. greedily uses client to update parameters of mode Kg- = arg max; 7i (see Eq. and Eq. I0 2) FedTEM (linear): our algorithm where the temporal prior is the periodic linear function P( (t) (t); 3) FedTEM (cosine): our algorithm where the temporal prior is the cosine function P(i E (t) = gc.((t); 4) FedTEM (soft): our algorithm with the soft temporal prior introduced in Section A.3] The results are shown in Figure 3 We find all three priors improves the "EM Only" baseline in most cases. "FedTEM (linear)" is better on image datasets, while "FedTEM (cosine)" is better on the language dataset. The soft prior has very weak assumptions about the prior, only requiring the estimated ratios to be non-increasing or non-dccreasing within certain intervals, but still improves the results.
EM Only:
Table 1: Stats of the datasets. We use |D.| to denote number of samples in each subset. On SO,
Classes is the vocabulary size.
Dataset
# Classes
Iz
[Ty]
[Pz,]
Iz
[Zl
[Dzal
[
EMNIST
62
Digits
3383
341.873
Characters
3400
329.712
77.483
CIFAR
110
CIFARIO
500
50,000
CIFAR100
500
50,000
20,000
Stack Overflow
10K
Questions
171K
67M
Answers
17IK
67M
16M
A.6 VERIFYING THE EFFECT OF MLE FOR EVALUATION
To quantize the effect the batch-level MLE, we compare the results of our method with without MLE on EMNIST and CIFAR. As shown in Figure this only decreases the accuracy of the baseline model, indicating our model is much better at distinguishing the two modes in expectation.
15

---

Published as a conference paper at ICLR 2022
Accuracy
40.0
67.5
77.5
FesTTe (VLE)
TaCTEM ONLIE
ML2
211
2.23
40 305
03
6.25
0.5
1.0
2.0
4.0
13.0
Figure 4: Comparing the effect of MLE on EMNIST and CIFAR.
Accuracy on EMMST
Accuriccy CIFAR
Accuracy Stack Dewrflom
69
AccumIci 88
rug lin STRV
lin still
Dit.at 21
mg ih:
82
24
1.0
0.23
20.0
Figure 5: Compare the effect of the label smoothing regularization on FEDTEM under linear and cosine distribution shifts. We use the linear prior all
ENNIST (Unear Shift!
Accuracy CIARR (Linear
00
Dayer joamping indl Boyes Jureddy tamolingg
bapm (gready)
Figure 6: Comparing the Bayesian and ML. based branch selection EMNIST and CIFAR under linear distribution shifts. We also compare the greedy branch selection vs. sampling based branch selection on EMNIST.
Table 2: Training hyperparameters on the datasets. The network parameters trained for one epoch on each client. In addition, on Stack Overflow, we also limit each client to use no more than 512 samples during training.
Dataset
Server Opt Server LR
Client Opt
Client LR
[[())
Batch Size
Total Rounds
EMNIST
Adam
10-2.5
10-4
SGD
10-1.5
10
20
2049
CIFAR
Adam
10-1
SGD
10-1.5
10
20
8196
Stack Overflow
Adam
0.01
10
SGD
10 10-0.5
50
16
2048
p 10.0
Lo
400
Sourds
Rounds
Touncs
Rounds
Rounds
Rourds
Beurds
cosine, 0.1
cosine, 0.25
cosinie 0.5
cosine, 1.0
cosine, 2.0
cosini, 4
cosine, 10.0
3.7
0.5
05
5
05
202
150
700
400
200
100
200
100
203
530
203
420
200
400
Tourcc
Raunds
Bicands
Figure 7: Probability of sampling of clients from In (4) in each round, with T 2256, under the linear and cosine transit functions and different values of
16

---

Published as a conference paper at ICLR 2022
Accuracy shift)
Accuracy (cosine shift)
90
88
82
80
T-128
T=512
T-512
Figure 8: Evaluating the effect of the underlying 7 on FEDTEM with lincar prior.
A.7 ADDITIONAL. RELATED WORKS
As mentioned above, Semi-cyelic SGD (Eichner et al. 2019) and (Ding et al.] (0200 are the only previous works we are aware of that explicitly consider periodical distribution shift in FL. We now review other related works that either consider other types of distribution shift, or have (weak) similarity as the proposed FEDTEM method.
Heterogeneity and client distribution shift. Client heterogeneity is an active topic in FI., and various methods have been proposed to address the distribution difference among clients. No- tably, FedProx (Li et al. 2020b) applies proximal regularizer when perform client updates; SCAF- FOLD (Karimireddy et al. (020) uses control variates as local states on clients for variance reduc- tion; FedPA (Al-Shedivat et al. 2021) provides a Bayesian view and adopts posterior sampling; FedGen (Zhu al.) 2021) learns generator that can be shared among clients; FedRobust (Rei-) sizadeh et al. (020)) applies gradient descent ascent to tackle distribution shifts in the form of affine transforms in the image domain. Transfer learning, multi-task learning, and meta learning intro- duced into FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al.) 2017; Khodak et al. 2019). Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al.) 102
More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al.] (2020); Ghosh et al.) (2020), while still assuming the clients can be uniformly accessed during the training process. We kindly ask interested readers to find more papers on het- erogeneity in few recent surveys (Kairouz et al. 2019 Li et al. 2020a Wang et al. (02T). All these methods consider the distributions shift among different clients, while the proposed FEDTEM considers periodic distribution shift of client population.
Clustering and mixture models. We do not assume strong control over the clients and the avail- able client population is always changing due to the periodical distribution shift. We also not require stateful clients, since our prior is applied globally to all clients. Our mixture model is based on the feature space, therefore we do not require labeled data for unseen clients to infer its mode. To our knowledge, existing works in clustered FL, including those personalization. mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting. Dennis et al. (202T) propose one-shot clustering approach based the raw client data before training. which implicitly assumes all representative clients are available simultaneously. Clustering on raw data can also be unreliable when the data demonstrates complicated distributions in the input space. Clustered Federated Learning (Sattler et al. (220)) applies an one-shot clustering based on a trained global model, and then train a personalized model for each cluster. To achieve this. it implicitly assumes strong control over population and clients sampling. Ghosh et al. (2020) and Mansour et al. (2020) propose a similar algorithm that alternatively perform clustering and updating the personal- ized models for corresponding clusters. Both of them require labled data for the clients to compute the loss for model selection. During the preparation of this draft, we notice a concurrent work, Fe- dEM (Marfoq al.) 2021), which proposes a Federated EM algorithm to learn mixture model for each client and weigh the predictions from multiple models. The modified EM algorithm requires computing the loss and therefore labeled data for every client. It maintains different prior distri- bution for every client, which requires stateful clients and strong control over client sampling. The hierarchical or bi- hi-partitioning clustering process of (Briggs al. 2020 Fu et al.! 2021) requires all
17

---

Published as a conference paper at ICLR 2022
Accuracy on EMNIST (Linear Shift)
Accuracy on EMNIST (Cosine Shift)
88
Acturaty
86
86
84
Vanilla
Vanilla
82
Vanilla
FedTEM (linaar)
FodTEM (linear)
No Dist. Shift
No Dist. Shift
0.5
0.25
4.0
Figure 9: The effect of using multi-branch networks for the haselines. "Vamilla (s)" is the single-branch network trained with FedAdam under various distribtuion shifts, "No Dist. Shift (s)" is the single-branch network trained with FedAdam with distributoin shift.
clients to be available for clustering simultancously. violating our assumption that the distribution of the population is constantly changing and is not practical in on-device FL in the real world. In addi- tion, for unseen clients, either additional communications are needed for determining their clusters, or the client has to download models for all clusters, which adds significantly to the communica- tion cost. (Xie et al., 202I) maintains one set of weights on each client while maintaining multiple weights cluster centers the server. It takes two communication rounds for each weight update, which can be impractical. remains unclear if this mechanism can be generalized to unseen clients. (Marfoq al. (021)) trains different set of weights for each node/client, which requires all the nodes/clients to be available all the time. It is more suitable for the cross-silo setting but not prac- tical for on-device setting. (Duan al. 202T) clusters the clients based the similarities between their gradients at the initial model weights. This is impractical in our setting since the training pop- ulation changes throughout the day and we cannot assume clients from both modes are available at the first round for clustering. In addition, there is no guarantee that the gradient at the chosen round can separate the clusters well. And to compute the gradients, still requires labeled data. (Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which not practical for on-device FL since BN has been widely observed cause instability in on-device FL., and is often replaced by Group Normalization which does not track running statis- ties (Hsieh et al. 2020 Hsu et al. 2020). Another limitation is it assumes every client is seen during training.
Multi-branch networks in FL.. Multi-branch networks have been explored in FL when person- alized model is preferred for each client: the branchs are either locally stored on clients (Arivazhagan al. 2019 Liang et al. 2020), or reconstructed based on client data (Singhal al. 2021). All the atorementioned three works require labeled data and additional training efforts to learn the weights of the branches. By comparison, our method does not need re-training the branches and does not need labeled data for unseen clients to select the branches. The branch selection is based on the mixture model the feature space, and the weights of the branches are fixed.
A.8 COMPARING B-GNCH AND
BASELINES WITHOUT CLUSTERING
For fair comparisons, we have used two branches for the baselines ("Vanilla" and "No Dist. Shift") so that their capacities are the same as our method. Since these two baselines are just training the networks with FedAdam, it is more natural to train single-branch networks directly. In Figure 9 show that the single-branch branch and multi-branch networks obtain similar results under various settings, with without distribution shift. Our method still outperforms these single-bran branch baselines.
A.9 THE EFFECT OF LOCAL TRAINING SET SIZES
When the number of training samples is small, the EM estimates may have high variance, caus- ing misspecifications and resulting in worse performance. Meanwhile, models trained smaller training sets tend to have worse generalization, so the test accuracy will drop for any method in general. To which factor has more significant effect the test accuracy, we compare with the
18

---

Published as a conference paper at ICLR 2022
Accuracy on EMNIST (Cosine Shift)
Accuracy on CIFAR (Cosine Shift]
Accuracy on Stack Overllow (Cosine Shift)
Vanill
Venila
TedTTN (inwar)
(inear)
Vanili No Disi, Stitt FosTEM (inca))
24.0
85
22.5
04
23.0
:20
sameles cient
clont
Figure 10:
The effect of local training set sizes of each client. x-axis represents the maximum number of training samples per client. For "inf", , we use all available training samples from the original dataset. To add the challenge, ensure the distribution shift is different from the temporal prior. On EMNIST and CIFAR, consider cosine data distribution shift with p 1 and linear prior for FedTEM. On Stack Overflow, we consider cosine distribution shift linear prior. We also compare with results for baseline model trained without distribution shift (No Dist. Shift).
Training Real
Vocabulary Real
Number of Tokens Real
6.700
0.375
0.100
620020
550000
0.323
0.000
500020
8.300
410020
0,075
Trained eny during rightimms ony during saytime
0.050
350000
0.525
250020
14
10
Figure 11: The training accuracy, out vocabulary rates and total number tokens at each round for training a next word prediction model in a real FL system. In general, the out of vocabulary rates hecome lower on nighttime clients, and the sentence lengths become longer nighttime clients. The plots also show training is faster during nighttime, since more rounds are finished during nighttime.
baselines trained with or without distribution shift in Figure I0 where we only change the maximum training samples per client while keeping other hyperparameters unchanged. From the plot, we can see FedTEM maintains the advantage over the baselines and oracle (No Dist. Shift) under various training set sizes, though its accuracy also decays as the baselines due to the decreased numbers of training samples.
A. STATS OF TRAINING A LANGUAGE MODEL IN A PRACTICAL FL SYSTEM
In Figure II we show some stats that could characterize the data distribution shifts in a real FL system. Generally, the data of daytime clients is more difficult to fit and generalize. Data during daytime has higher out of vocabulary rates, and lower sequence length, indicating these sentences might be more arbitrary or informal than data during nighttime. From the plots, we can also see more rounds are completed during nighttime. This is because the product is mainly deployed in one region, so during daytime, fewer devices are idle for training and the round completion rates are lower (Yang et al., 2018). These plots further justify that the data shifts smoothly rather than abruptly in practice.
19


Reviews:
[{"review_id": "h90QHQo5_Ho", "reviewer": "Reviewer_YzWj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The paper proposes a flexible, intuitive and practical algorithm for being able to model distributional shift and incorporate the temporal information through a mixture distribution.\n- The multi-branch based network ensures feature parity across different distributions being considered in this paper.\n- While the paper is centered around two different distributions, it can be easily extended to multiple distributions to finer granularities depending on the need.\n- Experimental results across all the vision and language tasks show an impressive improvement performance.\n\nCons:\n- It seems the authors make an implicit assumption that a client at any given time consists of data from either of the distributions at any given time. Also, at any given time how does one ensure clients are selected from only one of the distributions being considered? As with FL, one can't have any information about the distribution of data on one client.\n- The choice of the smooth transition function seems a little ad-hoc. What is the rationale behind selecting such functions? What other functions were considered?\n- The performance improvement comes at the cost of added communication complexity and extra hyper parameters to tune.\n- How much does the mis-specification of GMM affect the performance of FedTERM. Especially while dealing with clients with fewer samples, the updates for GMM can be high variance prone. How is the performance of FedTERM affected in such scenarios? Please consider adding additional experiments to consider aforementioned scenarios. Finally, misspecification of the prior at each step can lead to the algorithm underperforming severely.", "summary_of_the_paper": "This paper handles the distributional shift observed in data in clients in a federated learning production setting. In particular, the distributional shift is modeled as a mixture of distributions.Furthermore, a multi-branch network is used to encapsulate the shifting distribution and a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (FedTEM) is proposed. Experiments for image classification on EMNIST and CIFAR datasets, and next word prediction on the Stack Overflow dataset show the efficacy of the proposed algorithm.", "summary_of_the_review": "This paper handles the distributional shift observed in data in clients in a federated learning production setting. In particular, the distributional shift is modeled as a mixture of distributions. The innovations include learning a GMM in a federated fashion and the development of the algorithm FedTerm. Experiments show the performance gain with the proposed algorithm. However, there are some gaps in the paper which require additional clarification and experiments. Upon addition of the clarifications and experiments, the paper will be a strong paper to be accepted.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_7bl0snk8WQ", "reviewer": "Reviewer_QWqv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The problem setting is interesting and new. \n2. It is interesting studying the periodic distribution shift problem in FL. The authors model the distribution shift with the prior daytime and nighttime modes in daily life.\n3. The authors propose a novel federated expectation-maximization algorithm (FedTEM) to learn a branched neural network model. Experiments on simulated environments show the proposed method achieves better performance when distribution shifts in the federation.\n\n \nWeaknesses:\n\n1. The periodic distribution shift problem in FL is very interesting. The distribution shift problem usually expects uncertain changes in every shift step. However, the targeting problem in this paper is a shift between two fixed distributions in a periodic manner. This setting is much less challenging than other distribution shift problems.  \n2. The experiment cannot well support the mentioned application scenarios on “Diurnal or Nocturnal”. It simulates two distributions in FL benchmark datasets that are not designed to distinguish daytime and nighttime.\n3. The baseline selection and experimental comparison is weak. \n\n\nQuestions:\n\n1. From an application perspective, there are many straightforward solutions to solve the proposed problem in training two models for “Diurnal or Nocturnal” respectively. For example, we can train two global models in a totally separate mode. I am not convinced that the proposed solution is necessary to solve the targeting problem. It would be much help to give more detailed examples to enhance the motivation.\n2. Could you please provide more details about the difference between a daytime distribution and nighttime distribution? For example, how the distribution changes with respect to p(y), p(x) or p(y|x) ?\n3. The proposed FedTEM estimates the mode a client belongs to and later chooses the corresponding branch of a neural network model for it. However, a similar strategy has been reported in the clustered FL framework as mentioned in A.4 - Clustering and mixture model, and other related methods, such as the below papers. Please pick up some paper as baseline methods to compare with.\n\nhttps://arxiv.org/abs/2004.11791 \n\nhttps://arxiv.org/abs/2108.08647\n\nhttps://arxiv.org/abs/1705.10467 \n\nhttps://arxiv.org/abs/2103.00697 \n\nhttps://arxiv.org/abs/2108.09749\n\nhttps://link.springer.com/chapter/10.1007/978-3-030-73194-6_3 \n\nhttps://link.springer.com/chapter/10.1007/978-3-030-60548-3_13\n\n\n", "summary_of_the_paper": "The paper is to solve the distribution shift between daytime modes and nighttime modes in federated learning with a mixture of distributions. The proposed method could be viewed as a two-group clustered federated learning method. In particular, the client clustering is based on the model parameters of prediction layers.", "summary_of_the_review": "In this paper, the authors propose a novel setting for federated learning, and then design a solution to handle the periodic distribution shift problem in FL. However, the claim is not well supported by the contents.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "TZGttT-SYj4", "reviewer": "Reviewer_4Moy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Detailed comments:\n\n(1) The contribution is incremental and lacks sufficient novelty. The multi-branch idea has been fully explored in both multi-task FL and heterogenous/personalized FL. The idea of client clustering has also been studied by several recent works in FL. The proposed solution, i.e., the EM algorithm, is also common for client clustering. These works should be discussed and compared in the paper. Label-smoothing seems helpful here but is also not new. \n\n(2) The newly proposed heuristics, e.g., temporal prior, different shifting schedules (cosine, linear, soft), one-hot cluster assignment, etc., lack rigorous justification. They are too hand-waving. It is not clear why they can be directly integrated into the EM algorithm without hurting its convergence or causing any inconsistency. Moreover, the performance seems to be sensitive to the application of these heuristics. \n\n(3) The datasets/benchmark evaluated in the paper do not have real periodic shifting distributions. The periodic shifting is artificially added but it is not clear why the realistic shifting in practice should follow one of the three schedules and p. \n\n(4) In experiments, when generating the dataset with periodic shifting distributions, q(t) is used to define the periods. However, the algorithm cannot have access to the ground truth q(t) defining the shifting schedule in practice: the application of $\\tilde q(t)$ looks like cheating.\n\n(5) There exist many FL methods specifically designed for addressing non-IID distributions across clients and better-personalized models for different clients. However, none of them has been compared in the experiments. \"min-loss\" is a baseline modified by this paper. Hence, it is hard to justify the advantage of the proposed method.\n\n(6) Fig. 1 is hard to interpret: why do the curves for models trained only during daytime and the ones trained only during nighttime reach peak accuracy at the same time?\n\n(7) The assumption that day-time (night-time) data representation follows a single-mode isotropic Gaussian distribution is too strong and may not hold in practical data. ", "summary_of_the_paper": "his paper proposes a federated learning method to address a specific non-IID challenge over clients, i.e., when some clients are only available during the daytime and others are only available at the night, and these two types of clients' data are drawn from a mixture of two corresponding distributions. They study a multi-branch model, similar to multi-task learning, such that all clients share the same backbone network to extract data representations but a different branch is applied on top of the backbone network for different distributions. Only the clients with the same data distribution use the same branch to produce the final prediction. Under the assumption that the data representation produced by the backbone network follows a mixture of isotropic Gaussian distributions, this paper proposes an EM algorithm to assign clients to different clusters: the E-step determines the (soft) membership of each client belonging to each cluster and assigns it to the branches, while the M-step estimates the mean and variance of the representation for each Gaussian component in the mixture. Some additional heuristics are shown to be important for this method to achieve promising performance, e.g., label smoothing, a temporal prior of the ratio of clients from day and night, linear/cosine/soft schedule of this ratio, one-hot (hard) assignment of clients to clusters, etc. In experiments, they simulate the day-night distribution shift on EMNIST, CIFAR, and Stackflow datasets in FL by splitting each dataset into two parts with different data types. They evaluate the proposed method under different shifting settings (different p) with a few baselines proposed in this paper. They also empirically studied the effects of label smoothing and different period lengths.", "summary_of_the_review": "The problem of FL with periodically shifting distributions is well-motivated by practical needs. The idea of tackling the non-IID problem in this case by assigning different clients to different branches also makes sense. However, this paper then reduces the problem to FL with client clustering, which is a crowded topic with several similar methods published recently. Given these works, the EM algorithm in this paper is not novel. Moreover, there are a great number of works in the FL community studying the statistical heterogeneity problem and they are all likely to be directly applied here since the day-night shifting is a special case of it. However, the experimental comparison ignores all of them and only compares to baselines developed(modified) in this paper, which fails to provide convincing evidence of the advantage of the proposed method over existing heterogeneous/personalized FL methods. No real day-night periodic shifting exists in any of the three datasets. Knowing the shifting parameter p is somehow cheating. The proposed heuristics are not sufficiently justified whether they are compatible with the EM algorithm and the Gaussian mixture model. Hence, I think there are currently many non-trivial problems that need to be fixed for this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "APxA0IN8N4E", "reviewer": "Reviewer_JqZ9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n\n- Well written and clear paper, with a good motivation (fig 1)\n- The proposed method is novel to the best of my knowledge, and provides an empirical solution to the problem initially tackled by Eichner et al.\n- Although the authors focus on the case of 2 components in the experiments, the method is quite generic and could be applied in different settings\n- The experiments are run on different datasets, with multiple relevant ablation baselines, and the results demonstrate the impact of the proposed method as well as the resilience to errors in modelisation (e.g. cosine instead of linear)\n\n# Weaknesses\n\n- In the experiments, the baselines are not that natural insofar as they also have multiple heads. I think it would also be relevant to compare to a network with a single branch.\n- Fig 2 is difficult to read due to the large number of curves in each plot: I suggest adding different markers to make it more colorblind-friendly.\n- The temporal prior is not learnt but fixed. Do you think it would be possible to fit it as well from data?\n- Only 3 seeds are used: in addition to the means and stds, it would be great to provide the individual points, or, alternatively, to use more seeds.", "summary_of_the_paper": "This paper investigates the problem of periodic client distribution shift in cross-device FL, as previously investiged by Eichner et al (semi-cyclic SGD). The authors propose to split the NN model in a shared part, with different heads for each component. A GMM model is added on top of the output of the shared part to allocate each client to a given head during training, with a temporal prior for the shift between classes. This GMM is fit via a federated EM scheme, in addition to the federation of the NN weights themselves. Experiments are conducted on EMNIST, CIFAR10/100 and StackOverflow datasets, for 2 components. The results demonstrate the superiority of the proposed method with respect to baseline approaches, including the addition of the temporal prior.", "summary_of_the_review": "Given the technical contributions as well as the well-conducted experiments and its overall quality, I am in favour of accepting this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper considers FL with periodically shifting distributions, which is a very relevant and timely research question in the area of federated learning, and learning under distributions shifts. The paper proposed an interesting unsupervised way to learn grouping clients into different branches during training, using a federated version of the EM algorithm. Overall the paper contains several solid contributions in some novel combinations, but remained borderline in terms over overall scores. While reviewers were generally positive about the approach, technical soundness and the importance of the question, still some concerns remained.

Concerns included on the level of novelty relative to several recent similar related FL works, and them being included as baselines. Several of these are now discussed in the rebuttal and revisions, but not all in sufficient depth. While the datasets used seem to offer sufficiently hard task from the split between the 'day' and 'night' distribution, several questions were raised if the treatment of priors is realistic enough. This includes the question of fair hyperparameter tuning with respect to the temporal priors, as well as potential misspecification of the same, towards a more principled treatment of the priors. The authors have answered several of the concerns in the revision, and have added more baseline comparisons, raising the paper narrowly above the acceptance bar in my assessment.

Time-wise, the very related paper Marfoq et al 2021 "Federated Multi-Task Learning under a Mixture of Distributions" seems to have been available 6 weeks before ICLR deadline. We thank the authors for having included it in discussion and experiments, but the discussion of related contributions needs to be expanded (main difference seems to be supervised vs unsupervised group assignment).

My impression is these points can be addressed in a camera ready version, and I hope the detailed feedback here by all reviewers below will be incorporated.
Decision: Accept (Poster)