Krishna Pillutla

Yassine Laguel

Jérôme Malick

Zaid Harchaoui

Differentially Private Federated Quantiles with the Distributed Discrete Gaussian Mechanism

The computation of analytics in a federated environment plays an increasingly important role in data science and machine learning. We consider the differentially private computation of the quantiles of a distribution of values stored on a population of clients. We present two quantile estimation algorithms based on the distributed discrete Gaussian mechanism that are compatible with secure aggregation. Based on a privacy-utility analysis and numerical experiments, we delineate the regime under which each one is superior. We find that the algorithm with suboptimal asymptotic performance works the best on moderate problem sizes typical in federated learning with client sampling. We apply these algorithms to augment distributionally robust federated learning with differential privacy.

Introduction

Federated analytics consists of a number of client devices, such as smartphones, collaboratively computing population statistics under the orchestration of a central server [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF][START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF]. A key component of this paradigm is the preservation of clients' privacy: any deployed algorithms must ensure that no client data is leaked through the released statistics via rigorous differential privacy guarantees [START_REF] Dwork | Our Data, Ourselves: Privacy Via Distributed Noise Generation[END_REF][START_REF] Dwork | Calibrating Noise to Sensitivity in Private Data Analysis[END_REF].

In the federated setting, the recent paradigm of distributed differential privacy has been growing in popularity [START_REF] Kairouz | The Distributed Discrete Gaussian Mechanism for Federated Learning with Secure Aggregation[END_REF]. It allows one to simulate the effect of a trusted central aggregator without actually trusting the orchestrating server, via the use of cryptographic secure aggregation protocols [START_REF] Bonawitz | Practical Secure Aggregation for Privacy-Preserving Machine Learning[END_REF]. Practical implementations of such algorithms are based on cryptographic techniques such as secure multiparty computation [START_REF] Evans | A Pragmatic Introduction to Secure Multi-Party Computation[END_REF], which add two key limitations. First, they require interactions with the server to be limited to sums of per-client vectors. Second, they require each component of the per-client vectors to be discretized to the ring Z M of integers modulo M .

In this work, we are interested in the problem of federated quantile computation of a per-client scalar. Besides being fundamental to robust statistics, quantiles are central to adaptive clipping in federated averaging [START_REF] Andrew | Differentially private learning with adaptive clipping[END_REF] and to distributionally robust versions of federated learning [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF]. Prevailing approaches to differentially private quantile computation [START_REF] Smith | Privacy-preserving Statistical Estimation with Optimal Convergence Rates[END_REF][START_REF] Tzamos | Optimal private median estimation under minimal distributional assumptions[END_REF][START_REF] Gillenwater | Differentially Private Quantiles[END_REF] typically require operations such as sorting or dynamic programming. It is unclear how to extend these approaches to satisfy distributed differential privacy, particularly, limiting communication with the server to sums of per-client vectors.

We study a simple alternative based on classical techniques from cumulative distribution estimation that are compatible with distributed differential privacy. We quantize the per-client scalars into a per-client histogram, which we then aggregate with a simplified variant of the distributed discrete Gaussian mechanism [START_REF] Kairouz | The Distributed Discrete Gaussian Mechanism for Federated Learning with Secure Aggregation[END_REF]. We show that the quantile error to obtain an (ε, δ)-differentially private quantile estimate of n scalars with a histogram of b bins scales as √ b/(εn) ignoring log factors.

Figure 1: Illustration of our differentially private algorithm for the computation of a quantile using (flat) histograms.

Flat histograms are known to be suboptimal for cumulative distribution estimation; an approach known as the hierarchical histogram method or tree aggregation attains a better asymptotic dependence on the number of bins [START_REF] Hay | Boosting the Accuracy of Differentially Private Histograms Through Consistency[END_REF][START_REF] Dwork | Differential privacy under continual observation[END_REF][START_REF] Hubert Chan | Private and Continual Release of Statistics[END_REF][START_REF] Smith | Is Interaction Necessary for Distributed Private Learning[END_REF][START_REF] Cormode | Answering Range Queries Under Local Differential Privacy[END_REF]. We show that the hierarchical histogram method with distributed differential privacy can attain a quantile error of log 2 b/(εn).

While the hierarchical method attains a better asymptotic dependence on the number b of bins, the flat method achieves a smaller error for moderate b, as might be typical in a federated learning setting due to client sampling [START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF]. We corroborate this observation with synthetic numerical experiments in quantile estimation and distributionally robust federated learning.

Setting and Algorithms

Suppose we have n clients, with each client i hosting a bounded scalar i ∈ [0, B], i ∈ {1, • • • , n} for a fixed B > 0. We wish to estimate the p-quantile of the empirical distribution of i 's, defined as

q p := min t : n i=1 I( i ≤ t) n ≤ p .
Distributed differential privacy simulates a trusted central aggregator by using a secure summation oracle [START_REF] Bonawitz | Practical Secure Aggregation for Privacy-Preserving Machine Learning[END_REF], which enables the computation of summations n i=1 v i where v i ∈ R d is a privacysensitive vector residing with client i without revealing any further information to a privacy adversary. Practical implementations of such algorithms are based on cryptographic techniques such as secure multiparty computation [START_REF] Evans | A Pragmatic Introduction to Secure Multi-Party Computation[END_REF], which requires each component of the vectors v i to be discretized to the ring Z M of integers modulo M .

Quantiles from Flat Histograms. We start with by estimating the quantile from a flat histogram, as illustrated in Figure 1. Each client i first computes a local histogram x i ∈ {0, 1} b on [0, B] with b bins and given edges 0 ≤ l 0 < l 1 < • • • < l b = B. Each client then adds random discrete Gaussian with noise ξ i ∼ N Z (0, σ 2 I b ) and scale parameter σ 2 , that is a random variable ξ over Z satisfying

P(ξ = i) = C exp - (i -µ) 2 2σ 2 for all i ∈ Z ,
where C is an appropriate normalizing constant. We finally sum them up using a secure summation oracle. At the end of all these steps, the server has a histogram ĥ ∈ R b which approximates the true histogram h = n i=1 x i of per-client scalars. The cumulative distribution Fj = ( ĥ1 + • • • + ĥj )/( ĥ1 +• • •+ ĥb ) induced by the approximate histogram ĥ estimates the true cumulative distribution

F j = (1/n) n i=1 I( i ≤ l j ).
Since it may fail to be monotonic, we estimate the quantile using a bin edge l j such that the estimated cumulative mass Fj is as close to p as possible:

