Private and Personalized Frequency Estimation in a Federated Setting

Authors: amrith  setlur
vitaly  feldman
kunal  talwar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:59:00

Abstract: Motivated by the problem of next word prediction on user devices we introduce and study the problem of personalized frequency histogram estimation in a federated setting. In this problem, over some domain, each user observes a number of samples from a distribution which is specific to that user. The goal is to compute for all users a personalized estimate of the user's distribution with error measured in KL divergence. We focus on addressing two central challenges: statistical heterogeneity and protection of user privacy.
Our approach to the problem relies on discovering and exploiting similar subpopulations of users which are often present and latent in real-world data, while minimizing user privacy leakage at the same time. We first present a non-private clustering-based algorithm for the problem, and give a provably joint differentially private version of it with a private data-dependent initialization scheme. Next, we propose a simple data model which is based on a mixture of Dirichlet distributions, to formally motivate our non-private algorithm and demonstrate some properties of its components. Finally, we provide an extensive empirical evaluation of our private and non-private algorithms under varying levels of statistical and size heterogeneity on the Reddit, StackOverflow, and Amazon Reviews datasets. Our results demonstrate significant improvements over standard and clustering-based baselines, and in particular, they show that it is possible to improve over direct personalization of a single global model.

Main paper: 
 introduction: 
 Federated learning algorithms jointly learn from decentralized user data, addressing statistical challenges in local learning [49, 58] . At the same time it presents two key challenges, among few others [44, 52] . Firstly, data distributions can vary across users (often called statistical heterogeneity [27, 76] ), which reduces the effectiveness of a single, shared model for all users. Secondly, the data may be sensitive and mutually beneficial collaboration may compromise user privacy [31, 71] . An important practical problem where user data distributions are diverse, and user-privacy is paramount is next-word prediction for the keyboard input on user devices. Users often have diverse vocabulary, writing styles, and topics that leads to varied data distributions [32, 36, 65] . Naturally, frequency estimation is also one of the most basic statistical tasks with numerous other applications [14, 47] .Motivated by personalized next-word prediction, we introduce and study the problem of personalized frequency estimation in a federated setting. In this problem, each user has very few (O(d)) samples from an unknown, user-specific distribution over a large, finite domain of size d. The users interact with a server that needs to provide the user with a personalized estimate of their distribution. We focus on error measured in KL divergence that is common in language modeling [3, 38] , and is equivalent to minimizing the negative log-likelihood of samples from the user distribution. For more general loss minimization problems, numerous personalized federated learning (PFL) algorithms [67, 76] , tackle statistical heterogeneity by first learning a single global model across all users (for e.g., FedAvg [57] ), and then finetuning (FT) the global model locally for each user (FedAvg+FT [18, 20] ), thus balancing collaboration and personalization. The downside of this approach is that it is agnostic to the structure of the user population which is highly heterogeneous, but often consists of a number of concentrated subpopulations of similar users.We study the personalized frequency estimation problem under an intuitive model of the user-population which consists of well-concentrated clusters of users where users in the same cluster have "similar" token distributions. For this, we propose an iterative algorithm that adapts Lloyd's clustering [53] to distribution estimation in KL divergence. Relying on the optimal competitive estimation rates for local Good-Turing estimators, we propose to estimate global cluster centers by averaging user-level Good-Turing estimates for the current users in the cluster. Furthermore, since the performance of clustering is significantly affected by the initialization [5] of cluster centers, we also propose a data-dependent initialization approach, also specific to estimation in KL. We improve the practicality of our algorithm by separately handling data-poor users in cases where users vary significantly in their local data sizes (size-heteogeneity).Privacy: To address the challenge of user-level privacy we give a joint differentially private (DP) version of our iterative algorithm which requires that the estimator for user i is differentially private with respect to the data of all the other users (but may depend arbitrarility on user i's data) [48] . This definition is necessary due to the final personalization step which happens locally on user device and thus does not present a privacy risk. Specifically, we make each iteration of our clustering algorithm provably private by relying on adaptive clipping, secure aggregation, and common noise addition mechanisms [23] . Initialization of clustering algorithms typically uses data points themselves as initial centers and hence presents a significant challenge for privacy preserving analysis. Our provably private data-dependent initialization algorithm for cluster centers runs the exponential mechanism over candidates randomly sampled around the estimated population mean.Empirical evaluation: We validate both non-private and private versions of our algorithm on the real world data datasets: Reddit [16] , StackOverflow [6] , and Amazon Reviews [62] , where Reddit has a token vocabulary of size 10k and the others have 32k tokens. We find that our method furnishes significant gains over standard and clustering baselines, reducing the error by over 26% and 42% (averaged over datasets) in the non-private and private settings respectively. In the private, sizeheterogeneous setting, we improve over the clustering baseline IFCA [33] by over 30%. Finally, we show that our data model inspired adjustments are pivotal in yielding the performance improvements in practice. Particularly, we justify the use of our Good-turing based estimator, and our private initialization scheme with a favorable privacy/utility tradeoff.Formal guarantees. While our focus is on the empirical performance, to guide intuition and design of algorithms, we introduce a relatively simple generative model where user distributions are sampled from a mixture of Dirichlets, where Dirichlets are sufficiently separated globally and concentrated locally. We derive and analyze different estimators (e.g., Bayes optimal, FedAvg [57] , Good-Turing [34] ). Even when the cluster identities are known, we show that FedAvg (average of user empirical distributions) has poor guarantees compared to our proposed estimator (average of user Good-Turing estimates), where the minimax error rates for the latter do not scale with dimension, and is optimal for some regimes of problem parameters. Thus, using this model we demonstrate how additional structure in the user data can be exploited in the context of our more concrete frequency estimation problem while preserving the privacy of user data. 
problem definition: 
 We use P to denote a distribution, P [v] for the probability of event v, P for an estimate of P (but P if the estimate is computed without user collaboration), P for sets. P for matrices (P :,i /P i,: index into the i th column/row respectively), and {P i } for a collection of elements P i indexed by i.Setup. Let V be a finite vocabulary (set) of tokens, with vocabulary size d =: |V|, and ∆(V) is the set of all probability distributions over V. In a federated setup we have a collection of users U, and each user u ∈ U has an unknown distribution Q u ∈ ∆(V) over the tokens. The data available to each user u is a dataset S u comprising of m u i.i.d. samples from Q u . We will use Q u to refer to the empirical distribution estimated from S u . The combined collection of datasets is referred to as meta-dataset S =: {S u }, and similarly Q =: { Q u } u is the collection of all empirical distributions. Unless specified otherwise, we assume that we are in the size-homogeneous setting where all users have a dataset of the same size m = m u . We refer to the more general case as size-heterogeneous.Goal. Given the meta-dataset S, the goal is to compute estimates { Q u } u such that:1. utility:u E[KL(Q u ∥ Q u )] is minimized, where the expectation is taken over sampling of meta-dataset S and any randomness in computing { Q u } u . 2. privacy: Q u is computed privately with respect to all users except u, i.e. the estimate is joint differentially private (see Definition 2.2).Definition 2.1 (u-neighboring meta-datasets). Meta-datasets S and S ′ are u-neighbors, i.e., S ∼ u S ′ if they differ only in inclusion or deletion of user u's private data. For any algorithm M, we denote M -u as the output of the algorithm to all other users except u. Definition 2.2 (Joint Differential Privacy [48] ). Given ε ≥ 0, δ ∈ (0, 1], and neighbouring relation ∼ u , a randomized mechanism M : D → Y from the set of meta-datasets D to an output space Y is (ϵ, δ)-joint differentially private if for all u-neighboring S ∼ u S ′ , and all events E ⊆ Y,EQUATIONwhere probabilities are taken over the randomness of M. We say that M is ρ-zero-concentrated joint DP (ρ-zCJDP) if for all u-neighboring meta-datasets S ∼ u S ′ :EQUATIONwhere D α (M -u (S)∥M -u (S ′ )) is the α-Rényi Divergence between the outputs of all users except u. 
