Original: D IURNAL OR N OCTURNAL ? F EDERATED L EARNINGOF M ULTI - BRANCH N ETWORKS FROM P ERIODICALLY S HIFTING D ISTRIBUTIONS
Revised: D IURNAL OR N OCTURNAL ? F EDERATED L EARNING FROM P ERIODICALLY S HIFTING D ISTRIBUTIONS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Federated learning has been deployed to train machine learning models from decentralized client data on mobile devices in practice.
Revised: Federated learning has been deployed to train machine learning models from de - centralized client data on mobile devices in practice.

Original: 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.
Revised: 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.

Original: 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 ﬁnd this intuitive model to better match the observations in practical federated learning systems. 
Revised: 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 grad ually shifts between daytime and nighttime modes, and ﬁnd this intuitive model to better match the observations in practical federated learning systems. Furthermore

Original: 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.
Revised: , we propose to jointly train a clustering model and a multi -branch network to allocate lightweight specialized branches to clients from different modes A tem - poral prior is used to signiﬁcantly boost the training performance.

Original: Experiments for image classiﬁcation on EMNIST and CIFAR datasets, and next word prediction on the Stack Overﬂow dataset show that the proposed algorithm can effectively mitigate the impact of the distribution shift and signiﬁcantly improve the ﬁnal model performance.
Revised: Experiments for image classiﬁcation on EMNIST and CIFAR datasets, and next word predic- tion on the Stack Overﬂow dataset show that the proposed algorithm can counter the effects of the distribution shift and signiﬁcantly improve the ﬁnal model performance

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: Initially introduced for decentralized training on mobile devices (McMahan et al., 2017), FL has been widely applied for various different applications including ﬁnance, health, digital assistance and personalized recommendations (see a few recent surveys (Yang et al., 2019; Kairouz et al., 2019;
Revised: Initially introduced for decentralized training on mobile devices (McMahan et al., 2017), FL has been widely applied for various different applications including ﬁnance, health, digital assistance and personalized recommendations (see a few recent surveys (Yang et al., 2019; Kairouz et al., 2019;

Original: Li et al., 2020a; Lim et al., 2020; Wang et al., 2021)).
Revised: Li et al., 2020a; Lim et al., 2020; Wang et al., 2021)).

Original: Speciﬁcally, 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).
Revised: Speciﬁcally, cross-device FL has been used in practice to improve utility and privacy of 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).

Original: A typical communication round of FL starts with a server broadcasting a global model to clients.
Revised: A typical communication round of FL starts with a server broadcasting a global model to clients.

Original: Clients then perform local computation on private data and only send back aggregated model updates.
Revised: Clients then perform local computation on private data and only send back aggregated model updates.

Original: Finally, the server aggregates the client updates and apply them to the global model before beginning the next round.
Revised: Finally, the server aggregates the client updates and apply them to the global model before beginning the next round.

Original: 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.
Revised: 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.

Original: 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 signiﬁcantcultural differences, which  cause a periodically shifting data distribution that may degrade the training stability and ﬁnal model performance (Yang et al., 2018; Eichner et al., 2019).
Revised: 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 signiﬁcant differences, which can cause a periodically shifting data distribution that may degrade the training stability and ﬁnal model performance (Yang et al., 2018; Eichner et al., 2019).

Original: For centralized systems where client data can be collected, such a problem may be mitigated by caching and uniformly sampling from cached data.
Revised: For centralized systems where client data can be collected, such a problem may be mitigated by caching and uniformly sampling from cached data.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Both works assume a block-cyclic structure where daytime clients and nighttime clients alternately participate in training.
Revised: These works assume a block-cyclic structure where daytime clients and nighttime clients alternately participate in training.

Original: Eichner et al.
Revised: Eichner et al.

