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

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: 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 decentralized 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 blockcyclic pattern as previous works, we model it with a mixture of distributions that gradually changes between daytimemodes 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 gradually changes between daytime and nighttime modes, and ﬁnd this intuitive model to better match the observations in practical federated learning systems.

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

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

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 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).

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 will aggregate 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 cross-device FL, only a small fraction of clients (from dozens to thousands) out of a large population (possibly in the millions) will participate in a training round.

Original: Moreover, in practical FL systems (Bonawitz et al., 2019; Paulik et al., 2021), clients can only participate when  local criteria are met, such as when mobile devices are plugged into power, 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 this reason, the central server has limited control over the “sampling” process.

Original: In addition, clients that satisfy the local criteria and participate  at a speciﬁc time of a day can have intra-domain features compared to inter-domain clients from another time, possibly due to user mobile usage preference and time zone difference, which cause a periodically shifting  distributionof clients 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 cultural 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).

Original: Conventional centralized training often assumes independent identically distributed (IID) data. FL must deal with non-IID (heterogeneous) client data because the orchestrator (server) is unable to random sample or shufﬂe the decentralized private data. In cross-device FL, the training objective is the combination of client objectives based on an unknown population distribution, and even if we considered the empirical risk minimization formulation of the training set of clients, it is challenging because of the heterogeneity and lack of sampling control due to privacy and system constraints (Wang et al., 2021). Transfer learning, multi-task learning, and meta learning are introduced in
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: FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al., 2017; Khodak et al., 2019). Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al., 2021). More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al. (2020); Ghosh et al.
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: (2020), while still assuming the clients can be uniformly accessed during the training process.

Original: We provide further discussions of related works in Appendix A.4.

Original: To our knowledge, there are very few previous works (Eichner et al., 2019; Ding et al., 2020) discussed 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 will alternatively participate in training.
Revised: Both 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  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.

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: 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  inherits these issues.
Revised: (Ding et al., 2020) is a variant of semi-cyclic SGD that inherits these issues.

Revised: We provide further discussions of related works in Appendix A.5.

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

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

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 a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (F ED TEM) to train the multi-branch networks.

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 passes from the night mode to day mode, 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.
Revised: 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 such temporal prior, we can train models that are even more accurate than models trained with uniformly sampled clients.

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

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 consider the federated learning algorithm for a set of clients
Revised: We consider the federated learning algorithm for a set of clients

Original: I , and the i -th client has data samples D i .
Revised: I , and the i -th client has data samples D i .

