Original: D IURNAL OR N OCTURNAL ? F EDERATED L EARNING  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: A BSTRACT
Revised: A BSTRACT

Original: Federated learning has been deployed to train machine learning models from de - centralized 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 block cyclic pattern as previous works we model it with a mixture of distributions that grad ually shifts between daytime and nighttime modes, and ﬁnd this intuitive model to better match the observations in practical federated learning systems.
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 shifts between daytime and nighttime modes, and ﬁnd this intuitive model to better match the observations in practical federated learning systems.

Original: Furthermore , we propose to jointly train a clustering model and a multi -branch network to allocate lightweight specialized branches to clients from different modes
Revised: Furthermore, we propose to jointly train a clustering model and a multi-branch network to allocate lightweight specialized branches to clients from different modes.

Original: A tem - poral prior is used to signiﬁcantly boost the training performance.
Revised: A temporal prior is used to signiﬁcantly boost the training performance.

Original: Experiments for image classiﬁcation on EMNIST and CIFAR datasets, and next word predic- tion on the Stack Overﬂow dataset show that the proposed algorithm can counter the effects of the distribution shift and signiﬁcantly improve the ﬁnal model performance
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 counter the effects of the distribution shift and signiﬁcantly improve the ﬁnal model performance.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In Federated Learning (FL), many clients collaboratively train a machine learning model with decentralized data under the orchestration of a central server (Kairouz et al., 2019).
Revised: In Federated Learning (FL), many clients collaboratively train a machine learning model with decentralized data under the orchestration of a central server (Kairouz et al., 2019).

Original: FL is designed for privacy protection: the private data of local clients will never be directly transferred to the server or shared with other clients, which follows the principle of data minimization and keeps the attack surface of the system small (Wang et al., 2021).
Revised: FL is designed for privacy protection: the private data of local clients will never be directly transferred to the server or shared with other clients which follows the principle of data minimization and keeps the attack surface of the system small (Wang et al., 2021).

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

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

Original: Speciﬁcally, cross-device FL has been used in practice to improve utility and privacy of applications such as next word prediction (Hard et al., 2018), emoji suggestion (Ramaswamy et al., 2019), query suggestion (Yang et al., 2018), out-of-vocabulary word discovery (Chen et al., 2019), and keyword trigger models (Granqvist et al., 2020; Hard et al., 2020).
Revised: Speciﬁcally, cross-device FL has been used in practice to improve utility and privacy of applications such as next word prediction (Hard et al., 2018), emoji suggestion (Ramaswamy et al., 2019), query suggestion (Yang et al., 2018), out-of-vocabulary word discovery (Chen et al., 2019), and keyword trigger models (Granqvist et al., 2020; Hard et al., 2020).

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

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

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

Original: In practical FL systems (Bonawitz et al., 2019; Paulik et al., 2021), clients can only participate when the local criteria is met, such as when mobile devices are charging, idle, and connected to an unmetered network.
Revised: In practical FL systems (Bonawitz et al., 2019; Paulik et al., 2021), clients can only participate when the local criteria is met, such as when mobile devices are charging, idle, and connected to an unmetered network.

Original: For the server, clients that satisfy their local criteria and participate training at different times of the day are usually from different time zones that can have signiﬁcant differences, which can cause a periodically shifting data distribution that may degrade the training stability and ﬁnal model performance (Yang et al., 2018; Eichner et al., 2019).
Revised: For the server, clients that satisfy their local criteria and participate training at different times of the day are usually from different time zones that can have signiﬁcant differences, which can cause a periodically shifting data distribution that may degrade the training stability and ﬁnal model performance (Yang et al., 2018; Eichner et al., 2019).

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

Original: However, due to the privacy and system constraints (Wang et al., 2021), the orchestrator (server) in FL systems is not allowed to collect the raw user data, and must deal with such non-IID, heterogeneous data distribution.
Revised: However, due to the privacy and system constraints (Wang et al., 2021), the orchestrator (server) in FL systems is not allowed to collect the raw user data, and must deal with such non-IID, heterogeneous data distribution.

Original: To our knowledge, there are only a few previous works (Eichner et al., 2019; Ding et al., 2020) discussing periodical distribution shift of client population in federated learning.
Revised: To our knowledge, there are only a few previous works (Eichner et al., 2019; Ding et al., 2020) discussing periodical distribution shift of client population in federated learning.

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

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

Original: (2019) proposed the semi-cyclic SGD approach, where clients participating training at different time slots will contribute to and only use the corresponding model of the group.
Revised: (2019) proposed the semi-cyclic SGD approach, where clients participating training at different time slots will contribute to and only use the corresponding model of the group.

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

Original: non-cyclic setting under their assumptions.
Revised: non-cyclic setting under their assumptions.

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

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

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

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

Original: Instead of adopting the block-cyclic structure (Eichner et al., 2019) we assume a smooth transition between the daytime mode and nighttime mode, and empirically verify through simulation that its impact on training better matches the observation in practical FL systems.
Revised: Instead of adopting the block-cyclic structure (Eichner et al., 2019), we assume a smooth transition between the daytime mode and nighttime mode, and empirically verify through simulation that its impact on training better matches the observation in practical FL systems.

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

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

Original: Unlike (Mansour et al., 2020; Ghosh et al., 2020; Marfoq et al., 2021), the feature-based clustering model does not rely on labelled data, and can be easily applied for inference on new clients. 
Revised: Unlike (Mansour et al., 2020; Ghosh et al., 2020; Marfoq et al., 2021), the feature-based clustering model does not rely on labelled data, and can be easily applied for inference on new clients. We propose to use the temporal