Original: (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.
Revised: (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.

Revised: By learning separate models for different blocks, they can obtain the same guarantee as the i.i.d.

Original:  However, there are several caveats of semi-cyclic SGD that makes it difﬁcult 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.
Revised: non-cyclic setting under their assumptions. However, there are several caveats of semi-cyclic SGD that makes it difﬁcult 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.

Original: (2) It maintains a version of the full model for each clients group, which potentially increases the communication cost or privacy risk.
Revised: (2) It maintains a version of the full model for each clients group, which potentially increases the communication cost or privacy risk.

Original: The assumption of the abrupt switch from the daytime group to the nighttime group at a speciﬁc time of a day is unintuitive in practice.
Revised: The assumption of the abrupt switch from the daytime group to the nighttime group at a speciﬁc time of a day is unintuitive in practice.

Original: (Ding et al., 2020) is a variant of semi-cyclic SGD that inherits these issues.
Revised: (Ding et al., 2020) is a variant of semi-cyclic SGD that inherits these issues.

Original: We providefurther discussions of  related works  in Appendix A.5.
Revised: We provide discussions of more related works such as heterogeneity, clustering, and multi-branch networks in Appendix A.7.

Original: In this paper, we study periodical distribution shift of clients, and make the following contributions:
Revised: In this paper, we study periodical distribution shift of clients, and make the following contributions:

Original: We revisit the periodical distribution shift.
Revised: We revisit the periodical distribution shift.

Original: 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.
Revised: Instead of adopting the block-cyclic structure (Eichner 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.

Original: We propose to jointly train a multi-branch network and a Gaussian mixture model to select the branch that better ﬁts the client’s distribution based on the feature representations.
Revised: We propose to jointly train a multi-branch network and a clustering model to select the branch that better ﬁts the client’s distribution based on the feature representations.

Original: The lightweight branches for the day and night modes only slightly increase the communication cost, but signiﬁcantly improve the model performance.
Revised: The lightweight branches for the day and night modes only slightly increase the communication cost, but signiﬁcantly improve the model performance.

Original: 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 appliedcan be easily applied for inference on new clients.
Revised: Unlike (Mansour et al., 2020; Ghosh et al., 2020; Marfoq et al., 2021), the feature-based clustering model does not rely on labelled data, and can be easily applied for inference on new clients.

Original: We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution(F ED TEM) to train the multi-branch networks.
Revised: We propose to use the temporal prior of the client distribution to enhance the clustering models.

Original: 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.
Revised: 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 to daytime, and vice versa.

Original: We use such temporal prior to guide the EM algorithm in the M step to learn parameters that better distinguish the two modes.

Original: By exploiting such temporal prior, we can train models that are even more accurate than models trained with uniformly sampled clients.
Revised: By exploiting this prior, we can train models that are even more accurate than models trained with uniformly sampled clients in a conventional federated simulation.

Original: We provide simulations of the distribution shift on three benchmark datasets (EMNIST, CIFAR and Stack Overﬂow) to evaluate the empirical performance of FL algorithms under the periodic distribution shift with smooth transition.
Revised: We provide simulations of the distribution shift on three benchmark datasets (EMNIST, CIFAR and Stack Overﬂow) to evaluate the empirical performance of FL algorithms under the periodic distribution shift with smooth transition.

Original: We perform extensive experiments, where the multibranch networks trained by F ED TEM 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 Overﬂow dataset under various degrees of distribution shifts.
Revised: We perform extensive experiments, where the multibranch networks trained by our method 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 Overﬂow dataset under various degrees of distribution shifts.

Original: 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.
Revised: 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.

Original:  M ODELING P ERIODICAL D ISTRIBUTION S HIFT
Revised:  M ODELING P ERIODICAL D ISTRIBUTION S HIFT

Original: FL setting.
Revised: FL setting.

Original: We considerthe federated learning algorithm for a set of clients I , and the i -th client has datasamples D i . The clients are heterogeneous, while the data on each client are IID (independent and identically distributed)sampled from its own distribution.
Revised: We consider federated learning algorithms for a set of clients I , where the i -th client has data D i sampled IID (independent and identically distributed) from its own distribution, but the distribution of different clients can be heterogeneous.

Original: We minimize the expected loss on all clients, minimize w
Revised: We minimize the expected loss on all clients, minimize w

Original: L ( w ) = (cid:88) i ∈I p
Revised: L ( w ) = (cid:88) i ∈I p

Original: i L i ( w ) , where L i ( w ) =
Revised: i L i ( w ) , where L i ( w ) =

Original: 1 |D i | (cid:88) ξ ∈D i (cid:96) ( w , ξ ) , and (cid:88)
Revised: 1 |D i | (cid:88) ξ ∈D i (cid:96) ( w , ξ ) , and (cid:88)

Original: i ∈I p i
Revised: i ∈I p i

Original: = 1 , (1) where p i is the weight (probability) of client i , and ξ = ( x , y ) is a training sample pair of data and label on a client.
Revised: = 1 , (1) where p i is the weight (probability) of client i , and ξ = ( x , y ) is a training sample pair of data and label on a client.

Original: By setting p i =
Revised: By setting p i =

Original: |D i | / (cid:80) j ∈I |D
Revised: |D i | / (cid:80) j ∈I |D

Original: j | , the federated training loss recovers the empirical risk minimization (ERM) objective on all client samples.
Revised: j | , the federated training loss recovers the empirical risk minimization (ERM) objective on all client samples.

Original: For simplicity, we abuse notation and use x ∈ D i to denote a training sample (without label) for client i .
Revised: For simplicity, we abuse notation and use x ∈ D i to denote a training sample (without label) for client i .

Original: Periodic distribution shift.
Revised: Periodic distribution shift.

Original: A subset of clients I (cid:48) ( t ) ⊂
Revised: A subset of clients I (cid:48) ( t ) ⊂

Original: I ( t ) are available for training in a communication round t .
Revised: I ( t ) are available for training in a communication round t .

Original: I ( t ) periodically changes with most of the clients from a daytime (nighttime) client group in midday (midnight).
Revised: I ( t ) periodically changes with most of the clients from a daytime (nighttime) client group in midday (midnight).

Original: Figure 1 (left) shows that the training loss in a cross-device FL system has daily oscillation.
Revised: Figure 1 (left) shows that the training loss in a cross-device FL system has daily oscillation.

Original: Such oscillations was also observed in (Yang et al., 2018), and Yang et al.
Revised: Such oscillations was also observed by Yang et al.

Original: (2018) conjecture it is due to the domain differences between clients from different time zones and cultural backgrounds.
Revised: (2018), who conjecture it is due to the domain differences between clients from different time zones.

Original: Eichner et al.
Revised: Eichner et al.

Original: (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.
Revised: (2019) study the block-cyclic structure, where the model is trained for T rounds each day, and the clients are from the day mode and night mode alternately, each last continuously for T/ 2 rounds.

Original: 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.
Revised: We plot the training curves of block-cyclic structure in Figure 1 (right), and observe it is different from the (left) curves from a practical FL systems.

Original: Smooth transition.
Revised: Smooth transition.

Original: We also assume the distribution changes periodically with a period of T .
Revised: We also assume the distribution changes periodically with a period of T .

Original: Unlike (Eichner et al., 2019), we assume clients at round t are a mixture of daytime clients and nighttime clients, denoted as I 1 ( t ) and I 2 ( t ) , respectively.
Revised: Unlike (Eichner et al., 2019), we assume clients at round t are a mixture of daytime clients and nighttime clients, denoted as I 1 ( t ) and I 2 ( t ) , respectively.

Original: Intuitively, since the available population is usually large around the clock, the population distribution of available clients will gradually shift, rather than abruptlyjumping from one mode to another as in the block-cyclic structure.
Revised: Intuitively, since the available population is usually large around the clock, the population distribution of available clients should shift gradually, rather than abruptly from one mode to another as in the block-cyclic structure.

Original: To better approximate the behavior in practice, we assume that in each period, clients come from the day mode I 1 ( t ) with a probability q ( t ) that varies smoothly between 0 and 1.
Revised: To better approximate the behavior in practice, we assume that in each period, clients come from the day mode I 1 ( t ) with a probability q ( t ) that varies smoothly between 0 and 1.

Original: Speciﬁcally, we simulate I 1 ( t ) and I 2 ( t ) with two disjoint sets of clients with different data distributions, and deﬁne q : R + → [0 , 1] to be  periodic function with a period of T .
Revised: Speciﬁcally, we simulate I 1 ( t ) and I 2 ( t ) with two disjoint sets of clients with different data distributions, and deﬁne q : R + → [0 , 1] to be a periodic function with a period of T .

Original: At each round t , we sample the clients from the following distribution
Revised: At each round t , we sample the clients from the following distribution

Original: [Equation]
Revised: [Equation]

Original: 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,
Revised: We consider a periodic Linear function (L) and a smooth Cosine function (C) for q ( t ) , each further parameterized by an exponent factor p > 0 to control smoothness of the transition,

Original: [Equation]
Revised: [Equation]

Original: We visualize the transition probability q ( t ) in Figure 7 of the appendix.
Revised: We visualize the transition probability q ( t ) in Figure 7 of the appendix.

Original: When p < 1 , more daytime clients are available in T rounds, and when p > 1 , more nighttime clients are available.
Revised: When p < 1 , more daytime clients are available in T rounds, and when p > 1 , more nighttime clients are available.

Original: This simulates the difference in the number of completed training rounds during daytime and nighttime observed in practice (Yang et al., 2018).
Revised: This can simulate the difference in the number of completed training rounds during daytime and nighttime observed in practice (Yang et al., 2018).

Original: Observation and insight.
Revised: Observation and insight.

Original: Figure 1 (middle) simulates the training curve with q L, 1 ( t ) to control the probability for sampling from I 1 ( t ) , which more accurately approximates the curves in real FL systems.
Revised: Figure 1 (middle) simulates the training curve with q L, 1 ( t ) to control the probability for sampling from I 1 ( t ) , which more accurately approximates the curves from a practical FL system.

Original: 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.
Revised: 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 a large amount of shared weights.

Original: 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.
Revised: Semi-cyclic SGD (Eichner et al., 2019) provides the worst-case guarantees when the domains of each block are unrelated, and argues that training a single model on both modes is not optimal, which is different from our observations.

Original: 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 ﬁgure, from which we can infer the peak momentum of daytime clients and nighttime clients in practice and use it to deﬁne a strong prior in the proposed FedTEM algorithm.
Revised: Another important observation is that the training accuracy reaches its minima (maxima) when the client population is most biased towards the day mode or night mode, e.g., round 1024 and 1152 in the middle ﬁgure, from which we can infer the peak moment of daytime clients and nighttime clients in practice and use it to deﬁne a strong prior to improve the learning process.

Original:  M ITIGATING P ERIODICAL D ISTRIBUTION S HIFT
Revised:  L EARNING FROM P ERIODICAL D ISTRIBUTION S HIFT

Original: In this section, we propose F ED TEM to tackle the periodical distribution shift, in which we joint train amulti-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.
Revised: 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.

Original: 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 onthe participating clients based on the inferred latent variables, with each branch corresponding to a mode.
Revised: 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.

Original: We will discuss how F ED TEM  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.
Revised: We consider two methods for enforcing the temporal prior: F ED TEM (Section 3.1), which is based on a Gaussian Mixture Model (GMM), and F ED TKM ( Section 3.2), which is based on K-means clustering.

Original: Multi-branch network.
Revised: Multi-branch network.

Original: We consider applications in which clients are training a model with common input and output space, despite the distribution shift of clients.
Revised: Eichner et al. (2019) showed the merits of training separate models for different data distributions when the distributions can be identiﬁed during evaluation.

Original: 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 efﬁciency.
Revised: However, training a single model with shared feature extraction layers on all available data can usually improve the data efﬁciency of representation learning.

Original: 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.
Revised: 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.

Original: 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 ( w f , x ) and specialized prediction branches g k ( w k , f ( w f , x )) for clients from mode k .
Revised: To handle the distribution shift while learning shared feature representations and alleviating the communication overhead, we adopt the weight-sharing strategy from multi-task learning (Caruana, 1997) to train a multi-branch network with shared feature extraction layers f ( w f , x ) followed by one of the specialized prediction branches g k ( w k , f ( w f , x )) for clients from each cluster k ( 1 ≤ k ≤ K ).

Original: In this paper, we set eachof the prediction branches to be a single linear layer.
Revised: We set each prediction branch g k to be a single linear layer, which is more communication efﬁcient and data efﬁcient than having K versions of the same model. Temporal prior.

Original: A k -branch network is potentially more communication efﬁcient and data efﬁcient than having k copies of models.
Revised: The temporal prior ˜ q ( t ) is an estimate for the ratio of clients coming from the daytime cluster, q ( t ) . 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 training curve. In between these minima and maxima, we consider three types of ˜ q ( t ) in our current experiments: 1)

Original: 3. L EARNING THE M IXTURE M ODEL WITH T EMPORAL P RIORS
Revised: [Equation] Appendix A.3. Training Objective.

Original: We learn a mixture model together with the multi-branch network to infer the mode andselect 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. 
Revised: 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 k ∗ i be the branch index of client i chosen by the clustering model, and ¯ k ∗ i be the index of the other branch. Each client trains the network with the following objective minimize

Original: In addition, (Marfoq et al., 2021) also assumes stateful clients, which can suffer from staleness in cross-device setting (Wang et al., 2021).
Revised: [Equation] where s ( (cid:15), y ) = (cid:15) 1 n +(1 − (cid:15) ) y is the label smoothing function for one-hot vector y , (cid:15) ∈ [0 , 1] determines the amount of label smoothing, and λ > 0 is the regularization strength.

Original: Our approach has three key differences: (1) our k -branch network is more communication and data efﬁcient compared to previous k networks, (2) our model selects the  branches based onthe featuresinstead of  the loss, andtherefore does not need labeled data to selectbranches for new clients; (3) we propose the temporal prior to tackle periodic distribution shift when learning the mixture model.
Revised: The label smoothing regularization updates the other branch jointly with the feature extractor to prevent staleness, while encouraging the two branches to specialize in different feature subspaces: branch ¯ k ∗ i is trained to become less certain on features of samples from cluster k ∗ i . 3. L EARNING A G AUSSIAN M IXTURE M ODEL WITH P ER - CLIENT T EMPORAL S TATISTICS We propose Federated Expectation-Maximization with Temporal prior (F ED TEM) to learn a Gaussian Mixture Model (GMM) to infer the cluster a client is from, and select the corresponding branch in the multi-branch network.

Original: We deﬁne discrete latent variables z and ζ to represent which mode a sample and a client comes from, respectively, which will be inferred by our model.
Revised: We deﬁne discrete latent variables z and ζ to represent which cluster a sample and a client is from, respectively.

Original: Our model also assumesall 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
Revised: F ED TEM assumes samples on the same client are from the same cluster, so the GMM prior P ( ζ ) =

Original: P (  x | z = k ) as a Gaussian distribution N ( f ( w , x ) | µ k , σ2 k ) on the feature representation f ( w , x ) .
Revised: P ( z ) . We deﬁne P ( x | z = k ) as a Gaussian distribution N ( f ( w , x ) | µ k , σ k ) in the feature space.

Original: For efﬁciency, we assume and learn a diagonal covariance σ2 k for the Gaussian models.
Revised: For efﬁciency, we assume and learn a diagonal covariance σ k for the Gaussian models.

Original: We use θ to denote the parameter of the GMM for convenience.

Original: Algorithm 1 summarizes the training process, and the details of each step are provided in the following sections.
Revised: Algorithm 1 summarizes the training process, and the details of each step are provided in the following sections.

Original: 3.1. M ODELING THE C LIENT D ISTRIBUTIONS AND S ELECTING B RANCHES FOR T RAINING
Revised: 3.1. M ODELING THE C LIENT D ISTRIBUTIONS AND S ELECTING B RANCHES FOR T RAINING

Original: 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 set D i , and then select the corresponding branch. We assume all samples on one client are from the same mode, so P ( ζ ) = P ( z ) .
Revised: The prior P ( ζ ) of the client distribution are expected to change due to the periodic distribution shift. The i -th client estimates its probability of coming from the k -th cluster based on its data D i before the local update steps during training.

Original: There are two options for inferring the modes 1 Given the current mixture model, estimate P ( ζ|D i )[-from P ( x | z ) of all x ∈ D i ; 
 Given the training samples-] on  the client, ﬁnd the Maximum Likelihood Estimation (MLE) of P ( z ) for the client and use it for P ( ζ ) .
Revised: Speciﬁcally, given the GMM parameters, the local maximum likelihood estimation (MLE) of P ( ζ i )[-from P ( x | z ) of all x ∈ D i ; 
 Given the training samples-] on each new client i is estimated by the MLE of p ( z ) on client i .

Original: Option 1 exploits the current parameters of the mixture model to infer the mode of each client. With the assumption that samples in D i are independent, from Bayes’ theorem, we can get P ( z |D i ) as [Equation] Intuitively, since the client distribution is constantly shifting, we should not expect the prior P ( ζ ) to be unchanged from the last round.
Revised: Let π ∗

Original: Option 2 re-estimates the prior P ( z )of the mixture model onthe current client through MLE and use it as  P ( ζ |D i )  .
Revised: i be the MLE of P ( z ) on client i , where the k -th dimension of π ∗ i , denoted as π ∗ ik , is the MLE of P ( z = k ) on client i .

Original: As shown in Figure 6, we ﬁnd this achieves better empirical results than Option 1. Speciﬁcally, we use π ∗ i , the MLE of P ( z ) on the data of client i , to select the branch for training:
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: where π ∗ ik is the k -th dimension of π ∗ i , corresponding to the MLE of P ( z = k ) .

Original: For completeness, we give the derivation of Eq.
Revised: For completeness, we give the derivation of Eq.

Original: 5 in Appendix A.1.
Revised: 5 in Appendix A.1, where we also compare with the empirical results of using the posterior P ( ζ |D i ) instead of the MLE π ∗ i for branch selection.

Original: After estimating the distribution of ζ , we can either greedily select the branch k ∗ i corresponding to the maximum probability, orwe can sample from this distribution.
Revised: We train the network with the estimated P ( ζ i ) . We can either greedily select the branch k ∗ i = arg max k P ( ζ i = k ) and optimize Eq. 4, or sample from the distribution P ( ζ i ) .

Original: Our experiments in Figureshow 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
Revised: We show the greedy approach achieves better empirical results in Figure 6, and use it by default.

Original: [Equation]
Revised: [Table]

Original: where λ > 0 is the regularization strength, s ( (cid:15), y )

Original: = (cid:15) 1 n + (1 − (cid:15) )

Original: y is the label smoothing function with 0 ≤ (cid:15) ≤ 1 to determine the amount of label smoothing.

Original: The regularization enables the other branch to be updated jointly with the feature extractor to prevent mismatch.

Original: 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.

Original: 3.1. U PDATING THE M IXTURE M ODEL 
Revised: 3.1. U PDATING THE M IXTURE M ODEL P ARAMETERS

Original: 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 w t

Original: +1 i , and update the parameters of the

Original: GMM based on w t +1 i . 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).
Revised: We introduce a federated Expectation-Maximization (EM) (Dempster et al., 1977) enhanced by the temporal prior to update the parameters of the GMM.

Original: In the federated setting, this can be achieved privately byﬁrst running EM on each client i to get the locally updated GMM parameters (˜ µ t +1 i , ˜ σ t +1 i , ˜ π t +1 i ) , and then aggregating on the server.
Revised: The temporal prior is enforced via a bottomup approach: we start by running EM on each client i and send all possible locally optimal GMM parameters to the server, and then select and aggregate the GMM parameters with the temporal prior on the server.

Original: 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 +11 , π t +11 ) for mode 1, while the remaining clients will only update ( µ t +12 , σ t +12 , π t +12 ) for mode 2. The server decides which clients are from mode 1 and which are from mode 2 by their ˜ π t
Revised: We ﬁnd the optimal GMM parameters after the local update steps as described in

Original: +1 i , since ˜ π t +1 i estimates the probability of an arbitrary sample from client i to come from each mode.
Revised: Section 3.1.1, so that the GMM is updated in the feature space of the updated network.

Original: We give details of the steps below.
Revised: We give details of the process in the following.

Original: E step.
Revised: E step.

Original: 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 {+, based on the locally updated network with parameters w t +
 i+}
Revised: For each client i ∈ I ( t ) , evaluate the posterior γ ik ( x ) for each sample x to infer the probability of x coming from cluster k {+, based on the locally updated network with parameters w t +
 i+}

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: M step on clients.

Original: In between these minima and maxima, we consider three types of ˜ q ( t ) in our current experiments: 1) Linear: ˜ q ( t ) = q L, 1 ( t ) ; 2) Cosine: ˜ q
Revised: To learn a GMM that better distinguishes the daytime and nighttime clusters, we incorporate the temporal prior ˜ q ( t ) into the M step by considering the posterior P ( ζ = k

Original: ( t ) = q C, 1 ( t ) ; 3) Soft: see Appendix A.2. Conditioned on ˜ q ( t ) , our goal is to estimate the posterior P ( ζ |D i , θ t , ˜ q ( t )) for each client to decide which mode they are from.
Revised: |D i , θ t , ˜ q ( t )) , where θ t = ( µ t , σ t , π t ) is the collection of all GMM parameters.

