Tackling Distribution Shifts in Federated Learning with Superquantile Aggregation
Abstract
Federated learning has emerged as the predominant framework for distributed machine learning over decentralized data, e g. in mobile phones.
The usual approaches suffer from a distribution shift: the model is trained to fit the average population distribution but is deployed on individual clients, whose data distributions can be quite different.
We present a distributionally robust approach to federated learning based on a risk measure known as the superquantile and show how to optimize it by interleaving federated averaging steps with quantile computation.
We demonstrate experimentally that our approach is competitive with usual ones in terms of average error and outperforms them in terms of tail statistics of the error
Introduction
Federated learning is a distributed machine learning framework where many clients (e.g. mobile devices) collaboratively train a model under the orchestration of a central server (e.g. service provider), while keeping the training data private and local to the client throughout the training process [16, 10].
It has found widespread adoption across industry [1, 18] for applications ranging from smart device apps [23, 6] to healthcare [2, 9].
A key feature of federated learning is statistical heterogeneity, i.e., client data distributions are not identically distributed [10, 13].
Each client is a user who generates diverse data depending on their unique personal, cultural, regional, and geographical characteristics.
This data heterogeneity in federated learning manifests itself as a train-test distributional shift.
Indeed, the usual approach minimizes the prediction error of the model on average over the population of clients available for training [16] while at test time, the same model is deployed on individual clients.
This approach can be liable to fail on tail clients whose data distribution is far from most of the population or who may have less data than most of the population.
It is highly desirable, therefore, to have a federated learning method that can robustly deliver good predictive performance across a wide variety of natural distribution shifts posed by individual clients.
We present in this paper a robust approach to federated learning that guarantees a minimum level of predictive performance to all clients, even in situations where the population is heterogeneous.
The approach we develop addresses these issues by minimizing a learning objective based on the notion of a superquantile [21, 20], a risk measure that captures the tail behavior of a random variable.
Our algorithm relies on quantile statistics of the losses to filter out clients on which to run federated averaging steps.
Experimental results on benchmark datasets show that our approach yields improved performance on tail clients over a number of state-of-the-art baselines while maintaining competitive performance on the average error.
Proposed Objective and Optimization Algorithm
Suppose we have n clients such as mobile phones.
The loss incurred by the model w ∈ R d on this client i is F i ( w ) :=
E z ∼ p i [ f ( w ; z )] , where p i is the distinct data distribution on client i and f ( w ; ξ ) is the loss function e.g. cross entropy, on data point z .
The usual objective of federated learning [16] is simply the empirical risk minimization (ERM) approach
[Equation]
Owing to the natural statistical heterogeneity in the data, the data distribution p encountered at test time on an unseen test client might be different from the population training distribution p train =
(1 /n ) (cid:80) ni
=1 p i , leading to poor performance on such clients.
Our goal is to improve the performance on such tail clients.
To this end, we directly minimize the average loss across tail clients above a certain tail threshold.
We formalize this through the notion of a risk measure known as the superquantile , a tail summary statistic of random variables [21].
The (1 − α ) -superquantile is defined for a continuous random variable Z and α ∈ (0 , 1) as S α ( Z ) =
E [ Z | Z > Q α ( Z )] , where Q α ( Z ) is the (1 − α ) -quantile of
Z .
A similar interpretation holds for discrete distributions; it is formally defined as
[Equation]
This is an instance of the continuous knapsack problem and can be solved optimally by a greedy algorithm [4].
Assuming u 1 < · · · < u n and αn is an integer, the optimal solution π ⋆ above satisfies π ⋆i = 1 / ( αn ) for i ≥ (1 − α ) n or that the u i ’s larger than their (1 − α ) quantile are averaged.
The ∆ -FL Objective and Distributional Robustness.
Instead of minimizing the average loss as in (1), our proposed framework, called ∆ -FL, minimizes the tail loss across clients, as measured by the superquantile.
Concretely, at level α ∈ (0 , 1) , we minimize
[Equation]
If we have a test client whose distribution p π
=
(cid:80) ni
=1 π
i p i can be written as a mixture of the training distributions p i , . .
. , p n , then the ∆ -FL objective minimizes max π i ≤ 1 / ( αn ) E
z ∼ p
π
[ f ( w ; z )] , the worst-case loss over all mixture distributions with a weight constraint π i ≤ 1 / ( αn ) .
Federated Optimization of ∆ -FL.
To design a federated optimization algorithm to optimize the ∆ -FL objective, we must overcome two challenges: (i) nonsmoothness and (ii) biased gradient estimation.
The superquantile a (cid:55)→ S α ( u 1 , · · · , u n ) is a nonsmooth function, leading to potential difficulties in optimization.
We overcome this challenge by deriving an expression for the subgradient of the ∆ -FL objective.
Concretely, when αn is an integer, we have
[Equation]
, where π ∗
[Equation]
=
[Equation]
and Q α = Q α ( F 1 ( w ) , · · · , F n ( w )) is the (1 − α ) -quantile of the losses.
See Appendix B for a proof.
The second challenge stems from the lack of unbiased gradient estimators for the superquantile.
Given m i.i.d.
copies Z 1 , . .
, Z m of a random variable Z , the empirical mean ¯ Z m =
(1 /m ) (cid:80) mi
{+=
+} Z i is an unbiased estimate of the population mean, i.e., E [ ¯ Z m ] =
E [ Z ] .
This is no longer true for the superquantile, i.e., E [ S α
( Z 1 , · · · , Z m )]
̸ = S α
( Z ) .
As a result, we cannot access unbiased stochastic gradients (here, m is the batch size).
In federated learning, it is not reasonable to assume that we have access to all the clients due to a diurnal availability pattern of clients [10].
We overcome this issue by actually minimizing the expected minibatch superquantile instead, defined as
[Equation]
where U m is the uniform distribution over all subsets of { 1 , . .
. , n } of batch size m .
This is a uniform close surrogate of the original objective [11, Prop. 1]
[Equation]
Using this expression, we design a federated optimization algorithm that steps of the usual federated averaging algorithm [16] with quantile estimation steps.
Specifically, in each communication round, the local updates w + i from the subsample of m selected clients i ∈ S are aggregated to update the global model with the following two steps:
• estimate the quantile ˆ Q α ≈ Q
α ( F i ( w ) : i ∈ S ) of the per-client losses to the server, and • aggregate the updates from tail clients where F i ( w ) ≥
ˆ Q α to find the new global model w + as w + = | S α | (cid:88) i ∈ S
α w
+ i , where S α
= { i : F i ( w ) ≥ ˆ Q α } .
The complete algorithm is given in Appendix A. Similar to the standard FedAvg algorithm [16] for ERM objective (1), this aggregation rule enjoys a simplification in the case of a single local update per-client with a learning rate γ .
Specifically, under the assumption of full client participation (i.e., m = n ), if the local update w − w
+ i = γ ∇ F i ( w ) is a single gradient step and ˆ Q α
= Q
α
( F 1 ( w ) , · · · , F n ( w )) is the exact quantile of the per-client losses, the aggregated update is simply a subgradient step w − w + = γ ∇ F α ( w ) where we denote the subgradient as ∇ F α ( w ) ∈ ∂F
α ( w ) .
Like FedAvg, our algorithm reduces the overall communication cost, which is often the bottleneck in bandwidth-constrained edge devices, while incurring a larger computation cost for each client.
Assuming that each F i is G -Lipschitz and L -smooth, this algorithm exhibits a rate of
O (1 / √ T ) in the nonconvex case in T communication rounds and O (exp( − T/κ 3 / 2 )
+ κ/T ) in the strongly convex case with local condition number κ — we refer to the long version [19] for details.
Numerical Experiments
We demonstrate the effectiveness of ∆ -FL in handling natural distribution shifts in federated learning.
Setup.
We measure the 90 th percentile of the per-client misclassification errors as a measure of the tail performance.
We repeat all experiments 5 times and report the mean and standard deviation.
We consider two learning tasks.
(a) Character Recognition : We use the EMNIST dataset [3], where the input x is a 28 × 28 grayscale image of a handwritten character and the output y is its label (0-9, a-z, A-Z).
Each client is a writer of the character x .
We train both a linear model and a LeNet-type convolutional network.
(b) Sentiment Analysis : We use the Sent140 dataset [5] where the input x is a tweet, and the output y = ± 1 is its sentiment.
Each client is a distinct Twitter user.
We train both a logistic regression and a Long-Short Term Memory neural network architecture (LSTM).
The LSTM is built on the GloVe embeddings of the words of the tweet [8].
Baselines.
We compare ∆ -FL with the following baselines:
We consider two methods which attempt to minimize the usual objective (1): FedAvg [16] and FedProx [14].
The latter augments FedAvg with a proximal term for more stable optimization.
We also consider a few heterogeneity-aware objectives: Tilted-ERM [12], which is the analog of ∆ -FL but using the log-sum-exp function and AFL [17], whose objective is obtained as the limit lim α → 0 F α ( w ) of the ∆ -FL objective.
We also consider q -FFL [15], which raises the per-client loss F i to the ( q + 1) th power, for some q > 0 .
We optimize q -FFL and Tilted-ERM with the federated optimization algorithms proposed in their respective papers.
We use q -FFL with q = 10 in place of AFL, as it was found to have more stable convergence with similar performance.
Hyperparameters.
We fix the number of clients per round to be m = 100 for each dataset-model pair except for Sent140-RNN, for which we use m = 50 .
We fix an iteration budget and tune a learning rate for FedAvg.
The same iteration budget and learning rate schedule were used for all other methods, including ∆ -FL.
All hyperparameters were tuned to find the best tail error ( 90 th percentile).
Results.
The results are in Tables 1 and 2.
We visualize the distribution of test errors in Figure 1. ∆ -FL consistently achieves the smallest 90 th percentile error.
∆ -FL achieves a 3 . 3% absolute ( 12% relative) improvement over any ERM objective on EMNIST-ConvNet.
Among the heterogeneityaware objectives, ∆ -FL achieves 1 . 8% improvement over the next best objective, which is Tilted-
[Table]
[Table] [Table] [Table] [Table]
[Table]
∆ -FL is competitive at multiple values of α .
For EMNIST-ConvNet, ∆ -FL with α ∈ { 0 . 5 , 0 . 8 } is better in 90 th percentile error than all other methods we compare to, and ∆ -FL with α = 0 .
1 is tied with Tilted-ERM, the next best method.
We also empirically confirm that ∆ -FL interpolates between
FedAvg ( α → 1 ) and AFL ( α → 0 ).
Yet, ∆ -FL is competitive in terms of average error.
Perhaps surprisingly, ∆ -FL gets the best test error performance on EMNIST-ConvNet and Sent140-Linear.
This suggests that the average test distribution is shifted relative to the average training distribution p α .
In the other cases, we find that the reduction in mean error is small relative to the gains in the 90 th percentile error.
Minimizing superquantile loss over all clients performs better than minimizing worst error over all clients.
Specifically, AFL, which aims to minimize the worst error among all clients, as well as other objectives which approximate it ( ∆ -FL with α → 0 , q -FFL with q → ∞ ) tend to achieve poor performance.
∆ -FL offers a more nuanced and more effective approach via the constraint set π i ≤ 1 / ( nα ) than the straight pessimistic approach minimizing the worst error among all clients. ∆ -FL yields improved prediction on non-conforming clients.
This can be observed from the histogram of ∆ -FL in Figure 1, which exhibits thinner tails than FedAvg or Tilted-ERM.
We see that the ERM objective of FedAvg sacrifices performance on the nonconforming clients.
Tilted-ERM does improve over FedAvg in this regard, but ∆ -FL has a thinner right tail than Tilted-ERM, showing better handling of heterogeneity. ∆ -FL yields improved prediction on data-poor clients.
We observe in Figure 1 that Tilted-ERM and q -FFL mainly improve the performance on data-rich clients, that is, clients with lots of data.
On the other hand, ∆ -FL gives a more significant reduction in misclassification error on data-poor clients, that is, clients with little data ( < 200 examples per client).
[Figure] [Figure]
Acknowledgements
We acknowledge support from NSF DMS 2023166, DMS 1839371, CCF 2019844, the CIFAR program “Learning in Machines and Brains”, faculty research awards and a JP Morgan PhD fellowship.
This work has been partially supported by MIAI – Grenoble Alpes, (ANR-19-P3IA-0003).
This work was performed while Krishna Pillutla was at the University of Washington
A Pseudocode
The pseudocode of the proposed optimization algorithm is given in Algorithm 1.
[-Algorithm 
-]Proof of the Subgradient Expression (3) . We first give a general expression for the subgradient. Define the notation
[Equation] so that S α ( u 1 , · · · , u n ) = max π ∈P α π ⊤ u . Proposition 1. Fix a w ∈ R d and let π ⋆ ∈ arg max π ∈P α (cid:80) n i =1 π
i F i ( w ) . Then, we have, n (cid:88) [Equation]
where ∂F α ( w ) denotes the regular subdifferential of F α . Input: Initial iterate w (0) , number of communication rounds T , number of clients per round m , number of local updates τ , local step size γ 1: for t = 0 , 1 , · · · , T − 1 do 2:
Sample m clients from [ n ] without replacement in S 3:
Estimate the (1 − α ) -quantile of F i ( w ( t ) ) for i ∈ S ; call this Q ( t ) 4: for each selected client i ∈ S in parallel do 5:
Set ˜ π ( t ) i
= I (cid:0) F i ( w ( t ) ) ≥ Q ( t ) (cid:1) 6:
Initialize w ( t ) k, 0 = w ( t ) 7: for k = 0 , · · · , τ − 1 do 8: w ( t ) i,k +1 =
(1 − γλ ) w ( t ) i,k − γ ∇ F i ( w ( t ) i,k ) 9: end for 10: end for 11: w ( t +1) =
(cid:80) i ∈
S ˜ π
( t )
i w ( t ) i,τ / (cid:80) i ∈
S ˜ π
( t ) i 12: end for 13: return w T
Proof.
Let g n ( w ) =
( F 1 ( w ) , · · · , F n ( w )) denote the concatenation of the losses into a vector.
Then,
F α ( w ) =
S
α ◦ g n ( w ) .
Since S α is convex, we get that its (convex) subdifferential
[e.g., 7, Cor.
4.4.4] is
[Equation]
Since g n is smooth and S α is convex with full domain, we obtain the regular subdifferential of S α ◦ g n by the chain rule [22, Thm. 10.6] as
[Equation]
where ∇ g n ( w ) ∈ R d × n is the transpose of the Jacobian matrix of g n .
Let Z ( w ) be a discrete random variable which takes the value F i ( w ) with probability 1 /n for i = 1 , . . . , n , and let Q α ( Z ( w )) denote its (1 − α ) -quantile.
Consider the weights ˆ π ∈ ∆ n − 1 given by a hard-thresholding based on whether F i ( w ) is larger than its (1 − α )
-quantile:
[Equation]
The objective defined by these weights is ˆ F α ( w )
= (cid:80) ni =
ˆ π
i F i ( w ) .
The next proposition shows that ˆ F α ( w ) =
F α ( w ) when αn is an integer, or is a close approximation in general.
Proposition 2.
Assume F 1 ( w ) < · · · < F n ( w ) and let i ⋆ = ⌈
αn ⌉ .
Then, we have,
(a) π ⋆ = arg max π ∈P α (cid:80) ni
=1 π
i F i ( w ) is
unique,
(b) Q α ( Z ( w ))
=
F i ⋆ ( w ) , (c) if αn is an integer, then ˆ π = π ⋆ so that ˆ F α ( w ) =
F α ( w ) , and, (d) if αn is not an integer, then
[Equation]
Proof.
We apply the property that the superquantile is a tail mean for discrete random variables [21,
Proposition 8] to get
[Equation]
Comparing with the definition F α ( w ) = (cid:80)
[Equation]
1 π [Equation] F i ( w ) , this gives a closed-form expression for π ⋆ , which is unique because F i ⋆ − 1 ( w ) <
F i ⋆ ( w ) <
F i ⋆ +1 ( w ) .
For (b), note that Q α ( Z ( w ))
=
inf { η ∈ R : P
( Z ( w ) > η ) ≤ α } equals F i ⋆ ( w ) by definition of i ⋆ .
Therefore, if αn is an integer, π ⋆ coincides exactly with ˆ π .
When αn is not an integer, we have
[Equation]
⋆
F i [Equation] w ) . The bound on ˆ F α ( w ) − F α ( w ) follows from elementary manipulations.