Original: The clients are heterogeneous, while the data on each client are IID (independent and identically distributed) sampled from its own distribution.
Revised: The clients are heterogeneous, while the data on each client are IID (independent and identically distributed) sampled from its own distribution.

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  ( 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  thefederated training curves in a cross-device FL system have daily oscillation.
Revised: Figure 1 (left) shows that the training loss in a cross-device FL system has daily oscillation.

Original: Such oscillation suggests domain difference of daytime clients and nighttime clients was also observed in (Yang et al., 2018), and Yang et al.
Revised: Such oscillations was also observed in (Yang et al., 2018), and Yang et al.

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

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 arealternatively from the daytime group and the nighttime group fora continuous T/ 2 rounds.
Revised: (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.

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

Original: Smooth transition.
Revised: Smooth transition.

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: We also assume the distribution changes periodically with a period of T . Unlike (Eichner et al., 2019), we assume clients at round t are a mixture of daytime clients and nighttime clients, denoted as 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 abruptly jumping 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 will gradually shift, rather than abruptly jumping from one mode to another as in the block-cyclic structure.

Original: To better approximate the behavior in practice, we focus on periodic distribution shifts in which clients come from the day mode I 1 ( t ) with a probability q ( t ) that varies smoothly from the day mode to night mode.
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 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 ) :
Revised: We consider two types of functions for q ( t ) :

Original: Linear (L) and Cosine (C) , each further parameterized by an exponent factor p > 0 to control  the transition,
Revised: Linear (L) and Cosine (C) , 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 6 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.
Revised: This simulates 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 approximate the curves in real FL system.
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 in real FL systems.

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

Original: Semi-cyclic SGD (Eichner et al., 2019) provides the worst-case guarantees when the domains of each block are unrelated, which  does not fully exploit the data in practice.
Revised: 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.

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

Original:  M ITIGATING P ERIODICAL D ISTRIBUTION S HIFT
Revised:  M ITIGATING 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 a multi-branch network and a Gaussian mixture model (GMM).
Revised: In this section, we propose F ED TEM to tackle the periodical distribution shift, in which we joint train a multi-branch network and a Gaussian mixture model (GMM).

Original: The GMM can infer latent variables that determine a participating client is from daytime mode or nighttime mode.
Revised: The GMM can infer latent variables that determine a participating client is from daytime mode or nighttime mode.

Original: The performance of the inference is enhanced by priors about the probability of daytime (nighttime) clients at the current round.
Revised: The performance of the inference is enhanced by priors about the probability of daytime (nighttime) clients at the current round.

Original: A multi-branch network is simultaneously trained on the participating clients based on the inferred latent variables, with each branch corresponding to a mode.
Revised: A multi-branch network is simultaneously trained on the participating clients based on the inferred latent variables, with each branch corresponding to a mode.

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

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: We consider applications in which clients are training a model with common input and output space, despite the distribution shift of clients.

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

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

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

Original: In this paper, we set each of the prediction branches to be a single linear layer.
Revised: In this paper, we set each of the prediction branches to be a single linear layer.

Original: A k -branch network is potentially more communication efﬁcient and data efﬁcient than having k copies of models.
Revised: A k -branch network is potentially more communication efﬁcient and data efﬁcient than having k copies of models.

Original: 3. L EARNING THE M IXTURE M ODEL WITH T EMPORAL P RIORS
Revised: 3. L EARNING THE M IXTURE M ODEL WITH T EMPORAL P RIORS

Original: We learn a mixture model together with the multi-branch network to infer the mode and select the corresponding branch for prediction of each client.
Revised: We learn a mixture model together with the multi-branch network to infer the mode and select the corresponding branch for prediction of each client.

Original: 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.
Revised: 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: In addition, (Marfoq et al., 2021) also assumes stateful clients, which can suffer from staleness in cross-device setting (Wang et al., 2021).

Original: Our approach has three key difference: (1) our k -branch network is more communication and data efﬁcient compared to previous k networks, (2) our model selects the branches based on the features instead of the loss, and therefore does not need labeled data to select branches for new clients; (3) we propose the temporal prior to tackle periodic distribution shift when learning the mixture model.
Revised: 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 on the features instead of the loss, and therefore does not need labeled data to select branches for new clients; (3) we propose the temporal prior to tackle periodic distribution shift when learning the mixture model.

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 mode a sample and a client comes from, respectively, which will be inferred by our model.

Original: Our model also assumes all samples from the same client are from the same mode during training, so the prior P ( ζ ) = P ( z ) , but our model can still be used to estimate the posterior for each sample during inference.
Revised: Our model also assumes all samples from the same client are from the same mode during training, so the prior P ( ζ ) = P ( z ) , but our model can still be used to estimate the posterior for each sample during inference.

Original: For convenience, we also introduce a vector π with each entry π k = P ( z = k ) to denote the prior distribution of z .
Revised: For convenience, we also introduce a vector π with each entry π k = P ( z = k ) to denote the prior distribution of z .

Original: We model
Revised: We model

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 ( x | z = k ) as a Gaussian distribution N ( f ( w , x ) | µ k , σ 2 k ) on the feature representation f ( w , x ) .

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 σ 2 k for the Gaussian models.

Original: We use θ to denote the parameter of the GMM for convenience.
Revised: 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.
Revised: 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.

Original: We assume all samples on one client are from the same mode, so P ( ζ ) = P ( z ) .
Revised: We assume all samples on one client are from the same mode, so P ( ζ ) = P ( z ) .

Original: There are two options for inferring the modes
Revised: There are two options for inferring the modes

Original: 1 Given the current mixture model, estimate P ( ζ |D i ) from P ( x |
Revised: 1 Given the current mixture model, estimate P ( ζ |D i ) from P ( x |

Original: z ) of all x ∈ D i ; 

Revised: z ) of all x ∈ D i ;

Original: 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: Given the training samples on the client, ﬁnd the Maximum Likelihood Estimation (MLE) of P ( z ) for the client and use it for P ( ζ ) .

Original: Option 1 exploits the current parameters of the mixture model to infer the mode of each client.
Revised: Option 1 exploits the current parameters of the mixture model to infer the mode of each client.

Original: With the assumption that samples in D i are independent, from Bayes’ theorem, we can get P ( z |D i ) as
Revised: With the assumption that samples in D i are independent, from Bayes’ theorem, we can get P ( z |D i ) as

Original: [Equation]
Revised: [Equation]

Original: Intuitively, since the client distribution is constantly shifting, we should not expect the prior P ( ζ ) to be unchanged from the last round.
Revised: Intuitively, since the client distribution is constantly shifting, we should not expect the prior P ( ζ ) to be unchanged from the last round.

Original: Option 2 re-estimates the prior P ( z ) of the mixture model on the current client through MLE and use it as P ( ζ |D i ) .
Revised: Option 2 re-estimates the prior P ( z ) of the mixture model on the current client through MLE and use it as P ( ζ |D i ) .

Original: As shown in Figure 5, we ﬁnd this achieves better empirical results than Option 1.
Revised: As shown in Figure 6, we ﬁnd this achieves better empirical results than Option 1.

Original: Speciﬁcally, we use π ∗ i , the MLE of P ( z ) on the data of client i , to select the branch for training:
Revised: Speciﬁcally, we use π ∗ i , the MLE of P ( z ) on the data of client i , to select the branch for training:

Original: [Equation]
Revised: [Equation]

Original: where π ∗ ik is the k -th dimension of π ∗ i , corresponding to the MLE of P ( z = k ) .
Revised: 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.

Original: After estimating the distribution of ζ , we can either greedily select the branch k ∗ i corresponding to the maximum probability, or we can sample from this distribution.
Revised: After estimating the distribution of ζ , we can either greedily select the branch k ∗ i corresponding to the maximum probability, or we can sample from this distribution.

Original: Our experiments in Figureshow that the greedy approach is more effective.
Revised: Our experiments in Figureshow that the greedy approach is more effective.

Original: Since our experiments only consider the case of two modes, let ¯ k ∗
Revised: Since our experiments only consider the case of two modes, let ¯ k ∗

Original: i be the other, less-likely branch.
Revised: i be the other, less-likely branch.

Original: On each client, we train the model to optimize the following objective
Revised: On each client, we train the model to optimize the following objective

Original: [Equation]
Revised: [Equation]

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

Original: = (cid:15) 1 n + (1 − (cid:15) )
Revised: = (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.
Revised: 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.
Revised: 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.
Revised: 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

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
Revised: 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
Revised: +1 i , and update the parameters of the

Original: GMM based on w t +1 i .
Revised: GMM based on w t +1 i .

Original: In this way, the GMM is synchronized with the latest network parameters and ready for inference immediately after training.
Revised: In this way, the GMM is synchronized with the latest network parameters and ready for inference immediately after training.

Original: The GMM parameters are usually estimated with the Expectation-Maximization (EM) algorithm
Revised: The GMM parameters are usually estimated with the Expectation-Maximization (EM) algorithm

Original: (Dempster et al., 1977).
Revised: (Dempster et al., 1977).

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 , ˜
Revised: 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 , ˜

Original: π t +1 i ) , and then aggregating on the server.
Revised: π t +1 i ) , and then aggregating 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.
Revised: 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.