Original: To achieve this, similar as in Section 3.1.1, we ﬁrst obtain the MLE of the prior on D i using the updated network parameters w t +1 i as [Equation]
Revised: The M step optimizes the following objective

Original: [Equation]
Revised: [Equation]

Original: where γ ik ( x ) comes from the E step in Eq. Then, we sort theprior probability ¯ π ∗ i 1 of mode 1 for all sampled clients i ∈ I (cid:48) ( t ) , and assign (cid:98) ˜ q ( t ) · |I (cid:48) ( t ) | + 12 (cid:99) clients  with the highest ¯ π ∗ i 1 to mode 1, denoted as ˜I (cid:48)
Revised: Since the clients cannot share their data with the server, we ﬁrst ﬁnd locally optimal parameters ˜

Original: 1 ( t ) .
Revised: θ t +1 i

Original: The remaining clients, denoted as ˜ I (cid:48) 2 ( t ) , are assigned to mode 2. As a result, for i 1 ∈ ˜ I (cid:48) 1 ( t ) , P ( ζ = 1 |D i , θ t , ˜ q ( t )) = 1 , while for i 2 ∈ ˜ I (cid:48) 2 ( t ) , P ( ζ =2 |D i , θ t , ˜ q ( t )) = 1 .
Revised: =

Original: 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 ﬁrst executed locally on each client i , where its objective is to ﬁnd the GMM parameters to maximize the expectation of the log likelihood of data, weighted by the time-varying client-level posterior P ( ζ = k
Revised: argmax θ (cid:80)

Original:  |D i , θ t , ˜ q ( t ) [Equation]
Revised: Kk =1 P ( ζ = k |D i , θ t , ˜ q ( t )) log P

Original: For each client, ˜ q ( t ) affects the posterior distribution P ( ζ |D i ,  ˜ q ( t ))  , and the effect depends on other clients.
Revised: ( D i , ζ = k | θ , ˜ q ( t )) on each client i , and then average on the server.

Original: 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.
Revised: In cross-device FL, it is very difﬁcult to have multiple communication rounds between the server and client (Bonawitz et al., 2019; Wang et al., 2021).

Original: A viable approach is  to restrict P ( ζ |D
Revised: However, it is hard to evaluate the posterior P ( ζ = k

Original:  i ,  ˜ q ( t )) to have a ﬁnite number of possible values, so that the client can solve Eq. 9 for all possible P ( ζ |D i , ˜ q ( t )) andsend all feasible solutions to the server.
Revised: |D i , θ t , ˜ q ( t )) 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 on the estimates of other clients.

Original: To make ˜ q ( t ) effective, we should set a minimum of K different values for  P ( ζ |D
Revised: A viable approach is to assume the posterior to be one-hot, i.e., for a certain cluster k ∗ , P ( ζ = k

Original: i , ˜ q ( t )) so that each mode corresponds to a different value. To learn more distinctive representations for the modes, we let P ( ζ = k |D i , ˜ q ( t)) = 1 if the client is estimated to come from mode k , so P ( ζ = j |D i , ˜ q ( t ))
Revised: |D i , θ t , ˜ q ( t ))

Original: = 

Revised: =

Original: for all j (cid:54)
Revised: if k = k ∗ , and P ( ζ = k |D i , θ t , ˜ q ( t ))

Original: = {+
+}
Revised: = {+
+}

Original:  k , and Eq.
Revised: if k (cid:54) = k ∗ .

Original: 9 becomes
Revised: Under this modeling assumption, we can learn more distinctive features for the clusters at no communication overhead. When k ∗ is given, we only need to solve

Original: [Equation]
Revised: [Equation]

Original: 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 ( w t 
Revised: where its optimal mean and variance, and the MLE of the GMM prior for cluster k ∗ on client i are ˜ µ t +1 ik

Original: +1 i , x )]
Revised: [Equation]