Original: We propose to use the temporal prior of the client distribution to enhancethe clustering models.
Revised: prior of the client distribution to enhance

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 goes from nighttime to daytime, and vice versa.
Revised: the clustering models. We assume participating clients per communication round is a mixture of daytime and nighttime clients, and the number of participating clients from the daytime group will gradually increase as time goes from nighttime to daytime, and vice versa.

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

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

Original: We perform extensive experiments, where the multibranch networks trained by our method outperform the distribution-oblivious baselines by a large margin: 3-5% on EMNIST, 2-14% on CIFAR, and 0.4-1.35% on the challenging Stack Overﬂow dataset under various degrees of distribution shifts.
Revised: We perform extensive experiments, where the multibranch networks trained by our method outperform the distribution-oblivious baselines by a large margin: 3-5% on EMNIST, 2-14% on CIFAR, and 0.4-1.35% on the challenging Stack Overﬂow dataset under various degrees of distribution shifts.

Original: By leveraging the temporal priors, the proposed method can take advantage of the periodic distribution shift and beat the strong baselines of training with uniformly sampled clients by 4%, 4% and 0.45%, respectively.
Revised: By leveraging the temporal priors, the proposed method can take advantage of the periodic distribution shift and beat the strong baselines of training with uniformly sampled clients by 4%, 4% and 0.45%, respectively.

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

Original: FL setting.
Revised: FL setting.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: Such oscillations was also observed by Yang et al.
Revised: Such oscillations was also observed by Yang et al.

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

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

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

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

Original: Smooth transition.
Revised: Smooth transition.

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: We consider a periodic Linear function (L) and a smooth Cosine function (C) for q ( t ) , each further parameterized by an exponent factor p > 0 to control smoothness of the transition, 
Revised: We consider a periodic Linear function (L) and a smooth Cosine function (C) for q ( t ) , each further parameterized by an exponent factor p > 0 to control smoothness of the transition, [Equation] ,p

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: Figure 1 (middle) simulates the training curve with q L, 1 ( t ) to control the probability for sampling from I 1 ( t ) , which more accurately approximates the curves from a practical 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 from a practical FL system.

Original: The three curves show that training a single model around the clock achieves the best performance on both modes despite the domain differences, which motivates us to train a model with a large amount of shared weights.
Revised: The three curves show that training a single model around the clock achieves the best performance on both modes despite the domain differences, which motivates us to train a model with a large amount of shared weights.

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

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

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

Original: We consider applications in which clients collaboratively train a model with a common input and output space, and the distribution of clients changes with time.
Revised: We consider applications in which clients collaboratively train a model with a common input and output space, and the distribution of clients changes with time.

Original: We propose to jointly train a multi- branch network and a clustering model to assign specialized branches for prediction on clients from different modes, guided by the temporal prior of the client distributions.
Revised: We propose to jointly train a multibranch network and a clustering model to assign specialized branches for prediction on clients from different modes, guided by the temporal prior of the client distributions.

Original: We consider two methods for enforcing the temporal prior: F ED TEM (Section 3.1), which is based on a Gaussian Mixture
Revised: We consider two methods for enforcing the temporal prior: F ED TEM (Section 3.1), which is based on a Gaussian Mixture

Original: Model (GMM), and F ED TKM ( Section 3.2), which is based on K-means clustering.
Revised: Model (GMM), and F ED TKM ( Section 3.2), which is based on K-means clustering.

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

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

Original: (2019) showed the merits of training separate models for different data distributions when the distributions can be identiﬁed during evaluation.
Revised: (2019) showed the merits of training separate models for different data distributions when the distributions can be identiﬁed during evaluation.

Original: However, training a single model with shared feature extraction layers on all available data can usually improve the data efﬁciency of representation learning.
Revised: However, training a single model with shared feature extraction layers on all available data can usually improve the data efﬁciency of representation learning.

Original: For example, for vision tasks, it helps to learn to extract common low-level features, while for language tasks it helps to learn shared embeddings and grammars from the context.
Revised: For example, for vision tasks, it helps to learn to extract common low-level features, while for language tasks it helps to learn shared embeddings and grammars from the context.

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

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

Original: Temporal prior.
Revised: Temporal prior.

Original: The temporal prior ˜ q ( t ) is an estimate for the ratio of clients coming from the daytime cluster, q ( t ) .
Revised: The temporal prior ˜ q ( t ) is an estimate for the ratio of clients coming from the daytime cluster, q ( t ) .

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

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

Original: [Equation]
Revised: Linear: ˜ q ( t ) = q L, 1 ( t ) ; 2) Cosine: ˜ q ( t ) =

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

Original: Training Objective.
Revised: Training Objective.

Original: Our model assumes clients are from either the daytime or the nighttime cluster ( K = 2 ), so it has two branches, and each branch is used for one group of clients.
Revised: Our model assumes clients are from either the daytime or the nighttime cluster ( K = 2 ), so it has two branches, and each branch is used for one group of clients.

Original: Let k ∗ i be the branch index of client i chosen by the clustering model, and ¯ k ∗ i be the index of the other branch.
Revised: Let k ∗ i be the branch index of client i chosen by the clustering model, and ¯ k ∗ i be the index of the other branch.

Original: Each client trains the network with the following objectiveminimize
Revised: Each client trains the network with the following objective

Original: [Equation]
Revised: [Equation]

Original: where s ( (cid:15), y ) = (cid:15) 1 n +(1 − (cid:15) )
Revised: where s ( (cid:15), y ) = (cid:15) 1 n +(1 − (cid:15) )

Original: y is the label smoothing function for one-hot vector y , (cid:15) ∈ [0 , 1] determines the amount of label smoothing, and λ > 0 is the regularization strength.
Revised: y is the label smoothing function for one-hot vector y , (cid:15) ∈ [0 , 1] determines the amount of label smoothing, and λ > 0 is the regularization strength.