private algorithm for learning personalized histograms: 
 In this section, we present our algorithm for learning private and personalized histograms in both size-homogeneous and size-heterogeneous settings. Without assumptions on the distribution of users, in the worst case, the best estimator would simply estimate each user's histogram locally, i.e., estimate each user's Q u using their local dataset S u . Avoiding such worst-cases, we assume that the real-world problem is more structured. Specifically, we assume that each user distribution Q u is well clustered around K unknown cluster centers in the population of users. Note, we assume no knowledge of the cluster centers or number of clusters K. In Section 5, we empirically validate the presence of clusters in real-world datasets by evaluating the performance of our clustering based algorithm on the same datasets, and in Section 4 we theoretically analyze some of the algorithmic design choices we make, in a stylized Bayesian setting that simulates a user-distribution with clusters.As a warmup, we first discuss our local finetuning algorithm, where the setup is simpler. Say an algorithm runs a private collaborative protocol, and hands a distribution P to user u, with the guarantee that KL(Q u ∥P ) is "reasonably" small, then how should the user u finetune P using the local dataset S u available to them? For this question, we now present our local finetuning algorithm.Let P be the distribution over tokens returned by a collaborative algorithm that uses datasets across all users to estimate a single distribution for either all users in the population, or all users in a specific cluster. In the next section, we discuss how to estimate P . Here, we discuss how to personalize P for each user. For a user u, with empirical distribution Q u from m u samples, we use FT(P, Q u ) to denote the personalized, locally finetuned estimate of user distribution Q u (see Eq. 3). Note that P is moved closer to Q u in KL divergence (depending on choice of α > 0). When the population of users is well concentrated, FT lowers variance by relying more on the global estimate P , while adding some bias ( Q u is an unbiased estimate of Q u , but P is biased). In Section 4, we show that FT is the Bayes optimal estimator in KL divergence, under certain assumptions on the user distribution.EQUATIONNext, we discuss how to collaboratively estimate P . For a user-distribution with multiple clusters, a natural approach is to first identify the groups of clustered users. Then, we can use the distribution that is close in KL distance, and in expectation over the distribution of users in the cluser. We refer to these as the cluster centers. In the next Section, we present our algorithm to learn cluster centers, which we personalize for each user in the cluster using our local finetuning update (Equation 3).As mentioned in the start of Section 3, we base our approach on the belief that the user-population comprises of unknown clusters. Thus, given the collection of size-homogeneous users U, the goal is to learn a user partition {z u } u , and the matrix of K cluster centers P ∈ R d×K . To learn cluster centers P, akin to the popular clustering algorithms like Lloyds [53] , the objective we optimize is:EQUATIONwith the main difference being that we are concerned with the KL distance, instead of ℓ 2 2 . Given P, user u's assigned cluster zu is simply given by arg min k KL(Q u ∥P :,k ). Conversely, applying Lemma 3.1, given cluster memberships the cluster center is determined by the cluster average. Lemma 3.1. For users with assigned cluster i, i.e., Zi =: {u : zu = i}, the following is true:1 |Z i | u∈Zi Q u ∈ arg min P ∈∆(V) u∈ Zi KL(Q u ∥P ).Since, we do not have access to Q u , we replace Q u in Lemma 3.1 and the cluster assignment step with the local Good-Turing [34] estimate Q gt u . Typically, one would use the empirical distribution Q u , which is an unbiased estimate of Q u . But, as we will see in Section 5, most real-world distributions are long-tailed and the Good-Turing estimate's approximation of the user distribution does not suffer as heavily from the high-dimensional nature of the problem. We discuss the Good-Turing estimate and analyze its error more formally in Section 4. This simplification means that the cluster center estimator averages the local Good-Turing estimates from each user in the cluster. In theory, we can also estimate the cluster center P in a better way, e.g., by directly running the Good-Turing correction on the cluster level token counts. This is a possible direction of improvement in future works.To iteratively optimize the objective in Equation 4, we present our algorithm HistogramCluster (see Alg. 3). We start with initial center estimates P(0) , and run T iterations to refine them. In each iteration we peform two steps: (i) collect user outputs for each user via UserOutput (see Alg. 1); and (ii) re-center clusters using an algorithm, which for now is the NonPrivateCenter (see Alg. 2). Later, we see how we need only update the re-centering algorithm to satisfy user-level privacy contraints. From Lemma 3.1 and Alg. 1, it is not hard to see that NonPrivateCenter which computes the mean of the local estimates of users in a cluster, greedily reduces the objective value in Eq. 4 after each iteration. When the misclustering rate is sufficiently low to begin with, averaging local Good-Turing estimates brings the estimated center closer (in KL divergence) to the true centers that minimize Eq. 4, thereby reducing the number of iterations for convergence.An extension of the convergence analysis in Balakrishnan et al. [7] , gives us an upper bound on the misclustering error of Alg. 3 which goes down exponentially with each iteration of clustering (Theorem 3.2). Similar to most clustering analyses [8, 17, 33] , we make suitable conditions on the initialization of cluster centers P(0) and the separation of true clusters in the user population.Require: Centers P, cluster k, user u. Ensure: User output (bu, Qu).1: bu ← 1(arg min j KL( Q gt u ∥ P:,j) = k) 2: Qu ← bu • Qgt u + (1 -bu) • 0 d 3: Return (bu, Qu).Require: User data A =: {(bu, Qu)}. Ensure: Center P .1: Return P ← u Qu / u bu.Require: Set of users U, initial centers P(0) , iterations T . Ensure: Estimates of cluster memberships {zu}u and centers P.1: Initialize t ← 1. 2: while t ≤ T do 3: for j = 1 to K do 4: Collect: A (t)k ← {UserOutput( P(t-1) , k, u)}u 5:Re-center: P(t) Theorem 3.2 (Alg. 3 convergence). For the model in Section 4, if ∀c, v, n c = Θ(K 2 ), P c [v] = Ω(1), centers are sufficiently separated, i.e., ∆ =: min i̸ =j KL(P i ∥P j ) = Ω(k 2 + k 3 d /n), and λ =: max i̸ =j KL(Pi∥Pj ) /∆. Given an initialization with assignment error O( 1 /λ) for any cluster, after t = log(|U|) iterations, w.h.p. the assignment error is O(exp(-∆(α + 1))).First, we note that it suffices to compute all the cluster center estimates in the usual model of DP to ensure that the personalized estimates are differentially private in the billboard model. This is true since the assigment of the user to a cluster center and finetuning can be done locally by each user. Thus, if we ensure that Alg. 2 applied to all the current clusters is differentially private in each iteration we can then bound the total privacy leak across iterations via advanced composition for DP algorithms. In our algorithm PrivateCenter (Alg. 4) we first use the Laplace mechanism in step 1 to privately compute the number of users in the cluster. Then, we use the Gaussian mechanism in step 2 to get an initial estimate of the center since all local Good-Turing estimates are distributions in ∆(V) and thus bounded in ℓ 2 . This initial estimate B is used to identify the quantile B ± c B/b, where the mean lies with high probability [69] , so that we can then get a refined estimate in step 3 when we add noise (with std. deviation scaling with quantile interval c). This adaptive clipping procedure is fairly common for practical and private estimation [4, 12] . Finally, to ensure the noised output is still a distribution contained in ∆(V), we use ℓ 2 projection, Π ∆(V) (P ) =: arg min Q∈∆(V) ∥Q -P ∥ 2 2 . In practice, we find that projection with ℓ 2 is comparable with a KL projection.Privacy preserving noisy sum. As can be seen from the description of PrivateCenter, it only requires being able to compute a noisy sum over a subset of users. In the federated setting implementing this summation as part of algorithm HistogramCluster may seem to require the server to know the (private) cluster membership of each user. However, this can be avoided by computing each such sum as a sum over all the users, where a user submits value 0 if they do not belong to the cluster, i.e., b u = 0, and b u = 1 otherwise (Alg. 1). Similarly, Alg. 1 also emits the true local estimate if it is part of the cluster and 0 d (vector of d 0s) otherwise. Privacy-preserving computation of a noisy sum over all the users is one of key primitives in FL with a number of known implementations, e.g., Bonawitz et al. [13] . Crucially for the privacy analysis, while each user now participates in the sum computations for every cluster, the inclusion/deletion of a single user can only affect the sums for a single cluster.Require: Data collection {(bu, Qu)}, privacy parameter ρ, clipping parameter c, projection operator Π ∆(V) . Ensure: Private center P .1: b ← u bu + Lap 3 /2ρ 2: B ← 1 b u Qu + N (0, 3 /2ρI d ) 3: Σ ← 6c 2 /bρ • diag(B) 4: P ← 1 b u Qu + N (0, Σ) , Qu =: Clip [Qu]The performance of Alg. 3 is crucially determined by the proximity of initialized centers to each of the true centers (Theorem 3.2), and more critical for higher private noisy tolerance. In Pri-vateInit (Alg. 5), we provide a private data-dependent initialization technique, that takes as input a private center Q 0 (private FedAvg), and samples K 2 points around it as an initial candidate set. Then, it chooses K points iteratively from this set by sampling via the exponential mechanism. The utility function for each candidate Q is determined by the average reduction in the clustering objective when adding Q into the current set of initializations I, vs. not. The privacy budget ρ controls the temperature, with lower budget enforcing higher smoothing. In the non-private case, we replace the candidate set with the set of empirical user estimates, and select a high value of ρ. This iterative procedure of selecting candidates with highest regret of omission at each step is meant to ensure that each Q u is reasonably close to some initial center (akin to [5] ). We remark that the execution of the sampling step does not lend itself easily for implementation in the distributed setting. In practice, one would need to use a small dataset collected centrally for this.3.4 Putting it all together: Algorithm for private and personalized frequency estimationIn Alg. 6 we present our complete algorithm. We first compute a private center of all user histograms using Alg. 4 and privacy parameter ρ /3. We pass this to Alg. 5 with privacy parameter ρ /3 to get initial centers for Alg. 3, where the re-centering step applies Alg. 4 with privacy parameter ρ /3T . The privately learned cluster centers from Alg. 6 are then finetuned for each user locally, using Eq. 3. Theorem 3.3 (End-to-end privacy guarantee). Algorithms 4, 5 are both user-level ρ-zCJDP. Our end-to-end private and personalized estimation algorithm is (ρ + 2 ρ log( 1 /δ), δ)-JDP, ∀δ > 0.Require: Set of users U, dataset S, number of clusters K, number of clustering iterations T , finetuning parameter α, sampling parameter ω, privacy parameter ρ, clipping parameter c, projection operator Π ∆(V) . Ensure: Private and personalized estimates Qu for each u ∈ U .1:Q0 ← PrivateCenter(U, ρ, c, Π ∆(V) )# Estimate global center 2: P(0) ← Algorithm PrivateInit(Q0, S, K, ρ, ω, c) # Estimate initial cluster centers 3: {zu}u and centers P ← Algorithm HistogramCluster(U, P(0) , T ) # Estimate cluster memberships 4: ∀u : Qu ← FT( Pzu , α) # Local finetuning using Eq. 3 5: Return { Qu : u ∈ U}.When users vary in size, we need to weigh their local estimates while estimating the global center or the cluster center in the re-centering step of Alg. 3. For user u, the weight w u =:( 1 /σ 2 u ) / u∈Zc ( 1 /σ 2 u ).Here σ 2 u is the variance of the user's local estimate, i.e., Var[ Q u ] (from Lemma 3.4). When P zu [v] is known this weighted estimate is optimal under ℓ 2 [37] . Since P zu [v] is unknown we replace P zu [v] in Lemma 3.4 with a uniform average of Q u from heavier users, following Cummings et al. [21] .Lemma 3.4 (Variance of Q u ). For user u, Var[ Q u [v]] = Pz u [v] /α+1(1 -1 /mu) + Pz u [v](1-Pz u [v]) /mu.In the private setting though this approach cannot be used directly since the sensitivity of the weighted mean estimate is too high for users with a lot of data, and more so for small sized clusters. Cummings et al. [21] propose an algorithm to estimate the private heterogeneous mean of user data from different Bernoulli distributions, but their approach does not directly transfer to our setting due to the highdimensional nature of our user means and the relatively poor sensitivity of Alg. 3 to errors in the estimation of cluster centers at each iteration. Thus we consider a two-stage approach, split across data-rich and data-poor users. In the first stage we privately cluster users with sufficiently large datasets applying Alg. 3, and treating them as size-homogeneous data-rich users. Once cluster centers are learned privately, using Alg. 1, each data-poor user assigns itself to the closest center in KL divergence. Finally for each cluster, we apply Algorithm 2 from Cummings et al. [21] that re-centers each cluster privately, based on the newly added users. 