Original: j .
Revised: =

Original: The last equation comes from our modeling assumption that all x ∈ D i are from the same mode. From Eq.
Revised: [Equation] Note ˜ µ t +1 ik ∗ and ˜ σ t +1 ik ∗ are identical for different k ∗ ; see Appendix A.2 for derivations.

Original: 10, we can see the optimal solution is tokeep parameters for other modes unchanged, while only updating them for the  k -th mode as
Revised: 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 D i for each cluster k , i.e., ˜ π ∗ ik , to the server. The MLE is deﬁned in

Original: [Equation]
Revised: Eq.

Original: In this way, we can pack the all the feasible ˜ µ t +1 ik , [˜ σ t +1 ik ] 2 into matrices that have the same size as one set of GMM parameters, which does not increase the communication cost.
Revised: The total number of parameters sent to the server can be less than a single GMM model.

Original: M step (server).
Revised: M step on server with temporal prior.

Original: The server collects the MLE of the priors { ¯ π ∗ i | i ∈ I (cid:48)( t ) }
Revised: The temporal prior is enforced via a bottom-up approach: each client sends all possible solutions and MLEs { (˜ µ t +1 ik , ˜ σ t +1 ik , ˜ π ∗ ik ) | i ∈ I (cid:48)

Original: (Eq.
Revised: ( t ) , k ∈ { 1 , 2 }} to the server; the server then estimates the one-hot posterior P ( ζ |D i , θ t , ˜ q ( t )) based on ˜ π ∗ ik . A set ˜ I (cid:48) 1 ( t ) of (cid:98) ˜

Original: 8) and all the feasible solutions { (˜ µ t +1 i , ˜ σ t+1 i ) |  i ∈ I (cid:48) (  t ) } from the clients.
Revised: q ( t ) · |I (cid:48) ( t ) | + 12 (cid:99) clients with highest ˜ π ∗ i 1 will have P ( ζ = 1 |D i , θ t , ˜ q ( t ))

Original: Then, with the temporal prior, it divides the clients into two sets I (cid:48) 1 ( t ) and I (cid:48) 2  ( t ) asdescribed before.
Revised: = 1 , while the remaining clients, constituting a set ˜ I (cid:48) 2 ( t ) , will have P ( ζ = 2 |D i , θ t , ˜ q ( t )) = 1 . Then, similar as

Original: Similar as FedAvg (McMahan et al., 2017), it updates the  parameters for each mode k  weighted by the number ofsamples from each client as
Revised: 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 ∈ ˜ I (cid:48) k

Original: [Equation]
Revised: [Equation]

Original:  where M tk = (cid:80) i ∈I (cid:48) k ( t ) |D i | is the total number of samples of clients assigned to mode k .
Revised: t ) : [Equation] where M tk = (cid:80) i ∈I (cid:48) k ( t ) |D i | is the total number of samples from clients assigned to mode k .

Original: Notethe resulting σ t +1 is not guaranteed to be an unbiased estimate, but we found this process to give satisfying results in practice.
Revised: Note σ t +1 is not necessarily an unbiased estimate, but we ﬁnd this estimate gives good results in practice.

Original: 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 ofclients and samples from both modes are roughly the same.
Revised: We also use a running average of the GMM prior π t +1 k during training and set it to be a uniform distribution during inference; see Appendix A.4 for details. 3. L EARNING A C LUSTERING M ODEL WITH A GGREGATED T EMPORAL S TATISTICS

Original: Therefore, we usea 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.
Revised: F ED TEM collects per-client ˜ π ∗ i to use 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).

Original: During training, we estimate the prior π t in every step to match theshifting distribution.
Revised: We propose an alternative Federated K -Means algorithm augmented by the

Original: However, we ﬁnd it beneﬁcial to use a running average of π t 
Revised: Temporal prior (F ED TKM), where the temporal prior is enforced based on only aggregated results, described in Algorithm 2. F ED TKM clusters the i th client and selects the branches based on the averaged distance of the features to the k cluster centers, [Equation] where the superscript t denotes the communication round, and α k are distance scalar factors to control the “prior” of the clusters and enforce the temporal priors.

Revised: During training, a client i computes d ik on its training set to select the branch with minimal d ik , while at test time, we use a minibatch to estimate d ik .

Revised: In the case of two centers, we ﬁx α 1 = 1 , and learn a scalar α 2 > 0 to rescale the distances to the second cluster c 2 : when α 2 > 1 , it makes the F ED TKM algorithm more likely to assign the client to the ﬁrst cluster, and vice versa.

Revised: α 2 > 1 is updated by a quantile-based (private) estimator inspired by (Andrew et al., 2021).

Revised: The clients train the network with their selected branches by Eq.

Revised: 11 through the objective in Eq.

Revised: After the local training steps, each client estimates its cluster assignment again with the locally updated parameters ˜ w ti , and computes an indicator variable q

Original: {+i ( t ) = 
+}
Revised: {+i ( t ) = 
+}

Revised: [ d i 1 ( ˜ w ti ) <

Revised: d i 2 ( ˜ w ti )] for quantile estimation, and the feature means ¯ f ( ˜ w ti )

Revised: =

Revised: 1 |D i | (cid:80) x ∈D i f ( ˜ w ti , x ) for k -means centers.

Revised: Then q i ( t ) and ¯ f ( ˜ w i ) can be sent to a trusted, secure aggregator to get the qunatile of clients assigned to cluster 1, ¯ q ( t ) , and the averaged k -means centers c t +11 , c t + as

Revised: [Equation]

Revised: q

Revised: [Equation]

Revised: |

Original: +1 k= βπ tk
Revised: [Equation] i [Equation] k

Revised: [Equation]

Revised: ¯ f (

Revised: [Equation]

Revised: w ti ) , k

Revised: = 1 or 2 , (12) where ˜

Original: +(1 − β ) π k M tk / ( (cid:80) Kj =1 M tj )during training, which achieves better results than setting β = 0 or using uniform prior in our preliminary experiments.
Revised: I 1 ( t ) = { i | i ∈ I (cid:48) ( t ) , d i 1 ( ˜ w t i ) < d i 2 ( ˜ w t i ) } , and ˜ I 2 ( t ) = { i | i ∈ I (cid:48) ( t ) , d i 1 ( ˜ w t i ) ≥ d i 2 ( ˜ w t i ) } .

Original: E XPERIMENTS
Revised: The distance factor α 2 is updated by tracking the difference between the empirical quantile ¯ q ( t ) and the oracle quantile ˜ q ( t ) estimated by the temporal prior, [Equation] (13) where η t ≥ 0 is the step size of the geometric update at step t . In all experiments, we use the periodical linear function ˜ q ( t ) =

Revised: q L, 1 ( t ) as the temporal prior.

Original: [Table]
Revised: [Table]

Revised: 3. C OMPARISONS AND D ISCUSSIONS