Q p ( F ) := l j * p ( F ) where j * p ( F ) = arg min j∈[b] Fj -p . (1)
The full algorithm is given in Algorithm 1. We also consider a variant of the algorithm where the cumulative distribution Fj is estimated using the exact counts, i.e., the denominator ĥ1 + • • • + ĥb is replaced by the true sum

h 1 + • • • + h b = n (option 2 in Algorithm 1).
Algorithm 1 Quantile Computation with Distributed Differential Privacy: Flat Histograms p ( ĥ) corresponding to index j * p ( ĥ); cf. Eq. ( 1)

Input: Ring size M , n clients each with i ∈ [0, B], target quantile p ∈ (0, 1), discretization l 0 , l 1 , • • • , l b of [0, B], variance proxy σ 2 , scaling factor c ∈ Z + 1: Each client i computes local histogram x i = I(l j-1 ≤ i < l j ) b j=1 2: Each client i samples ξ i ∼ N Z (0, σ 2 I b ) and sets xi = (cx i + ξ i ) mod M 3: Compute s = (
Quantiles from Hierarchical Histograms. A hierarchical histogram H maintains the number of clients not only in every single bin but also in groups of bins organized as a binary tree. Concretely, H(r, j) maintains the number of clients whose losses lie between the bin edges l 2 r (j-1)+1 and l 2 r j for index j = 1, 1 The lower levels r = 0 and r = 1 correspond respectively to individual bins and pairs of bins, while the topmost level r = log 2 b -1 refers to two groups: the first b/2 bins and the last b/2 bins. We skip the topmost level in the tree because the count at this node is the publicly known number n of clients.

• • • , b/2 r and level r = 0, • • • , log 2 b -1.
Each client i first computes its local hierarchical histogram X i as

X i (r, j) = I l 2 r (j-1)+1 ≤ i < l 2 r j ,
such that the overall hierarchical histogram can be obtained as H = n i=1 X i . As previously, each client adds discrete Gaussian noise ξ i ∼ N Z (0, σ2 I) with scale parameter σ 2 . These noisy hierarchical histograms are summed up using a secure summation oracle so that the server receives an approximate hierarchical histogram Ĥ which approximates the true H = n i=1 X i . The benefit of hierarchical histograms for cumulative distribution estimation comes from utilizing counts at nodes higher up in the tree. Using a maximal dyadic partition P j of the range [1, j], we have F (j) = (r,o)∈Pj H(r, o) from summing up |P j | ≤ log 2 b terms. For instance, the dyadic partition for j = 15 is [START_REF] Andrew | Differentially private learning with adaptive clipping[END_REF][START_REF] Hay | Boosting the Accuracy of Differentially Private Histograms Through Consistency[END_REF] = [START_REF] Andrew | Differentially private learning with adaptive clipping[END_REF][START_REF] Dwork | Calibrating Noise to Sensitivity in Private Data Analysis[END_REF] ∪ [START_REF] Dwork | Differential privacy under continual observation[END_REF][START_REF] Freedman | On the histogram as a density estimator: L 2 theory[END_REF] ∪ [START_REF] Gillenwater | Differentially Private Quantiles[END_REF][START_REF] Guyon | Design of experiments of the neurips 2003 variable selection benchmark[END_REF] ∪ [START_REF] Hay | Boosting the Accuracy of Differentially Private Histograms Through Consistency[END_REF], where the counts of each range on the right side can be obtained from an intermediate node in the hierarchical histogram H.

Similarly, we estimate the cumulative distribution from the approximate hierarchical histogram Ĥ as Fj = (r,o)∈Pj Ĥ(r, o) where P j is a maximal dyadic partition of [1, j]. Similar to the previous approach, we estimate the quantile with [START_REF] Andrew | Differentially private learning with adaptive clipping[END_REF]. The full algorithm is given in Algorithm 2.

Utility and Privacy Analysis

We now analyze the privacy and utility of Algorithm 1. First, recall the definition of zero-concentrated differential privacy [START_REF] Bun | Concentrated Differential Privacy: Simplifications, Extensions, and Lower Bounds[END_REF]: a randomized algorithm A satisfies (1/2)ε 2 -concentrated differential privacy if the Rényi α-divergence D α (A(x) A(x )) ≤ αε 2 /2 for all α ∈ (0, ∞) and all sequences x, x of inputs which differ by the addition or removal of the data of one client. Intuitively, the addition or removal of one client should not change the output distribution of the randomized algorithm by much, as measured by the Rényi divergence. A smaller value of ε implies a stronger privacy guarantee.

Error Criterion. The bin edge l j corresponding to index j ∈ [b] approximates the p-quantile well if the cumulative mass F j ≈ p. We measure this error of approximation by the difference between the two sides. Formally, we define the error R p (F, j) of approximating the p-quantile of the cumulative function

F of a (hierarchical) histogram with index j ∈ [b] by R p (F, j) = |F j -p| . (2)
Algorithm 2 Quantile Computation with Distributed Differential Privacy: Hierarchical Histograms

Input: Ring size M , n clients each with i ∈ [0, B], target quantile p ∈ (0, 1), discretization l 0 , l 1 , • • • , l b of [0, B], variance proxy σ 2 , scaling factor c ∈ Z + 1: Each client i computes a hierarchical histogram X i (r, j) = I l 2 r (j-1)+1 ≤ i < l 2 r j for j = 1, • • • , b/2 r and r = 0, • • • , log 2 b -1 2: Each client i samples ξ i (r, j) ∼ N Z (0, σ 2 ) i.i.d. and sets Xi (r, j) = cX i (r, j) + ξ i (r, j)
mod M for each r, j 3: Compute S = ( n i=1 Xi ) mod M securely and set Ĥ = S/c 4: Define the cumulative distribution F ∈ R b as F (j) = (1/n) (r,o)∈Pj Ĥ(r, o) using a maximal dyadic partition P j of [1, j] 5: return Quantile estimate l j * p ( F ) corresponding to index j * p ( F ); cf. Eq. ( 1)

We define the best achievable error a R * p (F ) = min j∈[b] R p (F, j). Note that the best index j * p (F ) for estimating the p-quantile of the cumulative function F as defined in (1) satisfies j * p (F ) = arg min j∈[b] R p (F, j). Lastly, we define the quantile error ∆ p (F, F ) of estimating the p-quantile of the cumulative function F from that of F as

∆ p ( F , F ) = R p F, j * p ( F ) . (3)
Essentially, if the index j * p ( F ) computed from the estimate F corresponds to the p -quantile of F , the quantile error satisfies

∆ p ( F , F ) = |p -p |.
Privacy and Utility Analysis. We now analyze the differential privacy bound of the flat histogram approach of Algorithm 1 as well as the error in the quantile computation. Theorem 1 (Flat histogram). Fix a δ > 0. Suppose that σ ≥ 1/2 and c > 0 are given. We have that Algorithm 1 satisfies

(1/2)ε 2 -concentrated DP with ε = min c 2 nσ 2 + ψb 2 , c √ nσ + ψ √ b ,
where ψ = 10

n-1 i=1 exp -2π 2 σ 2 i/(i + 1) ≤ 10(n -1) exp(-2π 2 σ 2 ). Further, if the modular base satisfies M ≥ 2 + 2cn + 2n 2σ 2 log(8nb/δ), then we have with probability at least 1 -δ that the quantile error of cumulative function F returned by Algorithm 1 (option 2) is at most

∆ p ( F , F ) ≤ R * p ( F ) + 2σ 2 b c 2 n log 4 δ ,
where R * p ( F ) is the error in the estimation of p-quantile of histogram F .

Let us interpret the result. The effective noise scale is σ/c. Since the dominant term of the privacy error

is ε ≈ c/(σ √ n), we choose σ/c ≈ (ε √ n) -1
, so that the algorithm satisfies (1/2)ε 2 -concentrated DP. The role of c is to avoid degeneracy of the discrete Gaussian as σ → 0. In particular, the theorem requires σ ≥ 1/2. Ignoring constants and log factors, the error is

∆ p ( F , F ) R * p ( F ) + √ b εn .
While the quantile error increases with the number b of bins, the discretization error from reducing the scalars to a histogram typically reduces with b. For instance, for a uniform discretization of [0, B], the p-quantile of the histogram n), the conditions of the theorem require M n 3/2 . We give a similar bound for Algorithm 1 with estimated n (option 1) in Appendix A. Theorem 2 (Hierarchical Histogram). Fix a δ > 0. Suppose that σ ≥ 1/2 and c > 0 are given. We have that Algorithm 2 satisfies (1/2)ε 2 -concentrated DP with (ψ is given in Theorem 1) Further, if the modular arithmetic is performed on base M ≥ 2 + 2cn + 2n 2σ 2 log(16nb/δ), then we have with probability at least 1 -δ that the quantile error of cumulative function F returned by Algorithm 2 is at most