formal analysis in a stylized data model: 
 In Section 3, we presented a clustering based iterative algorithm (Alg. 3) for our frequency estimation problem introduced in Section 2, wherein we made several key design choices. In particular, we used average of Good-Turing estimators to estimate the cluster center, and used the update in Equation 3to locally finetune the learnt cluster centers. Now, we analyze these algorithms in a stylized model. Bayesian Model. Each user u belongs to a cluster z u ∈ [K] with cluster center P zu . The user histogram Q u ∼ Dir(αP zu ), and the user dataset S u ∼ Q mu u is sampled i.i.d. from Q u . We use Z c =: {u : z u = c}(Q u [v]) = O( Pz u [v] /1+α), and as α → ∞, Q u → P zu in weak topology [28] . Please note that the Dirichlet assumption is mainly for simplicity and our results (e.g., Theorem 4.3) only require each cluster's user distribution to be exponentially concentrated along each token, i.e., Pr(|Qu [v] -P zu [v]| ≥ t) = O(e -t ).Purely local learning. We present two local estimators that estimate Q u : empirical Q u . and Good-Turing [30] Q gt u . The naïve estimate Q u is the average of user data in S u . Next, we define the Good-Turing estimator. Let the frequency of token v in user's dataset S u (of size m u ) be cnt u,v . We denote the local frequency of the count j as ϕ u,j =: v∈V 1(cnt u,v = j), i.e., the number of token in S u with count j. Following Orlitsky and Suresh [63] , for a token v that appears j = cnt u,v times, the Good-Turing estimator for Q u [v] is given by:EQUATIONwhere N u is the normalization factor so that v Q gt u [v]=1. Let E u be the class of natural estimators that assign the same probability to tokens with equal counts in S u . Then, following Lemma 6 in Orlitsky and Suresh [63] and Theorem 2 in Acharya et al. [2] we conclude that with respect to E u , the Good-Turing estimate Q gt u has a worst case suboptimality gap of Õ(min( 1 /m, d /m) (see Lemma 4.1 for the full statement). The key point to note here is that when the user dataset is small, we do not suffer from the vocabulary size d, unlike the empirical estimate Q u . This is mainly because Good-Turing more accurately estimates the probability of unseen words [56] .Lemma 4.1. ( Q gt u suboptimality gap) For any Q u , the suboptimality gap of Q gt u with respect to E u is E[KL(Q u ∥ Q gt u )] -min Ã∈Eu E[KL(Q u ∥ Ã)] = Õ(min( 1 /m, d /m)), that matches minimax rates.Bayes optimal finetuning when the cluster center is given. When P zu is known, FT in Eq. 3 is Bayes optimal in KL and when only an estimate of the center is known, the error of the "plug-in" estimate scales linearly with the error in the center's estimate (Theorem 4.2).Theorem 4.2 (Bayes optimal local learning). Given P zu , α, the estimator FT(P zu , Q u ) is Bayes optimal in KL divergence, for the Dirichlet prior Dir(αP zu ). When Pzu is the estimated cluster center, andQopt u is the Bayes optimal estimate of Q u , then KL( Qopt u ∥FT( Pzu , Q u )) ≤ α /α+mu KL(P zu ∥ Pzu )).Estimating the cluster center given cluster members. Given all users in a cluster c, one estimate of the cluster center P c is the solution to the maximum log-likelihood objective in Eq. 6, which is typically solved with the FedAvg [49] algorithm. We denote this estimate as P fa c and is a simple weighted average of empirical estimates (see proof in Appendix E). Similary, for user u in cluster c, the personalized output from the PFL algorithm FedAvg+FT is given by FT( P fa c , Q u ). P fa c =:u∈Zc mu / u ′ ∈Zc m u ′ • Q u ∈ arg max Q∈∆(V) 1 n c u∈Zc log Pr(S u | Q) (FedAvg) (6)Our algorithms (Alg. 6, Alg. 2) use a different estimator. In the size-homogenous case, to estimate the cluster center we average the local Good-Turing estimates Q gt u for all users in the cluster:EQUATIONTo analyze the accuracy guarantees of our estimate P gt c , we first define the class of competing estimators for it. Let c nt v =: ⌊ u∈Z cntu,v / u mu⌋ be the average count of token v across all users in a cluster. We define Ē to be the class of estimators that assign the same mass to tokens v 1 , v 2 if c nt v1 = c nt v2 . Such a class is natural for estimators that rely on aggregated statisitics, and includes the Good-Turing correction applied to cluster level counts. In Theorem 4.3, for any cluster center P c , we compare the suboptimality gap in KL divergence for our estimate P gt c in Eq. 7 with FedAvg. Theorem 4.3 ( P gt c in Eq. 7 has lower suboptimality gap than FedAvg). For the competetive class of estimators with the same average count, i.e., Ē (defined above), and suboptimality gap from Lemma 4.1, the suboptimality gap for P gt in Eq. 7, w.r.t.Ē is Õ( α+m /α+1 • min( 1 /m, d /m)). The suboptimality gap over Ē for FedAvg P fa is Ω ( d(α+m) /m(α+1)). Õ, Ω hides polylog factors in m, n.When m is small and comparable to the intra-cluster std. deviation, i.e., m /α+1 = O(1), then the suboptimality gap for our estimator does not suffer from large dimension d . On the other hand, FedAvg guarantees are much weaker and degrade with dimension for small user datasets (common in practice). From, Theorem 4.2, we also conclude that our final estimate for Q u , given by the finetuned center: FT( P , Q u ) would also have stronger error guarantees than FedAvg+FT. Partitioning users into clusters. In practice, neither the cluster membership z u nor the center P c is known. The Bayes optimal estimate for P c , z u is the mean of the posterior Pr(Q u | S, α) (c) Amazon Reviews Figure 1 : Performance before finetuning: We compare the test NLL loss before local personalization (finetuning) for baselines FedAvg, MAML, IFCA with our approach. NLL is uniformly averaged over users and each value is averaged over 50 random runs (error bars indicate 95% confidence intervals).(see Theorem 4.2 proof), but since the posterior is intractable, we instead compute the maximum likelihood estimate (MLE). Asymptotically, MLE's accuracy is vindicated by Doob's theorem [59] , and its variance matches the Cramér-Rao bound [69] . Still, the joint MLE for P c , z u is a non-concave maximization problem. Akin to solving the MLE for mixture distributions [61] , we maximize an evidence lower bound with Expectation-Maximization (EM) [7] . Since each Q u belongs to only one cluster, it must be that z u ∈ arg min c∈[k] KL(Q u ∥P c ). Thus, the MLE for the cluster centers is given by objective in Eq. 4 from Section 3, which is optimized iteratively by our Algorithm 6.Takeaways. We note: (i) when cluster memberships are known, average of local Good-Turings has a lower suboptimality gap than FedAvg; and (ii) solutions to the clustering objective in Eq. 4 estimates cluster centers for each user; and (iii) if the center estimates are accurate then finetuning them with FT yields estimates that are equally close to the Bayes optimal solution. These findings validate our algorithmic design choices for the clustering and finetuning algorithms in Section 3. 
empirical evaluation: 
 We empirically evaluate our approach on real-world datasets and present: (i) contrary to popular belief [76] , we show there are real-world distributions where clustering based algorithms (e.g., ours) significantly outperform the popular PFL baseline FedAvg+FT; (ii) our method achieves better privacy-utility tradeoff than private versions of clustering-based baselines, and the gap amplifies in the size-heterogeneous case; and (iii) we show our method's performance improvements can be largely attributed to algorithmic design choices influenced by our data model and analysis in Section 4.Datasets. We evaluate methods on three real-world datasets: Reddit [16] , StackOverflow [6] , and Amazon Reviews [62] . For Reddit, we use the NLTK tokenizer [11] with a vocabulary of size 10k tokens, and for the other two datasets, we use the Huggingface (bert-case-uncased) tokenizer [74] with a vocabulary size of 32k tokens. For more details on datasets and hyperparameters see Appendix C.Baselines. We evaluate two baselines that learn a single global model: FedAvg [57] and MAML [29] ; and a clustering-based approach which learns multiple models, one for each cluster: IFCA [33] /HypCluster [54] . Given a single global model (or a cluster level model), each user can finetune this locally in different ways: full batch gradient descent (GD) initialized at global model, RTFA [19] , or with our problem specific method in Eq. 3. We compare these on FedAvg model (Figure 3(d) ), and fix the best one (ours) as the local finetuning approach for all algorithms. In the private setting, we adapt the baselines by using standard techniques [1, 25] that privatize l 2 bounded gradients for MAML/IFCA, and l 1 bounded probability distributions for FedAvg.Evaluation metric. For distribution Q u , estimate Qu , KL(Q u , Qu ) = NLL(Q u , Qu )-H(Q u ), where NLL(Q u , Qu ) =: -w∈V Q u,w log Q u,w is the negative log-likelihood loss, and H(Q u ) is the entropy of Q u . Following language modeling [3, 38] , we use the test set of the user to get an unbiased estimate of the NLL loss, by replacing Q u with the empirical distribution on the test set. Since the test set is insufficient to get an unbiased estimate of H(Q u ), as a reference point, we instead report the entropy of the global center or FedAvg estimate H( P f a ), which is a rough estimate of the "hardness" of estimating the user distribution. We refer to this difference in test NLL and H( P f a ) as NLL sub-optimality gap. In the private case, we ensure the algorithms satisfy (15, 10 -10 )-JDP.Our approach significantly reduces test NLL in private and non-private settings. In Figure 1 we report a significant reduction in test NLL sub-optimality gap even before local finetuning, by atleast 25-45% in the non-private setting and upto 50% in the private case. Non-privately MAML/IFCA perform better than FedAvg on StackOverflow and Amazon, but their private versions perform similar or worse than the private FedAvg. This is because the gradient-based optimization in MAML and IFCA can incur very high per-iteration privacy overheads which only accumulates more for the latter that additionally suffers from poor convergence due to imperfect cluster initialization. On the other hand, our method achieves a privacy-utility tradeoff that is comparable to the gradient-free FedAvg, and converges in fewer iterations when clustering is initialized with centers from Alg. 5. In Figure 2 we compare the performance after we finetune the global/cluster-level estimates locally for each user (using Eq. 3), and verify that the relative gains remain consistent with pre-finetuning. In particular we note that finetuning global FedAvg/MAML models does not do better than finetuning cluster-level models, vindicating the presence of concentrated subpopulations in an otherwise highly heterogeneous real-world user distribution.Algorithmic design decisions influenced by our model in Section 4. In Section 4 we identified the Good-Turing based estimator (Eq. 7) to suffer less from vocabulary size, compared to FedAvg (Theorem 4.3), and in practice too we observe a big improvement in the test NLL loss of Alg. 3 when it uses Good-Turing as the local estimate, vs. empirical (Figure 3(c) ). The relative gap is particularly wider on the large vocabulary datasets Amazon and StackOverflow. In Figure 3(b) , we show that our data-dependent private initialization (Alg. 5), that is also based on our mixture of Dirichlet model, plays a crucial role in lowering the final test loss, and we attribute this gain to a significant reduction in clustering iterations (from 200 to 20), thereby reducing privacy overhead. In Figure 3(d) we note that the Bayes optimal finetuning algorithm in our data model (Eq. 3) does better than typical gradient-based (GD) or proximal term based (RTFA) approaches. Finally, in Figure 3 (a), we plot test NLL as we vary the number of clusters assumed by Alg. 3. We note that for runs on all datasets we choose K=10 which was found to be optimal on the Reddit validation set (see Appendix C), even though it is clearly (slightly) suboptimal on the other two, suggesting that Alg. 3's performance is not too sensitive to the choice of K practice.Size-heterogeneous case. In Figure 4 plot test NLL in the size-heterogeneous case where our algorithm's gains (> 40% error reduction) over single global model baselines FedAvg and MAML are even more pronounced than the size-homogeneous results in Figures 1, 2 . For the better suited clustering baseline IFCA, the privacy utility tradeoff worsens (compared to size-homogenous) due to poorer higher sensitivity of gradient estimates when dataset sizes vary. On the other hand, our two-stage approach reduces privacy noise in the clustering stage by partitioning users into different stages and still improves performance for data scarce users in the second stage by incorporating them in the re-centering of clusters from previous stage. Compared to IFCA we observe > 50% reduction in test NLL suboptimality gap for private estimation on Amazon Reviews. 
related work: 
 Private and personalized federated learning. Multiple works propose and analyze clustering based algorithms that learn a diverse set of models for heterogeneous user distributions [26, 33, 35, 54, 55, 66, 73, 77] . With the exception of [66] that first learns a global model and then uses it to partition users based on their losses, most works learn diverse models from scratch using gradient based algorithms, as done by FedAvg [57] . The key difference being that they first partition clients based on their loss [33] or gradients [72] and then use the same gradients to update different models, one for each cluster. Additionally, their analysis also holds mainly for smooth/strongly convex loss functions, for e.g., least squares [33, 54] . This is a natural and promising approach but we are not aware of practical results showing that it can improve on the more direct combination of FedAvg and FT [19, 76] . In fact, Wu et al. [76] raise concerns about mode collapse with clustering iterations. Contrary to the works above, we focus on a non-gradient based approach specifically for frequency estimation in KL divergence, which can be ill-conditioned in practice, and also provide an algorithm for the more challenging size-heterogeneous setting. As we show in our work, this problem requires solutions that are different from the ones proposed for general loss families in machine learning. Moreover, unlike the above, we provide privacy guarantees for our clustering based personalization. For more discussion on related works please see Appendix B Distribution estimation in KL divergence. A multitude of works on mixture of Gaussians, and mixed linear regression propose and analyze distribution estimation algorithms [9, 22, 45, 61] , but their guarantees are mainly for parametric estimation errors in ℓ 1 /ℓ 2 metric. In contrast, we are concerned with histogram estimation in KL divergence which presents interesting challenges since this is not a proper distance metric (e.g., does not satisfy triangle inequality). On the other hand, [30, 34, 64] study Good-Turing estimators and give estimation error guarantees in KL divergence for categorical distributions over a fixed alphabet. Our work extends estimators of Acharya et al. [2] , Orlitsky and Suresh [63] to the federated setting where the goal is to estimate a full population of distributions that share a latent structure. Relevant to our objective and metric is [17] that analyzes guarantees for information theoretic clustering. Their analysis shows that one can adapt analysis for other metrics (e.g., Hellinger) to obtain worst-case approximation guarantees. In contrast, we investigate practical and private algorithms for the federated setting. 
conclusion: 
 We introduce the problem of private and personalized frequency estimation in KL divergence. For this, we propose an iterative algorithm that privately learns clusters in the population of all user frequencies. Each user in a cluster locally finetunes their corresponding cluster center to produce personalized and private frequency estimates with formal joint DP guarantees. We improve the privacy-utility tradeoff of our algorithm by proposing a novel data-dependent private initialization for clustering that empirically reduces number of clustering iterations. We also present a two-stage version of our approach to separately handle the harder size-heterogenous setting. In a Bayesian model where the user distributions are distributed as Dirichlets around well-separated centers, we reason about different collaborative and local estimators, and provide formal guarantees for some of our algorithmic design choices, like Good-Turing estimators, and the choice of the local finetuning algorithm. Empirically, we test our algorithm on three real-world datasets and show a significant reduction in test NLL, by 25-45% in the non-private setting and upto 50% in the private case. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: In Appendix A we discuss some limitations of our work. This paper presents work whose goal is to advance the field of Private Federated Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work is not focused on releasing any pretrained models or datasets.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We make sure to cite relevant work including references to the three publicly available federated learning datasets we use in our work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: We do not release any new assets as part of this work.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: We do not do any crowdsourcing experiments.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: We do not conduct any experiments with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: As a first step towards language modeling we only focus on privately estimating the marginal token distribution for each user, and an interesting future direction would be to extend this to conditional next-token distributions. Nevertheless, we find that existing algorithms in the federated learning literature when applied to our problem setting perform suboptimally, and require problem specific interventions we introduce, like Algorithms 3, 5, and Good-Turing based estimators. Our current work also does not handle a federated version of the private initialization algorithm we propose. We suggest handling the initialization at the central node that maintains a small dataset of clients. This is a known limitation for sampling from the exponential mechanism in a federated setting. At the same time we would like to point out, that all our other algorithms, including private ones can be implemented in the federated setting.In our work we consider joint, user-level privacy guarantees in the billboard model [39] , whereas some prior works at the intersection of privacy and personalization only provide weaker record-level privacy guarantees [40, 51] . While other works [50, 58] study the notion of user-level privacy in federated setups, they are mainly concerned with learning a single model. Cummings et al. [21] also study the problem of private heterogeneous mean estimation but do not consider the clustered setting, or the subtleties of estimatation in KL divergence. More recent works (e.g., [42, 43] ) propose personalization algorithms that are user-level private, but are mainly tailored for settings where the different models for each user share representations or lie in a hidden low rank subspace. In contrast, we explore another latent structure more suited for histogram estimation, i.e., one where histograms need only be concentrated around the vertices of a low-dimensional polytope. Other works [10, 41, 75] propose gradient-based algorithms for private multi-task learning where they learn different models for each user (task), regularized using inter-task prior relationship matrices. In contrast, our work does not assume any such prior knowledge.Datasets. We evaluate methods on three real-world datasets: Reddit [16] , StackOverflow [6] , and Amazon Reviews [62] . For the Reddit dataset, we use the NLTK tokenizer [11] with a vocabulary of size 10k tokens, and use the Huggingface (bert-case-uncased) tokenizer [74] with a size of 32k tokens for the other two. We have ≈ 10k users in Reddit, each with at least 1k tokens. For StackOverflow and Amazon Reviews we have ≈ 19k and 50k clients respectively, each with 500 tokens For all datasets, we partition the data for each client into 60 : 40 train/test splits. Additionally, we set aside 5% of users in each, as a validation, to tune cluster count K, privacy clip bound c, etc. In the private setting, we adapt the baselines to the private setting by using standard techniques [1, 25] to privatize l 2 bounded gradients for MAML/IFCA, and l 1 bounded probability distributions for FedAvg, so as to ensure all algorithms achieve JDP-(15, 10 -10 ). For more details see Appendix C.Hyperparameters. We validate hyperparameters for our algorithms and baselines using a hold out validation set of users. For clustering we use K = 10, as validated by Figure 5 , and fix this for all datasets and clustering baselines. We run clustering for T = 50 iterations non-privately and T = 20 iterations privately on all datasets. We tune the finetuning parameter λ by sweeping across {0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5} and find λ = 0.25, λ = 0.15 and λ = 0.1 to be optimal on Reddit, StackOverflow and Amazon respectively. We train MAML and IFCA using SGD with momentum, with learning rate 0.01 and momentum parameter 0.9. For RTFA, we find the proximal regularization parameter of 0.2 to be optimal on Reddit and StackOverflow, 0.1 to perform better on Amazon Reviews. For private training, we use a clipping threshold of c = 0.1 in Alg. 4, and use c = 4.0 for Alg. 4.Computational resources. None of our experiments require very high computational requirements and can be run with one 3090Ti card. All runs can be reproduced in approximately 500 GPU hours. EQUATIONand for x < λ,EQUATION)Lemma D.2 (Bernstein's Inequality). Let X 1 , X 2 , . . . , X n be n independent zero mean random variables such that with probability ≥ 1 -EQUATION)EQUATIONLemma D.3 (Banerjee et al. [8] ). The solution for the optimization problem:arg min Q (1 -λ) KL(P 0 || Q) + λ KL(P 1 || Q)is given by (1 -λ)P 0 + λP 1 .Lemma E.1. ( Q gt u achieves minimax suboptimality gap; restated) For any Q u , the suboptimality gap ofQ gt u with respect to E u is E[KL(Q u ∥ Q gt u )] -min Ã∈E E[KL(Q u ∥ Ã)] = Õ(min( 1 /m, d /m)), that matches minimax rates.Proof. From Lemma 8 in [63] we know that for the class of natural estimators E u , the suboptimality gap defined above is given by E[KL(M k ∥ Mk )], whereM k =: v∈V 1(cnt v = k)Q u [v], Mk =: v∈V 1(cnt v = k) Qu [v].We use the definition of ϕ u,t from [63] , where:ϕ u,t = v∈V 1(cnt u,v = t) (12)Now, from Theorem 2 in [63] we also know that KL(M k ∥ Mk ) = Õ( 1/m) with probability at least 1 -1/m over the randomness in draw of the user dataset S u . Since their estimate is lower bounded by 1 /2m, we can convert their high probability guarantees to guarantees in expectation. Further, we note that in their proof of Theorem 2, we can easily upper bound the followingt (M k -Mk ) 2 Mk ≤ k ϕ u,t m ≤ t ϕ u,t m = d min the worst case. Thus, in the worst case, E[KL(M k ∥ Mk )] is upper bounded by the minimum of O( d /m) and Õ( 1 /m). This completes the proof of Lemma 4.1.Theorem E.2 (Finetuning is Bayes optimal given center; restated). If cluster center P zu , and Dirichlet parameter α is known then FT(P zu , Q u ) is Bayes optimal in KL divergence.Proof. We note that the distribution of Q u given the center P cu is a Dirichlet-Multinomial distribution. Further, the Dirichlet distribution is a conjugate prior for the Multinomial.Hence, the posterior distribution for Q u | Q u , P cu is a Dirichlet Multinomial with mean: ( α α+m )P u + ( m α+m ) Q u . To see why, mean is the Bayes optimal estimator, we invoke Lemma D.3, which implies:(1 -λ)P + λ Q u ∈ arg min Q (1 -λ) KL(P || Q) + λ KL( Q u || Q)Now the Bayes risk of algorithm M, given the center estimate P is :BayesRisk(M) =: E Qu∼Pr(Qu| Qu,P ) KL(Q u ∥M( Q u ))Recursively applying Lemma D.3 on every measurable subspace (under measure defined by the posterior Pr(Q u | Q u , P )) of the set of token distributions ∆(V), we conclude that the optimal solution is E[Q u | Q u , P ] when Q u ∼ Pr(Q u | Q u , P). This concludes our proof. Note that the same proof also applies for Lemma 3.1, but applied iteratively on each element in the summation.Lemma E.3 (Accuracy of plugging center estimate into finetuning; restated). Let Pzu be the estimated cluster center for user u, and Qopt u be the Bayes optimal estimate of Q u . Then,KL( Qopt u ∥FT( Pzu , Q u )) ≤ α α + m u KL(P zu ∥ Pzu ).Proof. This result is a simple application of Jensen's inequality for KL divergence.KL( Qopt u ∥FT( Pzu , Q u )) = KL α α + m u P zu + α α + m u Q u ∥ α α + m u Pzu + α α + m u Q u ≤ α α + m u KL(P zu ∥ Pzu )Lemma E.4 (FedAvg estimate). The FedAvg model is given byQ fa = 1 n Qu∈S Q u .Proof. The log likelihood objective log Pr(S u ∥Q) is equivalent upto additive constant with the objective KL( Q u ∥Q u ). We can once again use the result from Lemma D.3 to conclude that the optimal is simply the average of local means.Theorem E.5 ( P gt c in Eq. 7 has lower suboptimality gap than FedAvg; restated). For the competetive class of estimators with the same average count, i.e., Ē (defined above), and suboptimality gap from Lemma 4.1, the suboptimality gap for P gt in Eq. 7, w.r.t. Ē is Õ( α+m /α+1 • min( 1 /m, d /m)). The suboptimality gap over Ē for FedAvg P fa is Ω ( d(α+m) /m(α+1)). Õ, Ω hides polylog factors in m, n.Proof. For notational convenience, we denote the cluster center with P (instead of P c ), and the number of users in the cluster as n (instead of n c ). To derive bounds on competetive estimation, where the competetive class can only use the average count for each token, we need to first ensure that our Good-Turing based estimator can actively estimate the following quantity:M t =: v∈V 1(⌊cnt v ⌋ = t) P v and F t =: v∈V 1(⌊cnt v ⌋ = t) P gt v (13)Similarly, the local equivalent of M t for each user is defined as:EQUATIONAs a starting point, we will first derive our results with the genie-aided estimator, where we have oracle access to the ration of expected frequencies for every user, and then switch the analysis to the more general version when we do not. The genie-aided local estimator Qgen u,v is:EQUATION)and the corresponding average of genie-aided Good-Turing estimates:P gen v = 1 n u Qgen u,v .Next, we derive results for the setting where m u ∼ Poi(m) so that the local frequency counts of words become independent. Then, using the results on Poisson sampling in Mitzenmacher and Upfal [60] , we can convert the bounds to the setting where each user samples exactly m tokens.We will use 1 t v as a shorthand for 1(⌊cnt v ⌋ = t) and similarly use 1 t u,v as a shorthand for 1(cnt u,v = t). The following quantities define the global and local frequencies of average and local counts respectively:EQUATIONFor F t defined using the genie-aided estimator, to bound the competetive KL estimation error we first need to bound the KL distance between ground truth M t and estimate F t under high probability over the draw of the meta-dataset S.EQUATIONUsing Lemma E.6 we lower bound the denominator in the KL upper bound in Eq. 17.Lemma E.6 (Lower bound on F t ). With probability atleast 1 -nm -2 , F t = Ω tϕt m .Proof. The proof relies on three results for local Good-Turing (per-user) estimates. We will combine the three results, and then apply a union bound over all the users. Let v be some token in the vocabulary that has an average count of t, i.e., 1 t v = 1. From Lemma 15, 16 and Claim 20 in Acharya et al. [2] , with probability at least 1 -m -3 , Q un u = Ω( cntu,v /m), where Q un u is the unnormalized local Good-Turing estimate. Additionally, from Lemma 17 in Acharya et al. [2] , the normalization factor N u for user u satisfies with probability at least 1 -10m -2 , N u = 1 + Õ( 1 /m 1/4 ). Since m u is a Poisson distribution from Lemma D.1, the number of data points m u for the user satisfies |m u -m| ≤ √ m, with constant probability. Putting these together, and applying a union bound over all n users we conclude:F t = Ω 1 n v∈V u 1 t v cnt u,v m = Ω cnt v ϕ t m = Ω tϕ t m ,with probability at least 1 -nm -2 .Next we try to upper bound (M t -F t ) 2 for some t ≥ 1. But before that we introduce some more helpful results.Lemma E.7 (User's Q u concentrates around center P ). With probability at least 1 -δ, ∀v ∈ V:1 n u Q u,v -P v = O P v log d /δ n(1 + α) Proof. Since Q u,v ∈ [0, 1]is a bounded random variable with variance Pv(1-Pv) /(1+α), we can apply the Bernstein inequality in Lemma D.2. This would give us a confidence interval for a single token v, which we can then union bound over all tokens v ∈ V to get the result.Lemma E.8 (variance of linear estimators; Claim 21 in Acharya et al. [2] ). For every distribution p,Var v t 1 t v f (x, t) ≤ x t E[1 t v ]f 2 (x, t).Proof. By Poisson sampling, the multiplicities are independent. Furthermore, the variance of the sum of independent random variables is the sum of their variances. Hence,Var v t 1 t v f (x, t) = v Var t 1 t v f (x, t) ≤ v E   t 1 t v f (x, t) 2   = v E t (1 t v f (x, t)) 2 (a) = v t E[1 t v ]f 2 (x, t). (b) For t ̸ = t ′ , E[1 t v 1 t ′v ] = 0 and hence (a). (b) uses the fact that 1 t v is an indicator random variable. Lemma E.9 (expected sensitivity of the local frequencies; Claim 20 in [2] ). For every user distribution p over V, from which we draw n i.i.d. samples, let t be the local count and ϕ t = v∈V 1 t v is the frequency of the local count,E[ϕ t ] -E[ϕ t+1 ] = O E[ϕ t ] max log m t + 1 , log n t + 1 + 1 n .Proof. We consider the two cases t + 1 ≥ log n and t + 1 < log n separately. Consider the case when t + 1 ≥ log n. We first show thatEQUATION)The first equality follows by substituting E[1 t x ] = e -npx (np x ) t /t!. For the inequality, note that if |np x -t -1| ≤ 25(t + 1) log n, then the inequality follows. If not, then by the Chernoff boundE[1 t x ] = Pr(t x = t) ≤ n -3 and hence E[1 t x ] -E[1 t+1 x ] ≤ E[1 t x ] + E[1 t+1 x ] ≤ 2/n 3 . By definition, E[Φ t ] -E[Φ t+1 ] = x E[1 t x ] -E[1 t+1 x ]. Substituting, |E[Φ t ] -E[Φ t+1 ]| ≤ x E[1 t x ] -E[1 t+1 x ] (a) = x e -npx (np x ) t t! 1 - np x t + 1 = x:npx≤1 e -npx (np x ) t t! 1 - np x t + 1 + x:npx>1 e -npx (np x ) t t! 1 - np x t + 1 (b) ≤ x:npx≤1 np x t! + x:npx>1 5e -npx (np x ) t t! log n t + 1 + 2 n 3 ≤ 1 n 2 + O E[Φ t ] log n t + 1 + 2n n 3 ≤ O E[Φ t ] log n t + 1 + 1 n .where (a) follows from the fact that E[1 t x ] = e -npx (np x ) t /t!. (b) follows from the fact that np x ≤ 1 in the first summation and Equation ( 4). The proof for the case t + 1 < log n is similar and hence omitted.Lemma E.10 (Upper bound on |M t -F t |). For any t ≥ 1, with high probability at least 1-npoly(m).EQUATIONOur proof mainly relies on the concentration properties of Dirichlet distribution and some properties of the local Good-Turing estimate. While the first term is fairly streaightforward to bound, the second term requires us to handle the fact that 1 t v is a random variable that depends on the global statistics for the token v, while Qgt u,v -Q u,v depends only on the local statistics. This requires us to analyze this differently from Acharya et al. [2] .Using Lemma E.8, followed by Lemma E.7 we can bound term A in Eq. 18 in the following way:EQUATIONThis completes our derivation for the upper bound on the first term (A) in Eq. 18. For the second term (B), we use the definition of the genie-aided local Good-Turing estimator and bound its variance in the folliowing way:Var 1 n v∈V u 1 t v (Q u,v -Qgt u,v ) ≤ v∈V E[1 t v ] 1 n u Q u,v -Qgt u,v 2 = v∈V E[1 t v ] 1 n u Q u,v - cnt u,v + 1 m u E[ϕ u,cntu,v+1 ] E[ϕ u,cntu,v ] 2 ≤ 2 v∈V E[1 t v ] 1 n u Q u,v - cnt u,v + 1 m u 2 (22) + 2 v∈V E[1 t v ] 1 n u (cnt u,v + 1)(E[ϕ u,cntu,v+1 -ϕ u,cntu,v ]) m u E[ϕ u,cntu,v ]The first inequality above uses Lemma E.8, and the second equality uses the definition of the genieaided estimator. In Eq. 22 we have two terms, and we bound the first term using Jensen inequality followed by the DKW inequality [70] for empirical estimators of the cumulative distribution. With high probability 1 -δ,u v∈V E[1 t v ] 1 n u Q u,v - cnt u,v + 1 m u 2 ≤ 1 n u v∈V Q u,v - cnt u,v + 1 m u 2 = 1 n u v∈V E[1 t v ]O cnt u,v log 2 ( d /δ) m 2 u ≤ v∈V E[1 t v ]O cnt v log 2 ( d /δ) m 2 u ≤ Õ tE[ϕ t ] m 2The final inequality above uses the definition of t, and the concentration of m u , i.e., with high probabibility:|m u -E[m u ]| = |m u -m| = Õ( √ m).Next we will bound the second term in Eq. 22, for which we use some properties of the expected sensitivity of the local frequency counts (Lemma E.9).v∈VEQUATIONApplying Jensen's inequality it is easu to see that we can bound the above term with:v∈VEQUATIONCombining the result in Equation 19 with Equation 23, and then applying Lemma 8 from Orlitsky and Suresh [63] completes the proof of the Lemma.Let us first consider the following simplified setting which deviates from our setup in Section 4 in two ways: (1) Q u = P , i.e., each user's true distribution matches exactly the cluster center (this matches the α → ∞ case in our setup); and (2) each user independently samples m u ∼ Poi(m), which is the number of samples in their local dataset, that is used to estimate the empirical distribution Q u (note that m u = m for our results in Section 5). We shall now derive results in this setting, and then consider the more general case.Without loss of generality, let us fix the center to be c, with some collection of users Z c . Then, we know that cnt v is distributed as a Poisson random variable with mean m u∈Zc Q u [v]. Thus, with probability at least 1 -δ:| cntu /n -u∈Zc mQu[v] /n| ≤ log( 2 /δ).We mainly rely on Lemma 13 for the Genie aided estimator from [2] . The key difference being that E[ϕ k ], which is words with total count k would differ from E[ϕ cntu,v ]. For, this we first bound ⌊ k /n⌋ -cnt u,v with high probability. We do this, by:cnt uv -E(cnt u,v ) = δ Õ( E[cnt u,v ]) = δ Õ( mQ u [v]) + m Õ( P [v] α + 1 )since the counts follow Poisson distribution under indpendent sampling. Thus, when α = Ω(m), the first term dominates, and we see that each this gives us:k n ∈ cnt u,v -cnt u,v , cnt u,v + cnt u,vwith high probability. Then from Claim 20 in [2] , qwe know that|E[ϕ k ] -E[ϕ cntu,v ]| = O(k)The rest of our analysis relies with the genie-aided estimator mainly relies on Theorem 2 in Acharya et al. [2] . Using the result in Lemma E.10, we conclude that with probability 1 -npoly(m) over the draw of the meta-dataset S, the KL distance:EQUATIONNext, we use Lemma 15 from Acharya et al. [2] to prove the same claim for the estimator that is not genie-aided. And finally, we reuse the conversion of the high probability result to one in expectation using arguments similar to our proof of Lemma 4.1. This completes our analysis of the upper bound on competetive regret for the average of local Good-Turings (our estimator).Finally, for the lower bound in Theorem 4.3, we invoke known lower bounds on the concentration of the Dirichlet distribution Ω( 1 /α+1), and minimax statistical lower bounds on the risk of the empirical estimator Ω( d m ), to conclude that the competetive risk of FedAvg is at least Ω ( d /m + 1 /α+1). This completes our proof of Theorem 4.3. [46] ). The Γ function is sandwiched as follows. There is a positive constant C, such that for all x > 0:Cx x-1 /2 e -x ≤ Γ(x) ≤ Cx x-1 /2 e -x e 1 /12x . This implies log Γ(x) = log C + (x -1 /2) log x -x + O( 1 /x). For large enough x, we substitute Γ(x) ≈ log C + (x -1 /2) log x -x.Theorem E.12 (Clustering approximates MLE objective; restated). The joint likelihood for P c and z u can be upperbounded with the clustering objective in Eq. 4, and this approximation is tight upto an additive term of O( 1 /α).Let V be a fixed set of words, d =: |V| is the size of the set, and ∆(V) is the set of all possible discrete probability distributions over V. Each user u in the federated setup has an unknown distribution Q u ∈ ∆(V). Further, each Q u belongs to one of k clusters, and the cluster membership is denoted by z u , where z u ∈ R k is a 1-hot k-dimensional vector. Each cluster c ∈ [k] is associated with a center P c ∈ ∆(V), and the matrix with column vectors as cluster centers is denoted as P ∈ R d×K , i.e., P =: [P 1 , P 2 , . . . , P k ].Independently for each user u, given the cluster membership z u , and some α > 0, the user's distribution Q u is sampled from a Dirichlet around the cluster center Pz u , i.e., Q u ∼ Dir(αPz u ). Further, each user samples an i.i.d. dataset S u of size n u . When all users sample dataset of fixed size n u = n, we refer to this setting as size-homogeneous, and the more general case as sizeheterogeneous. Given the set of datasets {S u } u and the distribution assumptions above, we can write down the maximum-likelihood estimate (MLE) for the cluster centers P and cluster memberships {z u } u :EQUATION= arg maxP,{zu}u u log DirMul(n u Q u | P, {z u } u ) = arg max P,{zu}u u log Γ(α) + log Γ(n u + 1) -log Γ(n u + α) + x log Γ( Q u [x] + αPz u [x]) -log Γ(αPz u [x]) -log Γ( Q u [x] + 1) , F.1 Proof of Theorem 3.2Theorem F.2 3 convergence; restated). For the model in Section 4, if ∀c, v, n c = Θ(K 2 ), P c [v] = Ω(1), centers are sufficiently separated, i.e., ∆ =: min i̸ =j KL(P i ∥P j ) = Ω(k 2 + k 3 d /n), and λ =: max i̸ =j KL(Pi∥Pj ) /∆. Given an initialization with assignment error O( 1 /λ) for any cluster, after t = log(|U|) iterations, w.h.p. the assignment error is O(exp(-∆(α + 1))).Lemma F.3 (Pinsker [68] ).For P, Q ∈ ∆ |V|-1 , the KL(P || Q) ≥ 1 2 ∥P -Q∥ 2 1 = 2TV 2 (P, Q). Lemma F.4 (KL upper bound). For P, Q ∈ ∆ |V|-1 , the KL divergence KL(P || Q) ≤ χ 2 (P Q)Proof Overview. The main technique we follow is to use Pinsker's inequality to lower bound KL divergence with TV 2 (P, Q) and upper bound KL with χ 2 . This allows us to then treat the clustering in the l 2 2 metric, and pay the d-dimensional penalty term 1 d . Based on this general principle, the following section presents the EM convergence analysis of Algorithm 3, extending results from [7] to our setting. While the steps below closely mirrors their analysis upto the final terms where we need to upper and lower bound KL, we present the full proof for completeness.Let us begin by introducing some notaion. For any S ⊆ [n], define W S = i∈S w i . Recall thatEQUATIONEQUATIONEQUATIONh .Then we have nEQUATIONhg . The mis-clustering rate at iteration s can be written asA s = 1 n n i=1 I{ẑ (s) i ̸ = z i } = 1 n g̸ =h∈[k] 2 n (s)gh .We define a cluster-wise mis-clustering rate at iteration s asG s = max h∈[k] g̸ =h∈[k] n (s) gh n (s) h , g̸ =h∈[k] n (s) hg n * h .The first term in the maximum operator of definition of G s can be understood as the false positive rate of cluster h and the second term is the true negative rate of cluster h. It is easy to see the relationship that A s ≤ G s .Let ∆ = min g̸ =h∈[k] ∥θ g -θ h ∥ be the signal strength. For h ∈ [k], let θ(s) h be the estimated center of cluster h at iteration s. Define our error rate of estimating centers at iteration s asΛ s = max h∈[k] 1 ∆ ∥ θ(s) h -θ h ∥. Lemma F.5. Based on our definition of M above, G 0 = O 1 2 -6 √ r k 1 λ . Proof. First we show λ = O √ d. Since P c [v] = 1 d, we can invoke the KL upper bound above to show that ∆ λ = O( 1 d ), and furthermore, the KL for each pair is Θ(M ). Thus, this tells us:EQUATIONIn other words the signal to noise ratio is ∆/σ ≥ k 2 + k 3 d n .Based on the above result, we also get:∆ σ 1/K 1 + kd/n ≥ C √ k.Recall σ for us is given by Thus, we satisfy all conditions needed for Theorem 3.2 in their paper, the proof of which we replicate below for reader's convenience.They assume the following initialization condition:EQUATIONLemma F. 6 .EQUATION) with probability greater than 1 -exp(-0.3n). Lemma F.7.EQUATION)with probability greater than 1 -exp(-0.5n). Lemma F.8. For any fixed i ∈[n], S ⊆ [n], t > 0 and δ > 0, Pr    w i , 1 |S| j∈S w j ≥ 3σ 2 (t |S| + d + log(1/δ)) |S|    ≤ exp -min t 2 4d , t 4 + δ.Lemma F.9.EQUATIONwith probability greater than 1 -n -3 . Lemma F.10. For any fixed θ 1 , • • • , θ k ∈ R d and a > 0, we haveEQUATIONfor all g ̸ = h ∈ [k] 2 with probability greater than 1 -n -3 .The following two lemmas give the iterative relationship between the error of estimating centers and the error of estimating labels. Let E be the intersection of high probability events in Lemma F.6, Lemma F.7 Lemma F.9, Lemma F.10 and the initialization condition (29) . Then we have P{E c } ≤ 3n -3 + ν. In the rest part of the proof, if not otherwise stated, we all condition on the event E and the following analysis are deterministic. Lemma F.11. On event E, if G s ≤ 1 2 , then we haveEQUATIONLemma F.12. On event E, if Λ s ≤ 1-ϵ 2 and r ≥ 36ϵ -2 , thenEQUATION)Proof of Lemma F.11. For any B ⊆ [n], define ȲB = 1|B| i∈B y i . The error of estimated centers at step s can be written asθ(s) h -θ h = 1 n h i∈S hh (y i -θ h ) + 1 n h a̸ =h i∈S ah (y i -θ h ) = 1 n h i∈S hh w i + a̸ =h n ah n h ȲS ah -θ hAccording to our label update step, we have ∥y i -θ(s-1)h ∥ ≤ ∥y i - θ(s-1)a ∥ for any i ∈ S ah . This means for any i ∈ S ah , y i is closer to θ(s-1)h than θ(s-1)a , so is the average of {y i , i ∈ S ah }. Thus, we have∥ ȲS ah - θ(s-1) h ∥ ≤ ∥ ȲS ah -θ(s-1) a ∥.Consequently, triangle inequality gives usȲS ah -θ h ≤ ȲS ah -θ a + ∥ θ(s-1) a -θ a ∥ + ∥ θ(s-1) h -θ h ∥,which, combined with Lemma F.6 and the definition of Λ s-1 , yieldsȲS ah -θ h ≤ σ 3(n + d)/n ah + 2Λ s-1 ∆.Taking a weighted sum over a ̸ = h ∈ [k], we geta̸ =h n ah n h ȲS ah -θ h ≤ σ 3(n + d) n h a̸ =h √ n ah + 2Λ s-1 ∆ a̸ =h n ah n h ≤ σ 3(n + d) √ n h (k -1)G s + 2G s Λ s-1 ∆,where the Last inequality is due to Cauchy-Schwartz and the fact that a̸ =h n ah ≤ G s n h . Note thatW S hh = W T * h -a̸ =h W S ha .Triangle inequality and Lemma F.10 implyEQUATIONCombining the pieces, we obtainEQUATIONTherefore, we get the first term in (34) . To prove the second term, we decompose θ(s) h differently.EQUATIONEQUATIONThen, the error of θ(s) h can be upper bounded asθ(s) h -θ h = k a=1 n ah n h (θ a -θ h ) + 1 n h W T h ≤ a̸ =h n ah n h (θ a -θ h ) + 1 n h W T h .By triangle inequality,EQUATIONThis, together with Lemma F.6 and (36), impliesEQUATIONfor all h ∈ [k]. The proof is complete.Proof of Lemma F.12. For anyEQUATIONTriangle inequality implies∥θ g - θ(s) h ∥ 2 ≥ ∥θ g -θ h ∥ -∥θ h - θ(s) h ∥ 2 ≥ (1 -Λ s ) 2 ∥θ g -θ h ∥ 2 .Using the fact that (1 -x) 2 -y 2 ≥ (1 -x -y) 2 when y(1 -x -y) ≥ 0, we obtainEQUATION)Denote by∆ h = θ(s) h -θ h for h ∈ [k].Then,EQUATIONi = h ≤ I ϵ 2 ∥θ g -θ h ∥ 2 ≤ 2 ⟨w i , θ h -θ g + ∆ h -∆ g ⟩ ≤ I ϵ 2 2 ∥θ g -θ h ∥ 2 ≤ 2 ⟨w i , θ h -θ g ⟩ + I ϵ 2 2 ∆ 2 ≤ 2 ⟨w i , ∆ h -∆ g ⟩ .Taking a sum over i ∈ T * g and using Markov's inequality on the second term, we obtainEQUATIONNote that I ϵ 2 4 ∥θ g -θ h ∥ 2 ≤ ⟨w i , θ h -θ g ⟩ are independent Bernoulli random variables. By Lemma F.10, the first term in RHS of ( 43) can be upper bounded byEQUATIONBy Lemma F.7, the second term in RHS of ( 43) can be upper bounded byEQUATION)Combining ( 43), ( 44) and ( 45) and using the fact that ∥∆ g -∆ h ∥ 2 ≤ 4Λ 2 s ∆ 2 , we getn (s+1) gh ≤ n * g exp - ϵ 4 ∆ 2 32σ 2 + 5n * g log n + 384(n * g + d)σ 2 ϵ 4 ∆ 2 Λ 2 s .Consequently,EQUATIONSince Λ s ≤ 1/2 and r ≥ 20ϵ -2 , the RHS of ( 46) is smaller that 1/2 when αn ≥ 32k 2 log n. Thus,n (s+1) h ≥ n (s+1) hh ≥ 1 2 n * h ≥ 1 2 αnfor all h ∈ [k] and we haveEQUATIONwhich, together with (46) , impliesG s+1 ≤ exp - ϵ 4 ∆ 2 32σ 2 + log(2/α) + 5k log n α 2 n + 768 ϵ 4 r 2 Λ 2 sUnder the assumptions that ϵ 4 α∆ 2 /σ 2 ≥ r 2 ϵ 4 ≥ 36, we have the desired result (35) .Proof. From Lemma F.11, a necessary condition for Proof. From the proof of Lemma F.11, the error of estimating θ h at iteration s can be written asΛ 0 ≤ 1 2 -4 √ r is G 0 ≤ ( 1 2 -6 √ r ) 1 λ . Setting ϵ = 7 √ r in Lemma F.EQUATIONIn addition, by Lemma F.11 and Lemma F.12, there is a constant C 1 such thatΛ s ≤ 3 r + G s + 2G s Λ s-1 ≤ C 1 r + C 1 r Λ s-1 + 0.7Λ s-1 + C 1 k log n α 2 n 1/4for all s ≥ 1. Therefore, when r is large enough, we haveΛ ≤ C 2 r -1 + C 2 k log n α 2 n 1/4for all s ≥ log n. Then by ( 42), we haveEQUATIONi = h ≤ I β 1 ∥θ g -θ h ∥ 2 ≤ 2 ⟨w i , θ h -θ g + ∆ h -∆ g ⟩ where (1 -2Λ s ) 2 ≥ β 1 := 1 -4C 2 r -1 -4C 2 k log n α 2 n 1/4.In order to prove that A s attains convergence rates, we first upper bound the expectation of A s and then derive the high probability bound using Markov's inequality. Similar to the two-mixture case, we need to upper bound the inner product ⟨w i , ∆ h -∆ g ⟩ more carefully. Note that {T * h , h ∈ [k]} are deterministic sets, we could use concentration equalities to upper bound W T * h and u h parts separately.EQUATIONi = h into three terms.EQUATIONi = h ≤ I β∥θ g -θ h ∥ 2 ≤ 2 ⟨w i , θ h -θ g ⟩ +I β 2 ∆ 2 ≤ 2 ⟨w i , u h -u g ⟩ +I β 4 ∆ 2 ≤ 2 ⟨w i , v h -v g ⟩ ,where β 2 and β 4 will be specified later andEQUATIONEQUATIONEQUATIONLet us first consider the expectation of J 1 . Using Chernoff's bound, we havePr β∥θ g -θ h ∥ 2 ≤ 2 ⟨w i , θ h -θ g ⟩ ≤ exp - β 2 ∥θ h -θ g ∥ 2 8σ 2 ≤ exp - β 2 ∆ 2 8σ 2 .Thus,EJ 1 ≤ k exp - β 2 ∆ 2 8σ 2 = exp - γ∆ 2 8σ 2 , with γ = β 2 -8σ 2 log k ∆ 2 ≥ β 2 -8/r 2 .We use Markov Inequality to upper bound J 2 . Markov's inequality and Lemma F.7 give us1 n n i=1 I β 2 ∆ 2 ≤ 2 ⟨w i , u h -u zi ⟩ ≤ 4 nβ 2 2 ∆ 4 g∈[k] i∈T * g (w ′ i (u h -u g )) 2 ≤ 24σ 2 nβ 2 2 ∆ 4 g∈[k] (n * g + d)∥u h -u g ∥ 2 .(48) impliesJ 2 ≤ 96σ 2 G s nβ 2 2 ∆ 2 h∈[k] g∈[k] (n * g + d) ≤ 96σ 2 k(n + kd) αnβ 2 2 ∆ 2 A s = 12 √ k r A s .Here the second inequality is due to the fact that G s ≤ A s /α. And we choose β 2 = 8k/r in the last equality.Finally, we upper bound the expectation of J 3 . Given z i = g, we haveEQUATIONEQUATION)), where we have used the assumption that n * g ≥ αn and αn ≥ 36r 2 . Thus,EJ 3 ≤ 2k exp - ∆ 2 σ 2 ,Combining the pieces, we haveEA s+1 ≤ E [J 1 ] + E [J 2 I{E}] + E [J 3 ] + P{E c } ≤ exp - γ∆ 2 8σ 2 + 12 √ k r EA s + 2k exp - ∆ 2 σ 2 , with γ = (β 1 -8k/r -64/r) 2 -8/r 2 = 1 -o(1).Here only prove the case that r → ∞. For the finite case, all the o(1) in the following proof can be substituted by a small constant.EA s ≤ 1 2 s-⌈log r⌉ + 2 exp -(1 -η) ∆ 2 8σ 2 + 2 n 3 ≤ 2 exp -(1 -η) ∆ 2 8σ 2 + 3 n 3Lemma F.16 (Composition of CDP [15] ). If M is an adaptive composition of differentially private algorithms M 1 , . . . , M T , then 1. if M 1 , . . . , M T are (ϵ 1 , δ 1 ), . . . , (ϵ T , δ T )-DP then M is ( t ϵ t , t δ t )-DP, and 2. if M 1 , . . . , M T are ρ 1 , . . . , ρ T -zCDP then M is ( t ρ t )-zCDP.We can achieve differential privacy via noise addition proportional to sensitivity [24] . Definition F.17 (Sensitivity). Let f : X n → R d be a function, its ℓ 2 -sensitivity is defined to be ∆ f,2 = max X∼X ′ ∈X n ∥f (X) -f (X ′ )∥ 2 , Here, X ∼ X ′ denotes that X and X ′ are neighboring datasets (i.e., those that differ in exactly one entry).For functions with bounded ℓ 1 -sensitivity, we can achieve ϵ-DP by adding noise from a Laplace distribution proportional to ℓ 1 -sensitivity. For functions taking values in R d for large d it is more useful to add noise from a Gaussian distribution proportional to the ℓ 2 -sensitivity, to get (ϵ, δ)-DP and ρ-zCDP. Lemma F.18 (Gaussian Mechanism). Let f : X n → R d be a function with ℓ 2 -sensitivity ∆ f,2 . Then the Gaussian mechanismM f (X) = f (X) + N 0, ∆ f,2 √ 2ρ 2 • I d×d satisfies ρ-zCDP. F.2.2 Proof of Theorem 3.3From the Lemmas above for Gaussian and Laplace mechanism, it is easy to see that the adaptive clipping step in Algorithm 4 satisfies 2ρ/3-zCDP, where as the Laplace mechanism is ρ/3-zCDP.Then, we can use composition lemma to complete show ρ-zCDP of Algorithm 4.For Algorithm 5, each step uses satisfies ρ/K-zCDP from the guarantees for exponential mechanism, since we clip the KL difference with c. Here too, composition gives us the result.For the final end-to-end algorithm, we use the private aggregation scheme, we describe, and changing the user data would only affect two cluster estimates, but since we use add-one, leave-one definition of privacy, the entire step is ρ/3T -zCDP. Finally, composing over the T rounds of clustering and with initial private center and private init we get the final result. We then use Lemma 3.6 from [15] to convert zCDP guarantees to (ε, δ)-DP.• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Given two probability distributions P, Q over Ω, Dα(P ∥Q) = 1 α-1 logx P (x) α Q(x) 1-α . 