Original: The server decides which clients are from mode 1 and which are from mode 2 by their ˜ π t
Revised: The server decides which clients are from mode 1 and which are from mode 2 by their ˜ π t

Original: +1 i , since ˜ π t
Revised: +1 i , since ˜ π t

Original: +1 i estimates the probability of an arbitrary sample from client i to come from each mode.
Revised: +1 i estimates the probability of an arbitrary sample from client i to come from each mode.

Original: We give details of the steps below.
Revised: We give details of the steps below.

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

Original: [Equation]
Revised: [Equation]

Original: Temporal prior.
Revised: Temporal prior.

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

Original: The temporal prior ˜ q ( t ) is a rough estimate for the ratio of clients coming from the day mode.
Revised: The temporal prior ˜ q ( t ) is a rough estimate for the ratio of clients coming from the day mode.

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

Original: In between these minima and maxima, we consider three types of ˜ q ( t ) in our current experiments: 1) Linear:
Revised: In between these minima and maxima, we consider three types of ˜ q ( t ) in our current experiments: 1) Linear:

Original: ˜ q ( t ) =
Revised: ˜ q ( t ) =

Original: q L, 1 ( t ) ; 2) Cosine: ˜ q
Revised: q L, 1 ( t ) ; 2) Cosine: ˜ q

Original: ( t ) =
Revised: ( t ) =