Original: Dataset.
Revised: Privacy.

Revised: In our simulation experiments (Section 4), F ED TEM often performs better than F ED TKM.

Revised: However, the server has to observe the GMM parameters for participating client in F ED TEM, which may require new techniques to satisfy the strong data minimization and data anonymization principles (Wang et al., 2021;

Revised: Bonawitz et al., 2021).

Revised: On the other hand, F ED TKM only uses aggregated results, which is compatible with the strong data minimization principle, and easier to be further protected by other privacy techniques such as secure aggregation (Bonawitz et al., 2019) and differential privacy (McMahan et al., 2018; Kairouz et al., 2021).

Revised: Differences from previous methods.

Revised: Compared with previous methods that jointly train a clustering model with multiple networks, our method has four key differences that better suits cross-device

Revised: FL under temporal distribution shift.

Revised: First, our method introduces the temporal prior to regularize the clustering model, which is missing in previous works.

Revised: Second, our k -branch network with shared feature extractor is more communication and data efﬁcient than using k networks.

Revised: Third, our clustering model selects the branches based on the feature, which eliminates the need for labeled data for clients not participate in training.

Revised: Clustering models of (Ghosh et al., 2020; Marfoq et al., 2021) are based on the loss, which is impractical for clients without labeled data.

Revised: Fourth, our method does not maintain states on clients.

Revised: (Marfoq et al., 2021) assumes stateful clients, which is impractical for cross-device setting where each client only participate limited times (Wang et al., 2021).

Revised: We give more detailed comparison with previous methods in Appendix A.7.

Original:  We consider two image classiﬁcation datasets, EMNIST and CIFAR, and one next word prediction task on Stack Overﬂow (SO).
Revised: Dataset. We consider two image classiﬁcation datasets, EMNIST and CIFAR, and one next word prediction task on Stack Overﬂow (SO).

Original: The split of the day mode ( I 1 ) and night mode ( I 2 ), and other statistics, are shown in Table 1.
Revised: The split of the day mode ( I 1 ) and night mode ( I 2 ), and other statistics, are shown in Table 1 in the Appendix.

Original: On SO, we use a vocabulary size of 10K and report the test accuracy without special tokens.
Revised: On SO, we use a vocabulary size of 10K and report the test accuracy without special tokens.

Original: We truncate each sentence to have no more than 20 tokens.
Revised: We truncate each sentence to have no more thantokens.

Revised: 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 .

Original: Architecture of the multi-branch networks.
Revised: Architecture of the multi-branch networks.

Original: For image classiﬁcation, on EMNIST, we train LeNet with 2 Conv layers and 2 FC layers, while on CIFAR, we train a ResNet-18 with the Batch 
Revised: On EMNIST, we train LeNet with 2 Conv layers and 2 FC layers. On CIFAR, we train a ResNet-18 with the Batch Norm replaced by Group Norm (Wu [Figure]

Original: Normalization replaced by Group Normalization (Wu & He, 2018) for stability in federated learning.
Revised: & He, 2018) for stability in federated learning.

Original: We use the last FC layer of these convolutional networks as the multi-branch part and share all the remaining layers for all modes.
Revised: We use the last FC layer of these convolutional networks as the multi-branch part and share all the remaining layers.

Original: We use the output from the shared feature extractor for GMM, which are 128 and 512 dimensional  respectively.
Revised: We use the output from the shared feature extractor for clustering, which are 128 and 512 dimensional for EMNIST and CIFAR respectively.

Original: On SO, we train a single-layer LSTM (Hochreiter & Schmidhuber, 1997) with a hidden size of 670 and embedding size of 96.
Revised: On SO, we train a single-layer LSTM (Hochreiter & Schmidhuber, 1997) with a hidden size of 670 and embedding size of 96.

Original: To alleviate the communication overhead and prevent overﬁtting, we deﬁne the branches to be the last projection layer before the ﬁnal prediction layer and share the remaining weights for both modes, which is 15x smaller than the ﬁnal prediction layer.
Revised: To alleviate the communication overhead and prevent overﬁtting, we deﬁne the branches to be the last projection layer before the ﬁnal prediction layer, which is 15x smaller than the ﬁnal prediction layer.

Original: We concatenate the average of all tokens’ features from both branches for GMM, which has dimensionality 192.

Original: With these conﬁgurations, the  models on EMNIST, CIFAR and SO increase the communication cost by 0.7%, 0.5% and 1.6%, respectively.
Revised: The communication overheads of the extra branch are 0.7%, 0.5% and 1.6% respectively for models on EMNIST, CIFAR and SO.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: We use FedAdam (Reddi et al., 2021) as the optimizer.
Revised: We use FedAdam (Reddi et al., 2021) as the optimizer.

Original: Table 2 shows the training hyperparameters.
Revised: Table 2 shows the training hyperparameters, with implementation details and practical notes given in Section A.4.

Original: 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.
Revised: For

Original: For F ED TEM and Min Loss, we also do a grid search on the label smoothing parameters (cid:15) and λ .
Revised: Min Loss, F ED TEM and F ED TKM, we also do a grid search on the label smoothing parameters (cid:15) and λ .

Original: For each set of hyperparameters, we run 3 experiments with different random seeds and report their mean and standard error.
Revised: For each hyperparameter setting, we run 3 experiments with different random seeds and report their mean and standard error.

Original: Methods to compare. First, we provide ablation studies on the temporal prior to quantify its efﬁcacy, and see which  prior (Linear, Cosine, or Soft) works better in practice.Note for Linear and
Revised: Ablation study and temporal prior. In Appendix A.5 and Figure 3, we provide ablation studies on the temporal prior to quantify its efﬁcacy, and see which temporal prior (Linear, Cosine, or Soft) works better in practice.

Original: Cosine, we always use the same periodical linear or cosine function with p = 1 under all types and p ’s for the distribution shift.

Original: 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 ﬁt the distribution shift and improves the results.
Revised: 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 ﬁt the distribution shift and improves the results.

Original: 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
Revised: With these observations, for F ED TEM, we choose the linear prior ˜ q L, 1 ( t ) on EMNIST and CIFAR and the cosine prior ˜ q C, 1 ( t ) on SO. For F ED TKM, we only consider the linear prior ˜ q L, 1 ( t ) .

Original: Figure 2.
Revised: Baselines.

Original: 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.
Revised: For fair comparisons, unless otherwise stated, we adapt all methods into our 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.

Original: 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.
Revised: 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.

Original: Such networks trained under shifting data distribution are denoted as “ Vanilla ”, while those trained without distribution shift are denoted as “ No Dist.
Revised: Such networks trained under shifting data distribution are denoted as “ Vanilla ”, while those trained without distribution shift are denoted as “ No Dist.

Original: Shift ” (NDS).In
Revised: Shift ” (NDS).In

Original: Appendix A.6, we also give results of training single-branch models in these two settings, which obtained almost the same results.
Revised: Appendix A.8, we also provide results of training single-branch models in these two settings, which obtained almost identical results.

Original: 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 ﬁrst half of each period, and the other model on the other half.
Revised: The baseline methods are: 1) “ SCSGD ”, where similar as SemiCyclic SGD (Eichner et al., 2019), we train one model during ﬁrst half of each period, and the other model on the other half.

Original: During test time, since there is no direct identiﬁer 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.
Revised: 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 a uniform distribution over all labels.

Original: 2) T-shot K-means , which is an enhancement ofthe 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.
Revised: 2) T-shot K-means , which is an enhanced variant of 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.

Original: For both training and test, it selects the branches for each sample according to the nearest cluster center.
Revised: For both training and test, it selects the branches for each sample according to the nearest cluster center.

Original: 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).
Revised: 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).

Original: Besides, we have also applied thesame label smoothing regularization (Eq. 6) to IFCA, which improves the results.
Revised: We applied the label smoothing regularization (Eq. 4) to IFCA, which improves the results.

Original: FedEM (Marfoq et al., 2021), where we use the same algorithm on participating clients during training.
Revised: FedEM (Marfoq et al., 2021), where we use the same algorithm on participating clients during training.

Original: For testing, we let it “cheat” and assume all test samples are labeled, so that it can execute itsmodiﬁed EM steps based on the loss to select the branches.
Revised: FedEM is implemented with privilege to access sample labels for test clients, to enable its EM steps which uses the loss.

Original: Main Results
Revised: Main Results.

Original: Thebest results of allcompared methods are shown in Figure 2.
Revised: The results of all methods are shown in Figure 2.

Original: By comparing results of Vanilla and NDS, we ﬁnd the temporal distribution shift indeed causes issues for normal [Figure] training, causing the worst-case decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and SO.
Revised: By comparing results of Vanilla and NDS, we ﬁnd the temporal distribution shift indeed degrades vanilla training, causing the worstcase decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and SO.