Original: The label smoothing regularization updates the other branch jointly with the feature extractor to prevent staleness, while encouraging the two branches to specialize in different feature subspaces: branch ¯ k ∗ i is trained to become less certain on features of samples from cluster k ∗ i .
Revised: The label smoothing regularization updates the other branch jointly with the feature extractor to prevent staleness, while encouraging the two branches to specialize in different feature subspaces: branch ¯ k ∗ i is trained to become less certain on features of samples from cluster k ∗ i .

Original: 3. L EARNING A G AUSSIAN M IXTURE M ODEL WITH P ER - CLIENT T EMPORAL S TATISTICS
Revised: 3. L EARNING A G AUSSIAN M IXTURE M ODEL WITH P ER - CLIENT T EMPORAL S TATISTICS

Original: We propose Federated Expectation-Maximization with Temporal prior (F ED TEM) to learn a Gaussian Mixture Model (GMM) to infer the cluster a client is from, and select the corresponding branch in the multi-branch network.
Revised: We propose Federated Expectation-Maximization with Temporal prior (F ED TEM) to learn a Gaussian Mixture Model (GMM) to infer the cluster a client is from, and select the corresponding branch in the multi-branch network.

Original: We deﬁne discrete latent variables z and ζ to represent which cluster a sample and a client is from, respectively.
Revised: We deﬁne discrete latent variables z and ζ to represent which cluster a sample and a client is from, respectively.

Original: F ED TEM assumes samples on the same client are from the same cluster, so the GMM prior P ( ζ  ) = P ( z  )  .
Revised: F ED TEM assumes samples on the same client are from the same cluster, so the GMM prior P ( ζ = k ) = P ( z = k ) for any k .

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

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

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: The prior P ( ζ ) of the client distribution are expected to change due to the periodic distribution shift.
Revised: The prior P ( ζ ) of the client distribution are constantly changing due to the periodic distribution shift.

Original: The i -th client estimates its probability of coming from the k -th cluster based on its data D i before the local update steps during training.
Revised: To stay up-to-date, the i -th client estimates its probability of coming from the k -th cluster based on its data D i before the local update steps during training.

Original: Speciﬁcally, given the GMM parameters,  the local maximum likelihood estimation (MLE) of P ( ζ i ) on each new client i is estimated by the MLE of p ( z ) on client i .
Revised: Speciﬁcally, given the GMM parameters, since P ( ζ ) = P ( z ) , the local maximum likelihood estimation (MLE) of P ( ζ i ) on each new client i is equal to the MLE of p ( z ) on client i .

Original: Let π ∗
Revised: Let π ∗

Original: i be the MLE of P ( z ) on client i , where the k -th dimension of π ∗ i , denoted as π ∗ ik , is the MLE of P ( z = k ) on client i .
Revised: i be the MLE of P ( z ) on client i , where the k -th dimension of π ∗ i , denoted as π ∗ ik , is the MLE of P ( z = k ) on client i .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

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

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

Original: We train the network with the estimated P ( ζ i ) .
Revised: We select the branch based on the MLE of P ( ζ i ) and train the network by optimizing Eq. 4.

Original: We can either greedily select the branch k ∗ i = arg max k P ( ζ i = k ) and optimize Eq. 4, or sample from the  distribution P ( ζ i) .
Revised: We can either greedily select the branch k ∗ i = arg max k π ∗ i , or sample from the discrete distribution π ∗ i .

Original: We show the greedy approach achieves better empirical results in Figure 6, and use it by default. 
Revised: We show the greedy approach achieves better empirical results in Figure 6, and use it by default. optimizing Eq. 4 and get ˜ w t

Revised: +1 i ; 9:

Revised: On D i , compute the MLE ˜ π ∗ i and all feasible Gaussian parameters (

Original: [Table]
Revised: ˜ µ t +1 i , ˜ σ t +1 i ) ; (cid:46) see Eq. Server aggregates the model update (cid:80) p i ( ˜ w t +1 i − w t ) 11: Update network parameters to w t +1 with the prescribed server optimizer.

Revised: Server collects GMM params { ˜ π ∗ i , ˜ µ t

Revised: +1 i , ˜ σ t +1 i | i ∈

Revised: I (cid:48) ( t ) } ; 13:

Revised: Update the GMM parameters to ( µ t +1 , σ t +1 , π t +1 ) with the temporal prior on ¯ π ∗ i ; (cid:46) see Eq.

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

Original: We introduce a federated Expectation-Maximization (EM) (Dempster et al., 1977) enhanced by the temporal prior to update the parameters of the GMM.
Revised: We introduce a federated Expectation-Maximization (EM) (Dempster et al., 1977) enhanced by the temporal prior to update the parameters of the GMM.

Original: The temporal prior is enforced via a bottomup approach: we start by running EM on each client i and send all possible locally optimal GMM parameters to the server, and then select and aggregate the GMM parameters with the temporal prior on the server.
Revised: The temporal prior is enforced via a bottomup approach: we start by running EM on each client i and send all possible locally optimal GMM parameters to the server, and then select and aggregate the GMM parameters with the temporal prior on the server.

Original: We ﬁnd the optimal GMM parameters after the local update steps as described in
Revised: We ﬁnd the optimal GMM parameters after the local update steps as described in

Original: Section 3.1.1, so that the GMM is updated in the feature space of the updated network.
Revised: Section 3.1.1, so that the GMM is updated in the feature space of the updated network.

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

Original: E step.
Revised: E step.

Original: For each client i ∈ I ( t ) , evaluate the posterior γ ik ( x ) for each sample x to infer the probability of x coming from cluster k , based on the locally updated network with parameters w t
Revised: For each client i ∈ I ( t ) , evaluate the posterior γ ik ( x ) for each sample x to infer the probability of x coming from cluster k , based on the locally updated network with parameters w t

Original: +

Revised: +

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: M step on clients.
Revised: M step on clients.