Original: q C, 1 ( t ) ; 3) Soft: see Appendix A.2.
Revised: q C, 1 ( t ) ; 3) Soft: see Appendix A.2.

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

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

Original: +1 i as
Revised: +1 i as

Original: [Equation] 
Revised: [Equation] [Equation]

Original: where γ ik ( x ) comes from the E step in Eq.
Revised: where γ ik ( x ) comes from the E step in Eq.

Original: Then, we sort the prior probability ¯ π ∗ i 1 of mode 1 for all sampled clients i ∈
Revised: Then, we sort the prior probability ¯ π ∗ i 1 of mode 1 for all sampled clients i ∈

Original: I (cid:48) ( t ) , and assign (cid:98) ˜ q ( t ) · |I (cid:48)
Revised: I (cid:48) ( t ) , and assign (cid:98) ˜ q ( t ) · |I (cid:48)

Original: ( t ) | + 12 (cid:99) clients with the highest ¯ π ∗ i 1 to mode 1, denoted as ˜ I (cid:48)
Revised: ( t ) | + 12 (cid:99) clients with the highest ¯ π ∗ i 1 to mode 1, denoted as ˜ I (cid:48)

Original: 1 ( t ) .
Revised: 1 ( t ) .

Original: The remaining clients, denoted as ˜ I (cid:48) 2 ( t ) , are assigned to mode 2.
Revised: The remaining clients, denoted as ˜ I (cid:48) 2 ( t ) , are assigned to mode 2.

Original: As a result, for i 1 ∈ ˜ I (cid:48) 1 ( t ) , P ( ζ = 1 |D i , θ t , ˜ q ( t ))
Revised: As a result, for i 1 ∈ ˜ I (cid:48) 1 ( t ) , P ( ζ = 1 |D i , θ t , ˜ q ( t ))

Original: = 1 , while for i 2 ∈ ˜ I (cid:48) 2 ( t ) , P ( ζ = 2 |D i , θ t , ˜ q ( t ))
Revised: = 1 , while for i 2 ∈ ˜ I (cid:48) 2 ( t ) , P ( ζ = 2 |D i , θ t , ˜ q ( t ))

Original: = 1 .
Revised: = 1 .

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

Original: M step (client).
Revised: M step (client).

Original: In our federated setting, the M step happens on both the server and the clients.
Revised: In our federated setting, the M step happens on both the server and the clients.

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

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

Original: [Equation]
Revised: [Equation]

Original: For each client, ˜ q ( t ) affects the posterior distribution P ( ζ |D
Revised: For each client, ˜ q ( t ) affects the posterior distribution P ( ζ |D

Original: i , ˜ q ( t )) , and the effect depends on other clients.
Revised: i , ˜ q ( t )) , and the effect depends on other clients.

Original: Therefore, the distributions in Eq.
Revised: Therefore, the distributions in Eq.

Original: 9 should be estimated on the server.
Revised: 9 should be estimated on the server.

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

Original: Wang et al., 2021) so that the server can send the re-estimated prior to the clients.
Revised: Wang et al., 2021) so that the server can send the re-estimated prior to the clients.

Original: A viable approach is to restrict P ( ζ |D
Revised: A viable approach is to restrict P ( ζ |D

Original: i , ˜ q ( t )) to have a ﬁnite number of possible values, so that the client can solve Eq.
Revised: i , ˜ q ( t )) to have a ﬁnite number of possible values, so that the client can solve Eq.

Original: 9 for all possible P ( ζ
Revised: 9 for all possible P ( ζ

Original: |D i , ˜ q ( t )) and send all feasible solutions to the server.
Revised: |D i , ˜ q ( t )) and send all feasible solutions to the server.

Original: To make ˜ q ( t ) effective, we should set a minimum of K different values for P ( ζ |D
Revised: To make ˜ q ( t ) effective, we should set a minimum of K different values for P ( ζ |D

Original: i , ˜ q ( t )) so that each mode corresponds to a different value.
Revised: i , ˜ q ( t )) so that each mode corresponds to a different value.

Original: To learn more distinctive representations for the modes, we let P ( ζ = k |D i , ˜ q ( t ))
Revised: To learn more distinctive representations for the modes, we let P ( ζ = k |D i , ˜ q ( t ))

Original: = 1 if the client is estimated to come from mode k , so P ( ζ = j |D i , ˜ q ( t ))
Revised: = 1 if the client is estimated to come from mode k , so P ( ζ = j |D i , ˜ q ( t ))