ε = min    c 2 log 2 2 b nσ 2 + ψb, c log 2 b √ nσ + ψ √ 2b    .
∆ p ( F , F ) ≤ R * p ( F ) + 4σ 2 c 2 n log 2 b log 4b δ .
Comparing the Two Approaches. Asymptotically, the flat histogram approach has an error of √ b/εn, while the hierarchical one has an error of poly log b/εn, showing the asymptotic optimality (up to log factors) of the latter. However, the behavior is more nuanced for intermediate values of b and n. In particular, disregarding the error term ψ in the privacy cost, we have that the flat approach is better for

2b log(4/δ) εn ≤ 4 log 3 2 (b) log(4b/δ) εn .
With a failure probability of δ = 0.05, this corresponds (approximately) to b ≤ 2000. Under the Freedman-Diaconis rule [START_REF] Freedman | On the histogram as a density estimator: L 2 theory[END_REF] where one chooses the number of bins as b = O(n 2/3 ), the flat histogram is better for n approximately below 2.5 × 10 6 . In federated learning scenarios where quantile statistics are computed along with learning (e.g. [START_REF] Andrew | Differentially private learning with adaptive clipping[END_REF][START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF]), n is the number of clients per round and is of the order of 10 2 to 10 4 in typical cross-device settings [START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF]. Our analysis suggests that we should prefer the flat histogram approach in this case although it is asymptotically suboptimal.

Numerical Experiments: Quantile Estimation

We test the privacy and utility of Algorithms 1 and 2 in synthetic examples. We consider n ∈ {128, 512} clients, each with a scalar that is distributed as Unif([0, B]) or χ 2 (4) distributions. We clip all scalars to [0, B] with B = 10. For all experiments, we consider the uniform quantization of [0, B] into b bins, where b ∈ {32, 64}. We fix the ring size at M = 2 18 , which is large enough to prevent any modular wraparound as required by the theorems.

Both Algorithms 1 and 2 require two further parameters: the variance-proxy σ 2 and the granularity parameter c, which together determine the privacy leakage. As per the discussion following Theorem 1, we take σ 2 = 2 and vary c. We plot the utility of each algorithm, as measured by the quantile error (3) versus the privacy parameter ε required for (ε, 10 -5 )-differential privacy.

Results. Figure 2 shows the numerical results. First, we note that all three approaches have small quantile errors. For the left plot, the flat histogram with estimated n (Algorithm 1, option 1) has a quantile error of 0.03 at ε = 1; this means we might find the 47 th or 53 rd percentile instead of the median. This error quickly falls below 0.01 at ε = 5. Second, we note that the flat histogram approaches tend to outperform the hierarchical approach at small ε despite their asymptotic suboptimality. For instance, the flat histogram with estimated n (error = 0.03) vastly outperforms the hierarchical approach (error = 0.09) at ε = 1 in the left plot. These values are 0.10 and 0.26 respectively for the right plot; the latter might be too inexact to even be useful depending on the application. At ε ≥ 10, all the methods are within one standard deviation of each other. Overall, these results corroborate the theoretical analysis of Section 3.

Application to Distributionally Robust Federated Learning

We develop an end-to-end differentially private version of ∆-FL [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF], a distributionally robust approach to federated learning. It minimizes the tail mean of the per-client losses, known formally as the superquantile [START_REF] Tyrrell | Conditional Value-at-Risk for General Loss Distributions[END_REF], which is defined for a continuous random variable Z as S p (Z) = E[Z|Z > Q p (Z)], where Q p (Z) is its p-quantile. In each round of their algorithm consists of two steps: (a) Quantile estimation: Compute the p-quantile of the per-client losses, and, (b) Tail Aggregation: Aggregate local updates from clients whose loss is larger than the quantile.

We propose an end-to-end differentially private version of this algorithm, called DP-∆-FL. We estimate the quantile of the losses clipped to a tuned bound B using Algorithm 1 or 2. We clip the weight updates to a norm bound C, and add Gaussian noise, similar to DP-FedAvg [START_REF] Mcmahan | Learning Differentially Private Recurrent Language Models[END_REF]. The total privacy loss is calculated by composing the privacy loss across both the quantile and weight updates, and the number of rounds together with amplification by subsampling using the bounds of [START_REF] Zhu | Poission Subsampled Rényi Differential Privacy[END_REF].

We calculate the noise scales σ q of the quantile and σ w of the weight update so that the overall algorithm satisfies (ε, 1/n)-differential privacy and tune the loss bound B, norm bound C, the number of bins b and the ratio of the privacy budget to assign to the quantile estimation.