Decision: The reviewers were reasonably enthusiastic about the paper. There were a couple of reviewers on the fence. We would recommend the authors to take care of the concerns raised in the reviews, in particular, w.r.t. the presentation (especially Section 4), and w.r.t. the baselines chosen for the experiments. We recommend the authors to meticulously incorporate the rebuttal discussion in a future version of the paper.

Reviews: review: paper_summary: paper_summary: The paper proposes a federated learning approach for frequency histogram estimation in a distributed setting. The proposed approach relies on first clustering users that have similar subpopulation distributions before performing the estimation in a privacy-preserving manner. Finally, the performance of the approach is also validated on three real-world datasets.
strength_weakness: strength_weakness: The paper focuses on next word prediction, which is a quite focused application. Instead, it would be more interesting to try to broaden the applicability of the approach, in particular by testing it on other type of datasets. The possibility to extend the framework developed beyond the KL-divergence should also be discussed.

The neighbouring notion considered for differential privacy seems quite strong in terms of the privacy guarantees provided but the authors failed to motivate why this is right notion to consider in the federated learning in comparison to record level differential privacy. The necessity of the Dirichlet assumption as well as the concentration of the users’ histograms should also be justified. 

A discussion on the method that should be used to estimate the number of clusters K is currently missing in the paper. This is a crucial point as this parameter is likely to have an important impact on the success of the approach. 