Original: = 

Revised: =

Original: for all j (cid:54)
Revised: for all j (cid:54)

Original: =
Revised: =

Original: k , and Eq.
Revised: k , and Eq.

Original: 9 becomes
Revised: 9 becomes

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 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

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

Original: j .
Revised: j .

Original: The last equation comes from our modeling assumption that all x ∈ D i are from the same mode.
Revised: The last equation comes from our modeling assumption that all x ∈ D i are from the same mode.

Original: From Eq.
Revised: From Eq.

Original: 10, we can see the optimal solution is to keep parameters for other modes unchanged, while only updating them for the k -th mode as
Revised: 10, we can see the optimal solution is to keep parameters for other modes unchanged, while only updating them for the k -th mode as

Original: [Equation]
Revised: [Equation]

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

Original: M step (server).
Revised: M step (server).

Original: The server collects the MLE of the priors { ¯ π ∗ i | i ∈ I (cid:48)
Revised: The server collects the MLE of the priors { ¯ π ∗ i | i ∈ I (cid:48)

Original: ( t ) }
Revised: ( t ) }

Original: (Eq.
Revised: (Eq.

Original: 8) and all the feasible solutions { (˜ µ t +1 i , ˜ σ t +1 i ) | i ∈
Revised: 8) and all the feasible solutions { (˜ µ t +1 i , ˜ σ t +1 i ) | i ∈

Original: I (cid:48) ( t ) } from the clients.
Revised: I (cid:48) ( t ) } from the clients.

Original: Then, with the temporal prior, it divides the clients into two sets I (cid:48) 1 ( t ) and I (cid:48) 2 ( t ) as described before.
Revised: Then, with the temporal prior, it divides the clients into two sets I (cid:48) 1 ( t ) and I (cid:48) 2 ( t ) as described before.

Original: Similar as FedAvg (McMahan et al., 2017), it updates the parameters for each mode k weighted by the number of samples from each client as
Revised: Similar as FedAvg (McMahan et al., 2017), it updates the parameters for each mode k weighted by the number of samples from each client as

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

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

Original: Setting the prior.
Revised: Setting the prior.

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

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

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

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

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

Original: +1 k = βπ tk
Revised: +1 k = βπ tk

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: +(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.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Revised: [Table]

Original: Dataset.
Revised: Dataset.

Original: We consider two image classiﬁcation datasets, EMNIST and CIFAR, and one next word prediction task on Stack Overﬂow (SO).
Revised: 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.

Original: On Stack Overﬂow, 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 thantokens.
Revised: We truncate each sentence to have no more than 20 tokens.

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

Original: For image classiﬁcation, on EMNIST, we train
Revised: For image classiﬁcation, on EMNIST, we train

Original: LeNet with 2 Conv layers and 2 FC layers, while on CIFAR, we train a ResNet-18 with the Batch
Revised: LeNet with 2 Conv layers and 2 FC layers, while on CIFAR, we train a ResNet-18 with the Batch

Original: Normalization replaced by Group Normalization (Wu & He, 2018) for stability in federated learning.
Revised: Normalization replaced by Group Normalization (Wu & 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[Table] 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 for all modes.

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 GMM, which are 128 and 512 dimensional respectively.

Original: On Stack Overﬂow, 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 and share the remaining weights for both modes, 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.
Revised: We concatenate the average of all tokens’ features from both branches for GMM, which has dimensionality 192.

Original: Methods to compare. The comparison with baselines are shown in Figure 2. For fair comparison, we maintain capacity of the network unchanged. For the baselines, we train the same multi-branch network without any branch-selection, and take the averaged output from both branches for prediction.
Revised: With these conﬁgurations, the models on

Original: We consider a baseline trained under the same distribution shift (denoted as Vanilla), and a strong baseline trained without any distribution shift where the clients are sampled uniformly from all clients (denoted as “No Dist.

Original: Shift” or NDS). For our method, F ED TEM, we compare the effect of the three different temporal priors (Linear, Cosine and Soft). We also compare with the version without temporal prior (“EM Only”) to demonstrate the efﬁcacy of the prior, 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. 8 and Eq.
Revised: EMNIST, CIFAR and SO increase the communication cost by 0.7%, 0.5% and 1.6%, respectively.

Original: We also compare with an approach (“Min Loss”) similar to IFCA (Ghosh et al., 2020), where same as IFCA, 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 during test when labels are unavailable, which is achieved by evaluating the

Original: KLD between the prediction and a uniform distribution on the labels.

Original: We have also applied the same label smoothing regularization (Eq. 6) to “Min Loss” and found it improves the baseline results.

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.

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

Original: We use β = 0 .
Revised: We use β = 0 .

Original: 99 for the moving average of π k in all experiments.
Revised: 99 for the moving average of π k in all experiments.

Original: For F ED TEM and Min Loss, we also do a grid search on the label smoothing parameters (cid:15) and λ .
Revised: For F ED TEM and Min Loss, 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 set of hyperparameters, we run 3 experiments with different random seeds and report their mean and standard error. Methods to compare.

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

Revised: Note for Linear and

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

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

Original: [Figure]
Revised: With these observations, we use the linear prior on EMNIST and CIFAR ( FedTEM (linear) ) and the cosine prior on SO ( FedTEM (cosine) ), to compare with other methods. Results are shown in Figure 2. For fair comparisons, we adapt all methods into our settings to maintain two important conditions unless otherwise stated: 1) training the same multi-branch network to keep model capacities the same; 2) labels are not available for test clients.

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.

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

