Diurnal or Nocturnal? Federated Learning of Multi-branch Networks from Periodically Shifting Distributions

Federated learning has been deployed to train machine learning models from decentralized client data on mobile devices in practice . The clients available for training are observed to have periodically shifting distributions changing with the time of day , which can cause instability in training and degrade the model performance . In this paper , instead of modeling the distribution shift with a block-cyclic pattern as previous works , we model it with a mixture of distributions that gradually changes between daytime and nighttime modes , and find this intuitive model to better match the observations in practical federated learning systems . We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution ( FedTEM ) , which jointly learns a mixture model to infer the mode of each client , while training a network with multiple light-weight branches specializing at different modes . Experiments for image classification on EMNIST and CIFAR datasets , and next word prediction on the Stack Overflow dataset show that the proposed algorithm can effectively mitigate the impact of the distribution shift and significantly improve the final model performance . 1 INTRODUCTION . In Federated Learning ( FL ) , many clients collaboratively train a machine learning model with decentralized data under the orchestration of a central server ( Kairouz et al. , 2019 ) . FL is designed for privacy protection : the private data of local clients will never be directly transferred to the server or shared with other clients , which follows the principle of data minimization and keeps the attack surface of the system small ( Wang et al. , 2021 ) . Initially introduced for decentralized training on mobile devices ( McMahan et al. , 2017 ) , FL has been widely applied for various different applications including finance , health , digital assistance and personalized recommendations ( see a few recent surveys ( Yang et al. , 2019 ; Kairouz et al. , 2019 ; Li et al. , 2020a ; Lim et al. , 2020 ; Wang et al. , 2021 ) ) . Specifically , cross-device FL has been used in practice to boost utility and privacy for applications such as next word prediction ( Hard et al. , 2018 ) , emoji suggestion ( Ramaswamy et al. , 2019 ) , query suggestion ( Yang et al. , 2018 ) , out-of-vocabulary word discovery ( Chen et al. , 2019 ) , and keyword trigger models ( Granqvist et al. , 2020 ; Hard et al. , 2020 ) . A typical communication round of FL starts with a server broadcasting a global model to clients . Clients then perform local computation on private data and only send back aggregated model updates . Finally , the server aggregates the client updates and apply them to the global model before beginning the next round . In practical FL systems ( Bonawitz et al. , 2019 ; Paulik et al. , 2021 ) , clients can only participate when the local criteria is met , such as when mobile devices are charging , idle , and connected to an unmetered network . For the server , clients that satisfy their local criteria and participate training at different times of the day are usually from different time zones that can have significant cultural differences , which cause a periodically shifting data distribution that may degrade the training stability and final model performance ( Yang et al. , 2018 ; Eichner et al. , 2019 ) . For centralized systems where client data can be collected , such a problem may be mitigated by caching and uniformly sampling from cached data . However , due to the privacy and system constraints ( Wang et al. , 2021 ) , the orchestrator ( server ) in FL systems is not allowed to collect the raw user data , and must deal with such non-IID , heterogeneous data distribution . To our knowledge , there are only a few previous works ( Eichner et al. , 2019 ; Ding et al. , 2020 ) discussing periodical distribution shift of client population in federated learning . Both works assume a block-cyclic structure where daytime clients and nighttime clients alternately participate in training . Eichner et al . ( 2019 ) proposed the semi-cyclic SGD approach , where clients participating training at different time slots will contribute to and only use the corresponding model of the group . However , there are several caveats of semi-cyclic SGD that makes it difficult to apply in practice : ( 1 ) It assigns models to clients based on their participation time , but not all clients will participate in federated learning , hence it is hard to decide the correct group for these clients . ( 2 ) It maintains a version of the full model for each clients group , which potentially increases the communication cost or privacy risk . ( 3 ) The assumption of the abrupt switch from the daytime group to the nighttime group at a specific time of a day is unintuitive in practice . ( Ding et al. , 2020 ) is a variant of semi-cyclic SGD that inherits these issues . We provide further discussions of related works in Appendix A.5 . In this paper , we study periodical distribution shift of clients , and make the following contributions : 1 . We revisit the periodical distribution shift . Instead of adoopting the block-cyclic structure ( Eichner et al. , 2019 ) , we assume a smooth transition between the day mode and night mode , and empirically verify through simulation that its impact on training better matches the observation in practical FL systems . 2 . We propose to jointly train a multi-branch network and a Gaussian mixture model to select the branch that better fits the client ’ s distribution based on the feature representations . The lightweight branches for the day and night modes only slightly increase the communication cost , but significantly improve the model performance . Unlike ( Mansour et al. , 2020 ; Ghosh et al. , 2020 ; Marfoq et al. , 2021 ) , the feature-based mixture model does not rely on labelled data , and can be easily applied can be easily applied for inference on new clients . 3 . We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution ( FEDTEM ) to train the multi-branch networks . We assume participating clients per communication round is a mixture of daytime and nighttime clients , and the number of participating clients from the daytime group will gradually increase as time passes from the night mode to day mode , and vice versa . We use such temporal prior to guide the EM algorithm in the M step to learn parameters that better distinguish the two modes . By exploiting such temporal prior , we can train models that are even more accurate than models trained with uniformly sampled clients . 4 . We provide simulations of the distribution shift on three benchmark datasets ( EMNIST , CIFAR and Stack Overflow ) to evaluate the empirical performance of FL algorithms under the periodic distribution shift with smooth transition . We perform extensive experiments , where the multibranch networks trained by FEDTEM outperform the distribution-oblivious baselines by a large margin : 3-5 % on EMNIST , 2-14 % on CIFAR , and 0.4-1.35 % on the challenging Stack Overflow dataset under various degrees of distribution shifts . By leveraging the temporal priors , the proposed method can take advantage of the periodic distribution shift and beat the strong baselines of training with uniformly sampled clients by 4 % , 4 % and 0.45 % , respectively . 2 MODELING PERIODICAL DISTRIBUTION SHIFT . FL setting . We consider the federated learning algorithm for a set of clients I , and the i-th client has data samples Di . The clients are heterogeneous , while the data on each client are IID ( independent and identically distributed ) sampled from its own distribution . We minimize the expected loss on all clients , minimize w L ( w ) = ∑ i∈I piLi ( w ) , where Li ( w ) = 1 |Di| ∑ ξ∈Di ` ( w , ξ ) , and ∑ i∈I pi = 1 , ( 1 ) where pi is the weight ( probability ) of client i , and ξ = ( x , y ) is a training sample pair of data and label on a client . By setting pi = |Di|/ ∑ j∈I |Dj | , the federated training loss recovers the empirical risk minimization ( ERM ) objective on all client samples . For simplicity , we abuse notation and use x ∈ Di to denote a training sample ( without label ) for client i . Periodic distribution shift . A subset of clients I ′ ( t ) ⊂ I ( t ) are available for training in a communication round t. I ( t ) periodically changes with most of the clients from a daytime ( nighttime ) client group in midday ( midnight ) . Figure 1 ( left ) shows that the training loss in a cross-device FL system has daily oscillation . Such oscillations was also observed in ( Yang et al. , 2018 ) , and Yang et al . ( 2018 ) conjecture it is due to the domain differences between clients from different time zones and cultural backgrounds . Eichner et al . ( 2019 ) study the block-cyclic structure , where training takes place over T rounds each day , and clients are from the day mode and night mode alternately , each last continuously for T/2 rounds . We simulate the training curves of block-cyclic structure in Figure 1 ( right ) , and observe it is different from the ( left ) curves from real FL systems . Smooth transition . We also assume the distribution changes periodically with a period of T . Unlike ( Eichner et al. , 2019 ) , we assume clients at round t are a mixture of daytime clients and nighttime clients , denoted as I1 ( t ) and I2 ( t ) , respectively . Intuitively , since the available population is usually large around the clock , the population distribution of available clients will gradually shift , rather than abruptly jumping from one mode to another as in the block-cyclic structure . To better approximate the behavior in practice , we assume that in each period , clients come from the day mode I1 ( t ) with a probability q ( t ) that varies smoothly between 0 and 1 . Specifically , we simulate I1 ( t ) and I2 ( t ) with two disjoint sets of clients with different data distributions , and define q : R+ → [ 0 , 1 ] to be periodic function with a period of T . At each round t , we sample the clients from the following distribution P ( i ∈ I1 ( t ) ) = q ( t ) , P ( i ∈ I2 ( t ) ) = 1− q ( t ) . ( 2 ) We consider two types of functions for q ( t ) : Linear ( L ) and Cosine ( C ) , each further parameterized by an exponent factor p > 0 to control smoothness of the transition , qL , p ( t ) = ∣∣∣∣2 t mod TT − 1 ∣∣∣∣p , qC , p ( t ) = [ 12 ( cos ( 2πt/T ) + 1 ) ] p . ( 3 ) We visualize the transition probability q ( t ) in Figure 7 of the appendix . When p < 1 , more daytime clients are available in T rounds , and when p > 1 , more nighttime clients are available . This simulates the difference in the number of completed training rounds during daytime and nighttime observed in practice ( Yang et al. , 2018 ) . Observation and insight . Figure 1 ( middle ) simulates the training curve with qL,1 ( t ) to control the probability for sampling from I1 ( t ) , which more accurately approximates the curves in real FL systems . All three curves show that training a single model around the clock achieves reasonable performance on both modes ( much better than random ) despite the domain differences , which motivates us to train a model with large amount of shared weights . Semi-cyclic SGD ( Eichner et al. , 2019 ) provides the worst-case guarantees when the domains of each block are unrelated , which argues that training a single model on both modes is not optimal , but does not match our observations . Another important observation is that the training accuracy reaches its minima ( maxima ) when the clients are most biased towards the day mode or night mode , e.g. , round 1024 and 1152 in the middle figure , from which we can infer the peak momentum of daytime clients and nighttime clients in practice and use it to define a strong prior in the proposed FedTEM algorithm .

Summary: his paper proposes a federated learning method to address a specific non-IID challenge over clients, i.e., when some clients are only available during the daytime and others are only available at the night, and these two types of clients' data are drawn from a mixture of two corresponding distributions. They study a multi-branch model, similar to multi-task learning, such that all clients share the same backbone network to extract data representations but a different branch is applied on top of the backbone network for different distributions. Only the clients with the same data distribution use the same branch to produce the final prediction. Under the assumption that the data representation produced by the backbone network follows a mixture of isotropic Gaussian distributions, this paper proposes an EM algorithm to assign clients to different clusters: the E-step determines the (soft) membership of each client belonging to each cluster and assigns it to the branches, while the M-step estimates the mean and variance of the representation for each Gaussian component in the mixture. Some additional heuristics are shown to be important for this method to achieve promising performance, e.g., label smoothing, a temporal prior of the ratio of clients from day and night, linear/cosine/soft schedule of this ratio, one-hot (hard) assignment of clients to clusters, etc. In experiments, they simulate the day-night distribution shift on EMNIST, CIFAR, and Stackflow datasets in FL by splitting each dataset into two parts with different data types. They evaluate the proposed method under different shifting settings (different p) with a few baselines proposed in this paper. They also empirically studied the effects of label smoothing and different period lengths.