In terms of the structure of the paper, it ends quite abruptly lacks a conclusion summarizing the key findings of the paper and also discussing possible future works.

A few typos :
-« Baysian prior » -> « Bayesian prior »
-« privacy preserving analysis » -> « privacy-preserving analysis »
-« estimatation in KL divergence » -> « estimation in KL divergence »
-« For competetive estimator » -> « For competitive estimator »
-« LLoyd’s K-means » -> « Llloyd’s K-means »
-« differentially private » -> « differentially-private »
-« Privacy preserving noisy sum » -«  « Privacy-preserving noisy sum »
questions: questions: What is the relationship of the privacy parameter rho compared to standard parameters epsilon and delta in differential privacy?

Please see the main points raised in the weaknesses section.
limitations: limitations: The authors have a specific section discussing the limitations of their work.
paper_summary: paper_summary: The paper presents both non-private and private per-user histogram estimation approaches that can be applied to problems such as next word predictions. The approach relies on iteratively clustering among different users to find similar user subgroups and then fine-tuning within each user to get user-specific estimations. The privacy guarantee is achieved by privately finding clusters using all users data and is quantified by the notion of joint-differential privacy. The paper theoretically analyzes the sub-optimality of their algorithm and empirically evaluate the performance by comparing to several baselines.
strength_weakness: strength_weakness: - The largest concern to me is the presentation clarity of the paper. 
    - In section 4, the authors seem to present many types of assumptions regarding possible user data distributions, but the sub-titles are somewhat confusing as it tries to explain model, learning, different clustering cases, histogram estimators and also theorems. It would be much clearer if more relevant parts are put together. 
    - It would be easier to understand why these assumptions are stated if in the experiment section the authors could include whether certain dataset fall into which category of the assumed data structure/distributions, and in the algorithm section how the algorithms are adapt to these cases of data structure/distributions.
    - It was also unclear to me regarding the final pipeline (the 'putting it all together' section). For example, to my understanding Algorithm 3 seems to be the 'main' algorithm where the cluster centers are determined, then to get user-specific estimations (the $\hat{Q}_u$s) do users fine-tune in their own ways or follow Algorithm 1 or something else?