Original: Surprisingly, with F ED TEM, the accuracy can be even higher than models trained in the NDS setting in most cases, indicating F ED TEM obtains more distinctive feature representations by learning to separate the two modes in the feature space and learning specialized prediction branches.
Revised: While with F ED TEM and F ED TKM, the accuracy can be even higher than models trained in the NDS setting in most cases, demonstrating the efﬁcacy of the temporal prior for learning more distinctive feature representations for the daytime and nighttime modes with specialized prediction branches.

Original: The improvement on the strong NDS baseline can be as high as 4%, 4% and 0.45% on the three datasets.
Revised: The improvement on the strong NDS baseline can be as high as 4%, 4% and 0.45% with

Original: F ED TEM  is not better than NDS when p deviates too much from 1, but the data distribution is extremely skewed in such settings, where F ED TEM still improves over most other methods in most cases.
Revised: F ED TEM on the three datasets. F ED TEM is not better than NDS when p deviates too much fromas the data distribution is extremely skewed in such settings, while our methods are still better than most other methods in most cases.

Original: For other methods, SCSGD performs surprisingly well on EMNIST when p is large, but the performance quickly drops on the more complicated 
Revised: For baseline methods, SCSGD performs surprisingly well on EMNIST when p is large, but the performance quickly drops on the more complicated CIFAR and

Original: CIFAR and SO.
Revised: SO.

Original: T-shot K-means achieves improvements for extreme p ’s on EMNIST, but it is not signiﬁcantly better than “Vanilla” on CIFAR, due to the difﬁculty in reliable clustering in more complicated image spaces.
Revised: T-shot K-means achieves improvements for extreme p ’s on EMNIST, but it is not signiﬁcantly better than “Vanilla” on CIFAR, due to the difﬁculty in reliable clustering in more complicated image spaces.

Original: Also note it is not straightforward to do clustering on raw data of SO so we have not implemented it.
Revised: We did not implement T-shot K-means for SO due to the difﬁculty of clustering the raw data of the language task.

Original: Min Loss (IFCA) alleviates the drop and sometimes catches up with NDS, but cannot achieve improvements over NDS.
Revised: Min Loss (IFCA) alleviates the drop and sometimes competes with NDS, but cannot achieve improvements over NDS.

Original: By contrast,results of FedEM is the worst even when it is “cheating” on the test clients.
Revised: By contrast, FedEM often achieves inferior results even with the privilege access to the test labels for EM estimation.

Original: 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.
Revised: 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 a few rounds on EMNIST and CIFAR, and at most one round on SO.

Original: Effect of label smoothing regularization.
Revised: Effect of label smoothing regularization.

Original: Shown in Figure 3.
Revised: Shown in Figure 5.

Original: Since EMNIST is relatively easy and the network on it is small, the label smoothing does not show signiﬁcant improvements.
Revised: Since EMNIST is relatively easy and the network on it is small, the label smoothing does not show signiﬁcant improvements.

Original: However, it is critical to the success on CIFAR and SO.
Revised: However, it is critical to the success on CIFAR and SO.

Original: Without the regularization, one branch will not be updated simultaneously with the feature extractor, resulting in instability and low test accuracy.
Revised: Without the regularization, one branch will not be updated simultaneously with the feature extractor, resulting in instability and low test accuracy.

Original: MLE for Evaluation.
Revised: MLE for Evaluation.

Original: During evaluation,for image classiﬁcation, we ﬁnd it beneﬁcial in most cases to use the MLE in Eq.
Revised: For F ED TEM, during evaluation, we ﬁnd it beneﬁcial for image classiﬁcation to use the MLE in Eq.

Original: 5 on the minibatches (typically of size no larger than 64) to select the best branch, compared with the sample-level inference.
Revised: 5 on the minibatches (typically of size no larger than 64) to select the best branch, compared with the sample-level inference.

Original: We show its beneﬁt in Figure 5.
Revised: We show its beneﬁt in Figure 4 in the appendix.

Original: In practice, similar approach can be realized through caching.
Revised: In practice, similar approaches can be realized through caching.

Original: We ﬁnd Min Loss does not show much difference when a similar approach is applied, where branches are chosen by the average of persample certainty.
Revised: We ﬁnd Min Loss does not show much difference when a similar approach is applied, where branches are chosen by the averaged certainty on minibatches.

Original: 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.
Revised: The baseline model even gets worse performance with this approach, as shown in Figure 4, indicating our model is much better at distinguishing two modes in expectation.

Original: However, this approach does not show improvements for F ED TEM on SO.
Revised: However, this approach does not show improvements for F ED TEM on SO.

Original: Effect of T .
Revised: Effect of T .

Original: As shown in Figure 8, for F ED TEM with linear prior on EMNIST, a smaller T tends to decrease the performance when p deviates too much from 1.
Revised: As shown in Figure 8, for F ED TEM with linear prior on EMNIST, a smaller T tends to decrease the performance when p deviates too much from 1.

Original: In such scenarios, the distribution changes frequently and abruptly.
Revised: In such scenarios, the distribution changes frequently and abruptly.

Original: However, the performance of F ED TEM is maintained when the distribution changes in a frequent but more balanced way.
Revised: However, the performance of F ED TEM is maintained when the distribution changes in a frequent but more balanced way.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this paper, we demonstrated that the smooth transition model for simulating the periodical distribution shift better matches  practical FL systems, and developed F ED TEM algorithm to train multibranch networks to tackle the distribution shift.
Revised: In this paper, we showed the inﬂuence 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.

Original: F ED TEM incorporated priors of the temporal distribution shifts into learning a GMM to guide the network to select corresponding branches for clients of different modes.
Revised: Our methods incorporates priors of the temporal distribution shifts to learn a mixture or clustering model to guide the network and select corresponding branches for clients from different modes during both training and inference.

Original: The GMM is deﬁned in the feature space and does not require labelled data for inference, hence is ready to use for new test clients.
Revised: The clustering models are deﬁned in the feature space and does not require labelled data for inference, hence is ready to be deployed on clients without labeled data.

Original: The branches are light-weight with little communication overhead, and the model demonstrates signiﬁcant improvements in test accuracy on the benchmark datasets.
Revised: The branches are lightweight with little communication overhead, and the model demonstrates signiﬁcant improvements in test accuracy on the benchmark datasets.

Original: F ED TEM also satisﬁes many other real-world constraints like single round communication and stateless clients (Bonawitz et al., 2019; Paulik et al., 2021;
Revised: Our methods also satisﬁes many other real-world constraints like single round communication and stateless clients (Bonawitz et al., 2019; Paulik et al., 2021;

Original: Wang et al., 2021), which makes it ready to be deployed in practical FL systems.
Revised: Wang et al., 2021), which makes it possible to be deployed in practical FL systems.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Revised: Our code is available on GitHub.

Original: For reproducibility, we have reported the mean and standard error for 3 experiments with different random seeds for all settings.
Revised: For reproducibility, we have reported the mean and standard error for 3 experiments with different random seeds for all settings.

Original: The dataset we used for simulations are all available in Tensorﬂow Federated and our implementation is largely based on the code from Adaptive Federated Optimization.
Revised: The dataset we used for simulations are all available in Tensorﬂow Federated, and our implementation is largely based on the code from Adaptive Federated Optimization (Reddi et al., 2021).

Original: 1 We have listed our hyperparameters in Table 2.
Revised: 2 We have listed our hyperparameters in Table 2.

Original: We have given the pseudo code in Algorithmand the detailed formulas and derivations of our algorithm in Section 3.1.
Revised: We have given the pseudo code in Algorithm 1 and the detailed formulas and derivations of our algorithm in Section 3.1.

Original: We are committed to opensource our implementation upon publication.
Revised: 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 program, and the Sloan Foundation.

Original: A A PPENDIX A. D ERIVATION OF O PTION
Revised: A A PPENDIX A. F URTHER DISCUSSIONS OF BRANCH SELECTION WITH GMM

Revised: Derivation of the MLE.

Original:  We consider maximizing the expectation of the complete-data log likelihood log P ( D i , Z | θt +1 ) under the posterior distribution P ( Z |D i , θt ) , where Z denotes the collection of latent variables associated with each of the sample x ∈ D i .
Revised: Recall θ is the parameters of the GMM model. We consider maximizing the expectation of the complete-data log likelihood log P ( D i , Z | θ ) under the posterior distribution P ( Z |D i , θ ) , where Z is the collection of latent variables associated with each of the sample x ∈ D i .

Original: Formally, it is solving the following constrained optimization problem
Revised: Formally, it is solving the following constrained optimization problem

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

Original: We introduce an indicator function 1 [ z = k ] which is 1 if z = k or 0 otherwise.
Revised: We introduce an indicator function 1 [ z = k ] which is 1 if z = k or 0 otherwise.

Original: In this way, the log likelihood can be represented as
Revised: In this way, the log likelihood can be represented as

Original: [Equation]
Revised: [Equation]

Original: Meanwhile,
Revised: Meanwhile,

Original: [Equation]
Revised: [Equation]

Original: Plug Eq.
Revised: Plug Eq.

Original: 14 and Eq. 15 back into Eq.13, and
Revised: 15 and Eq. 16 back into Eq.

