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: This paper investigates the problem of periodic client distribution shift in cross-device FL, as previously investiged by Eichner et al (semi-cyclic SGD). The authors propose to split the NN model in a shared part, with different heads for each component. A GMM model is added on top of the output of the shared part to allocate each client to a given head during training, with a temporal prior for the shift between classes. This GMM is fit via a federated EM scheme, in addition to the federation of the NN weights themselves. Experiments are conducted on EMNIST, CIFAR10/100 and StackOverflow datasets, for 2 components. The results demonstrate the superiority of the proposed method with respect to baseline approaches, including the addition of the temporal prior.