- The novelty seems not very strong as the clustering part seems to be the major contribution but the papers apply standard K-means like clustering algorithms except changes to KL as the target is to minimize some KL as in Equation 5.
questions: questions: - In section 4 one of the assumption is 'the user histograms are well concentrated along each token’s marginal', what does 'each token’s marginal' means and can it be described in mathematical language in terms of 'well concentrated'?
- Since the authors defined their own DP relaxation definition, it will be helpful if the authors could clarify:
    - Is there a conversion between scale of $\epsilon$ to the pure or approximate DP case? The $\epsilon=15$ value reported in the evaluation section seems to be a quite large number for meaningful privacy protection, how should it be interpreted under the $\rho$-zCJDP notion?
    - I did not completely understand the $\rho$ parameter in Definition 3.2. Is it another relaxation defined if it's defined as the divergence term $D_{\alpha} < \rho \alpha$?
limitations: limitations: see above
paper_summary: paper_summary: This work introduces a private and personalized histogram estimation approach based on Good-Turing estimates, when the clients are clustered. Theoretically, they provide performance guarantees for different settings when cluster center is known and cluster assignments are known. In practice, they note that neither cluster centers nor assignments are known, hence inspired by theoretical results from the simpler settings they provide practical algorithms. Later they also introduce differential private version of the proposed methods whose performance can increased further with a particular initialization method. For experiments they consider three tasks where the proposed method outperforms the competition.
strength_weakness: strength_weakness: - Lack of comparison to local only training and other personalized estimation methods.