Original: To learn a GMM that better distinguishes the daytime and nighttime clusters, we incorporate the temporal prior ˜ q ( t ) into the M step by considering the posterior P ( ζ = k
Revised: To learn a GMM that better distinguishes the daytime and nighttime clusters, we incorporate the temporal prior ˜ q ( t ) into the M step by considering the posterior P ( ζ = k

Original: |D i , θ t , ˜ q ( t )) , where θ t = ( µ t , σ t , π t ) is the collection of all GMM parameters.
Revised: |D i , θ t , ˜ q ( t )) , where θ t = ( µ t , σ t , π t ) is the collection of all GMM parameters.

Original: The M step optimizes the following objective
Revised: The M step optimizes the following objective

Original: [Equation]
Revised: [Equation]

Original: Since the clients cannot share their data with the server, we ﬁrst ﬁnd locally optimal parameters ˜
Revised: Since the clients cannot share their data with the server, we ﬁrst ﬁnd locally optimal parameters ˜

Original: θ t +1 i = argmax θ (cid:80)
Revised: θ t +1 i = argmax θ [Equation]

Original: Kk =1 P ( ζ = k |D i , θ t , ˜ q ( t )) log P ( D i , ζ = k | θ , ˜ q ( t )) on each client i , and then average on the server.
Revised: Kk = [Equation] P ( ζ = k |D i , θ t , ˜ q ( t )) log P ( D i , ζ = k | θ , ˜ q ( t )) on each client i , and then average on the server.

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

Original: However, it is hard to evaluate the posterior P ( ζ = k
Revised: However, it is hard to evaluate the posterior P ( ζ = k

Original: |D i , θ t , ˜ q ( t )) locally without additional communication rounds, since ˜ q ( t ) is an estimate of the ratio of clients from one cluster, and the posterior for one client depends on the estimates of other clients.
Revised: |D i , θ t , ˜ q ( t )) locally without additional communication rounds, since ˜ q ( t ) is an estimate of the ratio of clients from one cluster, and the posterior for one client depends on the estimates of other clients.