Revised: Shift ” (NDS).In

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

Revised: 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: 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: 2) T-shot K-means , which is an enhancement of the one-shot

Revised: 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: For both training and test, it selects the branches for each sample according to the nearest cluster center.

Revised: 3) “ Min Loss (IFCA) ”, where same as

Revised: 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: Besides, we have also applied the same label smoothing regularization (Eq. 6) to IFCA, which improves the results.

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

Revised: For testing, we let it “cheat” and assume all test samples are labeled, so that it can execute its modiﬁed EM steps based on the loss to select the branches.

Original: Main Results
Revised: Main Results

Original: The best results of all compared methods are shown in Figure 2.
Revised: The best results of all compared 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  training, causing the worst-case decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and Stack Overﬂow.
Revised: 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.

Original: The Min Loss approach alleviates the drop and sometimes catches up with NDS.

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

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% on the three datasets.

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  other methods in most cases. 
Revised: 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. For other methods, SCSGD performs surprisingly well on EMNIST when p is large, but the performance quickly drops on the more complicated

Original: Among different temporal priors for F ED TEM, we ﬁnd the linear prior toperform the best in the two image classiﬁcation tasks, while the smoother cosine prior performs the best on SO.
Revised: CIFAR and SO. 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. Also note it is not straightforward to do clustering on raw data of SO so we have not implemented it.

Revised: Min Loss (IFCA) alleviates the drop and sometimes catches up with

Revised: NDS, but cannot achieve improvements over NDS.

Revised: By contrast, results of FedEM is the worst even when it is “cheating” on the test clients.

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

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

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

Original: On EMNISTwhere the task is relatively easy and  network  is simple, the  smoothingregularization 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 of the branches will not be updated simultaneously with the feature extractor,causing instability and resulting in  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: During evaluation, for image classiﬁcation, we ﬁnd it beneﬁcial in most cases 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 4.
Revised: We show its beneﬁt in Figure 5.

Original: In practice, similar approach can be realized through caching.
Revised: In practice, similar approach 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 average of persample certainty.

Original: The baseline model even gets worse performance with this approach, as shown in
Revised: The baseline model even gets worse performance with this approach, as shown in

Original: Figure 4, indicating our model is much better at distinguishing two modes in expectation.
Revised: Figure 5, 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 7, 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 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.

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

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

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 light-weight 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: 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;

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

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

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
Revised: The dataset we used for simulations are all available in Tensorﬂow

Original: Federated, and our implementation is largely based on the code from Adaptive Federated Optimization.
Revised: Federated and our implementation is largely based on the code from Adaptive Federated Optimization.

Original: 1 We have listed our hyperparameters in Table 2.
Revised: 1 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 Algorithmand the detailed formulas and derivations of our algorithm in Section 3.1.

Original: We are committed to opensource our implementation upon publication.
Revised: We are committed to opensource our implementation upon publication.

Original: A A PPENDIX A. D ERIVATION OF O PTION
Revised: A A PPENDIX A. D ERIVATION OF O PTION

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

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.
Revised: 14 and Eq.

Original: 15 back into Eq. 13,
Revised: 15 back into Eq. 13,

Original: and
Revised: and

Original: ignore terms that do not depend on π t +

Revised: ignore terms that do not depend on π t +

