Diurnal or Nocturnal? Federated Learning from Periodically Shifting Distributions
Abstract
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 blockcyclic pattern as previous works, we model it with a mixture of distributions that gradually changes between daytime and nighttime modes, and ﬁnd this intuitive model to better match the observations in practical federated learning systems.
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 classiﬁcation on
EMNIST and CIFAR datasets, and next word prediction on the Stack Overﬂow dataset show that the proposed algorithm can effectively mitigate the impact of the distribution shift and signiﬁcantly improve the ﬁnal model performance.
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 ﬁnance, 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)).
Speciﬁcally, cross-device FL has been used in practice to boost utility and privacy for applications such as next word prediction (Hard et al., 2018), emoji suggestion (Ramaswamy et al., 2019), query suggestion (Yang et al., 2018), out-of-vocabulary word discovery (Chen et al., 2019), and keyword trigger models (Granqvist et al., 2020; Hard et al., 2020).
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 signiﬁcant cultural differences, which cause a periodically shifting data distribution that may degrade the training stability and ﬁnal model performance (Yang et al., 2018; Eichner et al., 2019).
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 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.
(2) It maintains a version of the full model for each clients group, which potentially increases the communication cost or privacy risk.
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.
(Ding et al., 2020) is a variant of semi-cyclic SGD that inherits these issues.
We provide further discussions of related works in Appendix C.5.
In this paper, we study periodical distribution shift of clients, and make the following contributions:
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.
We propose to jointly train a multi-branch network and a Gaussian mixture model to select the branch that better ﬁts the client’s distribution based on the feature representations.
The lightweight branches for the day and night modes only slightly increase the communication cost, but signiﬁcantly 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.
We propose a Federated Expectation-Maximization algorithm enhanced by Temporal priors of the shifting distribution (F ED TEM) 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.
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.
We perform extensive experiments, where the multi-branch networks trained by F ED TEM outperform the distribution-oblivious baselines by a large margin: 3-5% on EMNIST, 2-14% on CIFAR, and 0.4-1.35% on the challenging Stack Overﬂow dataset under various degrees of distribution shifts.
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.
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 D i .
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 ) = (cid:88) i ∈I p
i L i ( w ) , where L i ( w ) =
1 |D i | (cid:88) ξ ∈D i (cid:96) ( w , ξ ) , and (cid:88)
i ∈I p i
= 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.
By setting p i =
|D i | / (cid:80) j ∈I |D
j | , the federated training loss recovers the empirical risk minimization (ERM) objective on all client samples.
For simplicity, we abuse notation and use x ∈ D i to denote a training sample (without label) for client i .
Periodic distribution shift.
A subset of clients I (cid:48) ( 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(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 I 1 ( t ) and I 2 ( 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 I 1 ( t ) with a probability q ( t ) that varies smoothly between 0 and 1.
Speciﬁcally, we simulate I 1 ( t ) and I 2 ( t ) with two disjoint sets of clients with different data distributions, and deﬁne q : R + → [0 , 1] to be periodic function with a period of T .
At each round t , we sample the clients from the following distribution
[Equation]
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,
[Equation]
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 q L, 1 ( t ) to control the probability for sampling from I 1 ( 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 ﬁgure, from which we can infer the peak momentum of daytime clients and nighttime clients in practice and use it to deﬁne a strong prior in the proposed FedTEM algorithm.
Mitigating Periodical Distribution Shift
We propose F ED TEM to tackle the periodical distribution shift, in which we jointly train a multibranch network and a Gaussian mixture model (GMM).
The GMM can infer latent variables that determine a participating client is from daytime mode or nighttime mode.
The learning process of the GMM is enhanced by priors about the probability of daytime (nighttime) clients at the current round.
A multi-branch network is simultaneously trained on the participating clients based on the inferred latent variables, with each branch corresponding to a mode.
Our algorithm satisﬁes many desirable properties for cross-device FL, which are missing in many previous works in clustered
FL or personalization.
The server and clients only exchange information once in a communication round, and clients do not maintain local states.
Also, the GMM model does not rely on labeled data, facilitating deployment on new clients without labeled data.
We will discuss how F ED TEM is applied in cross-device FL in detail in Section A.
[Table]
Experiments
We give the main results in this section, and defer the details in Section B.
Methods to compare.
First, we provide ablation studies on the temporal prior to quantify its efﬁcacy, and see which prior (Linear, Cosine, or Soft) works better in practice.
Note for Linear and Cosine, we always use the same periodical linear or cosine function with p = 1 under all types and p ’s for the distribution shift.
The detailed comparisons are given in Appendix C.3 and Figure 4, from which we can see: 1) all three types of priors improve the baseline in most cases; 2) linear prior works better on image datasets, while cosine prior works better on SO; 3) with weakened assumptions, Soft prior is able to ﬁt the distribution shift and improves the results.
With these observations, we use the linear prior on EMNIST and CIFAR ( FedTEM (linear) ) and the cosine prior on SO ( FedTEM (cosine) ), to compare with other methods.
Results are shown in
Figure 2.
For fair comparisons, we adapt all methods into our settings to maintain two important conditions unless otherwise stated: 1) training the same multi-branch network to keep model capacities the same; 2) labels are not available for test clients.
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.
Such networks trained under shifting data distribution are denoted as “ Vanilla ", while those trained without distribution shift are denoted as “ No Dist.
Shift " (NDS).In Appendix C.6, we also give results of training single-branch models in these two settings, which obtained almost the same results.
For other methods addressing data heterogeneity, we consider: 1) “ SCSGD ", where similar as Semi-Cyclic SGD (Eichner et al., 2019), we train one model during ﬁrst half of each period, and the other model on the other half.
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.
T-shot K-means , which is an enhancement of the 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.
For both training and test, it selects the branches for each sample according to the nearest cluster center.
3) “ Min Loss (IFCA) ", where same as IFCA (Ghosh et al., 2020), we choose the branch with minimum loss on the local training set of each client during training, but different from
IFCA, we choose branches with highest certainty on the unlabeled test set clients, using the same criterion as 1).
Besides, we have also applied the same label smoothing regularization (Eq. 6) to
IFCA, which improves the results.
FedEM (Marfoq et al., 2021), where we use the same algorithm on participating clients during training.
For testing, we let it “cheat" and assume all test samples are labeled, so that it can execute its modiﬁed EM steps based on the loss to select the branches.
Main Results
The best results of all compared methods are shown in Figure 2.
By comparing results of Vanilla and NDS, we ﬁnd the temporal distribution shift indeed causes issues for normal training, causing the worst-case decrease in accuracy by more than 7%, 20% and 1% respectively on EMNIST, CIFAR and SO.
Surprisingly, with F ED TEM, the accuracy can be even higher than models trained in the NDS setting in most cases, indicating F ED TEM obtains more distinctive feature representations by learning to separate the two modes in the feature space and learning specialized prediction branches.
The improvement on the strong NDS baseline can be as high as 4%, 4% and 0.45% on the three datasets.
F ED TEM is not better than NDS when p deviates too much from 1, but the data distribution is extremely skewed in such settings, where F ED TEM still improves over most other methods in most cases.
For other methods, SCSGD performs surprisingly well on EMNIST when p is large, but the performance quickly drops on the more complicated CIFAR and SO.
T-shot
K-means achieves improvements for extreme p ’s on EMNIST, but it is not signiﬁcantly better than “Vanilla" on CIFAR, due to the difﬁculty in reliable clustering in more complicated image spaces.
Also note it is not straightforward to do clustering on raw data of SO
so we have not implemented it.
Min Loss (IFCA) alleviates the drop and sometimes catches up with NDS, but cannot achieve improvements over NDS.
By contrast, results of FedEM is the worst even when it is “cheating" on the test clients.
This is probably due to the staleness of the priors, since we simulate the on-device setting on a large population where clients only participate training for a few rounds on EMNIST and
CIFAR, and at most one round on SO.
Conclusions
In this paper, we demonstrated that the smooth transition model for simulating the periodical distribution shift better matches practical FL systems, and developed F ED TEM algorithm to train multi-branch networks to tackle the distribution shift.
F ED TEM incorporated priors of the temporal distribution shifts into learning a GMM to guide the network to select corresponding branches for clients of different modes.
The GMM is deﬁned in the feature space and does not require labelled data for inference, hence is ready to use for new test clients.
The branches are light-weight with little communication overhead, and the model demonstrates signiﬁcant improvements in test accuracy on the benchmark datasets.
F ED TEM also satisﬁes many other real-world constraints like single round communication and stateless clients (Bonawitz et al., 2019; Paulik et al., 2021;
Wang et al., 2021), which makes it ready to be deployed in practical FL systems.
Acknowledgements
The authors would like to thank Galen Andrew and Brendan McMahan for helpful discussions.
A Details of F ED TEM
Multi-branch network.
We consider applications in which clients are training a model with common input and output space, despite the distribution shift of clients.
Though there exists distribution shift, training a model that shares common features by leveraging data from all clients can enable representation learning with better data efﬁciency.
For vision tasks, for example, a common model can use data across domains to learn common low-level features, while for language tasks it helps with learning common embedding and grammatical rules from the context.
To alleviate the communication overhead, we adopt the weight-sharing strategy from multi-task learning (Caruana, 1997) to train a multi-branch network with shared feature representation layers f ( w f , x ) and specialized prediction branches g k ( w k , f ( w f , x )) for clients from mode k .
In this paper, we set each of the prediction branches to be a single linear layer.
A k -branch network is potentially more communication efﬁcient and data efﬁcient than having k copies of models.
A.1 Learning the Mixture Model with Temporal Priors
We learn a mixture model together with the multi-branch network to infer the mode and select the corresponding branch for prediction of each client.
Similar joint learning strategies are applied in (Ghosh et al., 2020; Marfoq et al., 2021) to learn a clustering or mixture model for selection from multiple networks, but both works have the impractical assumption that labeled data is available for all clients.
In addition, (Marfoq et al., 2021) also assumes stateful clients, which can suffer from staleness in cross-device setting (Wang et al., 2021).
Our approach has three key differences: (1) our k -branch network is more communication and data efﬁcient compared to previous k networks, (2) our model selects the branches based on the features instead of the loss, and therefore does not need labeled data to select branches for new clients; (3) we propose the temporal prior to tackle periodic distribution shift when learning the mixture model.
We deﬁne discrete latent variables z and ζ to represent which mode a sample and a client comes from, respectively, which will be inferred by our model.
Our model also assumes all samples from the same client are from the same mode during training, so the prior P ( ζ ) = P ( z ) , but our model can still be used to estimate the posterior for each sample during inference.
For convenience, we also introduce a vector π with each entry π k = P ( z = k ) to denote the prior distribution of z .
We model
P ( x | z = k ) as a Gaussian distribution N ( f ( w , x ) | µ k , σ 2 k ) on the feature representation f ( w , x ) .
For efﬁciency, we assume and learn a diagonal covariance σ 2 k for the Gaussian models.
We use θ to denote the parameter of the GMM for convenience.
Algorithm 1 summarizes the training process, and the details of each step are provided in the following sections.
A.1.1 Modeling the Client Distributions and Selecting Branches for Training
During training, on each client, before proceeding to the local update steps, we estimate the probabilities of client i coming from each mode k given its training set D i , and then select the corresponding branch.
We assume all samples on one client are from the same mode, so P ( ζ ) = P ( z ) .
There are two options for inferring the modes
1 Given the current mixture model, estimate P ( ζ |D i ) from P ( x |
z ) of all x ∈ D i ;
Given the training samples on the client, ﬁnd the Maximum Likelihood Estimation (MLE) of P ( z ) for the client and use it for P ( ζ ) .
Option 1 exploits the current parameters of the mixture model to infer the mode of each client.
With the assumption that samples in D i are independent, from Bayes’ theorem, we can get P ( z |D i ) as
[Equation]
Intuitively, since the client distribution is constantly shifting, we should not expect the prior P ( ζ ) to be unchanged from the last round.
Option 2 re-estimates the prior P ( z ) of the mixture model on the current client through MLE and use it as P ( ζ |D i ) .
As shown in Figure 6, we ﬁnd this achieves better empirical results than Option 1.
Speciﬁcally, we use π ∗ i , the MLE of P ( z ) on the data of client i , to select the branch for training:
[Equation]
where π ∗ ik is the k -th dimension of π ∗ i , corresponding to the MLE of P ( z = k ) .
For completeness, we give the derivation of Eq.
5 in Appendix C.1.
After estimating the distribution of ζ , we can either greedily select the branch k ∗ i corresponding to the maximum probability, or we can sample from this distribution.
Our experiments in Figure 6 show that the greedy approach is more effective.
Since our experiments only consider the case of two modes, let ¯ k ∗
i be the other, less-likely branch.
On each client, we train the model to optimize the following objective
[Equation]
where λ > 0 is the regularization strength, s (
[Equation]
, y ) =
[Equation]
1 n +
[Equation]
1 −
[Equation]
) y is the label smoothing function with 0 ≤ (cid:15) ≤ 1 to determine the amount of label smoothing.
The regularization enables the other branch to be updated jointly with the feature extractor to prevent mismatch.
Meanwhile, it encourages learning distinctive features for the two clusters through label smoothing: the other branch treats samples from the current mode as outliers and is trained to become less certain on these samples.
A.1.2 Updating the Mixture Model
After the local updates of the network parameters on client i with the chosen branch as described in
Section A.1.1, we obtain the updated network parameters w t
+1 i , and update the parameters of the
GMM based on w t +1 i .
In this way, the GMM is synchronized with the latest network parameters and ready for inference immediately after training.
The GMM parameters are usually estimated with the
Expectation-Maximization (EM) algorithm
(Dempster et al., 1977).
In the federated setting, this can be achieved privately by ﬁrst running EM on each client i to get the locally updated GMM parameters (˜ µ t +1 i , ˜ σ t +1 i , ˜
π t +1 i ) , and then aggregating on the server.
Further, to target the temporal distribution shift, we integrate our priors ˜ q ( t ) into the aggregation step on the server, so that a ratio of ˜ q ( t ) clients will be used to update ( µ t +11 , σ t +11 , π t +11 ) for mode 1, while the remaining clients will only update ( µ t +12 , σ t +12 , π t +12 ) for mode 2.
The server decides which clients are from mode 1 and which are from mode 2 by their ˜ π t
+1 i , since ˜ π t
+1 i estimates the probability of an arbitrary sample from client i to come from each mode.
We give details of the steps below.
E step.
For each client i ∈ I ( t ) , compute the posterior γ ik ( x ) for each sample x to infer the probability of x coming from mode k
[Equation]
Temporal prior.
We integrate the temporal prior into the M step, so that we can estimate GMM parameters that better capture the differences in the modes behind the shifting distributions.
The temporal prior ˜ q ( t ) is a rough estimate for the ratio of clients coming from the day mode.
From the observations in Section 2, we can locate the time when q ( t ) is most likely to be 0 or 1 by observing when the minima and maxima occur from the train curve.
In between these minima and maxima, we consider three types of ˜ q ( t ) in our current experiments: 1) Linear:
˜ q ( t ) =
q L, 1 ( t ) ; 2) Cosine: ˜ q
( t ) =
q C, 1 ( t ) ; 3) Soft: see Appendix C.2.
Conditioned on ˜ q ( t ) , our goal is to estimate the posterior P ( ζ |D i , θ t , ˜ q ( t )) for each client to decide which mode they are from.
To achieve this, similar as in Section A.1.1, we ﬁrst obtain the MLE of the prior on D i using the updated network parameters w t
+1 i as
[Equation]
where γ ik ( x ) comes from the E step in Eq.
Then, we sort the prior probability ¯ π ∗ i 1 of mode 1 for all sampled clients i ∈
I (cid:48) ( t ) , and assign (cid:98) ˜ q ( t ) · |I (cid:48)
( t ) | + 12 (cid:99) clients with the highest ¯ π ∗ i 1 to mode 1, denoted as ˜ I (cid:48)
1 ( t ) .
The remaining clients, denoted as ˜ I (cid:48) 2 ( t ) , are assigned to mode 2.
As a result, for i 1 ∈ ˜ I (cid:48) 1 ( t ) , P ( ζ = 1 |D i , θ t , ˜ q ( t ))
= 1 , while for i 2 ∈ ˜ I (cid:48) 2 ( t ) , P ( ζ = 2 |D i , θ t , ˜ q ( t ))
= 1 .
Note it is possible to design soft assignments where the posteriors are not one-hot, but this one-hot assignment has smaller communication overhead, as we will discuss in the M steps below.
M step (client).
In our federated setting, the M step happens on both the server and the clients.
It is ﬁrst executed locally on each client i , where its objective is to ﬁnd the GMM parameters to maximize the expectation of the log likelihood of data, weighted by the time-varying client-level posterior
P ( ζ = k |D i , θ t , ˜ q ( t ))
[Equation]
For each client, ˜ q ( t ) affects the posterior distribution P ( ζ |D
i , ˜ q ( t )) , and the effect depends on other clients.
Therefore, the distributions in Eq.
9 should be estimated on the server.
However, we cannot send client data to the server, and it is almost impossible in practice to have multiple communication rounds between the server and client (Bonawitz et al., 2019;
Wang et al., 2021) so that the server can send the re-estimated prior to the clients.
A viable approach is to restrict P ( ζ |D
i , ˜ q ( t )) to have a ﬁnite number of possible values, so that the client can solve Eq.
9 for all possible P ( ζ
|D i , ˜ q ( t )) and send all feasible solutions to the server.
To make ˜ q ( t ) effective, we should set a minimum of K different values for P ( ζ |D
i , ˜ q ( t )) so that each mode corresponds to a different value.
To learn more distinctive representations for the modes, we let P ( ζ = k |D i , ˜ q ( t ))
= 1 if the client is estimated to come from mode k , so P ( ζ = j |D i , ˜ q ( t ))
=
for all j (cid:54)
=
k , and Eq.
9 becomes
[Equation]
where d is the dimension of the features, µ kj , σ kj denotes the j -th dimension of µ k , σ k , and we use similar notation for π k and [ f ( w t
+1 i , x )]
j .
The last equation comes from our modeling assumption that all x ∈ D i are from the same mode.
From Eq.
10, we can see the optimal solution is to keep parameters for other modes unchanged, while only updating them for the k -th mode as
[Equation]
In this way, we can pack the all the feasible ˜ µ t +1 ik , [˜ σ t +1 ik ] 2 into matrices that have the same size as one set of GMM parameters, which does not increase the communication cost.
M step (server).
The server collects the MLE of the priors { ¯ π ∗ i | i ∈ I (cid:48)
( t ) }
(Eq.
8) and all the feasible solutions { (˜ µ t +1 i , ˜ σ t +1 i ) | i ∈
I (cid:48) ( t ) } from the clients.
Then, with the temporal prior, it divides the clients into two sets I (cid:48) 1 ( t ) and I (cid:48) 2 ( t ) as described before.
Similar as FedAvg (McMahan et al., 2017), it updates the parameters for each mode k weighted by the number of samples from each client as
[Equation]
where M tk = (cid:80) i ∈I (cid:48) k ( t ) |D i | is the total number of samples of clients assigned to mode k .
Note the resulting σ t +1 is not guaranteed to be an unbiased estimate, but we found this process to give satisfying results in practice.
Setting the prior.
During test, to compare the model performance with established baselines and ensure fairness, we consider a static test set where the number of clients and samples from both
4.0 10.0 p
A cc u r a c y
Accuracy on EMNIST reg (lin.
shift)
no reg (lin.
shift) reg (cos shift) no reg (cos shift) 0.1 0.25 0.
4.0 10.0 p
A cc u r a c y
Accuracy on CIFAR reg (lin.
shift)
no reg (lin.
shift) reg (cos shift) no reg (cos shift) 0.1 0.25 0.
4.0 10.0 p
A cc u r a c y
Accuracy on Stack Overflow reg (lin.
shift)
no reg (lin.
[Figure]
[Figure] [Figure]
shift)
no
[Figure] [Figure]
modes are roughly the same.
Therefore, we use a uniform distribution for the priors on the test set.
This does not make the GMM worse if the underlying P ( x | z ) is the same for the training and test sets, and our estimate of P ( x | z ) is accurate.
During training, we estimate the prior π t in every step to match the shifting distribution.
However, we ﬁnd it beneﬁcial to use a running average of π t
+1 k = βπ tk
+ (1 − β ) π k M tk / ( (cid:80) Kj =1 M tj ) during training, which achieves better results than setting β = 0 or using uniform prior in our preliminary experiments.
B Additional Details of Experiments
Dataset
[Table]
Dataset.
We consider two image classiﬁcation datasets, EMNIST and CIFAR, and one next word prediction task on Stack Overﬂow (SO).
The split of the day mode ( I 1 ) and night mode ( I 2 ), and other statistics, are shown in Table 1.
On SO, we use a vocabulary size of 10K and report the test accuracy without special tokens.
We truncate each sentence to have no more than 20 tokens.
Architecture of the multi-branch networks.
For image classiﬁcation, on EMNIST, we train
LeNet with 2 Conv layers and 2 FC layers, while on CIFAR, we train a ResNet-18 with the Batch
Normalization replaced by Group Normalization (Wu & He, 2018) for stability in federated learning.
We use the last FC layer of these convolutional networks as the multi-branch part and share all the remaining layers for all modes.
We use the output from the shared feature extractor for GMM, which are 128 and 512 dimensional respectively.
On SO, we train a single-layer LSTM (Hochreiter & Schmidhuber, 1997) with a hidden size of 670 and embedding size of 96.
To alleviate the communication overhead and prevent overﬁtting, we deﬁne the branches to be the last projection layer before the ﬁnal prediction layer and share the remaining weights for both modes, which is 15x smaller than the ﬁnal prediction layer.
We concatenate the average of all tokens’ features from both branches for GMM, which has dimensionality 192.
With these conﬁgurations, the models on
EMNIST, CIFAR and SO increase the communication cost by 0.7%, 0.5% and 1.6%, respectively.
Hyperparameters.
We use FedAdam (Reddi et al., 2021) as the optimizer.
Table 2 shows the training hyperparameters.
We use T = 256 for the majority of our results, and we evaluate both linear and cosine distribution shifts under p ∈ { 0 . 1 , 0 . 25 , 0 . 5 , 1 , 2 , 4 , 10 } .
We use β = 0 .
99 for the moving average of π k in all experiments.
For F ED TEM and Min Loss, we also do a grid search on the label smoothing parameters (cid:15) and λ .
For each set of hyperparameters, we run 3 experiments with different random seeds and report their mean and standard error.
Effect of label smoothing regularization.
Shown in Figure 3.
Since EMNIST is relatively easy and the network on it is small, the label smoothing does not show signiﬁcant improvements.
However, it is critical to the success on CIFAR and SO.
Without the regularization, one branch will not be updated simultaneously with the feature extractor, resulting in instability and low test accuracy.
MLE for Evaluation.
During evaluation, for image classiﬁcation, we ﬁnd it beneﬁcial in most cases to use the MLE in Eq.
5 on the minibatches (typically of size no larger than 64) to select the best branch, compared with the sample-level inference.
We show its beneﬁt in Figure 5.
In practice, similar approach can be realized through caching.
We ﬁnd Min Loss does not show much difference when a similar approach is applied, where branches are chosen by the average of per-sample certainty.
The baseline model even gets worse performance with this approach, as shown in Figure 5, indicating our model is much better at distinguishing two modes in expectation.
However, this approach does not show improvements for F ED TEM on SO.
Effect of T .
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.
In such scenarios, the distribution changes frequently and abruptly.
However, the performance of F ED TEM is maintained when the distribution changes in a frequent but more balanced way.
C Appendix C.1 Derivation of Option
We consider maximizing the expectation of the complete-data log likelihood log P ( D i , Z | θ t +1 ) under the posterior distribution P ( Z |D i , θ t ) , where Z denotes the collection of latent variables associated with each of the sample x ∈ D i .
Formally, it is solving the following constrained optimization problem
[Equation]
We introduce an indicator function 1 [ z = k ] which is 1 if z = k or 0 otherwise.
In this way, the log likelihood can be represented as
[Equation] [Equation]
Meanwhile,
[Equation]
Plug Eq.
14 and Eq.
15 back into Eq. 13,
and
ignore terms that do not depend on π t +
, we ﬁnd we are solving the following problem for π t + maximize
[Equation]
to
[Equation]
The Lagrangian multiplier of this problem is
[Equation]
Let the derivative w.r.t.
π k be 0,
[Equation]
Multiply both sides of Eq.
18 with π k , sum over 1 ≤ k ≤ K , and apply the constraint that (cid:80)
Kk =1 π k = 1 , we have
[Equation]
Plug this back into Eq.
18, we get the equation as desired.
[Equation]
C.2 The Soft Prior
The soft prior ˜ q S ( t ) gives strong signals only at time steps where we believe q ( t ) = 0 or q ( t )
= 1 .
In between, it only requires the ratio to be non-increasing or non-decreasing, where the ratio is estimated using the posterior from Eq.
7 as
[Equation]
i.e., the ratio of samples whose posterior has higher probability on mode 1.
In this way, the model can still learn to distinguish the modes but we do not enforce a strong prior at every time step.
Speciﬁcally, it is deﬁned as
[Equation]
C.3 Ablation studies: the effect of temporal prior
To quantify the effect of temporal prior on training, we compare four versions of our algorithm: 1)
EM Only : only apply the EM part of our algorithm without the temporal priors, where the only difference from F ED TEM is that the server, in the M step, greedily uses client i to update parameters of mode k ∗ i = arg max k ¯ π ∗ ik (see Eq.
8 and Eq.
12); 2) FedTEM (linear) :
our algorithm where the temporal prior is the periodic linear function P ( i ∈ I 1 ( t )) =
q L, 1 ( t ) ; 3) FedTEM (cosine) : our algorithm where the temporal prior is the cosine function P ( i ∈ I 1 ( t )) =
q C, 1 ( t ) ; 4) FedTEM
[Figure]
[Figure]
(soft) : our algorithm with the soft temporal prior introduced in Section C.2.
The results are shown in
Figure 4.
We ﬁnd all three priors improves the “EM Only" baseline in most cases.
“FedTEM (linear)" is better on image datasets, while “FedTEM (cosine)" is better on the language dataset.
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.
C.4 Verifying the effect of MLE for evaluation
To quantize the effect the batch-level MLE, we compare the results of our method with or without
MLE on EMNIST and CIFAR.
As shown in Figure 6, this only decreases the accuracy of the baseline model, indicating our model is much better at distinguishing the two modes in expectation.
[Table]
C.5 Additional Related Works
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.
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.
Heterogeneity and client distribution shift.
Client heterogeneity is an active topic in FL, and various methods have been proposed to address the distribution difference among clients.
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; FedGen
[Figure] [Figure]
(Zhu et al., 2021) learns a generator that can be shared among clients; FedRobust (Reisizadeh et al., 2020) applies gradient descent ascent to tackle distribution shifts in the form of afﬁne transforms in the image domain.
Transfer learning, multi-task learning, and meta learning are introduced into FL to explicitly handle distribution shifts among clients assuming heterogeneous clients are from different domains or tasks (Smith et al., 2017; Khodak et al., 2019).
Continual learning is introduced to handle distribution shift due to streaming tasks on each client (Yoon et al., 2021).
More recently, distribution shift in clusters of clients instead of each individual client are studied in Mansour et al.
Ghosh et al.
(2020), while still assuming the clients can be uniformly accessed during the training process.
We kindly ask interested readers to ﬁnd more papers on heterogeneity in a few recent surveys (Kairouz et al., 2019;
Li et al., 2020a; Wang et al., 2021).
All these methods consider the distributions shift among different clients, while the proposed F ED TEM considers periodic distribution shift of client population.
Clustering and mixture models.
We do not assume strong control over the clients and the available client population is always changing due to the periodical distribution shift.
We also do not require stateful clients, since our prior is applied globally to all clients.
Our mixture model is based on the feature space, therefore we do not require labeled data for unseen clients to infer its mode.
To our knowledge, existing works in clustered FL, including those on personalization, mostly fail to satisfy at least one of the three properties and therefore not applicable in our setting.
Dennis et al.
(2021) propose a one-shot clustering approach based on the raw client data before training, which implicitly assumes all representative clients are available simultaneously.
Clustering on raw data can also be unreliable when the data demonstrates complicated distributions in the input space.
Clustered
Federated Learning (Sattler et al., 2020) applies an one-shot clustering based on a trained global model, and then train a personalized model for each cluster.
To achieve this, it implicitly assumes strong control over population and clients sampling.
Ghosh et al.
and Mansour et al.
(2020) propose a similar algorithm that alternatively perform clustering and updating the personalized models for corresponding clusters.
Both of them require labled data for the clients to compute the loss for model selection.
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.
The modiﬁed EM algorithm requires computing
[Figure]
the loss and therefore labeled data for every client.
It maintains a different prior distribution for every client, which requires stateful clients and strong control over client sampling.
The hierarchical or bi-partitioning clustering process of (Briggs et al., 2020; Fu et al., 2021) requires all 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.
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.
(Xie et al., 2021) maintains one set of weights on each client while maintaining multiple weights as cluster centers on the server.
It takes two communication rounds for each weight update, which can be impractical.
It remains unclear if this mechanism can be generalized to unseen clients.
(Marfoq et al., 2021) trains a different set of weights for each node/client, which requires all the nodes/clients to be available all the time.
It is more suitable for the cross-silo setting but not practical for on-device setting.
(Duan et al., 2021) clusters the clients based on the similarities between their gradients at the initial model weights.
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.
In addition, there is no guarantee that the gradient at the chosen round can separate the clusters well.
And to compute the gradients, it still requires labeled data.
(Andreux et al., 2020) advocates using different sets of BNs to track different running statistics for different silos, which is not practical for on-device FL since BN has been widely observed to cause instability in on-device
FL, and is often replaced by Group Normalization which does not track running statistics (Hsieh et al., 2020; Hsu et al., 2020).
Another limitation is it assumes every client is seen during training.
Multi-branch networks in FL.
Multi-branch networks have been explored in FL when a personalized model is preferred for each client:
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).
All the aforementioned three works require labeled data and additional training efforts to learn the weights of the branches.
By comparison, our method does not need re-training the branches and does not need labeled data for unseen clients to select the branches.
The branch selection is based on the mixture model in the feature space, and the weights of the branches are ﬁxed.
C.6 Comparing single-branch and multi-branch baselines without clustering
For fair comparisons, we have used two branches for the baselines (“Vanilla" and “No Dist.
Shift") so that their capacities are the same as our method.
Since these two baselines are just training the networks with FedAdam, it is more natural to train single-branch networks directly.
In Figure 9, we show that the single-branch and multi-branch networks obtain similar results under various settings, with or without distribution shift.
Our method still outperforms these two single-branch baselines.
C.7 The effect of local training set sizes
When the number of training samples is small, the EM estimates may have high variance, causing misspeciﬁcations and resulting in worse performance.
Meanwhile, models trained on smaller training sets tend to have worse generalization, so the test accuracy will drop for any method in general.
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.
From the plot, we can see FedTEM maintains the advantage over the baselines and oracle (No Dist.
Shift) under various training set sizes, though its accuracy also decays as the baselines due to the decreased numbers of training samples.