Original: A viable approach is to assume the posterior to be one-hot, i.e., for a certain cluster k ∗ , P ( ζ = k
Revised: A viable approach is to assume the posterior to be one-hot, i.e., for a certain cluster k ∗ , P ( ζ = k

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

Original: = 

Revised: =

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

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

Original: = 

Revised: =

Original: if k (cid:54) =
Revised: if k (cid:54) =

Original: k ∗ .
Revised: k ∗ .

Original: Under this modeling assumption, we can learn more distinctive features for the clusters at no communication overhead.
Revised: Under this modeling assumption, we can learn more distinctive features for the clusters at no communication overhead.

Original: When k ∗ is given, we only need to solve
Revised: To see this, when k ∗ is given, we only need to solve

Original: [Equation]
Revised: [Equation]

Original: where its optimal mean and variance, and the MLE of the GMM prior for cluster k ∗ on client i are ˜µ t +1 ik
Revised: where its optimal mean and variance, and the MLE of the GMM prior for cluster k ∗ on client i are ˜

Original: [Equation]
Revised: [Equation]

Original: = [Equation] Note ˜ µ t +1 ik ∗ and ˜ σ t +1 ik ∗ are identical for different k ∗ ; see Appendix A.2 for derivations.
Revised: Note ˜ µ t +1 ik ∗ and ˜ σ t +1 ik ∗ are identical for different k ∗ ; see Appendix A.2 for derivations.

Original: Clients send the optimal mean and variance to the server.
Revised: The clients send the optimal mean and variance to the server.

Original: To enforce temporal priors on the server, clients also send the MLE of the prior on D i for each cluster k , i.e., ˜ π ∗ ik , to the server.The MLE is deﬁned in Eq.
Revised: To enforce temporal priors on the server, clients also send the MLE of the prior on D i for each cluster k , i.e., ˜ π ∗ ik , to the server.

Original: The total number of parameters sent to the server can be less than a single GMM model.
Revised: In this way, the total number of parameters sent to the server can be less than a single GMM model.

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

Original: The temporal prior is enforced via a bottom-up approach: each client sends all possible solutions and MLEs { (˜ µ t +1 ik , ˜ σ t +1 ik , ˜
Revised: The temporal prior is enforced via a bottom-up approach: each client sends all possible solutions and MLEs { (˜ µ t +1 ik , ˜ σ t +1 ik , ˜

Original: π ∗ ik ) | i ∈ I (cid:48)
Revised: π ∗ ik ) | i ∈ I (cid:48)

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

Original: A set ˜ I (cid:48) 1 ( t ) of (cid:98) ˜
Revised: A set ˜ I (cid:48) 1 ( t ) of (cid:98) ˜

Original: q ( t ) · |I (cid:48)
Revised: q ( t ) · |I (cid:48)

Original: ( t ) | + 12 (cid:99) clients with highest ˜ π ∗ i 1 will have P ( ζ = 1 |D
Revised: ( t ) | + 12 (cid:99) clients with highest ˜ π ∗ i 1 will have P ( ζ = 1 |D

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

Original: = 1 , while the remaining clients, constituting a set ˜ I (cid:48) 2 ( t ) , will have P ( ζ = 2 |D i , θ t , ˜ q ( t ))
Revised: = 1 , while the remaining clients, constituting a set ˜ I (cid:48) 2 ( t ) , will have P ( ζ = 2 |D i , θ t , ˜ q ( t ))

Original: = 1 .
Revised: = 1 .

Original: Then, similar as
Revised: Then, similar as

Original: FedAvg (McMahan et al., 2017), the server updates the GMM parameters for each mode k as the weighted average of the optimal GMM parameters of each client i ∈ ˜ I (cid:48)k [Equation]
Revised: FedAvg (McMahan et al., 2017), the server updates the GMM parameters for each mode k as the weighted average of the optimal GMM parameters of each client i ∈ ˜ I (cid:48)

Original:  t ) :
Revised: k ( t ) :

Original: [Equation]
Revised: [Equation]

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

Original: Note σ t +1 is not necessarily an unbiased estimate, but we ﬁnd this estimate gives good results in practice.
Revised: Note σ t +1 is not necessarily an unbiased estimate, but we ﬁnd this estimate gives good results in practice.

Original: We also use a running average of the GMM prior π t +1 k during training and set it to be a uniform distribution during inference; see Appendix A.4 for details.
Revised: We also use a running average of the GMM prior π t +1 k during training and set it to be a uniform distribution during inference; see Appendix A.4 for details.

Original: 3. L EARNING A C LUSTERING M ODEL WITH A GGREGATED T EMPORAL S TATISTICS
Revised: 3. L EARNING A C LUSTERING M ODEL WITH A GGREGATED T EMPORAL S TATISTICS

Original: F ED TEM collects per-client ˜
Revised: F ED TEM collects per-client ˜

Original: π ∗ i to use the temporal prior and update the GMM parameters, which may not satisfy the strong aggregation-only data minimization principle (Bonawitz et al., 2021;
Revised: π ∗ i to use the temporal prior and update the GMM parameters, which may not satisfy the strong aggregation-only data minimization principle (Bonawitz et al., 2021;

Original: Wang et al., 2021).
Revised: Wang et al., 2021).

Original: We propose an alternative Federated K -Means algorithm augmented by the
Revised: We propose an alternative Federated K -Means algorithm augmented by the

Original: Temporal prior (F ED TKM), where the temporal prior is enforced based on only aggregated results, described in Algorithm 2.
Revised: Temporal prior (F ED TKM), where the temporal prior is enforced based on only aggregated results, described in Algorithm 2.

Original: F ED TKM clusters the i th client and selects the branches based on the averaged distance of the features to the k cluster centers,
Revised: F ED TKM clusters the i th client and selects the branches based on the averaged distance of the features to the k cluster centers,

Original: [Equation]
Revised: [Equation]

Original: where the superscript t denotes the communication round, and α k are distance scalar factors to control the “prior” of the clusters and enforce the temporal priors.
Revised: where the superscript t denotes the communication round, and α k are distance scalar factors to control the “prior” of the clusters and enforce the temporal priors.

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

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

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

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

Original: 11 through the objective in Eq.
Revised: 11 through the objective in Eq.

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

Original: i ( t ) = 

Revised: i ( t ) =

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

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

Original: =
Revised: =

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

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

Original: |

Original: [Equation]
Revised: [Equation]

Original:  i [Equation] k [Equation]
Revised: where r i 1 =

Original: ¯ f (

Original: [Equation]

Original: w ti ) , k = 1 or 2 , (12) where ˜ I 1 ( t ) = { i| i ∈ I (cid:48) ( t ) , d i  1  (˜ w ti ) <
Revised: q i ( t ) , r i 2 = 1 − q i ( t ) .

Original: d i 2 ( ˜ w t i ) } , and ˜ I 2 ( t ) = { i | i ∈ I (cid:48)

Original: ( t ) , d i 1 ( ˜ w t i ) ≥ d i 2 ( ˜ w t i ) } .

Original: The distance factor α 2 is updated by tracking the difference between the empirical quantile ¯ q ( t ) and the oracle quantile ˜ q ( t ) estimated by the temporal prior,
Revised: The distance factor α 2 is updated by tracking the difference between the empirical quantile ¯ q ( t ) and the oracle quantile ˜ q ( t ) estimated by the temporal prior,

Original: [Equation]
Revised: [Equation]

Original: (13) where η t ≥ 0 is the step size of the geometric update at step t .
Revised: 13) where η t ≥ 0 is the step size of the geometric update at step t .

Original: In all experiments, we use the periodical linear function ˜ q ( t ) =
Revised: In all experiments, we use the periodical linear function ˜ q ( t ) =

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

Original: [Table]
Revised: [Table]

Original: 3. C OMPARISONS AND D ISCUSSIONS
Revised: 3. C OMPARISONS AND D ISCUSSIONS

Original: Privacy.
Revised: Privacy.

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

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

Original: Bonawitz et al., 2021).
Revised: Bonawitz et al., 2021).

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

Original: Differences from previous methods.
Revised: Differences from previous methods.

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

Original: FL under temporal distribution shift.
Revised: FL under temporal distribution shift.

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

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

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

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

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

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

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

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

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 in the Appendix.
Revised: The split of the day mode ( I 1 ) and night mode ( I 2 ), and other statistics, are shown in Table 1 in the Appendix.

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

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

Original: For the distribution shift, we use T = 256 for the majority of our results, and we compare results under both linear and cosine distribution shifts with p ∈ { 0 . 1 , 0 . 25 , 0 .
Revised: For the distribution shift, we use T = 256 for the majority of our results, and we compare results under both linear and cosine distribution shifts with p ∈ { 0 . 1 , 0 . 25 , 0 .

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

Original: On EMNIST, we train LeNet with 2 Conv layers and 2 FC layers.
Revised: On EMNIST, we train LeNet with 2 Conv layers and 2 FC layers.

Original: On CIFAR, we train a ResNet-18 with the Batch Norm replaced by Group Norm (Wu
Revised: On CIFAR, we train a ResNet-18 with the Batch Norm replaced by Group Norm (Wu

Original: [Figure]
Revised: [Figure]

Original: & He, 2018) for stability in federated learning.
Revised: & He, 2018) for stability in federated learning.

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

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

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

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

Original: The communication overheads of the extra branch are 0.7%, 0.5% and 1.6% respectively for models on EMNIST, CIFAR and SO.
Revised: The communication overheads of the extra branch are 0.7%, 0.5% and 1.6% respectively for models on EMNIST, CIFAR and SO.

Original: Hyperparameters.
Revised: Hyperparameters.

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

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

Original: For
Revised: For

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

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

Original: Ablation study and temporal prior.
Revised: Ablation study and temporal prior.

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

Original: From the results we can see: 1) all three types of priors improve the baseline in most cases; 2) linear prior works better on image datasets, while cosine prior works better on SO; 3) with weaker assumptions, Soft prior is able to ﬁt the distribution shift and improves the results.
Revised: From the results we can see: 1) all three types of priors improve the baseline in most cases; 2) linear prior works better on image datasets, while cosine prior works better on SO; 3) with weaker assumptions, Soft prior is able to ﬁt the distribution shift and improves the results.

