Year: 2021

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 changes between daytime and nighttime modes, and find this intuitive model to better match the observations in practical federated learning systems. We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (FedTEM), which jointly learns a mixture model to infer the mode of each client, while training a network with multiple light-weight branches specializing at different modes. Experiments for image classification on EMNIST and CIFAR datasets, and next word prediction on the Stack Overflow dataset show that the proposed algorithm can effectively mitigate the impact of the distribution shift and significantly improve the final model performance.

text: 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 changes between daytime and nighttime modes, and find this intuitive model to better match the observations in practical federated learning systems. We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (FedTEM), which jointly learns a mixture model to infer the mode of each client, while training a network with multiple light-weight branches specializing at different modes. Experiments for image classification on EMNIST and CIFAR datasets, and next word prediction on the Stack Overflow dataset show that the proposed algorithm can effectively mitigate the impact of the distribution shift and significantly improve the final model performance.

heading: 1 INTRODUCTION
text: In Federated Learning (FL), many clients collaboratively train a machine learning model with decentralized 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 to 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 applications including finance, health, digital assistance and personalized recommendations (see a few recent surveys (Yang et al., 2019; Kairouz et al., 2019; Li et al., 2020a; Lim et al., 2020; Wang et al., 2021)). Specifically, cross-device FL has been used in practice to boost utility and privacy for applications such as next word prediction (Hard 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 (Granqvist et al., 2020; Hard et al., 2020).
A typical communication round of FL starts with a server broadcasting a global model to clients. Clients then perform local computation on private data and only send back aggregated model updates. 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; Paulik et al., 2021), clients can only participate when the local criteria is met, such as when mobile devices are charging, idle, and connected to an unmetered network. For the server, clients that satisfy their local criteria and participate training at different times of the day are usually from different time zones that can have significant cultural differences, which cause a periodically shifting data distribution that may degrade the training stability and final model performance (Yang et al., 2018; Eichner et al., 2019). For centralized systems where client data can be collected, such a problem may be mitigated by caching and uniformly sampling from cached data. However, due to the privacy and system constraints (Wang et al., 2021), the orchestrator (server) in FL systems is not allowed to collect the raw user data, and must deal with such non-IID, heterogeneous data distribution.
To our knowledge, there are only a few previous works (Eichner et al., 2019; Ding et al., 2020) discussing periodical distribution shift of client population in federated learning. Both 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 to and only use the corresponding model of the group. 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 it is hard to decide the correct group for these clients. (2) It maintains a 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 a specific time of a day is unintuitive in practice. (Ding et al., 2020) is a variant of semi-cyclic SGD that inherits these issues. We provide further discussions of related works in Appendix A.5.
In this paper, we study periodical distribution shift of clients, and make the following contributions:
1. We revisit the periodical distribution shift. Instead of adoopting the block-cyclic structure (Eichner et al., 2019), we assume a smooth transition between the day mode and night 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 a multi-branch network and a Gaussian mixture 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., 2021), the feature-based mixture model does not rely on labelled data, and can be easily applied can be easily applied for inference on new clients.
3. We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (FEDTEM) to train the multi-branch networks. 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 passes from the night mode to day mode, and vice versa. We use such temporal prior to guide the EM algorithm in the M step to learn parameters that better distinguish the two modes. By exploiting such temporal prior, we can train models that are even more accurate than models trained with uniformly sampled clients.
4. We provide simulations of 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 multibranch networks trained by FEDTEM outperform the distribution-oblivious baselines by a large margin: 3-5% on EMNIST, 2-14% 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 proposed 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.