We experiment with a synthetic dataset contains k = 10 classes in d = 20 dimensions, and n = 2500 training clients, each with 100 examples. We consider robustness to a label shift, while the classconditional distribution is fixed. The marginal across classes is a Dirichlet distribution with parameter 0.5 for the training clients and 0.01 for the test clients. For further details, we refer to Appendix B.

Results. We compare the distributional robustness in terms of the 90 th percentile of the per-client errors in Figure 3; we also show the mean test error. First, we see that both the flat or hierarchical histograms lead to similar performance for DP-∆-FL (equal up to one std.).

We also observe that DP-∆-FL has a much smaller tail error than DP-FedAvg at each ε; the gap is 4.3 pp at ε = 5. Its performance does not degrade much with privacy and is within 1pp of ∆-FL without DP at ε = 5. Importantly, the 90 th percentile error of DP-∆-FL is better than FedAvg without DP. In terms of the mean error, DP-∆-FL is within 0.2-0.4 pp to DP-FedAvg, comparable to the gap between the non-private versions (0.24 pp).

We conclude from the experiments that (a) DP-∆-FL has a privacy-utility tradeoff similar to DP-FedAvg compared to their respective non-private versions, and (b) the distributional robustness of ∆-FL is also enjoyed by its private variant. Theorem 3. Consider the setting of Theorem 1. If the modular base satisfies M ≥ 2 + 2cn + 2n 2σ 2 log(8nb/δ), then we have with probability at least 1-δ that the quantile error of cumulative function F returned by Algorithm 1 (option 1, estimated count) is at most

∆ p ( F , F ) ≤ R * p ( F ) 1 + 2σ 2 b c 2 n log 4 δ + (1 + p) 2σ 2 b c 2 n log 4 δ ,
where R * p ( F ) is the error in the estimation of p-quantile of histogram F .

Compared to Theorem 1, this only has an additional factor on the irreducible error term R * p ( F ).

A.1 Privacy-Utility with Flat Histograms

We now prove Theorem 1 and Theorem 3.

Proof of Theorem 1 and Theorem 3. We start by defining and controlling the probabilities of some events. Throughout, let δ > 0 be fixed. Define the event

E mod = n i=1 b j=1 - M -2 2n ≤ cx i,j + ξ i,j ≤ M -2 2n . ( 4
)
Note that under E mod , no modular wraparound occurs in the algorithm, i.e., xi = cx i + ξ i and ĥ =

n i=1 xk c = n i=1 x k + ξ k c .
We will show later that E mod holds with high probability; for now, we assume that it holds.

Privacy Analysis. We start by establishing the sensitivity of the sum query over x k 's as 1. Define the input space X to be the canonical basis vectors in R b , i.e., the set of all vectors in {0, 1} b with only one 1, and let X * = ∪ ∞ r=1 X r denote the set of all sequences of elements of X . We consider the rescaled sum query A((x 1 , • • • , x N )) = n i=1 cx i . The L 2 sensitivity S(A) of this query A is supremum over all X ∈ X * and X which is obtained by concatenating x to X:

S(A) = sup X,X A(X) -A(X ) 2 = sup x ∈X c x 2 = c .
We invoke the privacy bound of sums of discrete Gaussians (see e.g. [START_REF] Kairouz | The Distributed Discrete Gaussian Mechanism for Federated Learning with Secure Aggregation[END_REF]) to claim that an algorithm A returning A(x) + n i=1 ξ k satisfies (1/2)ε 2 -concentrated DP where ε is as in the theorem statement. The fact that the quantile and all further functions of it remains private follows from the postprocessing property of DP (also known as the data-processing inequality).

Utility Analysis with Option 1 (estimated count, Theorem 3). Define n = b j=1 ĥj , as the noisy analogue to n = b j=1 h j . We bound the quantile error as

∆ p F , F = R p F, j * p ( F ) = 1 n j * p ( F ) j=1 h j -p ≤ 1 n j * p ( F ) j=1 h j -ĥj + 1 n j * p ( F ) j=1 ĥj -np + p n |n -n| ≤ max j ∈[b] 1 cn j j=1 n i=1 ξ i,j + 1 + |n -n| n R * p ( F ) + p n |n -n| .
Let us define an event E sum under which the first term and last terms are bounded:

E sum =    max j∈[b] j j =1 n i=1 ξ i,j ≤ 2σ 2 nb log(4/δ)    . ( 5
)
Under E sum , we also have

|n -n| = 1 c b j=1 n i=1 ξ i,j ≤ 2σ 2 nb c 2 log 4 δ .
Plugging this back into ∆ p ( ĥ, h) gives us the desired bound, provided E sum holds.

Utility Analysis with Option 2 (exact count, Theorem 1). Similar to the previous case, we have

∆ p F , F = R p F, j * p ( F ) = 1 n j * p ( F ) j=1 h j -p ≤ 1 n j * p ( F ) j=1 h j -ĥj + 1 n j * p ( F ) j=1 ĥj -p ≤ max j ∈[b] 1 cn j j=1 n i=1 ξ i,j + R * p ( F ) .
Under E sum , the desired bound holds.

Bounding the Failure Probability. The algorithm fails when at least one of E mod or E sum fail to hold. We have from Claim 4 that P(E mod ) ≥ 1 -δ/4 under the given assumptions. From, Claim 5, we have P(E sum |E mod ) ≥ 1 -δ/2. We bound the total failure probability of the algorithm with a union bound as

P( Ēsum ∪ Ēmod ) ≤ P( Ēsum |E mod ) P(E mod ) + P( Ēsum | Ēmod ) P( Ēmod ) + P( Ēmod ) ≤ P( Ēsum |E mod ) + 2 P( Ēmod ) ≤ δ .
We state and prove bounds on probabilities of the events E mod , E sum defined above. Claim 4. If M ≥ 2 + 2cn + 2n 2σ 2 log(8nb/δ), then P(E mod ) ≥ 1 -δ/4.

Proof. Each discrete Gaussian random variable ξ i,j is centered and sub-Gaussian with variance proxy σ 2 . A Cramér-Chernoff bound (see e.g. Lemma 10) gives us the exponential tail bound

P |ξ i,j | > 2σ 2 log(8nb/δ) ≤ δ 4nb .
Using a union bound for i ∈ [n], j ∈ [b] and x i,j ∈ {0, 1} completes the proof.

Claim 5. We have that

P(E sum |E mod ) ≥ 1 -δ/2.
Proof. Under E mod , no modular wraparound occurs, so we can ignore the modulo operations. Each discrete Gaussian random variable ξ i,j is centered and sub-Gaussian with variance proxy σ 2 , i.e., E[ξ i,j ] = 0 and E[exp(λξ i,j )] ≤ exp(λ 2 σ 2 /2) for all λ ∈ R. Therefore, ζ j := n i=1 ξ i,j is centered and sub-Gaussian with variance proxy nσ 2 , since E[ζ j ] = 0, and