Original: With these observations, for F ED TEM, we choose the linear prior ˜ q L, 1 ( t ) on EMNIST and CIFAR and the cosine prior ˜ q C, 1 ( t ) on SO.
Revised: With these observations, for F ED TEM, we choose the linear prior ˜ q L, 1 ( t ) on EMNIST and CIFAR and the cosine prior ˜ q C, 1 ( t ) on SO.

Original: For F ED TKM, we only consider the linear prior ˜ q L, 1 ( t ) .
Revised: For F ED TKM, we only consider the linear prior ˜ q L, 1 ( t ) .

Original: Baselines.
Revised: Baselines.

Original: For fair comparisons, unless otherwise stated, we adapt all methods into our settings to maintain two important properties: 1) training the same multi-branch network to keep model capacities the same; 2) labels are not available for test clients.
Revised: For fair comparisons, unless otherwise stated, we adapt all methods into our settings to maintain two important properties: 1) training the same multi-branch network to keep model capacities the same; 2) labels are not available for test clients.

Original: For the baseline without any branch selection technique, we train the same multi-branch network by taking the averaged output from both branches for prediction.
Revised: For the baseline without any branch selection technique, we train the same multi-branch network by taking the averaged output from both branches for prediction.

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

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

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

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

Original: During test time, since there is no indicator of which mode the clients come from in practice, we select the models according to the certainty of their predictions, measured by evaluating the KLD between the prediction and a uniform distribution over all labels.
Revised: During test time, since there is no indicator of which mode the clients come from in practice, we select the models according to the certainty of their predictions, measured by evaluating the KLD between the prediction and a uniform distribution over all labels.

Original: 2) T-shot
Revised: 2) T-shot

Original: K-means , which is an enhanced variant of one-shot K-means (Dennis et al., 2021), collecting the cluster centers on raw data from all participating clients during a whole period of T rounds before training.
Revised: K-means , which is an enhanced variant of one-shot K-means (Dennis et al., 2021), collecting the cluster centers on raw data from all participating clients during a whole period of T rounds before training.

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

Original: 3) “ Min Loss (IFCA) ”, where same as IFCA (Ghosh et al., 2020), we choose the branch with minimum loss on the local training set of each client during training, but different from
Revised: 3) “ Min Loss (IFCA) ”, where same as IFCA (Ghosh et al., 2020), we choose the branch with minimum loss on the local training set of each client during training, but different from

Original: IFCA, we choose branches with highest certainty on the unlabeled test set clients, using the same criterion as 1).
Revised: IFCA, we choose branches with highest certainty on the unlabeled test set clients, using the same criterion as 1).

Original: We applied the label smoothing regularization (Eq. 4) to IFCA, which improves the results.
Revised: We applied the label smoothing regularization (Eq. 4) to IFCA, which improves the results.

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

Original: FedEM is implemented with privilege to access sample labels for test clients, to enable its EM steps which uses the loss.
Revised: FedEM is implemented with privilege to access sample labels for test clients, to enable its EM steps which uses the loss.

Original: Main Results.
Revised: Main Results.

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

Original: By comparing results of Vanilla and NDS, we ﬁnd the temporal distribution shift indeed degrades vanilla training, causing the worstcase decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and
Revised: By comparing results of Vanilla and NDS, we ﬁnd the temporal distribution shift indeed degrades vanilla training, causing the worstcase decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and

Original: SO.
Revised: SO.

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

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

Original: F ED TEM on the three datasets.
Revised: F ED TEM on the three datasets.

Original: F ED TEM is not better than NDS when p deviates too much fromas the data distribution is extremely skewed in such settings, while our methods are still better than most other methods in most cases.
Revised: F ED TEM is not better than NDS when p deviates too much fromas the data distribution is extremely skewed in such settings, while our methods are still better than most other methods in most cases.

Original: For baseline methods, SCSGD performs surprisingly well on
Revised: For baseline methods, SCSGD performs surprisingly well on

Original: EMNIST when p is large, but the performance quickly drops on the more complicated CIFAR and
Revised: EMNIST when p is large, but the performance quickly drops on the more complicated CIFAR and

Original: SO.
Revised: SO.

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

Original: We did not implement T-shot K-means for SO due to the difﬁculty of clustering the raw data of the language task.
Revised: We did not implement T-shot K-means for SO due to the difﬁculty of clustering the raw data of the language task.

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

Original: By contrast, FedEM often achieves inferior results even with the privilege access to the test labels for EM estimation.
Revised: By contrast, FedEM often achieves inferior results even with the privilege access to the test labels for EM estimation.

Original: This is probably due to the staleness of the priors in the challenging periodic distribution shift setting, as we simulate cross-device FL on a large population where clients only participate training for a few rounds on EMNIST and CIFAR, and at most one round on SO.
Revised: This is probably due to the staleness of the priors in the challenging periodic distribution shift setting, as we simulate cross-device FL on a large population where clients only participate training for a few rounds on EMNIST and CIFAR, and at most one round on SO.

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

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

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

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

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

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

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

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

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

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

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