Original: , we ﬁnd we are solving the following problem for π t +
Revised: , we ﬁnd we are solving the following problem for π t +

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

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

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: 18 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]
Revised: [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: 18, we get the equation as desired.

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Figure]

Original: [Figure]
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: 7 as

Original: [Equation]

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. A BLATION STUDIES : THE EFFECT OF TEMPORAL PRIOR

Revised: 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: 8 and Eq.

Revised: 12); 2) FedTEM (linear) :

Revised: our algorithm where the temporal prior is the periodic linear 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 )) =

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

Revised: The results are shown in Figure 4.

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

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

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.

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, 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 5, 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. [Figure] [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure] [Figure]
Revised: [Figure] [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) is the only previous work 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. Transfer learning, multi-task learning, and meta learning are introduced into FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al., 2017; Khodak et al., 2019). Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al., 2021). More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al.

Revised: Ghosh et al.

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: Several methods have proposed tostudy the  distribution shift under clustering or mixture models of clients  in FL.
Revised: We do not assume strong control over the clients and the available client population is always changing due to the periodical distribution shift. We also do not require stateful clients, since our prior is applied globally to all clients. Our mixture model is based on the feature space, therefore we do not require labeled data for unseen clients to infer its mode. To our knowledge, existing works in clustered FL, including those on personalization, mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting.

Original: Clustered Federated Learning (Sattler et al., 2020) applies an one shot clustering  based on trained global model, and then train a personalized model for each cluster.
Revised: Dennis et al. (2021) propose a one-shot clustering approach based on the raw client data before training, which implicitly assumes all representative clients are available simultaneously.

Original: Sattler et al.
Revised: Clustering on raw data can also be unreliable when the data demonstrates complicated distributions in the input space. Clustered Federated Learning (Sattler et al., 2020) applies an one-shot clustering based on a trained global model, and then train a personalized model for each cluster.

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

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 preprint (Marfoq et al., 2021) proposes a Federated 
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:  EM algorithm to learn a mixture model for each client and weigh the predictions from multiple models.
Revised: The modiﬁed EM algorithm requires computing the loss and therefore labeled data for every client. It maintains a different prior distribution for every client, which requires stateful clients and strong control over client sampling.

Original: By enforcing the probability in mixture models of a client to be an one-hot vector, (Marfoq et al., 2021) can recover the clustering algorithm in (Ghosh et al., 2020; Mansour et al., 2020).
Revised: The hierarchical or bi-partitioning clustering process of (Briggs et al., 2020; Fu et al., 2021) requires all [Figure] 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: 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: (Xie et al., 2021) maintains one set of weights on each client while maintaining multiple weights as cluster centers on the server.

Original: However, these methods do not consider the periodic distribution shift.
Revised: It takes two communication rounds for each weight update, which can be impractical. It remains unclear if this mechanism can be generalized to unseen clients. (Marfoq et al., 2021) trains a different set of weights for each node/client, which requires all the nodes/clients to be available all the time. It is more suitable for the cross-silo setting but not practical for on-device setting.

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

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: Moreover, these methods are designed for personalization when labelled data can be accessed during inference time to calculate loss and select the corresponding personalized models.
Revised: In addition, there is no guarantee that the gradient at the chosen round can separate the clusters well. And to compute the gradients, it still requires labeled data. (Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which is not practical for on-device FL since BN has been widely observed to cause instability in on-device FL, and is often replaced by Group Normalization which does not track running statistics (Hsieh et al., 2020; Hsu et al., 2020). Another limitation is it assumes every client is seen during training.

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).

Revised: All the aforementioned three works require labeled data and additional training efforts to learn the weights of 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.

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

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

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

Original:  In this paper, we design multi-branch networks to mitigate periodical distribution shift by training each branch with corresponding clients.
Revised: Shift”) so that their capacities are the same as our method. Since these two baselines are just training the networks with FedAdam, it is more natural to train single-branch networks directly. In Figure 9, we show that the single-branch and multi-branch networks obtain similar results under various settings, with or without distribution shift. Our method still outperforms these two single-branch baselines.

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

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

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

Revised: To see which factor has more signiﬁcant effect on the test accuracy, we compare with the

Revised: [Figure]

Revised: 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: From the plot, we can see FedTEM maintains the advantage over the baselines and oracle (No Dist.

Revised: Shift) under various training set sizes, though its accuracy also decays as the baselines.

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

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

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

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.

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

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).

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