E[exp(λζ j )] = n i=1 E[exp(λξ i,j )] ≤ exp(λ 2 σ 2 n/2)
by independence. We get a bound on the partial sums from the maximal inequality (see e.g. Lemma 11); this involves constructing a martingale (

j j =1 ζ j ) b
j=1 and applying the maximal inequality. The bound we get is

P   max j∈[b] j j =1 ζ j > t   ≤ exp - t 2 2σ 2 nb .
Plugging in t = 2σ 2 nb log(2/δ) completes the proof.

Algorithm 3

The ∆-FL Algorithm [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF] 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:
Find the p-quantile of f i (w (t) )) i∈S ; call this Q (t) 4:

for each selected client i ∈ S in parallel do

: Set π(t) i = I f i (w (t) ) ≥ Q (t) 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: In this expression, the quantile Q p (Z) is obtained as the left end-point of the argmin set over η.

w (t+1) = i∈S π(t) i w (t) i,τ / i∈S π (t)
The ∆-FL objective is to minimize the superquantile of the per-client losses:

g p (w) := S p f 1 (w), • • • , f n (w) , (8)
where w ∈ R d denotes the model parameters.

Algorithm. The algorithm to optimize the ∆-FL objective proposed by [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF] can be interpreted more transparently as below. We take this alternate interpretation as opposed to the one used in [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF] as it directly allows us to derive a version with end-to-end differential privacy.

The superquantile is a nonsmooth function, and so is the ∆-FL objective g p . However, leveraging duality, the following subdifferential expression can be derived. Proposition 12. Suppose the per-client objective f i is L-smooth for each i, and let the parameter p be such that np is an integer. For all w such that f i (w) = f j (w) for i = j, a subgradient of the objective (8) is given by

∂g p (w) n i=1 π i f i (w) , where π * i = I(f i (w) ≥ Q p ) n j=1 I(f j (w) ≥ Q p ) , and
Q p = Q p (f 1 (w), • • • , f n (w))
is the p-quantile of the losses and ∂g p refers to the regular subdifferential of g p .

Proof. Without loss of generality, we assume that f 1 (w) < • • • < f n (w). Due to the dual expression for the supqerquantile [START_REF] Föllmer | Convex measures of risk and trading constraints[END_REF], we have the equivalent expression

g p (w) = max n i=1 π i f i (w) : 0 ≤ π i ≤ 1 n(1 -p) , n i=1 π i = 1 . (9)
By Danskin's theorem, we have that

∂g p (w) n i=1 π i ∇f i (w) , (10)
where π attains the argmax in the expression above. Next, we invoke the a closed form expression of the superquantile for discrete random variables [START_REF] Tyrrell | Conditional Value-at-Risk for General Loss Distributions[END_REF]Proposition 8] (this is analogous to the expression

S p (Z) = E[Z | Z > Q p (Z)] for continuous random variables Z) to get g p (w) = 1 (1 -p)n n i=i +1 f i (w) + 1 - (1 -p)n (1 -p)n f i (w) ,
where i = (1 -p)n . Comparing with [START_REF] Dwork | Differential privacy under continual observation[END_REF], this gives a closed-form expression for π , which is unique because F i -1 (w) < F i (w) < F i +1 (w). By the definition of the quantile, also observe that Q p = f i (w). Plugging in this closed form expression of π into (10) completes the proof.

Using this expression, the algorithm of [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF] interleaves federated averaging steps with quantile estimation. Specifically, 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:

• compute Q p = Q p (f i (w) : i ∈ S),
• aggregate the updates from the tail clients where f i (w) ≥ Q p to find the new global model w + as

w + = 1 |S p | i∈Sp w + i , where S p = {i : f i (w) ≥ Q p } .
Similar to FedAvg, 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, the aggregated update is simply a subgradient step w -w + = γ∇g p (w) where we denote the subgradient as ∇g p (w) ∈ ∂g p (w).

The overall algorithm is summarized in Algorithm 3.

B.2 End-to-End Differential Privacy with DP-∆-FL

To obtain an end-to-end differentially private version of ∆-FL, we make two modifications to Algorithm 3. First, we estimate the quantile with distributed differential privacy, using either Algorithm 1 or Algorithm 2. Second, we modify the weight aggregation step (line 11) by clipping the weight updates and add Gaussian noise to obtain differential privacy via the Gaussian mechanism. The overall algorithm is given in Algorithm 4.

Privacy Accounting. We now discuss the privacy spent in each communication round. For simplicity, we assume the number m (t) = i∈S I(f i (w (t) ) ≥ Q (t) ) of selected clients is publicly known. Claim 13. Consider the setting of Algorithm 4 with noise scale σ w , norm bound C and Algorithm 2 with b bins and noise scale σ = σ q . Each round of Algorithm 4 satisfies (1/2)ε 2 -concentrated DP where

1 2 ε 2 = 1 2 ε 2 q + σ 2 w 2C 2
, where ε q is obtained from Theorem 1 or 2 for Algorithm 1 or 2 respectively.

Proof. The (1/2)ε 2 q -concentrated DP of the quantile computation comes from Theorem 1 or 2. Since the contribution δ (t) i of each client has 2 norm δ (t) i ≤ C and we add Gaussian noise N (0, σ 2 w I d ), the weight update step satisfies σ 2 w /(2C 2 )-concentrated DP. The proof is completed by noting that concentrated differential privacy composes additively.

To obtain a bound on the concentrated DP of the entire algorithm, we rely on generic upper bounds of [START_REF] Zhu | Poission Subsampled Rényi Differential Privacy[END_REF] for privacy amplification by subsampling.

B.3 Experimental Setup

We consider a synthetic classification dataset and train a linear model on it.

Dataset Description. We create a 10-class classification dataset in d = 20 dimensions, inspired by [START_REF] Guyon | Design of experiments of the neurips 2003 variable selection benchmark[END_REF]. The input x for each class k is drawn from a Gaussian of mean µ i and identity covariance in R 15 . The means µ i 's are the corners of a random polytope in R 15 . We add 2 features that are linear combinations of the 15 informative ones and 3 features that are pure noise. Overall, the dataset can be generated using the make classification function of scikit-learn [START_REF] Pedregosa | Scikit-learn: Machine learning in Python[END_REF] as Estimate the p-quantile of f i (w (t) ) for i ∈ S with distributed differential privacy (Algorithm 1 or 2); call this Q (t)

4:

Set