Original: The baseline model even gets worse performance with this approach, as shown in Figure 4, indicating our model is much better at distinguishing two modes in expectation.
Revised: The baseline model even gets worse performance with this approach, as shown in Figure 4, indicating our model is much better at distinguishing two modes in expectation.

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

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

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

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

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

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this paper, we showed the inﬂuence of a smoothly shifting distribution between daytime and nighttime modes better matches the phenomenon in practical FL systems, and developed algorithms to train multi-branch networks to tackle the distribution shift.
Revised: In this paper, we showed the inﬂuence of a smoothly shifting distribution between daytime and nighttime modes better matches the phenomenon in practical FL systems, and developed algorithms to train multi-branch networks to tackle the distribution shift.

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

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

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

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

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

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

Original: Our code is available on GitHub.
Revised: Our code is available on GitHub.

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

Original: The dataset we used for simulations are all available in Tensorﬂow Federated, and our implementation is largely based on the code from
Revised: The dataset we used for simulations are all available in Tensorﬂow Federated, and our implementation is largely based on the code from

Original: Adaptive Federated Optimization (Reddi et al., 2021).
Revised: Adaptive Federated Optimization (Reddi et al., 2021).

Original: 2 We have listed our hyperparameters in
Revised: 3 We have listed our hyperparameters in

Original: Table 2.
Revised: Table 2.

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

Original:  A CKNOWLEDGEMENTS
Revised:  A CKNOWLEDGEMENTS

Original: The authors would like to thank Galen Andrew and Brendan McMahan for helpful discussions.
Revised: The authors would like to thank Galen Andrew and Brendan McMahan for helpful discussions.

Original: Goldstein was supported by the National Science Foundation #1912866, the DARPA GARD program, and the Sloan Foundation.
Revised: Goldstein was supported by the National Science Foundation #1912866, the DARPA GARD program, and the Sloan Foundation.

Original: A A PPENDIX A. F URTHER DISCUSSIONS OF BRANCH SELECTION WITH GMM
Revised: A A PPENDIX A. F URTHER DISCUSSIONS OF BRANCH SELECTION WITH GMM

Original: Derivation of the MLE.
Revised: Derivation of the MLE.

Original: Recall θ is the parameters of the GMM model.
Revised: Recall θ is the parameters of the GMM model.

Original: We consider maximizing the expectation of the complete-data log likelihood log P ( D i , Z | θ ) under the posterior distribution
Revised: We consider maximizing the expectation of the complete-data log likelihood log P ( D i , Z | θ ) under the posterior distribution

Original: P ( Z |D i , θ ) , where Z is the collection of latent variables associated with each of the sample x ∈ D i .
Revised: P ( Z |D i , θ ) , where Z is the collection of latent variables associated with each of the sample x ∈ D i .

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

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

Original: Meanwhile,
Revised: Meanwhile,

Original: [Equation]
Revised: [Equation]

Original: Plug Eq.
Revised: Plug Eq.

Original: 15 and Eq.
Revised: 15 and Eq.

Original: 16 back into Eq.
Revised: 16 back into Eq.

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

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

Original: An alternative for branch selection.
Revised: An alternative for branch selection.

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

Original: =
Revised: =

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

Original: Speciﬁcally,
Revised: Speciﬁcally,

Original: [Equation]
Revised: [Equation]

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

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

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

Original: = 

Revised: =

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

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

Original: = 

Revised: =

Original: if k (cid:54) =
Revised: if k (cid:54) =

Original: k ∗ , for some given k ∗ .
Revised: k ∗ , for some given k ∗ .

Original: Under this assumption, we need to solve
Revised: Under this assumption, we need to solve

Original: [Equation]
Revised: [Equation]

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

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

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

Original: +1 i , x ) , respectively.
Revised: +1 i , x ) , respectively.

Original: From Eq.
Revised: From Eq.

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

Original: A. I MPLEMENTATION DETAILS AND PRACTICAL NOTES
Revised: A. I MPLEMENTATION DETAILS AND PRACTICAL NOTES

Original: Setting the GMM prior.
Revised: Setting the GMM 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 it estimates P ( x | z ) accurately on the test set.
Revised: This does not make the GMM worse if it estimates P ( x | z ) accurately on the test set.

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
Revised: +1 k

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Setting η t for F ED TKM.
Revised: Setting η t for F ED TKM.

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

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

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

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

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

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

Original: To quantify the effect of temporal prior on training, we compare four versions of F ED TEM: 1)
Revised: To quantify the effect of temporal prior on training, we compare four versions of F ED TEM: 1)