heading: 2 MODELING PERIODICAL DISTRIBUTION SHIFT
text: FL setting. We consider the federated learning algorithm for a set of clients I , and the i-th client has data samples Di. The clients are heterogeneous, while the data on each client are IID (independent and identically distributed) sampled from its own distribution. We minimize the expected loss on all clients,
minimize w L(w) = ∑ i∈I piLi(w), where Li(w) = 1 |Di| ∑ ξ∈Di `(w, ξ), and ∑ i∈I pi = 1, (1)
where pi is the weight (probability) of client i, and ξ = (x,y) is a training sample pair of data and label on a client. By setting pi = |Di|/ ∑ j∈I |Dj |, the federated training loss recovers the empirical risk minimization (ERM) objective on all client samples. For simplicity, we abuse notation and use x ∈ Di to denote a training sample (without label) for client i. Periodic distribution shift. A subset of clients I ′(t) ⊂ I(t) are available for training in a communication round t. I(t) periodically changes with most of the clients from a daytime (nighttime)
client group in midday (midnight). Figure 1 (left) shows that the training loss in a cross-device FL system has daily oscillation. Such oscillations was also observed in (Yang et al., 2018), and Yang et al. (2018) conjecture it is due to the domain differences between clients from different time zones and cultural backgrounds. Eichner et al. (2019) study the block-cyclic structure, where training takes place over T rounds each day, and clients are from the day mode and night mode alternately, each last continuously for T/2 rounds. We simulate the training curves of block-cyclic structure in Figure 1 (right), and observe it is different from the (left) curves from real FL systems.
Smooth transition. We also assume the distribution changes periodically with a period of T . Unlike (Eichner et al., 2019), we assume clients at round t are a mixture of daytime clients and nighttime clients, denoted as I1(t) and I2(t), respectively. Intuitively, since the available population is usually large around the clock, the population distribution of available clients will gradually shift, rather than abruptly jumping from one mode to another as in the block-cyclic structure. To better approximate the behavior in practice, we assume that in each period, clients come from the day mode I1(t) with a probability q(t) that varies smoothly between 0 and 1. Specifically, we simulate I1(t) and I2(t) with two disjoint sets of clients with different data distributions, and define q : R+ → [0, 1] to be periodic function with a period of T . At each round t, we sample the clients from the following distribution
P (i ∈ I1(t)) = q(t), P (i ∈ I2(t)) = 1− q(t). (2)
We consider two types of functions for q(t): Linear (L) and Cosine (C), each further parameterized by an exponent factor p > 0 to control smoothness of the transition,
qL,p(t) = ∣∣∣∣2 t mod TT − 1 ∣∣∣∣p , qC,p(t) = [12(cos (2πt/T ) + 1) ]p . (3)
We visualize the transition probability q(t) in Figure 7 of the appendix. When p < 1, more daytime clients are available in T rounds, and when p > 1, more nighttime clients are available. This simulates 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 qL,1(t) to control the probability for sampling from I1(t), which more accurately approximates the curves in real FL systems. All three curves show that training a single model around the clock achieves reasonable performance on both modes (much better than random) 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, which argues that training a single model on both modes is not optimal, but does not match our observations. Another important observation is that the training accuracy reaches its minima (maxima) when the clients are most biased towards the day mode or night mode, e.g., round 1024 and 1152 in the middle figure, from which we can infer the peak momentum of daytime clients and nighttime clients in practice and use it to define a strong prior in the proposed FedTEM algorithm.

heading: 3 MITIGATING PERIODICAL DISTRIBUTION SHIFT
text: In this section, we propose FEDTEM to tackle the periodical distribution shift, in which we joint train a multi-branch network and a Gaussian mixture model (GMM). The GMM can infer latent variables that determine a participating client is from daytime mode or nighttime mode. The performance of the inference is enhanced by priors about the probability of daytime (nighttime) clients at the current round. A multi-branch network is simultaneously trained on the participating clients based on the inferred latent variables, with each branch corresponding to a mode. We will discuss how FEDTEM is applied in cross-device FL in detail, where server and client will only exchange information once in a communication round, and clients do not maintain local states.
Multi-branch network. We consider applications in which clients are training a model with common input and output space, despite the distribution shift of clients. Though there exists distribution shift, training a model that shares common features by leveraging data from all clients can enable representation learning with better data efficiency. For vision tasks, for example, a common model can use data across domains to learn common low-level features, while for language tasks it helps with learning common embedding and grammatical rules from the context. To alleviate the communication overhead, we adopt the weight-sharing strategy from multi-task learning (Caruana, 1997) to train a multi-branch network with shared feature representation layers f(wf ,x) and specialized prediction branches gk(wk, f(wf ,x)) for clients from mode k. In this paper, we set each of the prediction branches to be a single linear layer. A k-branch network is potentially more communication efficient and data efficient than having k copies of models.

heading: 3.1 LEARNING THE MIXTURE MODEL WITH TEMPORAL PRIORS
text: Algorithm 1 FEDTEM: Federated EM with Temporal Prior (Training) 1: Input: A stream of clients I(t) with periodical distribution shift; Number of communication rounds N ;
Number of rounds per day T . 2: Output: Network parameters wN+1 = (wN+1f ,w N+1 1 ,w N+1 2 ), mixture model parameters θ N+1 =
(µN+1,σN+1,πN+1). 3: for t = 1 to N do 4: A set of m clients I′(t) ⊂ I(t) is available for training; 5: Server broadcasts parameters of the networkwt and the GMM (µt,σt,πt) to I′(t); 6: for clients i ∈ I′(t) in parallel do 7: Estimate πi on Di, and choose the branch k∗i = arg maxk π∗ik; . see Eq. 5 8: Given k∗i , run local updates for the network by optimizing Eq. 6 and get w̃ t+1 i ; 9: On Di, compute the MLE π̄∗i and all feasible Gaussian parameters (µ̃t+1i , σ̃ t+1 i ); . see Eq. 8,11
10: Server aggregates the network and GMM params {w̃t+1i , π̄ ∗ i , µ̃ t+1 i , σ̃ t+1 i |i ∈ I ′(t)} from the clients; 11: Update the GMM parameters to (µt+1,σt+1,πt+1) with the temporal prior on π̄∗i ; . see Eq. 12 12: Update network parameters towt+1 with the prescribed server optimizer.
We learn a mixture model together with the multi-branch network to infer the mode and select the corresponding branch for prediction of each client. Similar joint learning strategies are applied in (Ghosh et al., 2020; Marfoq et al., 2021) to learn a clustering or mixture model for selection from multiple networks, but both works have the impractical assumption that labeled data is available for all clients. In addition, (Marfoq et al., 2021) also assumes stateful clients, which can suffer from staleness in cross-device setting (Wang et al., 2021). Our approach has three key differences: (1) our k-branch network is more communication and data efficient compared to previous k networks, (2) our model selects the branches based on the features instead of the loss, and therefore does not need labeled data to select branches for new clients; (3) we propose the temporal prior to tackle periodic distribution shift when learning the mixture model.
We define discrete latent variables z and ζ to represent which mode a sample and a client comes from, respectively, which will be inferred by our model. Our model also assumes all samples from the same client are from the same mode during training, so the prior P (ζ) = P (z), but our model can still be used to estimate the posterior for each sample during inference. For convenience, we also introduce a vector π with each entry πk = P (z = k) to denote the prior distribution of z. We model P (x|z = k) as a Gaussian distribution N (f(w,x)|µk,σ2k) on the feature representation f(w,x). For efficiency, we assume and learn a diagonal covariance σ2k for the Gaussian models. We use θ to
denote the parameter of the GMM for convenience. Algorithm 1 summarizes the training process, and the details of each step are provided in the following sections.

heading: 3.1.1 MODELING THE CLIENT DISTRIBUTIONS AND SELECTING BRANCHES FOR TRAINING
text: During training, on each client, before proceeding to the local update steps, we estimate the probabilities of client i coming from each mode k given its training setDi, and then select the corresponding branch. We assume all samples on one client are from the same mode, so P (ζ) = P (z). There are two options for inferring the modes
1 Given the current mixture model, estimate P (ζ|Di) from P (x|z) of all x ∈ Di; 2 Given the training samples on the client, find the Maximum Likelihood Estimation (MLE)
of P (z) for the client and use it for P (ζ).
Option 1 exploits the current parameters of the mixture model to infer the mode of each client. With the assumption that samples in Di are independent, from Bayes’ theorem, we can get P (z|Di) as
P (ζ = k|Di) = πk ∏ x∈Di P (x|z = k)∑K
j=1 πj ∏ x∈Di P (x|z = j) . (4)
Intuitively, since the client distribution is constantly shifting, we should not expect the prior P (ζ) to be unchanged from the last round. Option 2 re-estimates the prior P (z) of the mixture model on the current client through MLE and use it as P (ζ|Di). As shown in Figure 6, we find this achieves better empirical results than Option 1. Specifically, we use π∗i , the MLE of P (z) on the data of client i, to select the branch for training:
π∗ik = 1 |Di| ∑ x∈Di P (z = k|x) = 1 |Di| ∑ x∈Di πkN (f(wt,x)|µtk,σtk)∑K j=1 πjN (f(wt,x)|µtj ,σtj) , (5)
where π∗ik is the k-th dimension of π ∗ i , corresponding to the MLE of P (z = k). For completeness, we give the derivation of Eq. 5 in Appendix A.1.
After estimating the distribution of ζ, we can either greedily select the branch k∗i corresponding to the maximum probability, or we can sample from this distribution. Our experiments in Figure 6 show that the greedy approach is more effective. Since our experiments only consider the case of two modes, let k̄∗i be the other, less-likely branch. On each client, we train the model to optimize the following objective
minimize w
Li(w) = 1 |Di| ∑ ξ∈Di `CE(gk∗i (w,x),y) + λ`CE(gk̄∗i (w,x), s( ,y)), (6)
where λ > 0 is the regularization strength, s( ,y) = 1n + (1 − )y is the label smoothing function with 0 ≤ ≤ 1 to determine the amount of label smoothing. The regularization enables the other branch to be updated jointly with the feature extractor to prevent mismatch. Meanwhile, it encourages learning distinctive features for the two clusters through label smoothing: the other branch treats samples from the current mode as outliers and is trained to become less certain on these samples.

heading: 3.1.2 UPDATING THE MIXTURE MODEL
text: After the local updates of the network parameters on client i with the chosen branch as described in Section 3.1.1, we obtain the updated network parameters wt+1i , and update the parameters of the GMM based on wt+1i . In this way, the GMM is synchronized with the latest network parameters and ready for inference immediately after training. The GMM parameters are usually estimated with the Expectation-Maximization (EM) algorithm (Dempster et al., 1977). In the federated setting, this can be achieved privately by first running EM on each client i to get the locally updated GMM parameters (µ̃t+1i , σ̃ t+1 i , π̃ t+1 i ), and then aggregating on the server.
Further, to target the temporal distribution shift, we integrate our priors q̃(t) into the aggregation step on the server, so that a ratio of q̃(t) clients will be used to update (µt+11 ,σ t+1 1 , π t+1 1 ) for mode
1, while the remaining clients will only update (µt+12 ,σ t+1 2 , π t+1 2 ) for mode 2. The server decides which clients are from mode 1 and which are from mode 2 by their π̃t+1i , since π̃ t+1 i estimates the probability of an arbitrary sample from client i to come from each mode. We give details of the steps below.
E step. For each client i ∈ I(t), compute the posterior γik(x) for each sample x to infer the probability of x coming from mode k
γik(x) = P (z = k|x) = πtkN (f(w t+1 i ,x)|µtk,σtk)∑K
j=1 π t jN (f(w t+1 i ,x)|µtj ,σtj)
. (7)
Temporal prior. We integrate the temporal prior into the M step, so that we can estimate GMM parameters that better capture the differences in the modes behind the shifting distributions. The temporal prior q̃(t) is a rough estimate for the ratio of clients coming from the day mode. From the observations in Section 2, we can locate the time when q(t) is most likely to be 0 or 1 by observing when the minima and maxima occur from the train curve. In between these minima and maxima, we consider three types of q̃(t) in our current experiments: 1) Linear: q̃(t) = qL,1(t); 2) Cosine: q̃(t) = qC,1(t); 3) Soft: see Appendix A.2.
Conditioned on q̃(t), our goal is to estimate the posterior P (ζ|Di,θt, q̃(t)) for each client to decide which mode they are from. To achieve this, similar as in Section 3.1.1, we first obtain the MLE of the prior on Di using the updated network parameters wt+1i as
π̄∗ik = 1 |Di| ∑ x∈Di γik(x), (8)
where γik(x) comes from the E step in Eq. 7. Then, we sort the prior probability π̄∗i1 of mode 1 for all sampled clients i ∈ I ′(t), and assign bq̃(t) · |I ′(t)|+ 12c clients with the highest π̄ ∗ i1 to mode 1, denoted as Ĩ ′1(t). The remaining clients, denoted as Ĩ ′2(t), are assigned to mode 2. As a result, for i1 ∈ Ĩ ′1(t), P (ζ = 1|Di,θt, q̃(t)) = 1, while for i2 ∈ Ĩ ′2(t), P (ζ = 2|Di,θt, q̃(t)) = 1. Note it is possible to design soft assignments where the posteriors are not one-hot, but this one-hot assignment has smaller communication overhead, as we will discuss in the M steps below.
M step (client). In our federated setting, the M step happens on both the server and the clients. It is first executed locally on each client i, where its objective is to find the GMM parameters to maximize the expectation of the log likelihood of data, weighted by the time-varying client-level posterior P (ζ = k|Di,θt, q̃(t))
θ̃t+1i = argmax θ K∑ k=1 P (ζ = k|Di,θt, q̃(t)) logP (Di, ζ = k|θ, q̃(t)). (9)
For each client, q̃(t) affects the posterior distribution P (ζ|Di, q̃(t)), and the effect depends on other clients. Therefore, the distributions in Eq. 9 should be estimated on the server. However, we cannot send client data to the server, and it is almost impossible in practice to have multiple communication rounds between the server and client (Bonawitz et al., 2019; Wang et al., 2021) so that the server can send the re-estimated prior to the clients. A viable approach is to restrict P (ζ|Di, q̃(t)) to have a finite number of possible values, so that the client can solve Eq. 9 for all possible P (ζ|Di, q̃(t)) and send all feasible solutions to the server. To make q̃(t) effective, we should set a minimum of K different values for P (ζ|Di, q̃(t)) so that each mode corresponds to a different value. To learn more distinctive representations for the modes, we let P (ζ = k|Di, q̃(t)) = 1 if the client is estimated to come from mode k, so P (ζ = j|Di, q̃(t)) = 0 for all j 6= k, and Eq. 9 becomes
θ̃t+1i = argmax θ logP (Di, ζ = k|θ, q̃(t))
= argmax θ logP (ζ = k|θ, q̃(t)) + ∑ x∈Di logP (x|ζ = k,θ, q̃(t))
= argmax θ πk − ∑ x∈Di  d∑ j=1 ([f(wt+1i ,x)]j − µkj)2 2σ2kj + log σkj √ 2π  (10)
where d is the dimension of the features, µkj , σkj denotes the j-th dimension of µk,σk, and we use similar notation for πk and [f(wt+1i ,x)]j . The last equation comes from our modeling assumption that all x ∈ Di are from the same mode. From Eq. 10, we can see the optimal solution is to keep parameters for other modes unchanged, while only updating them for the k-th mode as
µ̃t+1ik = 1 |Di| ∑ x∈Di f(wt+1i ,x), [σ̃ t+1 ik ] 2 = 1 |Di| ∑ x∼Di (f(wt+1i ,x)− µ̃ t+1 ik ) 2, (11)
In this way, we can pack the all the feasible µ̃t+1ik , [σ̃ t+1 ik ] 2 into matrices that have the same size as one set of GMM parameters, which does not increase the communication cost.
M step (server). The server collects the MLE of the priors {π̄∗i |i ∈ I ′(t)} (Eq. 8) and all the feasible solutions {(µ̃t+1i , σ̃ t+1 i )|i ∈ I ′(t)} from the clients. Then, with the temporal prior, it divides the clients into two sets I ′1(t) and I ′2(t) as described before. Similar as FedAvg (McMahan et al., 2017), it updates the parameters for each mode k weighted by the number of samples from each client as
µt+1k = ∑
i∈Ĩ′k(t)
|Di| M tk µ̃t+1ik , [σ t+1 k ]
2 = ∑
i∈Ĩ′k(t)
|Di| M tk [σ̃t+1ik ] 2, πt+1k = M tk∑K j=1M t j , (12)
where M tk = ∑
i∈I′k(t) |Di| is the total number of samples of clients assigned to mode k. Note
the resulting σt+1 is not guaranteed to be an unbiased estimate, but we found this process to give satisfying results in practice.
Setting the prior. During test, to compare the model performance with established baselines and ensure fairness, we consider a static test set where the number of clients and samples from both modes are roughly the same. Therefore, we use a uniform distribution for the priors on the test set. This does not make the GMM worse if the underlying P (x|z) is the same for the training and test sets, and our estimate of P (x|z) is accurate. During training, we estimate the prior πt in every step to match the shifting distribution. However, we find it beneficial to use a running average of πt+1k = βπ t k +(1−β)πkM tk/( ∑K j=1M t j ) during training, which achieves better results than setting β = 0 or using uniform prior in our preliminary experiments.

heading: 4 EXPERIMENTS
text: Dataset. We consider two image classification datasets, EMNIST and CIFAR, and one next word prediction task on Stack Overflow (SO). The split of the day mode (I1) and night mode (I2), and other statistics, are shown in Table 1. On SO, we use a vocabulary size of 10K and report the test accuracy without special tokens. We truncate each sentence to have no more than 20 tokens.
Architecture of the multi-branch networks. For image classification, on EMNIST, we train LeNet with 2 Conv layers and 2 FC layers, while on CIFAR, we train a ResNet-18 with the Batch Normalization replaced by Group Normalization (Wu & He, 2018) for stability in federated learning. We use the last FC layer of these convolutional networks as the multi-branch part and share all the remaining layers for all modes. We use the output from the shared feature extractor for GMM, which are 128 and 512 dimensional respectively. On SO, we train a single-layer LSTM (Hochreiter & Schmidhuber, 1997) with a hidden size of 670 and embedding size of 96. To alleviate the communication overhead and prevent overfitting, we define the branches to be the last projection layer before the final prediction layer and share the remaining weights for both modes, which is 15x smaller than the final prediction layer. We concatenate the average of all tokens’ features from both branches for GMM, which has dimensionality 192. With these configurations, the models on EMNIST, CIFAR and SO increase the communication cost by 0.7%, 0.5% and 1.6%, respectively.
Hyperparameters. We use FedAdam (Reddi et al., 2021) as the optimizer. Table 2 shows the training hyperparameters. We use T = 256 for the majority of our results, and we evaluate both
linear and cosine distribution shifts under p ∈ {0.1, 0.25, 0.5, 1, 2, 4, 10}. We use β = 0.99 for the moving average of πk in all experiments. For FEDTEM and Min Loss, we also do a grid search on the label smoothing parameters and λ. For each set of hyperparameters, we run 3 experiments with different random seeds and report their mean and standard error.
Methods to compare. First, we provide ablation studies on the temporal prior to quantify its efficacy, and see which prior (Linear, Cosine, or Soft) works better in practice. Note for Linear and Cosine, we always use the same periodical linear or cosine function with p = 1 under all types and p’s for the distribution shift. The detailed comparisons are given in Appendix A.3 and Figure 4, from which 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 weakened assumptions, Soft prior is able to fit the distribution shift and improves the results. With these observations, we use the linear prior on EMNIST and CIFAR (FedTEM (linear)) and the cosine prior on SO (FedTEM (cosine)), to compare with other methods. Results are shown in Figure 2. For fair comparisons, we adapt all methods into our settings to maintain two important conditions unless otherwise stated: 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.6, we also give results of training single-branch models in these two settings, which obtained almost the same results. For other methods addressing data heterogeneity, we consider: 1) “SCSGD”, where similar as Semi-Cyclic SGD (Eichner 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 direct identifier 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 a uniform distribution over all labels. 2) T-shot K-means, which is an enhancement of the one-shot K-means (Dennis et al., 2021), collecting the cluster centers on raw 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 as IFCA (Ghosh et al., 2020), we choose the 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 as 1). Besides, we have also applied the same label smoothing regularization (Eq. 6) to IFCA, which improves the results. 4) FedEM (Marfoq et al., 2021), where we use the same algorithm on participating clients during training. For testing, we let it “cheat” and assume all test samples are labeled, so that it can execute its modified EM steps based on the loss to select the branches.
Main Results The best results of all compared methods are shown in Figure 2. By comparing results of Vanilla and NDS, we find the temporal distribution shift indeed causes issues for normal
training, causing the worst-case decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and SO. Surprisingly, with FEDTEM, the accuracy can be even higher than models trained in the NDS setting in most cases, indicating FEDTEM obtains more distinctive feature representations by learning to separate the two modes in the feature space and learning specialized prediction branches. The improvement on the strong NDS baseline can be as high as 4%, 4% and 0.45% on the three datasets. FEDTEM is not better than NDS when p deviates too much from 1, but the data distribution is extremely skewed in such settings, where FEDTEM still improves over most other methods in most cases. For other 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. Also note it is not straightforward to do clustering on raw data of SO so we have not implemented it. Min Loss (IFCA) alleviates the drop and sometimes catches up with NDS, but cannot achieve improvements over NDS. By contrast, results of FedEM is the worst even when it is “cheating” on the test clients. This is probably due to the staleness of the priors, since we simulate the on-device setting on a large population where clients only participate training for a few rounds on EMNIST and CIFAR, and at most one round on SO.
Effect of label smoothing regularization. Shown in Figure 3. 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 on 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. During evaluation, for image classification, we find it beneficial in most cases 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 5. In practice, similar approach can be realized through caching. We find Min Loss does not show much difference when a similar approach is applied, where branches are chosen by the average of persample certainty. The baseline model even gets worse performance with this approach, as shown in Figure 5, indicating our model is much better at distinguishing two modes in expectation. However, this approach does not show improvements for FEDTEM on SO.
Effect of T . As shown in Figure 8, for FEDTEM with linear prior on EMNIST, a 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 a frequent but more balanced way.

heading: 5 CONCLUSIONS
text: In this paper, we demonstrated that the smooth transition model for simulating the periodical distribution shift better matches practical FL systems, and developed FEDTEM algorithm to train multibranch networks to tackle the distribution shift. FEDTEM incorporated priors of the temporal distribution shifts into learning a GMM to guide the network to select corresponding branches for clients of different modes. The GMM is defined in the feature space and does not require labelled data for inference, hence is ready to use for new test clients. The branches are light-weight with little communication overhead, and the model demonstrates significant improvements in test accuracy on the benchmark datasets. FEDTEM also satisfies many other real-world constraints like single round communication and stateless clients (Bonawitz et al., 2019; Paulik et al., 2021; Wang et al., 2021), which makes it ready to be deployed in practical FL systems.

heading: 6 REPRODUCIBILITY STATEMENT
text: For reproducibility, we have reported the mean and standard error for 3 experiments with different random seeds for all 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.1 We have listed our hyperparameters in Table 2. We have given the pseudo code in Algorithm 1 and the detailed formulas and derivations of our algorithm in Section 3.1. We are committed to opensource our implementation upon publication.

heading: A APPENDIX

heading: A.1 DERIVATION OF OPTION 2
text: We consider maximizing the expectation of the complete-data log likelihood logP (Di,Z|θt+1) under the posterior distribution P (Z|Di,θt), whereZ denotes the collection of latent variables associated with each of the sample x ∈ Di. Formally, it is solving the following constrained optimization problem
maximize θ EZ∼P (Z|Di,θt) [logP (Di,Z|θ)] ,
subject to K∑ j=1 πj = 1, πj ≥ 0, for all 1 ≤ j ≤ K. (13)
We introduce an indicator function 1[z=k] which is 1 if z = k or 0 otherwise. In this way, the log likelihood can be represented as
logP (Di,Z|θ) = log |Di|∏ n=1 K∏ k=1 [P (xn, zn|θ)]1[z=k]
= |Di|∑ n=1 K∑ k=1 1[zn=k] [log πk + logP (xn|zn,θ)] .
(14)
Meanwhile,
EZ∼P (Z|Di,θt) [ 1[zn=k] ] = P (zn = k|xn,θt) =
πtkP (xn|zn = k,θt)∑K j=1 π t jP (xn|zn = j,θt) . (15)
Plug Eq. 14 and Eq. 15 back into Eq. 13, and ignore terms that do not depend on πt+1, we find we are solving the following problem for πt+1
maximize π |Di|∑ n=1 K∑ k=1 P (zn = k|xn,θt) log πk,
subject to K∑ j=1 πj = 1, πj ≥ 0, for all 1 ≤ j ≤ K.
(16)
The Lagrangian multiplier of this problem is
h(π, λ) = |Di|∑ n=1 K∑ k=1 P (zn = k|xn,θt) log πk + λ(1− K∑ j=1 πj). (17)
Let the derivative w.r.t. πk be 0,
∂h
∂πk = |Di|∑ n=1 P (zn = k|xn,θt) 1 πk − λ = 0. (18)
Multiply both sides of Eq. 18 with πk, sum over 1 ≤ k ≤ K, and apply the constraint that∑K k=1 πk = 1, we have
λ = |Di|∑ n=1 K∑ k=1 P (zn = k|xn,θt) = |Di|. (19)
Plug this back into Eq. 18, we get the equation as desired.
πtk = 1
|Di| |Di|∑ n=1 P (zn = k|xn,θt) (20)

heading: A.2 THE SOFT PRIOR
text: The soft prior q̃S(t) gives strong signals only at time steps where we believe q(t) = 0 or q(t) = 1. In between, it only requires the ratio to be non-increasing or non-decreasing, where the ratio is estimated using the posterior from Eq. 7 as
q̃′S(t) = |{x|x ∈ Di, γi1(x) > γi2(x)}|
|Di| , (21)
i.e., the ratio of samples whose posterior has higher probability on 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
q̃S(t) =  1, if t mod T = 1 0, if t mod T = T2 min(q̃S(t− 1), q̃′S(t)), if 1 < t mod T < T2 max(q̃S(t− 1), q̃′S(t)). if t mod T > T2
(22)

heading: A.3 ABLATION STUDIES: THE EFFECT OF TEMPORAL PRIOR
text: To quantify the effect of temporal prior on training, we compare four versions of our algorithm: 1) EM Only: only apply the EM part of our algorithm without the temporal priors, where the only difference from FEDTEM is that the server, in the M step, greedily uses client i to update parameters of mode k∗i = arg maxk π̄ ∗ ik (see Eq. 8 and Eq. 12); 2) FedTEM (linear): our algorithm where the temporal prior is the periodic linear function P (i ∈ I1(t)) = qL,1(t); 3) FedTEM (cosine): our algorithm where the temporal prior is the cosine function P (i ∈ I1(t)) = qC,1(t); 4) FedTEM (soft): our algorithm with the soft temporal prior introduced in Section A.2. The results are shown in Figure 4. 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-decreasing within certain intervals, but still improves the results.
A.4 VERIFYING THE EFFECT OF MLE FOR EVALUATION
To quantize the effect the batch-level MLE, we compare the results of our method with or without MLE on EMNIST and CIFAR. As shown in Figure 6, this only decreases the accuracy of the baseline model, indicating our model is much better at distinguishing the two modes in expectation.

heading: A.5 ADDITIONAL RELATED WORKS
text: As mentioned above, Semi-cyclic SGD (Eichner et al., 2019) and (Ding et al., 2020) 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 FL, and various methods have been proposed to address the distribution difference among clients. Notably, FedProx (Li et al., 2020b) applies proximal regularizer when perform client updates; SCAFFOLD (Karimireddy et al., 2020) uses control variates as local states on clients for variance reduction; FedPA (Al-Shedivat et al., 2021) provides a Bayesian view and adopts posterior sampling; FedGen (Zhu et al., 2021) learns a generator that can be shared among clients; FedRobust (Reisizadeh et al., 2020) 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 are introduced 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., 2021). 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 heterogeneity in a few recent surveys (Kairouz et al., 2019; Li et al., 2020a; Wang et al., 2021). 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 available client population is always changing due to the periodical distribution shift. We also do 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 on personalization, mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting. Dennis et al. (2021) propose a one-shot clustering approach based on 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., 2020) 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 personalized 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, FedEM (Marfoq et al., 2021), which proposes a Federated EM algorithm to learn a 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 a different prior distribution for every client, which requires stateful clients and strong control over client sampling. The hierarchical or bi-partitioning clustering process of (Briggs et al., 2020; Fu et al., 2021) requires all
clients to be available for clustering simultaneously, 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 addition, 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 communication cost. (Xie et al., 2021) maintains one set of weights on each client while maintaining multiple weights as cluster centers on the server. It takes two communication rounds for each weight update, which can be impractical. It remains unclear if this mechanism can be generalized to unseen clients. (Marfoq et al., 2021) trains a 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 practical for on-device setting. (Duan et al., 2021) clusters the clients based on the similarities between their gradients at the initial model weights. This is impractical in our setting since the training population 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, it still requires labeled data. (Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which is not practical for on-device FL since BN has been widely observed to cause instability in on-device FL, and is often replaced by Group Normalization which does not track running statistics (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 a personalized model is preferred for each client: the branchs are either locally stored on clients (Arivazhagan et al., 2019; Liang et al., 2020), or reconstructed based on client data (Singhal et al., 2021). All the aforementioned 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 in the feature space, and the weights of the branches are fixed.

heading: A.6 COMPARING SINGLE-BRANCH AND MULTI-BRANCH BASELINES WITHOUT CLUSTERING
text: 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, we show that the single-branch and multi-branch networks obtain similar results under various settings, with or without distribution shift. Our method still outperforms these two single-branch baselines.

heading: A.7 THE EFFECT OF LOCAL TRAINING SET SIZES
text: When the number of training samples is small, the EM estimates may have high variance, causing misspecifications and resulting in worse performance. Meanwhile, models trained on smaller training sets tend to have worse generalization, so the test accuracy will drop for any method in general. To see which factor has more significant effect on the test accuracy, we compare with the
baselines trained with or without distribution shift in Figure 10, 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.

heading: A.8 STATS OF TRAINING A LANGUAGE MODEL IN A REAL FL SYSTEM
text: In Figure 11, 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.

Review summaries:
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.

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.

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.

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.