m (t) = i∈S I f i (w (t) ) ≥ Q (t)
5:

for each selected client i ∈ S in parallel do 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:

Define the norm-clipped update contributed by the client We now split this dataset into a federated dataset with n = 2500 training clients and n = 500 validation and n = 500 test clients. The data distribution q i (x, y) = q i (y)q i (x|y) across the clients is designed to exhibit a label shift, i.e., the distribution q i (y) over labels for each client is different while the class-conditional distribution q i (x|y = k) = N (µ k , I d ) is the same across clients. The class distribution q i (y) on each training client i is drawn from a Dirichlet distribution Dir(0.5), while that for a validation or test client is drawn from Dir(0.01). We sample 100 input-output pairs for each training, validation, and test client.

δ (t) i =    C (w (t) i,τ -w (t) ) max C, w (t) i,τ -w (t) 2 , if f i (w (t) ) ≥ Q (t)
Model and Per-Client Objective. We use a linear model (with intercept) on each client and the multinomial logistic loss, also known as the cross entropy loss, to define the per-client objective.

Algorithms and Hyperparameters. We compare Algorithm 4 with DP-FedAvg [START_REF] Mcmahan | Learning Differentially Private Recurrent Language Models[END_REF], a version of FedAvg with differential privacy.

Both algorithms used a single full gradient step per client with a fixed learning rate of 0.1. For each algorithm, we sample 100 clients per round and run the training for a total of 1000 rounds. We vary the privacy budget ε ∈ {3, 5, 10, 15, 20} and tune the following hyperparameters for each algorithm.

For DP-FedAvg, we tune the 2 norm bound (analogous to C in Algorithm 4) and set the noise scale σ w depending on the target privacy budget ε and the norm bound C. For Algorithm 4, we allocate r-times the privacy budget of the weight updates to the quantile updates. In addition, we also tune:

• the loss upper bound B, so that all losses are truncated to [0, B],

• the number of bins b in the hierarchical histogram, • the 2 norm bound C for the weight update.

We tune all 4 hyperparameters with a grid search and set the noise scale σ w for the weight update, and σ q /c for the quantile update depending on the selected hyperparameters and the privacy budget ε. The objective of the grid search was to minimize the 90 th percentile of the misclassification errors across all validation clients. The ranges of the hyperparameters considered are quantile privacy ratio r ∈ {0.1, 0.25, 0.5, 0.75}, loss upper bound B ∈ {0.7, 0.9, 1.1, 1.3, 1.5},2 , number of bins b ∈ {16, 32, 64}, and update norm C ∈ {0.9, 1.1, 1.3, 1.5}. 3

n

  i=1 xi ) mod M securely and set histogram ĥ = s/c 4: Estimate the cumulative function with F ∈ R b as • Option 1 (estimated count): Fj = ( ĥ1 + • • • + ĥj )/n where n = ĥ1 + • • • + ĥb • Option 2 (exact count): Fj = ( ĥ1 + • • • + ĥj )/n 5: return Quantile estimate l j *