Original: 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: EM Only : only apply the EM part of our algorithm without the temporal priors, where the only difference from F ED TEM is that the server, in the M step, greedily uses client i to update parameters of mode k ∗ i = arg max k ˜ π ∗ ik (see Eq.

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

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

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

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

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

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

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

Original: The soft prior has very weak assumptions about the prior, only requiring the estimated ratios to be non-increasing or non-decreasing within certain intervals, but still improves the results
Revised: The soft prior has very weak assumptions about the prior, only requiring the estimated ratios to be non-increasing or non-decreasing within certain intervals, but still improves the results

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

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

Original: To quantize the effect the batch-level MLE
Revised: To quantize the effect the batch-level MLE

Original: [Table]
Revised: [Table]

Original: we compare the results of our method with or without
Revised: we compare the results of our method with or without

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

Original: As shown in Figure 6, this only decreases the accuracy of the baseline model, indicating our model is much better at distinguishing the two modes in expectation
Revised: As shown in Figure 6, this only decreases the accuracy of the baseline model, indicating our model is much better at distinguishing the two modes in expectation

Original: [Table]
Revised: [Table]

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

Revised: [Table]

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

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

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

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

Original: [Table]

Original: [Figure]
Revised: [Figure]

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

Original: [Table]

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

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

Original: [Figure]
Revised: [Figure]

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

Original: As mentioned above, Semi-cyclic SGD (Eichner et al., 2019) and (Ding et al., 2020) are the only previous works we are aware of that explicitly consider periodical distribution shift in FL.
Revised: As mentioned above, Semi-cyclic SGD (Eichner et al., 2019) and (Ding et al., 2020) are the only previous works we are aware of that explicitly consider periodical distribution shift in FL.

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

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

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

Original: Notably, FedProx (Li et al., 2020b) applies proximal regularizer when perform client updates; SCAFFOLD (Karimireddy et al., 2020) uses control variates as local states on clients for variance reduction; FedPA (Al-Shedivat et al., 2021) provides a Bayesian view and adopts posterior sampling;
Revised: Notably, FedProx (Li et al., 2020b) applies proximal regularizer when perform client updates; SCAFFOLD (Karimireddy et al., 2020) uses control variates as local states on clients for variance reduction; FedPA (Al-Shedivat et al., 2021) provides a Bayesian view and adopts posterior sampling;

Original: FedGen (Zhu et al., 2021) learns a generator that can be shared among clients; FedRobust (Reisizadeh et al., 2020) applies gradient descent ascent to tackle distribution shifts in the form of afﬁne transforms in the image domain.
Revised: FedGen (Zhu et al., 2021) learns a generator that can be shared among clients; FedRobust (Reisizadeh et al., 2020) applies gradient descent ascent to tackle distribution shifts in the form of afﬁne transforms in the image domain.

Original: Transfer learning, multi-task learning, and meta learning are introduced into FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al., 2017; Khodak et al., 2019).
Revised: Transfer learning, multi-task learning, and meta learning are introduced into FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al., 2017; Khodak et al., 2019).

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

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

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

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

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

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

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

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

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

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

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

Original: To our knowledge, existing works in clustered FL, including those on personalization, mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting.
Revised: To our knowledge, existing works in clustered FL, including those on personalization, mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting.

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

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

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

Original: Clustered Federated Learning (Sattler et al., 2020) applies an one-shot clustering based on a trained global model, and then train a personalized model for each cluster.
Revised: Clustered Federated Learning (Sattler et al., 2020) applies an one-shot clustering based on a trained global model, and then train a personalized model for each cluster.

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

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

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

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

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

Original: During the preparation of this draft, we notice a concurrent work, FedEM (Marfoq et al., 2021), which proposes a Federated EM algorithm to learn a mixture model for each client and weigh the predictions from multiple models.
Revised: During the preparation of this draft, we notice a concurrent work, FedEM (Marfoq et al., 2021), which proposes a Federated EM algorithm to learn a mixture model for each client and weigh the predictions from multiple models.

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

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

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

Original: [Figure]
Revised: [Figure]

Original: clients to be available for clustering simultaneously, violating our assumption that the distribution of the population is constantly changing and is not practical in on-device FL in the real world.
Revised: clients to be available for clustering simultaneously, violating our assumption that the distribution of the population is constantly changing and is not practical in on-device FL in the real world.

Original: In addition, for unseen clients, either additional communications are needed for determining their clusters, or the client has to download models for all clusters, which adds signiﬁcantly to the communication cost.
Revised: In addition, for unseen clients, either additional communications are needed for determining their clusters, or the client has to download models for all clusters, which adds signiﬁcantly to the communication cost.

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

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

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

Original: (Marfoq et al., 2021) trains a different set of weights for each node/client, which requires all the nodes/clients to be available all the time.
Revised: (Marfoq et al., 2021) trains a different set of weights for each node/client, which requires all the nodes/clients to be available all the time.

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

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

Original: This is impractical in our setting since the training population changes throughout the day and we cannot assume clients from both modes are available at the ﬁrst round for clustering.
Revised: This is impractical in our setting since the training population changes throughout the day and we cannot assume clients from both modes are available at the ﬁrst round for clustering.

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

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

Original: (Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which is not practical for on-device FL since BN has been widely observed to cause instability in on-device FL, and is often replaced by Group Normalization which does not track running statistics (Hsieh et al., 2020; Hsu et al., 2020).
Revised: (Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which is not practical for on-device FL since BN has been widely observed to cause instability in on-device FL, and is often replaced by Group Normalization which does not track running statistics (Hsieh et al., 2020; Hsu et al., 2020).

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

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

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

Original: the branchs are either locally stored on clients (Arivazhagan et al., 2019; Liang et al., 2020), or reconstructed based on client data (Singhal et al., 2021).
Revised: the branchs are either locally stored on clients (Arivazhagan et al., 2019; Liang et al., 2020), or reconstructed based on client data (Singhal et al., 2021).

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

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

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

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

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

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

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

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

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

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

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

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

Original: To see which factor has more signiﬁcant effect on the test accuracy, we compare with the baselines trained with or without distribution shift in Figure 10, where we only change the maximum training samples per client while keeping other hyperparameters unchanged.
Revised: To see which factor has more signiﬁcant effect on the test accuracy, we compare with the baselines trained with or without distribution shift in Figure 10, where we only change the maximum training samples per client while keeping other hyperparameters unchanged.

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

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

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

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

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

Original: Data during daytime has higher out of vocabulary rates and lower sequence length indicating these sentences might be more arbitrary or informal than data during nighttime.
Revised: Data during daytime has higher out of vocabulary rates and lower sequence length indicating these sentences might be more arbitrary or informal than data during nighttime.

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

Original: This is because the product is mainly deployed in one region, so during daytime, fewer devices are idle for training and the round completion rates are lower (Yang et al., 2018).
Revised: This is because the product is mainly deployed in one region, so during daytime, fewer devices are idle for training and the round completion rates are lower (Yang et al., 2018).

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