- The main results are rather in simpler settings, the analysis mostly builds on top of previous work; hence, theoretical contribution might be limited (compared to methodological).

- From the current experiments, performance dependence on variables such as number of clusters, clients, data samples per clients, heterogeneity etc. is not clear. I would suggest doing more ablations studies possibly using e.g. synthetic datasets to understand these.
questions: questions: See weaknesses.
limitations: limitations: Yes
paper_summary: paper_summary: This paper studies the problem of personalized frequency estimation in federated learning. The authors propose both private and non-private algorithms based on clustering and Good-Turing estimators and demonstrate promising empirical results on real-world datasets.
strength_weakness: strength_weakness: 1. The cluster assumption requires choosing the cluster number $K$.
2. KL divergence may not work well for sparse distributions, and this is possibly why the authors need to assume $P_c[v]$ is constant for all $v$ in Theorem 5.1. The authors motivate this work by estimating the frequency of words, but they are usually very concentrated over a very small fraction of words.
3. In Theorem 5.1, the minimum KL divergence is at least order $d$ which seems restrictive at first glance, especially considering the number of words could be of order $10^4-10^5$. Also, the $k^2$ in line 266 seems to be $K^2$.
4. The $\log(1/\delta)$ in the privacy guarantee in Theorem 5.2 can be pretty large when $\delta$ is very small. Since normally we require $\delta\sim o(1/n)$ where $n$ is the number of users, this introduces a $\log n$ factor to the privacy guarantee.
5. In the experiments, the authors only show results for $\varepsilon=15$.
6. Relevant papers to cite and discuss:

[1] Ozkara, Kaan, et al. "A statistical framework for personalized federated learning and estimation: Theory, algorithms, and privacy." International Conference on Learning Representations (ICLR), 2023. OpenReview. net, 2023.

[2] Liu Y, Suresh A T, Yu F X X, et al. Learning discrete distributions: user vs item-level privacy[J]. Advances in Neural Information Processing Systems, 2020, 33: 20965-20976.

[3] Liu Y, Suresh A T, Zhu W, et al. Algorithms for bounding contribution for histogram estimation under user-level privacy[C]//International Conference on Machine Learning. PMLR, 2023: 21969-21996.

[4] Huang, Ziyue, Yuting Liang, and Ke Yi. "Instance-optimal mean estimation under differential privacy." Advances in Neural Information Processing Systems 34 (2021): 25993-26004.

=====Update=====

Increased my score to 5 after the rebuttal
questions: questions: 1. How does the choice of $K$ affect the performance? Do the authors have a suggestion on how to choose $K$ in practice?
2. In line 256 ``Without any assumptions on the relationship between the user data distributions the best one can do is to locally estimate each user’s distribution'' I believe this should be true, but could the authors provide a brief justification for this claim?
3. Why did the authors not consider using privacy accounting methods such as RDP accountant which may give much tighter privacy bounds than the composition theorem of approximate DP?
limitations: limitations: The authors addressed the limitations and negative impact of their work.

Comments: title: Request for discussion
Request for discussion
Reply to Author Rebuttal
New experiments with private frequency estimation baselines
Choice of Baselines and New experiments with private frequency estimation baselines
Reply to Author Rebuttal
Request for discussion
comment: Thank you for increasing the score. We are committed to include the above discussion and clarifications in the final version. Thank you again for your thoughtful feedback and comments.
Thanks for your rebuttal. I appreciate the clarifications that have been provided and as well as the addition experiments and I will increase slightly my score.
Dear reviewer,

Apologies for bothering you! Since we only have one day left in the discussion period, we would be grateful and would sincerely appreciate if you could respond to our rebuttal, leaving us enough time to address any remaining questions.

Thanks, 
Authors
Dear reviewer,

Thank you for your acknowledgement and positive assessment of our work. We will definitely add the new ablations on data samples per client, different privacy budgets and related discussion to the main paper.

Since there is still one more day, we are also wondering if there would be some other discussion or evidence that we can provide in this period to help improve your evaluation of our paper further. Please let us know. Thanks a lot!

Thanks, 

Authors
Dear reviewer,

Apologies for bothering you! Since we only have one day left in the discussion period, we would be grateful and would sincerely appreciate if you could respond to our rebuttal, leaving us enough time to address any remaining questions.

Thanks, 
Authors
Thank you for raising the score. In the final version, using the extra page, we are certainly committed to provide more clarity on our theoretical motivation for algorithm design in Section 4. We will also add results obtained with the new baselines in the discussion above, further clarify the implementation details of existing baselines in the paper and add missing citations to the related work section. Thank you again for your detailed feedback and comments.
I thank the authors for their quick response and for adding new experiments. I have raised my score. However, given the discussion above I think the paper needs major revision to clarify the theoretical motivations for designing algorithms in Section 4, implementation details to compare baselines that are not closely targeting the same problem, and potentially add a clearer related work section and compare to work that are more closely related to the field of private histogram release. I think that is the best score I am confident to raise given the current version.
**[NEW EXPERIMENTS] Fair comparison with suggested private frequency estimation approaches [1, 2]:**

We compare our approach with the suggested baselines [1, 2], after we do the following:

1) tune the sketch length hyper-parameter $L$ (over 2, 4, 8, 16, 32) and set $L=8$ and tune sketch width $W$ (over 10,  20, 50, 100, 200, 500, 1000) and set $W=500$.

2) adapt the estimated global frequency (distribution) for each user using our finetuning algorithm (Eq. 2), that we also use to adapt cluster centers/FedAvg estimate for each user separately. 

We use the same setup (number of clients and local data points per client) as our main experiments in Figure 2, except for the privacy budget $\varepsilon$, which we additionally vary from 4 to 16.

___

| StackOverflow Test NLL Loss          |$\varepsilon=4$ | $\varepsilon=8$ | $\varepsilon=16$ |
|--------------------------------------|-------|-------|--------|
| CountSketch + Secure Aggregation [1] | 9.935 | 9.399 | 9.017  |
| 2-stg CountSketch + Secure Aggregation [2] | 9.783 | 9.132 | 8.755  |
| Ours                                 | **9.008** | **8.544** | **8.107**  |

___

| Reddit Test NLL Loss                 | $\varepsilon=4$ | $\varepsilon=8$ | $\varepsilon=16$ |
|--------------------------------------|-------|-------|--------|
| CountSketch + Secure Aggregation [1] | 7.618 | 7.012 | 6.806  |
| 2-stg CountSketch + Secure Aggregation [2] | 7.427 | 6.914 | 6.785  |
| Ours                                 | **7.013** | **6.589** | **6.312**  |
Thank you for your comments and suggestions for private global frequency estimation baselines. We add new experiments comparing our method with suggested baselines [1, 2], and find that for all ranges of the privacy budget our approach outperforms [1, 2], highlighting the need for learning clusters and cluster-level frequency counts (as opposed to global counts) to reduce the estimation error in KL, per-client. We also clarify that the baselines (FedAvg/MAML/IFCA) we used in our paper were originally proposed for more general settings, and we adapt each of these baselines in a fair manner for a meaningful comparison. We will add these experiments and discussion to the final version. **Please let us know if this addresses all outstanding concerns, and if so, we would be grateful if you might raise your score.**

___

**Adapting FedAvg/MAML/IFCA, which are proposed for convex optimization (superset of our setting which is to minimize KL distance):** We will make this more clear, but please note that all three baselines have been proposed and analyzed as general algorithms to solve convex objectives (our metric of KL divergence is also convex) in the federated setting. We apologize if this is not clear, but for a fair comparison, we adapt the baselines specifically to our setting in the following way:

- FedAvg and MAML both optimize the KL distance (as opposed to squared loss) to learn a single distribution that is close to the empirical distributions of each user in KL distance (Eq. 3 in our paper). Note that we do not use gradient-based approaches for FedAvg, and instead use the closed-form optimal (in Eq. 3, proof in Appendix C). 

- IFCA (Ghosh et al. [35]) is one of the most popular clustering baselines in FL, that was proposed for convex losses, and analyzed for k-means (squared loss), which we adapt to our setting by: 1) using the KL distance to the current cluster centers for determining the cluster affiliation of each user in a clustering round; and 2) estimating the new cluster centers by optimizing the KL objective, which is same as what we did for FedAvg/MAML but now only for clients in the cluster (Eq. 5 in our paper).  

- Estimates from all three baselines are personalized (adapted) for each user using the same method as ours (Eq. 2), which is not gradient-based. Additionally, we also try other popular approaches like RTFA, gradient-based (Fig. 3d) and find our non-gradient based approach to work best for all three baselines. The local adaptation step for MAML also uses the same non-gradient based method.    

___

**Comparison with private frequency estimation approaches [1, 2]**: Thank you for suggesting these baselines. Wu et al. [1] is based on the same algorithm (Count Sketch+Secure Aggregation) as Chen et al. [2], and mainly proposes a two-stage approach that adapts to the hardness of the problem instance (heterogeneity of the frequency counts). Both papers [1, 2] differ from our setting in the following ways:

1) They estimate a single global frequency count (FedAvg estimate in our case) in $l_1$ or $l_\infty$ norms, but our problem setting requires estimating personalized distributions for each user (and that too in KL divergence which is not a proper distance metric). 

2) They mainly use CountSketch to encode the local frequencies to reduce the per-round communication cost. 

_[1, 2] are concerned with communication cost:_  In the best case, when there is no communication constraint, their approach resorts to only using secure aggregation (which uses Gaussian mechanism in [1] and Poisson Binomial in [2]), and is thus same as ours which uses Gaussian Mechanism (Alg. 4). We are not concerned with reducing the communication cost, but if we were, it is straightforward to use the CountSketch algorithm at each round of clustering. Further, note that while CountSketch reduces the number of bits communicated by each client, the error suffered on each item's frequency (due to privacy noise) would not change at all. **So CountSketch can only improve communication complexity over our method and not privacy-utility tradeoff**. 

_Similar to [2], our approach also adapts to the range/hardness:_ In algorithm 4, for private center estimation for the global center (in FedAvg) or cluster-level center (in Alg. 3), we use the common adaptive clipping technique which first uses some privacy budget to estimate the range of the true center to a small confidence interval. Then, the user-level estimates are clipped to this small confidence interval before running secure aggregation with the rest of the privacy budget. Thus, **the adaptive clipping we use is similar to the two-stage approach in [2], in the sense that both adapt to the hardness of the underlying instance.** We find this interesting and will add this comparison to the main paper.   
___
**PLEASE SEE OUR CONTINUATION BELOW FOR OUR EXPERIMENTAL RESULTS**
Thanks for the explanations especially on the DP guarantees, it makes sense to me. 
- Regarding clarity on Section 4: I think ``cluster center known -> cluster center unknown but cluster membership known -> both cluster center and membership unknown'' makes sense, but the current section 4 lacks such structure which makes it difficult to understand which result/findings is linked to which setting. Since the authors mention Section 4 is for theoretical motivation and Section 6 tests the most general setting only, I would expect the authors to highlight and elaborate on the key findings that lead to the design of the actual algorithm.
- Regarding novelty and experiments: I appreciate the new experiments added, but my main concern still remains which the paper proposes an algorithm that targets to minimize the KL term between users' histogram and the estimated cluster center histogram distributions, which is also the metric to evaluate in experiments, but the baselines are not designed for the same problem. For example, FedAvg, MAML and IFCA are federated/distributed deep learning frameworks that train better models in federated setting on regression/classification tasks, are gradient-based, and are not particularly privacy-preserving designed. I would expect authors to compare to more similar approaches in private frequency estimation such as [1, 2].

For the above reasons I would keep my score.

[1] Private Federated Frequency Estimation: Adapting to the Hardness of the Instance. Wu et al. NIPS 2023. https://arxiv.org/abs/2306.09396.
[2] The communication cost of security and privacy in federated frequency estimation. Chen et al. AISTATS 2023. https://arxiv.org/abs/2211.10041.
Dear reviewers,

Apologies for bothering you! Since we are getting close to the end of the discussion period, we would be grateful and would sincerely appreciate if you could respond to our rebuttal, leaving us enough time to address any remaining questions.

Thanks,
Authors
Thanks for the rebuttal, I acknowledge reading the rebuttal and will keep my score.