n

  i=1 x i is at most B/b away from the true p-quantile of the scalars 1 , • • • , n . Finally, if we take σ = O(1) and c = O(ε √

Figure 2 :

 2 Figure 2: Distributed quantile estimation with Algorithms 1 and 2. Here, n is the number of clients, b is the number of bins, and bit width refers to log 2 M where M is the ring size. The scalar in each client is drawn from either Unif([0, B]) or χ 2 (4) distribution clipped to [0, B], where B = 10 in both cases. We plot the worst quantile error (3) in computing the p-quantile for p ∈ {0.1, 0.2, • • • , 0.9}, versus ε for (ε, 10 -5 )-DP. We plot the mean values across 10 runs, while the error bars denote the standard deviation.

  DP--FL + flat hist. DP--FL + hier. hist.

Figure 3 :

 3 Figure 3: ∆-FL vs. FedAvg with (ε, 1/n)-DP on a synthetic classification task in R 20 with 10 classes and n = 2500 clients and a tail parameter p = 0.75. The error bars denote the standard deviation across 5 runs.

5

 5

i 12 :

 12 end for 13: return w T

Algorithm 4 2 w 1 : 2 :

 4212 DP-∆-FL: ∆-FL with End-to-End Differential Privacy Input: Initial iterate w (0) , number of communication rounds T , number of clients per round m, number of local updates τ , local step size γ, 2 norm bound C for weight updates, noise variance σ for t = 0, 1, • • • , T -1 do Sample m clients from [n] without replacement in S 3:

14 :

 14 noise ξ (t) ∼ N (0, σ 2 w I d ) and update w (t+1) = w (t) return w T x , y = m a k e c l a s s i f i c a t i o n ( n s a m p l e s = i n t ( 5 e5 ) , n f e a t u r e s =20 , n i n f o r m a t i v e =15 , n r e d u n d a n t =2 , n r e p e a t e d =0 , n c l a s s e s =10 , n c l u s t e r s p e r c l a s s =1 , c l a s s s e p = 5 . 0 , h y p e r c u b e = F a l s e , r a n d o m s t a t e =2345 )

We assume for simplicity that b is a power of

so that log 2 b is an integer.

The loss at convergence was around 0.7, while that at random guessing is log 10 ≈ 2.3.

These correspond approximately to the 0.3, 0.5, 0.7, 0.9 quantiles of the update norms of FedAvg without differential privacy, during the latter half of training.

Acknowledgements

The authors thank Peter Kairouz and Lun Wang for fruitful discussions. 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 Additional Results and Proof Details

In this section, we give some additional results and present the proofs of the privacy-utility bounds.

For ease of handling negative integers, we perform modular arithmetic throughout over the ring

Flat Histogram with Estimated Counts. We first give the analogue of Theorem 1 for the option of estimating the counts.

A.2 Privacy-Utility with Hierarchical Histograms

We now prove Theorem 2.

Proof of Theorem 2. We start by defining and controlling the probabilities of some events. Throughout, let δ > 0 be fixed. Define the event

Note that under E mod , no modular wraparound occurs in the algorithm. Thus, for all valid levels r and indices j, we have Xi (r, j) = cX i (r, j) + ξ i (r, j) and

Next, we define the event

We will show later that E mod and E diff holds with high probability; for now, we assume that they hold.

Privacy Analysis. The privacy analysis follows similar to Theorem 1, with the difference that each hierarchical histogram has a total 2b entries (as opposed to b for the flat histogram) and an 2 sensitivity of log 2 b (as opposed to 1 for the flat histogram), since it contains one non-zero entry for each level of the tree.

Utility Analysis. We denote F j as F (j) for ease of reading. Using the triangle inequality, we get,

The first term is bounded under E diff and this gives the utility bound.

Bounding the Failure Probability. Follows similar to Theorem 1 by using Claim 6 and Claim 7 instead of Claim 4 and Claim 5 respectively.

We state and prove bounds on probabilities of the events E mod , E diff defined above.

Proof. Each discrete Gaussian random variable ξ i (r, j) is centered and sub-Gaussian with variance proxy σ 2 (cf. Property 8). A Cramér-Chernoff bound (cf. Lemma 10) gives us the exponential tail bound

.

Applying the union bound over i = 1, • • • , n and the 2b -2 nodes in each hierarchical histogram x i (each node corresponding to one (r, j) pair) completes the proof.

Claim 7. We have

Proof. Under E mod , we have that Ĥ(j) = H(j)+ n i=1

(r,o)∈Pj ξ i (r, o), where P j is the maximal dyadic partition of [1, j] with |P j | ≤ log 2 b. Thus, ζ j := Ĥ(j) -H(j) is sub-Gaussian with variance proxy n|P j |σ 2 ≤ nσ 2 log 2 b. A Cramér-Chernoff bound (cf. Lemma 10) gives us

.

Applying a union bound over j = 1, • • • , b completes the proof.

A.3 Useful Results

We state the sub-Gaussian property of the discrete Gaussians. Property 8. Let ξ be distributed according to N Z (µ, σ 2 ). Then, E[ξ] = µ. Furthermore, if µ = 0, then ξ is sub-Gaussian with variance proxy σ 2 , i.e., E[exp(λξ)] ≤ exp(λ 2 σ 2 /2) for all λ > 0.

The distributed discrete Gaussian mechanism gets privacy guarantees by adding a sum of discrete Gaussian random variables. We give a bound on its privacy. The following lemma is due to [START_REF] Kairouz | The Distributed Discrete Gaussian Mechanism for Federated Learning with Secure Aggregation[END_REF]. Lemma 9 (Privacy of Sum of Discrete Gaussians). Fix σ ≥ 1/2. Let A : X → d be a deterministic algorithm with 2 -sensitivity S for some input space X . Define a randomized algorithm A, which when given an input x ∈ X , samples

where ψ = 10

Next, we record standard concentration results (see e.g. [START_REF] Boucheron | Concentration inequalities: A nonasymptotic theory of independence[END_REF]). Lemma 10 (Cramér-Chernoff). Let ξ be a real-valued and centered sub-Gaussian random variable with variance proxy σ 2 , i.e., E[ξ] = 0 and E[exp(λξ)] ≤ exp(λ 2 σ 2 /2) for all λ > 0. Then, we have for any t > 0, P(|ξ| > t) ≤ 2 exp -t 2 2σ 2 . Lemma 11 (Maximal Inequality). Let ξ 1 , ξ 2 , • • • be i.i.d. centered sub-Gaussian random variables with variance proxy σ 2 , i.e., E[ξ j ] = 0 and E[exp(λξ j )] ≤ exp(λ 2 σ 2 /2) for all λ ∈ R and j = 1, 2, • • • . Then, it holds for any t > 0 and integer n ≥ 1 that

B DP-∆-FL: Background and Details

We describe the proposed DP-∆-FL approach, a variant of ∆-FL [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF] under end-to-end differential privacy.

B.1 ∆-FL Review

We review the ∆-FL objective and the corresponding optimization algorithm [START_REF] Laguel | A Superquantile Approach to Federated Learning with Heterogeneous Devices[END_REF].

Suppose we have n clients, where the objective on client i, denoted f i : R d → R, is the expected loss on client i under its data distribution q i for i = 1, • • • , n.

The ∆-FL objective minimizes the tail mean of the per-client losses, which can be formalized using a notion called superquantile [START_REF] Tyrrell | Optimization of Conditional Value-at-Risk[END_REF]. The p-superquantile of a continuous random variable Z is defined as S p (Z) = E[Z | Z > Q p (Z)], where Q p (Z) is the p-quantile of Z. More generally, the following definition is applicable to both discrete and continuous random variables:


Krishna Pillutla

Yassine Laguel

Jérôme Malick

Zaid Harchaoui

Tackling Distribution Shifts in Federated Learning with Superquantile Aggregation

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 [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF][START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF]. It has found widespread adoption across industry [START_REF] Kallista | Towards Federated Learning at Scale: System Design[END_REF][START_REF] Paulik | Federated Evaluation and Tuning for On-Device Personalization: System Design & Applications[END_REF] for applications ranging from smart device apps [START_REF] Yang | Applied Federated Learning: Improving Google Keyboard Query Suggestions[END_REF][START_REF] Hard | Federated Learning for Mobile Keyboard Prediction[END_REF] to healthcare [START_REF] Brisimi | Federated learning of predictive models from federated Electronic Health Records[END_REF][START_REF] Huang | Patient Clustering Improves Efficiency of Federated Machine Learning to Predict Mortality and Hospital stay time using Distributed Electronic Medical Records[END_REF].

A key feature of federated learning is statistical heterogeneity, i.e., client data distributions are not identically distributed [START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF][START_REF] Li | Federated Learning: Challenges, Methods, and Future Directions[END_REF]. 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 [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF] 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 [START_REF] Tyrrell | Conditional Value-at-Risk for General Loss Distributions[END_REF][START_REF] Tyrrell Rockafellar | Risk tuning with generalized linear regression[END_REF], 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 shows 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.

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∼pi [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 [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF] is simply the empirical risk minimization (ERM) approach

min w∈R d 1 n n i=1 F i (w) . (1)
Owing 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) n i=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 [START_REF] Tyrrell | Conditional Value-at-Risk for General Loss Distributions[END_REF]. 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

S α (u 1 , • • • , u n ) := max n i=1 π i u i : 0 ≤ π i ≤ 1 αn ∀ i ∈ [n], n i=1 π i = 1 .
This is an instance of the continuous knapsack problem and can be solved optimally by a greedy algorithm [START_REF] George B Dantzig | Discrete-variable extremum problems[END_REF]. 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

F α (w) := S α F 1 (w), • • • , F n (w) . (2)
If we have a test client whose distribution p π = n i=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. In order 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 → 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

∂F α (w) ∋ n i=1 π ⋆ i F i (w) , where π * i = I(F i (w) ≥ Q α ) n j=1 I(F j (w) ≥ Q α ) , (3)
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 Zm = (1/m) m i=1 Z i is an unbiased estimate of the population mean, i.e., E[ Zm ] = E[Z]. This is no longer true for the superquantile, i.e., E[S α (Z 1 , • • • , Z m )] ̸ = S α (Z).
As a result, we do not have access to 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 [START_REF] Peter Kairouz | Advances and Open Problems in Federated Learning[END_REF]. We overcome this issue by actually minimizing the expected minibatch superquantile instead, defined as

Fα,m (w) = E (i1,••• ,im)∼Um S α F i1 (w), . . . , F im (w) ,
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]