Original:  ignore terms that do not depend on π[-t +
-] , we ﬁnd we are solving the following problem for πt +
Revised: 14, and ignore terms that do not depend on π[-t +
-] , we ﬁnd we are solving the following problem for π

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The Lagrangian multiplier of this problem is
Revised: The Lagrangian multiplier of this problem is

Original: [Equation]
Revised: [Equation]

Original: Let the derivative w.r.t.
Revised: Let the derivative w.r.t.

Original: π k be 0,
Revised: π k be 0,

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: Multiply both sides of Eq.
Revised: Multiply both sides of Eq.

Original: 18 with π k , sum over 1 ≤ k ≤ K , and apply the constraint that (cid:80)
Revised: 19 with π k , sum over 1 ≤ k ≤ K , and apply the constraint that (cid:80)

Original: Kk =1 π k = 1 , we have
Revised: Kk =1 π k = 1 , we have

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Plug this back into Eq.
Revised: Plug this back into Eq.

Original: 18, we get the equation as desired.
Revised: 19, we get the equation as desired.

Original: [Equation]
Revised: [Equation]

Revised: An alternative for branch selection.

Revised: Instead of computing the MLE, with the assumption that samples in D i are IID, we can compute P ( ζ i |D i )

Revised: =

Revised: P ( z |D i ) using Bayes’ theorem and use P ( ζ i |D i ) for branch selection.

Revised: Speciﬁcally,

Original: [Equation]
Revised: [Equation]

Revised: As shown in Figure 6, we ﬁnd MLE achieves better empirical results than this alternative.

Revised: A. D ERIVATION OF THE LOCALLY OPTIMAL GMM S

Revised: We give the derivation of the optimal mean and variance of the GMM on each client, under the one-hot assumption that P ( ζ = k |D i , θ t , ˜ q ( t ))

Original: {+= 
+}
Revised: {+= 
+}

Revised: if k = k ∗ , and P ( ζ = k

Revised: |D i , θ t , ˜ q ( t ))

Original: {+= 
+}
Revised: {+= 
+}

Revised: if k (cid:54) =

Revised: k ∗ , for some given k ∗ .

Revised: Under this assumption, we need to solve

Revised: [Equation]

Revised: where d is the dimension of the features, [ µ k ∗ ] j , [ σ k ∗ ]

Revised: j and [ f ( w t +1 i , x )]

Revised: j denote the j -th dimension of µ k , σ k , f ( w t

Revised: +1 i , x ) , respectively.

Revised: From Eq.

Revised: 23, we can see the optimal solution for client i is to set the mean and variance for the k ∗ -th mode as

Revised: [Equation]

Original: A. T HE S OFT P RIOR
Revised: A. T HE S OFT P RIOR

Original: The soft prior ˜ q S ( t ) gives strong signals only at time steps where we believe q ( t ) = 0 or q ( t )
Revised: The soft prior ˜ q S ( t ) gives strong signals only at time steps where we believe q ( t ) = 0 or q ( t )

Original: = 1 .
Revised: = 1 .

Original: In between, it only requires the ratio to be non-increasing or non-decreasing, where the ratio is estimated using the posterior from Eq.
Revised: In between, it only requires the ratio to be non-increasing or non-decreasing, where the ratio is estimated using the posterior from Eq.

Original: 7 as
Revised: 6 as

Original: [Equation]
Revised: [Equation]

Original: i.e., the ratio of samples whose posterior has higher probability on mode 1.
Revised: i.e., the ratio of samples whose posterior has higher probability on mode 1.

Original: In this way, the model can still learn to distinguish the modes but we do not enforce a strong prior at every time step.
Revised: In this way, the model can still learn to distinguish the modes but we do not enforce a strong prior at every time step.

Original: Speciﬁcally, it is deﬁned as
Revised: Speciﬁcally, it is deﬁned as

Original: [Equation]
Revised: [Equation]

Revised: A. I MPLEMENTATION DETAILS AND PRACTICAL NOTES

Revised: Setting the GMM prior.

Revised: 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.

Revised: Therefore, we use a uniform distribution for the priors on the test set.

Revised: This does not make the GMM worse if it estimates P ( x | z ) accurately on the test set.

Revised: During training, we estimate the prior π t in every step to match the shifting distribution.

Revised: However, we ﬁnd it beneﬁcial to use a running average of π t

Revised: +1 k

Revised: [Equation]

Revised: training, which achieves better results than setting β = 0 or using uniform prior in our preliminary experiments.

Revised: We use β = 0 .

Revised: 99 for the moving average in all experiments.

Revised: Setting η t for F ED TKM.

Revised: Although we can identify the modes through the local minima and maxima of the training curve, the distribution shift can happen in arbitrary ways in between.

Revised: As a result, the certainty of the temporal prior should also be lower in between the modes.

Revised: To incorporate this heuristic, instead of using a constant η t , we let

Revised: [Equation]

Revised: [Equation]

Revised: We ﬁnd this to give better results when p (cid:54) = 1 , i.e., when the underlying distribution shift is biased and different from the linear temporal prior.

Original: A. A BLATION STUDIES : THE EFFECT OF TEMPORAL PRIOR
Revised: A. A BLATION STUDIES : THE EFFECT OF TEMPORAL PRIOR

Original: 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 F ED TEM is that the server, in the M step, greedily uses client i to update parameters of mode k ∗ i = arg max k ¯ π ∗ ik (see Eq.
Revised: To quantify the effect of temporal prior on training, we compare four versions of F ED TEM: 1) EM Only : only apply the EM part of our algorithm without the temporal priors, where the only difference from F ED TEM is that the server, in the M step, greedily uses client i to update parameters of mode k ∗ i = arg max k ˜ π ∗ ik (see Eq.

Original: 8 and Eq. 12); 2) FedTEM (linear) :
Revised: 9 and Eq. 10); 2) FedTEM (linear) :

Original: our algorithm where the temporal prior is the periodic linear function P ( i ∈ I 1 ( t )) =
Revised: our algorithm where the temporal prior is the periodic linear function P ( i ∈ I 1 ( t )) =

Original: q L, 1 ( t ) ; 3) FedTEM (cosine) : our algorithm where the temporal prior is the cosine function P ( i ∈ I 1 ( t )) =
Revised: q L, 1 ( t ) ; 3) FedTEM (cosine) : our algorithm where the temporal prior is the cosine function P ( i ∈ I 1 ( t )) =

Original: q C, 1 ( t ) ; 4) FedTEM (soft) : our algorithm with the soft temporal prior introduced in Section A.2.
Revised: q C, 1 ( t ) ; 4) FedTEM (soft) : our algorithm with the soft temporal prior introduced in Section A.3.

Original: The results are shown in Figure 4.
Revised: The results are shown in Figure 3.

Original: We ﬁnd all three priors improves the “EM Only” baseline in most cases.
Revised: We ﬁnd all three priors improves the “EM Only” baseline in most cases.

Original: “FedTEM (linear)” is better on image datasets, while “FedTEM (cosine)” is better on the language dataset.
Revised: “FedTEM (linear)” is better on image datasets, while “FedTEM (cosine)” is better on the language dataset.

Original: 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.
Revised: 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 [Table] [Table]

Original: A. V ERIFYING THE EFFECT OF MLE FOR EVALUATION
Revised: A. V ERIFYING THE EFFECT OF MLE FOR EVALUATION

Original: To quantize the effect the batch-level MLE, we compare the results of our method with or without
Revised: To quantize the effect the batch-level MLE [Table] we compare the results of our method with or without

Original: MLE on EMNIST and CIFAR.
Revised: MLE on EMNIST and CIFAR.

Original: 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.
Revised: 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 [Table] [Figure]

Original: [Figure] 
Revised: [Figure] [Figure] [Figure] [Table]

Original: [Figure] 
Revised: [Figure] [Table] [Figure] [Table]

Original: [Figure] 
Revised: [Figure] [Table] [Figure] [Table]

Revised: [Figure]

Original: [Table] 
Revised: [Table] [Figure] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: [Figure] 
Revised: [Figure] [Table] [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: A. A DDITIONAL R ELATED W ORKS
Revised: A. A DDITIONAL R ELATED W ORKS

Original: 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.
Revised: 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.

Original: We now review other related works that either consider other types of distribution shift, or have (weak) similarity as the proposed F ED TEM method.
Revised: We now review other related works that either consider other types of distribution shift, or have (weak) similarity as the proposed F ED TEM method.

Original: Heterogeneity and client distribution shift.
Revised: Heterogeneity and client distribution shift.

Original: Client heterogeneity is an active topic in FL, and various methods have been proposed to address the distribution difference among clients.
Revised: Client heterogeneity is an active topic in FL, and various methods have been proposed to address the distribution difference among clients.

Original: 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;
Revised: 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;

Original: 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 afﬁne transforms in the image domain.
Revised: 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 afﬁne transforms in the image domain.

Original: 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).
Revised: 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).

Original: Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al., 2021).
Revised: Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al., 2021).

Original: More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al.
Revised: More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al.

Original: Ghosh et al.
Revised: Ghosh et al.

Original: (2020), while still assuming the clients can be uniformly accessed during the training process.
Revised: (2020), while still assuming the clients can be uniformly accessed during the training process.

Original: We kindly ask interested readers to ﬁnd more papers on heterogeneity in a few recent surveys (Kairouz et al., 2019;
Revised: We kindly ask interested readers to ﬁnd more papers on heterogeneity in a few recent surveys (Kairouz et al., 2019;

Original: Li et al., 2020a; Wang et al., 2021).
Revised: Li et al., 2020a; Wang et al., 2021).

Original: All these methods consider the distributions shift among different clients, while the proposed F ED TEM considers periodic distribution shift of client population.
Revised: All these methods consider the distributions shift among different clients, while the proposed F ED TEM considers periodic distribution shift of client population.

Original: Clustering and mixture models.
Revised: Clustering and mixture models.

Original: We do not assume strong control over the clients and the available client population is always changing due to the periodical distribution shift.
Revised: We do not assume strong control over the clients and the available client population is always changing due to the periodical distribution shift.

Original: We also do not require stateful clients, since our prior is applied globally to all clients.
Revised: We also do not require stateful clients, since our prior is applied globally to all clients.

Original: Our mixture model is based on the feature space, therefore we do not require labeled data for unseen clients to infer its mode.
Revised: Our mixture model is based on the feature space, therefore we do not require labeled data for unseen clients to infer its mode.

Original: 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.
Revised: 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.

Original: Dennis et al.
Revised: Dennis et al.

Original: (2021) propose a one-shot clustering approach based on the raw client data before training, which implicitly assumes all representative clients are available simultaneously.
Revised: (2021) propose a one-shot clustering approach based on the raw client data before training, which implicitly assumes all representative clients are available simultaneously.

Original: Clustering on raw data can also be unreliable when the data demonstrates complicated distributions in the input space.
Revised: Clustering on raw data can also be unreliable when the data demonstrates complicated distributions in the input space.

Original: 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.
Revised: 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.

Original: To achieve this, it implicitly assumes strong control over population and clients sampling.
Revised: To achieve this, it implicitly assumes strong control over population and clients sampling.

Original: Ghosh et al.
Revised: Ghosh et al.

Original: and Mansour et al.
Revised: and Mansour et al.

Original: (2020) propose a similar algorithm that alternatively perform clustering and updating the personalized models for corresponding clusters.
Revised: (2020) propose a similar algorithm that alternatively perform clustering and updating the personalized models for corresponding clusters.

Original: Both of them require labled data for the clients to compute the loss for model selection.
Revised: Both of them require labled data for the clients to compute the loss for model selection.

Original: 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.
Revised: 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.

Original: The modiﬁed EM algorithm requires computing the loss and therefore labeled data for every client.
Revised: The modiﬁed EM algorithm requires computing the loss and therefore labeled data for every client.

Original: It maintains a different prior distribution for every client, which requires stateful clients and strong control over client sampling.
Revised: It maintains a different prior distribution for every client, which requires stateful clients and strong control over client sampling.

Original: The hierarchical or bi-partitioning clustering process of (Briggs et al., 2020; Fu et al., 2021) requires all
Revised: The hierarchical or bi-partitioning clustering process of (Briggs et al., 2020; Fu et al., 2021) requires all

Original: [Figure]
Revised: [Figure]

Original: 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.
Revised: 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.

Original: 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 signiﬁcantly to the communication cost.
Revised: 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 signiﬁcantly to the communication cost.

Original: (Xie et al., 2021) maintains one set of weights on each client while maintaining multiple weights as cluster centers on the server.
Revised: (Xie et al., 2021) maintains one set of weights on each client while maintaining multiple weights as cluster centers on the server.

Original: It takes two communication rounds for each weight update, which can be impractical.
Revised: It takes two communication rounds for each weight update, which can be impractical.

Original: It remains unclear if this mechanism can be generalized to unseen clients.
Revised: It remains unclear if this mechanism can be generalized to unseen clients.

Original: (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.
Revised: (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.

Original: It is more suitable for the cross-silo setting but not practical for on-device setting.
Revised: It is more suitable for the cross-silo setting but not practical for on-device setting.

Original: (Duan et al., 2021) clusters the clients based on the similarities between their gradients at the initial model weights.
Revised: (Duan et al., 2021) clusters the clients based on the similarities between their gradients at the initial model weights.

Original: 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 ﬁrst round for clustering.
Revised: 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 ﬁrst round for clustering.

Original: In addition, there is no guarantee that the gradient at the chosen round can separate the clusters well.
Revised: In addition, there is no guarantee that the gradient at the chosen round can separate the clusters well.

Original: And to compute the gradients, it still requires labeled data.
Revised: And to compute the gradients, it still requires labeled data.

Original: (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).
Revised: (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).

Original: Another limitation is it assumes every client is seen during training.
Revised: Another limitation is it assumes every client is seen during training.

Original: Multi-branch networks in FL.
Revised: Multi-branch networks in FL.

Original: Multi-branch networks have been explored in FL when a personalized model is preferred for each client:
Revised: Multi-branch networks have been explored in FL when a personalized model is preferred for each client:

Original: 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).
Revised: 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).

Original: All the aforementioned three works require labeled data and additional training efforts to learn the weights of the branches.
Revised: All the aforementioned three works require labeled data and additional training efforts to learn the weights of the branches.

Original: By comparison, our method does not need re-training the branches and does not need labeled data for unseen clients to select the branches.
Revised: By comparison, our method does not need re-training the branches and does not need labeled data for unseen clients to select the branches.

Original: The branch selection is based on the mixture model in the feature space, and the weights of the branches are ﬁxed.
Revised: The branch selection is based on the mixture model in the feature space, and the weights of the branches are ﬁxed.

Original: A. C OMPARING SINGLE - BRANCH AND MULTI - BRANCH BASELINES WITHOUT CLUSTERING
Revised: A. C OMPARING SINGLE - BRANCH AND MULTI - BRANCH BASELINES WITHOUT CLUSTERING

Original: For fair comparisons, we have used two branches for the baselines (“Vanilla” and “No Dist.
Revised: For fair comparisons, we have used two branches for the baselines (“Vanilla” and “No Dist.

Original: Shift”) so that their capacities are the same as our method.
Revised: Shift”) so that their capacities are the same as our method.

Original: Since these two baselines are just training the networks with FedAdam, it is more natural to train single-branch networks directly.
Revised: Since these two baselines are just training the networks with FedAdam, it is more natural to train single-branch networks directly.

Original: In Figure 9, we show that the single-branch and multi-branch networks obtain similar results under various settings, with or without distribution shift.
Revised: In Figure 9, we show that the single-branch and multi-branch networks obtain similar results under various settings, with or without distribution shift.

Original: Our method still outperforms these two single-branch baselines.
Revised: Our method still outperforms these two single-branch baselines.

Original: A. T HE EFFECT OF LOCAL TRAINING SET SIZES
Revised: A. T HE EFFECT OF LOCAL TRAINING SET SIZES

Original: When the number of training samples is small, the EM estimates may have high variance, causing misspeciﬁcations and resulting in worse performance.
Revised: When the number of training samples is small, the EM estimates may have high variance, causing misspeciﬁcations and resulting in worse performance.

Original: Meanwhile, models trained on smaller training sets tend to have worse generalization, so the test accuracy will drop for any method in general.
Revised: Meanwhile, models trained on smaller training sets tend to have worse generalization, so the test accuracy will drop for any method in general.

Original: To see which factor has more signiﬁcant 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.
Revised: To see which factor has more signiﬁcant 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.

Original: From the plot, we can see FedTEM maintains the advantage over the baselines and oracle (No Dist.
Revised: From the plot, we can see FedTEM maintains the advantage over the baselines and oracle (No Dist.

Original: Shift) under various training set sizes, though its accuracy also decays as the baselines due to the decreased numbers of training samples.
Revised: Shift) under various training set sizes, though its accuracy also decays as the baselines due to the decreased numbers of training samples.

Original: A. S TATS OF TRAINING A LANGUAGE MODEL IN A REAL FL SYSTEM
Revised: A. S TATS OF TRAINING A LANGUAGE MODEL IN A PRACTICAL FL SYSTEM

Original: In Figure we show some stats that could characterize the data distribution shifts in a real FL system.
Revised: In Figure we show some stats that could characterize the data distribution shifts in a real FL system.

Original: Generally, the data of daytime clients is more difﬁcult to ﬁt and generalize.
Revised: Generally, the data of daytime clients is more difﬁcult to ﬁt and generalize.

Original: 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.
Revised: 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.

Original: From the plots, we can also see more rounds are completed during nighttime.
Revised: From the plots, we can also see more rounds are completed during nighttime.

Original: 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).
Revised: 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).

Original: These plots further justify that the data shifts smoothly rather than abruptly in practice.
Revised: These plots further justify that the data shifts smoothly rather than abruptly in practice.