|F α (w) -Fα,m (w)| ≤ 3 √ αm max i=1,...,n |F i (w)| .
Using this expression, we design a federated optimization algorithm that steps of the usual federated averaging algorithm [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF] 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 + = 1 |S α | i∈Sα w + i , where S α = {i : F i (w) ≥ Qα } .
The full algorithm is given in Appendix A. Similar to the standard FedAvg algorithm [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF] for ERM objective [START_REF] Kallista | Towards Federated Learning at Scale: System Design[END_REF], 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). Similar to FedAvg, our algorithm reduces the overall communication cost, which is often the bottleneck in bandwidthconstrained edge devices, while incurring a larger computation cost at each client.

Numerical Experiments

In this section, 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 [START_REF] Cohen | EMNIST: an extension of MNIST to handwritten letters[END_REF], 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 [START_REF] Go | Twitter Sentiment Classification using Distant Supervision[END_REF] 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 [START_REF] Hochreiter | Long Short-Term Memory[END_REF].

Baselines. We compare ∆-FL with the following baselines: We consider two methods which attempt to minimize the usual objective (1): FedAvg [START_REF] Mcmahan | Communication-Efficient Learning of Deep Networks from Decentralized Data[END_REF] and FedProx [START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF]. The latter augments FedAvg with a proximal term for more stable optimization. We also consider a few heterogeneity-aware objectives: Tilted-ERM [START_REF] Li | Tilted Empirical Risk Minimization[END_REF], which is the analogue of ∆-FL but using the log-sum-exp function and AFL [START_REF] Mohri | Agnostic Federated Learning[END_REF], whose objective is obtained as the limit lim α→0 F α (w) of the ∆-FL objective. We also consider q-FFL [START_REF] Li | Fair Resource Allocation in Federated Learning[END_REF], 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 and2. We visualize in Figure 1 the distribution of test errors.

∆-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 heterogeneity aware objectives, ∆-FL achieves 1.8% improvement over the next best objective, which is Tilted- ERM. We note that q-FFL marginally outperforms ∆-FL on Sent140-Linear, but the difference 0.05% is much smaller than the standard deviation across runs. ∆-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 greater reduction in misclassification error on data-poor clients, that is clients with little data (< 200 examples per client).

A Pseudocode

The pseudocode of the proposed optimization algorithm is given in Algorithm 1. 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

: Set π(t) i = I F i (w (t) ) ≥ Q (t) 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

B Proofs

Proof of the Subgradient Expression (3). We first give a general expression for the subgradient. Define the notation

P α = π i ∈ R n : 0 ≤ π i ≤ 1 αn ∀ i ∈ [n], n i=1 π i = 1 , so that S α (u 1 , • • • , u n ) = max π∈Pα π ⊤ u. Proposition 1. Fix a w ∈ R d and let π ⋆ ∈ arg max π∈Pα n i=1 π i F i (w). Then, we have, n i=1 π ⋆ i F i (w) ∈ ∂F α (w) ,
where ∂F α (w) denotes the regular subdifferential of F α .

Proof. Let g n (w) = (F 1 (w),

(S α • g n ) = ∇g n (w)∂S α (u) ,
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: πi = I F i (w) ≥ Q α (Z(w)) , and, πi = πi

n i ′ =1 πi ′ . ( 4
)
The objective defined by these weights is Fα (w) = n i=1 π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α n i=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 0 ≤ F α (w) -Fα (w) ≤ max i=1,...,n |F i (w)| αn .

Proof. We apply the property that the superquantile is a tail mean for discrete random variables [START_REF] Tyrrell | Conditional Value-at-Risk for General Loss Distributions[END_REF]Proposition 8] to get

F α (w) = 1 αn n i=i ⋆ +1 F i (w) + 1 - ⌊αn⌋ αn F i ⋆ (w) .
Comparing with the definition F α (w) = n i=1 π ⋆ 1 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

Fα (w) = 1 n -i ⋆ + 1 n i=i ⋆ F i (w) .
The bound on Fα (w) -F α (w) follows from elementary manipulations.

Figure 1 :

 1 Figure 1: Per-client test misclassification error on EMNIST. Left: histogram of per-client errors. Right two: Scaltter plot of dataset size versus test error.

Algorithm 1 2 :

 12 The ∆-FL Algorithm 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 Sample m clients from [n] without replacement in S 3:

5

 5

Table 1 :

 1 90 th percentile of the distribution of test misclassification errors (in %). = 0.8 49.10 0.24 26.23 1.15 46.44 0.38 46.46 4.39 ∆-FL, α = 0.5 48.44 0.38 23.69 0.94 46.64 0.41 50.48 8.24 ∆-FL, α = 0.1 50.34 0.95 25.46 2.77 51.39 1.07 86.45 10.95

		EMNIST	Sent140
		Linear	ConvNet	Linear	RNN
	FedAvg	49.66 0.67 28.46 1.07 46.83 0.54 49.67 3.95
	FedProx	49.15 0.74 27.01 1.86 46.83 0.54 49.86 4.07
	q-FFL	49.90 0.58 28.02 0.80 46.39 0.40 48.66 4.68
	Tilted-ERM	48.59 0.62 25.46 1.49 46.69 0.49 46.54 3.27
	AFL	51.62 0.28 45.08 1.00 47.52 0.32 57.78 1.19
	∆-FL, α

Table 2 :

 2 Mean of the distribution of test misclassification errors (in %). .49 0.26 16.09 0.40 34.41 0.22 30.31 0.33 ∆-FL, α = 0.5 35.02 0.20 15.49 0.30 35.29 0.25 33.59 2.44 ∆-FL, α = 0.1 38.33 0.48 16.37 1.03 37.79 0.89 51.98 11.81

		EMNIST	Sent140
		Linear	ConvNet	Linear	RNN
	FedAvg	34.38 0.38 16.64 0.50 34.75 0.31 30.16 0.44
	FedProx	33.82 0.30 16.02 0.54 34.74 0.31 30.20 0.48
	q-FFL	34.34 0.33 16.59 0.30 34.48 0.06 29.96 0.56
	Tilted-ERM	34.02 0.30 15.68 0.38 34.70 0.31 30.04 0.25
	AFL	39.33 0.27 33.01 0.37 35.98 0.08 37.74 0.65
	α = 0.8 34

  • • • , 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 ∂S 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 [21, Thm. 10.6] as ∂

α (u) = arg max π∈Pα π ⊤ u .

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.