Private and Personalized Frequency Estimation in a Federated Setting
Authors: Amrith Setlur
Vitaly Feldman
Kunal Talwar
Venue: NeurIPS 2024 poster (2024)
Keywords: federated learning
personalization
privacy
Primary area: privacy

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.

Private and Personalized Frequency Estimation in a Federated Setting
Amrith Setlur' Carnegie Mellon University asetlurēcs. cmu,
Vitaly Feldman Apple vitaly edu@gmail com
Kunal Talwar Apple ktalwar@apple.com com
Abstract
Motivated by the problem of next word prediction user devices we introduce and study the problem of personalized frequency histogram estimation in a federated setting. In this problem, samples from a distribution which is specific to that user. The goal is to compute for all users personalized estimate of the user's 's distribution with error measured in KL divergence. We focus on addressing two central challenges: statistical heterogeneity and protection of user privacy. Our approach 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 non non- privariive clustering based algorithm for the problem, and give a provably joint differentially private version of with a private data- dependent initialization scheme. Next, we propose simple data model which is based on a mixture Dirichlet distributions, formally motivate our non- 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 the Reddit, StackOverflow,
some domain, each user observes number of
varying levels of statistical and size heterogeneity and Amazon Reviews datasets. Our results demonstrate significant improvements over standard and flustering-based baselines, and in particular, they show that it is possible to improve over direct personalization single global model.
Introduction
Federated learning algorithms jointly learn from decentralized user data, addressing statistical chal- lenges 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 heterogene- ity [27, 76]), which reduces the effectiveness of a single, shared model for all users. Secondly, the data may sensitive and mutually beneficial collaboration may compromise user privacy [31, 71]. important practical problem where user data distributions are diverse, and user-privv user-privacy is paramount 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, frequeecyy estimation is also one the most basic statistical tasks with numerous other applications [14, 47].
Motivated by personalized next-word prediction, we introduce and study the problem of per- sonalized 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 server that needs to provide the user with personalized estimate of their distribution. We focus on error measured in KI. divergence that is common in language model- ing [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
Part of the work was done during an internship at Apple.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(PFL) algorithms [67, 76], tackle statistical heterogeneity by first learning 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 agnostic to the structure of the user population which is highly heterogeneous, but often consists of a number of concentrated subpopulations of similar users.
Statistical heterogeneity: 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-Turine 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 significantly affected by the initialization [5] of cluster centers, we also propose a data-dependent initialization approach, also specific to estimation in KI.. We improve the practicality of our algorithm by separately handling data-poor users cases where users vary significantly in their local data sizes (size-heteogeneitt))
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 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, size- heterogeneous 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 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-Turee 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 Turing estimates), where the minimax eтror 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.
2 Problem definition
We use P to denote a distribution, P(u] for the probability of event u 1 for estimate of P (but P if the estimate is computed without user collaboration), P for sets. P for matrices (P.,/P.,, index into the 8th column/row respectively), and [P?} for a collection of elements P; indexed by i.
Sctup. Let V be a finite vocabulary (set) of tokens, with vocabulary size d =: [V], and A(V) is the set of all probability distributions over V. In a federated setup we have a collection of users 24 and each user E U has an unknown distribution Qu E A(V) over the tokens. The data available to each user u is dataset S, comprising of m. i.i.d. samples from Q.. We will use Q, refer to the empirical distribution estimated from S.. The combined collection of datasets is referred to as meta-dataset& S =; ) and similarly 2 =: {Qu, Ju 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 myy We refer to the more general case as size-heteroggneous.
Goal. Given the meta-dataset S, the goal is compute estimates {Quta such that:
1. ntility: Σn [[.L(Q.)) is minimized, where the expectation is taken over sampling of meta-dataset S and any randomness in computing {Qu)u.
2. privacy: Qu 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-datases 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 Dt, we denote W-,, as the output of the algorithm all other users except u.
Definition 2.2 (Joint Differential Privacy [48]). Given 2 0, 8 E (0, 1), and neighbouring relation
-us a randomized mechanism DR: DD + y from the set of meta-datasets 'D to an output space y is
(c, jojnn differentially private if for all u-neighboring S ~~ S', and all events E S )
Pr Pr(MN_.(S) E EJ S e Pr(3R_a(S") E EJ] +
(I)
where probabilities are taken over the randomness of 27. We say that M is p -e-eec-concenttrtee joint DP (p-zCJDP) if for all u-neighboring meta-datasets S ~u S':
Da (S)__()) S pa,
(2)
where Dr ( is the a-Rényi Divergence between the outputs of all users except
3 Private algorithm for learning personalized histograms
In this section, we present our algorithm for learning private and personalized histograms in both hiz-ehomogencous and bize-heterogeneous settings. Without assumptions 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 Qu using their local dataset S.,. Avoiding such worst-cases, we assume that the real- world problem is more structured. Specifically, we assume that each user distribution Qu 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 stylized Bayesian setting that simulates a distrdistribution with clusters.
As a warmup. we first discuss our local finetuning algorithm, where the setup is simpler. Say an algorithm runs private collaborative protocol, and hands distribution P to user u, with the guarantee that (T.(Qu.|PP (P is "reasonably" small, then how should the user u finetune using the local dataset S., available to them? For this question, we now present our local finetuning algorithm.
3.1 User-level local finetuning
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 specific cluster. In the next section, we discuss how to estimate P. Here, we discuss how to personalize P for each user. For user u. with empirical distribution Q. from m., samples, we use (T Q.) to denote the personalized, locally finetuned estimate of user distribution Q., (see Eq. 3). Note that P is moved closer to Qu 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., is an unbiased estimate of Q.. 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.
FT Q.) =: a P+ The
Qu (local finetuning)
+
a a+m,
Next, we discuss how to collaboratively estimate P. For a user-distribution with multiple clusters, 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).
3

---

3.2 Learning clusters in the population of size-homogeneous users
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 { Ju. and the matrix of K cluster centers P E R To learn cluster centers P, akin to the popular clustering algorithms like Lloyds [53], the objective we optimize is:
min (Po)
Σ wEt
min cE[K]
KL(QuulPo).
with the main difference being that we are concerned with the KL distance, instead of Ez. Given P. user u's assigned cluster &n is simply given by arg ming KL(Q=I|P... 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, ie., 2, =: {u: &n i). the following is true:
Σ uEZ;
Qu E arg min Σ KL(QuIP).
22
PEA(V) u
Since, do not have access Que replace Qu in Lemma 3.1 and the cluster assignment step Good-Turing
with the local Good Turing [34] estimate Q Typically, one would use the empirical distribution Q.. which is an unbiased estimate of Qu. But, as we will see in Section 5, most real-world distributions are long-tailed and the Good-Turing estimate's '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 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 possible direction of improvement in future works.
estimator averages the local
To iteratively optimize the objective in Equation 4, we present our algorithm Histogram Cluster (see Alg. 3). We start with initial center estimates P()), and run T' iterations refine them. In each iteration we peform two steps: (i) collect user outputs for each user via UserQutput (see Alg. 1); and (ii) re-center clusters using an algorithm, which for now 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 not hard to see that NonPrivateCenter which computes the mean of the local estimates of users in cluster, greedily reduces the objective value in Eq. 4 after each iteration. When the misclustering rate sufficiently low to begin with, averaging local Good- Sorr-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 (0) and the separation of true clusters in the user population.
Algorithm UserOutput
Algorithm 3 HistogramCluster
Require: Centers P. cluster k. user u.
Require: Set of users U4 initial centers p(0), iterations
Ensure: User output (b.., Q.).
Ensure: Estimates of cluster memberships (2,h, and centers P.
I: b, + 1 (arg min, KL [(:))) k)
2: Q.. + b. Qg + (1 b.)) O2
1: Initialize + 1.
2: while t < T' do
3: Return (bn,, Q.).
3: for = 1 K do
4: Collect: A
(UserOOtpu(((-)))) Ai u)]u 5: Re-center: ((()) (Non))rrvvteeeenter())
Algorithm 2 NonPrivateCenter
end for
Require: User data A =: {(bu, Q.)}. Ensure: Center P. I: Return P+ E.. Qa/2iaha. hu-
<- L+1
8: end while
9: Return: (2(T) P(T).
Theorem 3.2 (Alg. 3 convergence). For the model in Section 4, if Vc. v, ne O(K2) P(K3),P.l1
Q(1), centers are sufficiently separated, i.e., A =: min (L(P.))P;) 22k22 + ddd/)), and =: maxiz; KL(P.IP35/a. Given an initialization with assignment error (1/1/)) for any cluster, after t log(((() iterations, w.h.p. the assignment error is ((exp((exx)))) + 1)))

---

3.3 Learning clusters with user-level privacy
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. applied all the current clusters 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 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 A(V)
and thus bounded in lz. This initial estimate B is used to identify the quantile B = /B/b, where the mean lies with high probability [69], so that can then get a refined estimate in step 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 A(V), we use E2 projection, ( =: arg min)) ||Q PIl In practice, we find that projection with lg is comparable with a KL projection.
Privacy preserving noisy sum. As can be seen from the description of PrivateCenter, it only re- quires 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 sum over all the users, where a user submits value O if they do not belong to the cluster, ie., by = O, and by = otherwise (Alg. D). Similarly, Alg. also emits the true local estimate if it is part of the cluster and O, (vector of Os) otherwise. crivacyprreeevvng, computation of noisy sum over all the users is one of key primitives 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.
Algorithm 4 PrivateCenter
Algorithm Privatelnit
Require:
Data collection {(b.. Q.))).
Require: Center Qo. meta-datasel 8, # clusters K. privacy
privacy parameter p. clipping param- eter c, projection operator I[
parameter p. sampling parameter a, clipping threshold
Ensure: Matrix of initial centers ((0) R
1: Sample K² candidates: C ' Dir(aQo)
Ensure: Private center P.
I: b < Σ, by + Lap 3/2
2: Initialize set : I () and [Q].
3: while [Z] K do
2: B + : En Qu (N, 3/2ola))
4: I + IUQ Q where we sample Q from C with.
3: f 6e hop dBag(B)
Pr(() x exp
VSp
Σ elÂu.. Q) where i(Qu, Q)
4: P + (Eu Qu + +N00, E)).
2hV
Q.ES
@u =: Clip 1Q B-oy/
B-o 1
=: Clip min KL (T - min KL(@Q.IP))
PEE
PETU(()
5: P < AMM P
5: end while
6: Return P
6: Return p(0) with columns as entries I.
3.3.1 Private initialization for our clustering algorithm
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- vatelnit (Alg. 5), we provide a private data-dependent initialization technique, that takes input private center Qu (private FedAvg), and samples K2 points around 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 p 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 p. This iterative procedure of selecting candidates with highest regret of omission at each step is meant to ensure that each Q,, 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 small dataset collected centrally for this.
5

---

3.4 Putting it all together: Algorithm for private and personalized frequency estimation
In Alg. we present our complete algorithm. We first compute a private center of all user histograms using Alg. 4 and privacy parameter p/3. We pass this to Alg. 5 with privacy parameter P/3 to get initial centers for Alg. 3, where the re-centering step applies Alg. with privacy parameter 9/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 p-rCJDP. Our end-to-end private and personalized estimation algorithm is (p + 2 (0)g()))))))JD))))) vo > O.
Algorithm Tnd Algorithm
Require: Set of users U. dataset S, number of clusters K. number clustering iterations T. finetuning
parameter a, sampling parameter w. privacy parameter clipping parameter c, projection operator (V))
Ensure: Private and personalized estimates Qu for each u E Ut.
I: Qo + PrivateCenter((4, P. c, (v)) # Estimate global center
2: p( <- Algorithm Privatelnit(Qo. p,w,,)
# Estimate initial cluster centers
3: [āu]a and centers P <- Algorithm Manogrramlluste((4.. ,(O), # Estimate cluster memberships
4: Vu: Qu + FT (.....a) a)
Local finetuning using Eq. 3
5: Return {Qu E U).
3.5 Extending our private clustering algorithm to the size-heterogeneous setting
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/, =: (//4)/E (4/22)-
Here of is the variance of the user's local estimate, ie., Var [u]] (from Lemma 3.4). When Pz, [c] is
known this weighted estimate is optimal under &2 [37]. Since Pa. [c] is unknown we replace Ps..
Lemma 3.4 with uniform average of Qu. from heavier users, following Cummings et al. [21].
Lemma 3.4 (Variance of Qw). For user Var Var(Qu//u [v]] = P.., 111la/11111111 - 1/m.) + Ps., Pt.,
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 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 high- dimensional 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 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. 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, apply Algorithm 2 from Cummings et al. [21] that re-centers each cluster privately, based on the newly added users.
4 Formal analysis in 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 3 to locally finetune the learnt cluster centers. Now, we analyze these algorithms in stylized model.
Bayesian Model. Each user u belongs to a cluster 2u E [K] with cluster center Pz.. The user histogram Qn ~ Dir(a and the user dataset S, Q is sampled i.i.d. from Q.. We use Z. =: {u: : zu = c} to denote the set of users from cluster and n. = [Z is the number of users in c. Higher value of a implies more concentrated clusters since Var(Qu.v) [v]) [ol.,Ivl/1+o) and as a + 0o, Qu + Pz., 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..., Pr([Qu[u] - Pz.. [u]] 2 t) O(e "
Purely local learning. We present two local estimators that estimate Qu: empirical Q.. and Good- Turing [30] Qg. The naive estimate Qu is the average of user data in S,. Next, we define the

---

Good-Tuuing estimator. Let the frequency of token u in user's dataset S. (of size m.) be enty,v. We denote the local frequency of the count j as Pu. =: -.vey (cnt,. = j), f.e., the number of token in Sn with count j. Following Orlitsky and Suresh [63], for a token that appears j = enly, times, the Good-Turing estimator for Qu is given by:
j Qu.j+1:
Q [1] =:
ma No
1+1
su
otherwise,
(5)
No
Ou.z
where No. is the normalization factor so that En [%[[o]- Let E., be the class of natural estimators that assign the same probability to tokens with equal counts in S.. Then, following Lemma 6 in Orlitsky and Suresh [63] and Theorem 2 in Acharya et al. [2] we conclude that with respect to En, the Good-Turing estimate Q has worst case suboptimality gap of Ó(iin( (///
(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 Qu. This is mainly because Good- Seed-Turingg Turing more accurately estimates the probability of unseen words [56].
Lemma 4.1. (Qat suboptimality gap) For any Qu. the suboptimality gap of Q% with respect to & is Ommin( /1/m, 4/m)). that matches minimax rates.
E]KI(Q.]IQ min AEE, E[KL(Qu
Bayes optimal finetuning when the cluster center is given. When Pz. is known, FT in Eq. is Bayes optimal in KL and when only an estimate of the center 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 a, the estimator FT Qa) is Bayes opti- mal in KL divergence, for the Dirichlet prior Dir Dir(aP: s..). When Pz, the estimated cluster center, and Qurt is the Bayes optimal estimate of Qu. then KL, [[TT[](( (P.((P..Qu.) S <//+m. KL (P=,. (P..)) ))-
Estimating the cluster center given cluster members. Given all users in a cluster c, one estimate of the cluster center Pa 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 pfa 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(Pl) Qu).
Pfa
Σ (ma/Enrez, m..) Qu arg max uEZe
Σ log Pr(8., Q))
(FedAvg) (6)
QEA(V) ne uEZe
Our algorithms (Alg. 6, Alg. 2) use a different estimator. In the hize-homogenous case, to estimate the cluster center we average the local Good -urrin estimates Q for all users in the cluster:
QRt Σ
(Good-Turing based; ours).
The uEZ.
To analyze the accuracy guarantees of our estimate Pst, we first define the class of competing estimators for it. Let ent,, =: [Enez entu.=/E., m.. be the average count of token v across all users in cluster. We define E to be the class of estimators that assign the same mass to tokens 11, U2 ent., ent Such a class is natural for estimators that rely aggregated statisitics, and includes the -ood-n correction applied cluster level counts. In Theorem 4.3, for any cluster center Pa. we compare the suboptimality gap in KL divergence for our estimate Pgt in Eq. 7 with FedAvg.
Theorem 4.3 (P&t 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 pst in Eq. w.r.t. E is (1+tm/a+1 miní suboptimality gap over E for FedAvg pla is á (dla+m)/m(a+t)). Ö, Ω hides colylogfacton in m, n.
min(v/ / m. The
When m is small and comparable to the intra-cluster std. deviation, i.e., m/a+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.. given by the finetuned center: FT Q..) would also have stronger error guarantees than FedAvg+FT.
Partitioning users into clusters. In practice, neither the cluster membership 2n nor the center P. is known. The Bayes optimal estimate for P., Zu is the mean of the posterior Pr(Qu. S,a)
7

---

Roddit fineturing)
StackOverlow (wyo Instuning)
Amazon Reviows (w/o fineturing)
Test Negulive Luy, Liallooo
Text Negetine Lug Lailioad
Text Negettv Long ballbood
9.5
9.5
8.0
9.0
8.5
8.5
8.00
62
0.0
7.5
-..5
Fed/An
MAMAI
IFCA
Ours
FedAug
MAMI
IFCA
FedAug
MAMI
IFCA
Dur
(a) Reddit
(b) StackOverflow
(c) Amazon Reviews
Figure 1: Performance before finenmmin:: 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 hars 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 's theorem [59], and its variance matches the Cramér-Rao bound [69]. Still, the joint MLE for Pe: Zu is a non-concave maximization problem. Akin to solving the MLE for mixture distributions [61]. we maximize an evidence lower bound with Expectation- Expeecaaion-Maxiimizaiinnnnnnnnn (EM) [7]. Since each 2u belongs to only one cluster, it must be that 2n E arg cen] (L(Q.. Pa). 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 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.
5 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-
baselines, and the gap amplifies the size-heterogencouss case; and (iii) we show 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 vocabulary of size 10k tokens, and for the other two datasets, we use the Huggingface (bert-case-uncaset 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-basee approach which learns multiple models, one for each clus- ter: IFCA [33)/Hyp
[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
Reddit (with finctuning]
StackOverlow (with finetuning)
Amszann Rasinon (with finshunig)
Text Nageive Log Lhalhood
Tex Meggbbe Lag Liaiihed
8.0
Rarariatt
Tel: Negeiie log Liedtteod
0.50
Noephata
mP?
MMMM Phoae
8.5
8.00
8.0
7.75
7.50
7.5
FedAAg
MARA
FadlAwg
MANII
OHr
FecAug
IFCA
Oum
(a) Reddit
(b) StackOverflov
(c) Amazon Reviews
Figure 2: Performance finetuiin: the hozogeoomgessuus (a-c), and size-heterogennenous (d-f) settings. we compare the test NLL loss for baselines FF, MAML+FT, IFCA+FT with our Alg. 3+FT, where FT is implemented by Eq. 3. Uniformly averaged over users, each value is averaged over 50 random runs (error bars indicate 95% confidence intervals).
8

---

E.6
AL (isthraaerr
2.50
ming
3.3
525
11
120
270
$15
8.0
0.2
10
g:
10
12
24
12
(a)
(b)
(d)
Figure 3: Algortrhmic design choices: We evaluate test NLL for Alg. 3 we: (a) vary the number of clusters
K; (b) use Privatelnit or randomly initialize cluster centers; and (c) use average of Good-Turie Turing or empirical
average to estimate cluster centers. In (d) we evaluate different finetuning methods applied to the FedAvg model.
(Figure 3()),, 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 l2 bounded gradients for MAML/IFCA, and I bounded probability distributions for FedAvg.
Evaluation metric. For distribution Qu. estimate Qu. KL(Qu,, Q.) = NLL(Qu, HQQ,..
where NL Qw) =: Ewey Qu, log Qu, w is the negative log- likeliikoood loss, and H(Q,) is the entropy of Qu- Following language modeling [3, 38], we use the test set of the user to get an unbiased estimate of the NLL loss, by replacing Qu with the empirical distribution on the test set. Since the test set is insufficient to get an unbiased estimate of H(Q,), as a reference point, we instead report the entropy of the global center or FedAvg estimate H H(Pf)), which is a rough estimate of the "hardness" " of estimating the user distribution. We refer to this difference in test NLL and H ((.)))) as NLL sub-optimality gap. In the private case, we ensure the algorithms satisfy (15, 10 100 -JDP.
Our approach significantly reduces test NLL in private and non-private settings. In Figure 1 we report 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 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 we compare the performance after we finetune the giobal/cluster-leww estimates locally for each user (using Eq. 3). and verify that the relative gains remain consistent with pre- 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.
A Algorithmic design decisions influenced by our model in Section 4. In Section 4 we identified the Good-Turing -uring 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 uses Good Toriig as the local estimate, vs. empirical (Figure 3())). 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- hneereetroogs case. In Figure 4 plot test NLL in the algorithm's 's gains (> 40% error reduction) over single global model baselines FedAvg and MAML are even more pronounced than the size- size-homogeneou 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
case where our
9

---

Reddit (sne-beteenggreous))
StackOverlow (dize beterogentous)
Amazon Reviews (size- oetenogennooss
Tex Negutee Log Litelood
MPP")
Tex Nogatine Log Likeloood
Tee Negetive Log, Likehood
www Prints
10.5
9.5
10.0
-0
8.0
9.5
9.0
8.5
7.0
FeclAng
MAMAL
IFCA
8.0
IFCA
FodAAy
MAML
IFCA
Ours
(a) Reddit
(b) StackOverflov
(c) Amazon Reviews
Figure 4: Size-heterogenous setting: Comparison of test loss for baselines FedAvg+FT. MAML+FT, IFCA+FT
with our Alg. 3+FT, in the size-heterngeneous setting where each user's dataset vastly differs size. Uniformly averaged over users, each value is averaged over 50 random runs (error bars indicate 95% confidence intervals).
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.
6 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 global model and then uses it to partition users based on their losses, most works learn diverse models from scratch using gradient based algorithms,
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 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- 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 hize-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. 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 f metric. In contrast, we are concerned with histogram estimation in KL. divergence which presents interesting challenges since this is not proper distance metric (e.g., does not satisfy triangle inequality). On the other hand, [30, 34, 64] study Good-Tuuing estimators and give estimation error guarantees in KL divergence for categorical distributions over 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 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 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 two-stage version of our approach to separately handle the harder hize-heterogemous 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 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.
10

---

Acknowledgements
Part of this work was done when AS was intern at Apple, hosted by VF and KT. AS would like to thank Virginia Smith, Moshe Shenfeld, Hilal Asi, Abhishek Shetty, Aadirupa Saha, Steven Wu, Shengyuan Hu, Pratiksha Thaker, Tian Li, and Don Dennis for helfpful feedback and discussions. AS also thanks JP Morgan AI PhD Fellowship for their generous support.
References
[1] Martin Abadi, Andy Chu, Ian Goodfellow, H. Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC Conference on Computer and Communications Security, CCS' 16. ACM, October 2016. doi: (0.1145/97774492978883 URL http: /ett::/dxx.de doi org/10 /org/00.. 1..0..1455 1145/2976749.... 2978318.
[2] Jayadev Acharya, Ashkan Jafarpour, Alon Orlitsky, and Ananda Theertha Suresh. Optimal probability estimation with applications to prediction and classification. In Conference on Learning Theory, pages 764-796. PMLR, 2013.
[3] Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300- 22312, 2022.
[4] Galen Andrew, Om Thakkar, Brendan McMahan, and Swaroop Ramaswamy. Differentially private learning with adaptive clipping. Advances in Neural Information Processing Systems, 34:17455-17466, 2021.
[5] David Arthur and Sergei Vassilvitskii. K-means++ the advantages of careful seeding. In Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithms, pages 1027-1035, 2007.
[6] The TensorFlow Federated Authors. TensorFlow Federated Stack Overflow dataset, 2019. URL wwwg//eeeaaee//aaeeeeeeee aptaaapiiasaappaaaaaaaaaaasaaass httpaappccccpppo/////eeeeeeeeeee pttpppytttttssspppeppppp http://tt.//mmmamiiioa.o aataaetsssstaccovrrrll///lltt _att..
[7] Sivaraman Balakrishnan, Martin J. Wainwright, and Bin Yu. Statistical guarantees for the EM
algorithm: From population to sample-based analysis. The Annals of Statistics, 45(1):77 120,
2017. doi: 10.121411-AAOS 14355 URL https: / //doi otg/iiorg1.... 1214/1111111 A0S1435.
[8] Arindam Banerjee, Srujana Merugu, Inderjit Dhillon, Joydeep Ghosh, and John Lafferty. Clustering with bregman divergences. Journal of machine learning research, 6(10), 2005.
[9] Mikhail Belkin and Kaushik Sinha. Polynomial learning distribution families. In 2010 IEEE 5 Annual Symposium on Foundations of Computer Science, pages 103-112, 112. IEEE, 2010.
[10] Alberto Bietti, Chen-Yo Yu Wei, Miroslav Dudik, John Langford, and Steven Wu. Personalization improves privacy-accuracy tradeoffs in federated learning. In International Conference on Machine Learning, pages 1945-1962. PMLR, 2022.
[11] Steven Bird, Ewan Klein, and Edward Loper. Natural Language Processing with Python: Analyzing Text with the Natural Language Toolkit. O'Reilly Media, Inc., June 2009. ISBN 9780596516499 16499, URL httph::t///// :/wp:://www nltk. org/book/.
[12] Sourav Biswas, Yihe Dong, Gautam Kamath, and Jonathan Ullman. Coinpress: Practical private mean and covariance estimation. Advances in Neural Information Processing Systems, 33: 14475-14485, 2020.
[13] K. A. Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H. Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for federated learning on user-held data. In NIPS Workshop on Private Multi-Party Machine Learning, 2016. URL https:://// //ttri///.iiva// org/ abs/ 1611. 04482.
[14] Joseph Bonneau. The science of guessing: analyzing an anonymized corpus of 70 million passwords. In 2012 IEEE symposium on security and privacy, pages 538-552. IEEE, 2012.
11

---

[15] Mark Bun and Thomas Steinke. Concentrated differential privacy: Simplifications, extensions, and lower bounds. Theory of Cryptography Conference, pages 635-658 Springer. 2016.
[16] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprini arXiv: /8/2.0009// 2018.
[17] Kamalika Chaudhuri and Andrew McGregor. Finding metrie structure in information theoretic clustering. In COLT, volume 8, page 10, 2008.
[18] Gary Cheng, Karan Chadha, and John Duchi. Fine-tuning is fine in federated learning, arXiv preprint arXiv:2108.07313. 3, 2021.
[19] Gary Cheng, Karan Chadha, and John Duchi. Federated asymptotics: model to compare federated learning algorithms. In International Conference on Artificial Intelligence and Statistics, pages 10650-10689. PMLR, 2023.
[20] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. Advances in Neural Information Processing Systems, 35:10572-10586 2022.
[21] Rachel Cummings, Vitaly Feldman, Audra McMillan, and Kunal Talwar. Mean estimation with user-level privacy under data heterogeneity. Advances in Neural Information Processing Systems, 35:29 15:29139-291511 2022.
[22] Sanjoy Dasgupta. Learning mixtures of gaussians. 40th Annual Symposium on Foundations of Computer Science (Cat. No. 99CB37039), pages 634 644. IEEE, 1999.
[23] Cynthia Dwork. Differential privacy in new settings. Proceedings of the twenty-fist annual ACM-SIAM symposium on Discrete Algorithms, pages 174-183. SIAM, 2010.
[24] Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography; Third Theory of Cryptography Conference, TCC 2006, New York, NY, USA, March 4-7, 2006. Proceedings 3, pages 265-284. Springer, 2006.
200n.
[25] Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. Journal of Privacy and Confidentiality, 7(3):17-51, 2016.
[26] Mathieu Even, Laurent Massoulié, and Kevin Scaman. On sample optimality in personalized collaborative and federated learning. Advances in Neural Information Processing Systems, 35: 212-225, 2022.
[27] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprini arXiv:2002.07948, 2020.
[28] Thomas S Ferguson. A bayesian analysis of some nonparametric problems. The annals of statistics, pages 209-230, 1973.
[29] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126-16 1135. PMLR, 2017.
[30] William A Gale and Geoffrey Sampson. Good-turing frequency estimation without tears. Journal of quantitative linguistics, 2(3):217-377, 1995.
[31] Jonas Geiping, Hartmut Bauermeister, Hannah Dröge, and Michael Moeller. Inverting gradients- how easy is it to break privacy federated learning? Advances Neural Information Processing Systems, 33:16937-16947, 2020.
[32] Robin Geyer, Tassilo Klein, and Moin Nabi. Differentially private federated learning: client level perspective. arXiv preprint arXiv:1712.075577 2017.
12

---

[33] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33: 19586-19597, 2020.
[34] Irving Good. The population frequencies of species and the estimation of population parame- ters. Biometrika, 40(3-4):237-264, 1953.
[35] Dennis Grinwald, Philipp Wiesner, and Shinichi Nakajima. Solution simplex clustering for heterogeneous federated learning. arXiv preprini arXiv:2403.03383 2024.
[36] Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: 181 1eee11111100000 1.03604, 2018.
[37] Joachim Hartung, Guido Knapp, and Bimal K Sinha. Statistical meta-anala meta-analysis with applications. John Wiley & Sons, 2011.
[38] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[39] Justin Hsu, Zhiyi Huang, Aaron Roth, and Zhiwei Steven Wu. Jointly private convex pro- gramming. In Proceedings of the twenty-sevntth annual ACM-SIAM symposium on Discrete algorithms, pages 580-599, SIAM, 2016.
[40] Rui Hu, Yuanxiong Guo, Hongning Li, Qingqi Pei, and Yanmin Gong. Personalized federated learning with differential privacy. IEEE Internet of Things Journal, 7(10):9530-9539 2020.
[41] Shengyuan Hu, Zhiwei Steven Wu, and Virginia Smith. Private multi-task learning: Formulation and applications to federated learning. arXiv preprint arXiv:2108. /2978, 2021.
[42] Prateek Jain, Om Dipakbhai Thakkar, and Abhradeep Thakurta. Differentially private matrix completion revisited. In International Conference on Machine Learning, pages 2215-2224. PMLR, 2018.
[43] Prateek Jain, John Rush, Adam Smith, Shuang Song, and Abhradeep Guha Thakurta. Differen- tially private model personalization. Advances in Neural Information Processing Systems, 34: 29723-29735, 2021.
[44] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji. Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings. et al. Advances and open problems in federated learning, Foundations and Trends in Machine Learning, 14(1-2) 2):1-210, 2021.
[45] Adam Tauman Kalai, Ankur Moitra, and Gregory Valiant. Efficiently learning mixtures two gaussians. In Proceedings of the forty-second ACM symposium Theory of computing, pages 553 553-562,22 2010.
[46] Ekatherina A Karatsuba. On the asymptotic representation of the euler gamma function by ramanujan. Journal of computational and applied mathematics, 135(2):225-240, 2001.
[47] Vishesh Karwa and Aleksandra Slavković. Differentially private graphical degree sequences and synthetic graphs. In Privacy in Statistical Databases: UNESCO Chair Data Privacy, International Conference. PSD 2012, Palermo, Italy, September 26-28, 2012. Proceedings, pages 273-285. Springer, 2012.
[48] Michael Kearns, Mallesh Pai, Aaron Roth, and Jonathan Ullman. Mechanism design in large games: Incentives and privacy. In Proceedings of the 5th conference on Innovations in theoretical computer science, pages 403-410, 2014.
[49] Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint
2016.
13

---

[50] Daniel Levy. Ziteng Sun, Kareem Amin, Satyen Kale, Alex Kulesza, Mehryar Mohri, and Ananda Theertha Suresh. Learning with user-level privacy. Advances in Neural Information Processing Systems, 34: 14:12466-22479 2021.
[51] Jeffrey Li, Mikhail Khodak, Sebastian Caldas, and Ameet Talwalkar. Differentially private meta-learning, arXiv preprint 190.0:1900.058300, 2019.
[52] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Chal- lenges, methods, and future directions. IEEE signal processing magazine, 37(3):50-60, 2020.
[53] Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28 12:119-137, 1982.
[54] Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002. 10619, 2020.
[55] Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under mixture of distributions. Advances in Neural Information Processing Systems, 34:15434-15441 2021.
[56] David A McAllester and Robert E Schapire. On the convergence rate of good-turing estimators. In COLT, pages 1-6, 2000.
[57] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effieee
learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[58] H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private language models without losing accuracy. arXiv preprint 1rXiv:1710.06963, page 84, 2017.
[59] Jeffrey W Miller. A detailed treatment of doob's theorem. arXiv preprint arXiv: arXiv:1801 1801,0311032 2018.
[60] Michael Mitzenmacher and Eli Upfal. Probability and computing: Randomization and proba- bilistic techniques in algorithms and data analysis. Cambridge university press, 2017.
[61] Ankur Moitra and Gregory Valiant. Settling the polynomial learnability of mixtures of gaussians. In 2010 IEEE 51st Annual Symposium on Foundations of Computer Science, pages 93-102. IEEE, 2010.
[62] Jianmo Ni, Jiacheng Li, and Julian McAuley. Justifying recommendations using distantly- labeled reviews and fine- fine-graing grained aspects. In Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference natural language processing (EMNLP-IJCNLP), pages 188-197, 2019.
[63] Alon Orlitsky and Ananda Theertha Suresh. Competitive distribution estimation: Why is good-turing good. Advances in Neural Information Processing Systems, 28, 2015.
[64] Liam Paninski. Variational minimax estimation of discrete distributions under kl loss. Advances in Neural Information Processing Systems, 17, 2004.
[65] Alireza Salemi, Sheshera Mysore, Michael Bendersky, and Hamed Zamani. Lamp: When large language models meet personalization. arXiv preprint arXiv:2304. /406, 2023.
[66] Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural networks and learning systems, 32(8):371033722 2020.
[67] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[68] Ramon Van Handel. Probability in high dimension. Lecture Notes (Princeton University). 2014.
14

---

[69] Roman Vershynin. High-dimensional probability: An introduction with applications in data
science, volume 47. Cambridge university press, 2018.
[70] Martin Wainwright. Iligh-dimensional statistics; non-asymptotic viewpoint, volume 48.
Cambridge university press, 2019.
[71] Zhibo Wang. Mengkai Song, Zhifei Zhang. Yang Song, Qian Wang, and Hairong Qi. Beyond inferring class representatives: User-level privacy leakage from federated learning. In /EEE INFOCOM 2019-IEEE conference computer communications, pages 2512 25200. IEEE, 2019.
[72] Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Towards provably personalized federated learning via threshold-chustering of similar clients. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022), 2022.
[73] Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Provably personalized and robust federated learning. arXiv preprint arXiv:2306.0083 08393, 2023.
[74] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue. Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's 's transformers: State-of-the-an natural language processing, 2020.
[75] Huiwen Wu, Cen Chen, and Li Wang. A theoretical perspective on differentially private federated multi-task learning. arXiv preprint arXiv:201 1.00179, 2020.
[76] Shanshan Wu, Tian Li, Zachary Charles, Yu Xiao, Ziyu Liu, Zheng Xu, and Virginia Smith. Motley: Benchmarking heterogeneity and personalization in federated learning. arXiv preprint arXiv:2206.09262, 2022.
[77] Dun Zeng, Xiangjing Hu, Shiyu Liu, Yue Yu, Qifan Wang. and Zenglin Xu. Stochastic clustered federated learning, arXiv preprint arXiv:2303. 00880 2023.
15

---

Appendices
A Limitations of our work
As a first step towards language modeling we only focus on privately estimating the marginal token distribution for each user, and interesting future direction would be 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 federated version of the private initialization algorithm we propose. We suggest handling the initialization at the central node that maintains small dataset of clients. This 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.
B Additional related work
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 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 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 now-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.
C Additional details for empirical evaluation
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 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 x 10k users in Reddit, each with at least 1k tokens. For StackOverflow and Amazon Reviews we have s 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, adapt the baselines to the private setting by using standard techniques [1, 25] privatize l2 bounded gradients for MAML/IFCA, and l bounded probability distributions for FedAvg, so as to ensure all algorithms achieve JDP-(15, 10 10-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- 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 X = 0.25, À = 0.15 and = 0.1 to be optimal on Reddit, StackOverllow 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 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 4.0 for Alg.4. 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.
16

---

1.025
Test Nagative Log Likelhood
Reddt
1.000
0.975
0.950
0.925
12
16
Figure 5: Validating hyperparameter choice of K 10 Reddit dataset. We use the same value of K for the other two datasets as well.
D Useful Lemmas
Lemma D.1 (Chernoff bound). X Poi(A), then for 2 X,
(x (x- 2]2
Pr(X 2 x) S exp
(8)
2r
for x < X,
(x A)2
Pr(X S x) I exp
(9)
2A
Lemma D.2 (Bernstein's Inequality). Let X1, X2,.... Xn be independent zero mean random variables such that with probability 2 6o Xil M. Then,
1
Pr Σ
X; 21 S 2exp
61-
(10)
+ Mt/33
1=1
Ift = 2 (( log 1/2 + M log J. then
Pr Σ
2 2 Σ E[X?] [EEX?]]
log
M Mlog log
S 28 +
Gi-
(11)
Lemma D.3 (Banerjee et al. [8]). The solution for the optimization problem: arg arg min (1 A) KL Q) AKL(P) (1
is given by (1- AP1
E Missing proofs from Section 4
Lemma E.1. (Q% achieves minimax suboptimality gap: restated) For any Qu, the suboptimality gap of Q with respect to En, is [E[KT[QullQE)) min Aes ([.TT(Q..IAA)) Ó(min( /1/mm, d/m)), that matches minimax rates.
Proof. From Lemma 8 in [63] we know that for the class of natural estimators E,, the suboptimality gap defined above given by E[KL.[M& M.) where
M Σ 1ent, k Qu(u),
Mr Σ
1(ent, [)Q_[v).
VEV
We use the definition of ou, from [63], where:
Du,t = Σ
1(cnt = t)
(12)
vEV
17

---

Now, from Theorem in [63] we also know that KL(M& = -((1m/m with probability at least 1 1/m over the randomness in draw of the user dataset S,, Since their estimate lower bounded by 1/2m, we can convert their high probability guarantees to guarantees in expectation. Further, we note that their proof of Theorem 2, we can easily upper bound the following
MM 22
Pus
Ou.t
M
m
m
in the worst case. Thus, in the worst case, [[ [M [M)] is upper bounded by the minimum of (4/m) and O( ((7m))) This completes the proof of Lemma 4.1.
Theorem E.2 (Finetuning is Bayes optimal given center, restated). If cluster center Pe.., and Dirichlet parameter known then FT(P:.. Qu) Bayes optimal in KL divergence.
Proof. We note that the distribution of Qu given the center Pc. is a Dirichlet- Diriclet-Multinomial Multinomial distribution. Further, the Dirichlet distribution is conjugate prior for the Multinomial.
Hence, the posterior distribution for Q Qu: Po. a Dirichlet Multinomial with mean: ( )
(o4m Qu. To why, mean is the Bayes optimal estimator, we invoke Lemma D.3, which implies:
(1- + XQu arg min (1 KL (1 (KK
Now the Bayes risk of algorithm 21, given the center estimate P is
BayesRisk(DDU =: @...-P((..1...)) (I(Q..19TTQ..)))
Recursively applying Lemma D.3 on every measurable subspace (under measure defined by the posterior Pr(Qu Qw: P)) of the set of token distributions A(V), we conclude that the optimal solution is [E Qu. P] when Qu Pr(Q.. Qu.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 Ps. be the estimated cluster center for user u, and Qope be the Bayes optimal estimate of Q.. Then,
[FT(PP @w)) <
a KI(P=,, IP2.).
my
Proof. This result a simple application of Jensen's 's inequality for KL divergence.
KL
FT((:... Qw))
(
Pz. + a a my @..) ( Pa. at mu 2..)
= KL KL
a 1-+my my
a my
KL PP..)
+ m.,
Lemma E.4 (FedAvg estimate). The FedAvg model is given by Qia n Ee Q..
Proof. The log likelihood objective log Pr(S.ll))) is equivalent upto additive constant with the again use the result from Lemma D.3 conclude that the
objective ( Qu). We optimal is simply the average of local means.
L
Theorem E.5 (Pat in Eq. 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 ps in Eq. 7, w.r.t E is (aam//++1 min( / 1/mm, d/mx). The suboptimality gap over & for FedAvg pla is Q (Hl+tm)/m(att)) Ô, Ω hides factors in m,
18

---

Proof. For notational convenience, we denote the cluster center with P (instead of Pa), and the number of users in the cluster as n (instead n,). 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:
Mr
> = P. and F; =; Σ NEV
I([ent.] =t) Pe
(13)
VEV
Similarly, the local equivalent of Mr for each user is defined as:
Mu,e= = Σ (cntt,, = t) Qu,v and Fu,1 =: - = t) Q 0,00
(14)
DEV
DEV
As 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 Qge is:
cnty,r +
E (lou.ent....
Quen
X
(15)
Y7 E (u
and the corresponding average of cenie-aided Good Good-Turing Turing estimates:
pgen
u.p
Next, we derive results for the setting where my Poi(m) so that the local frequency counts of words become independent. Then, using the results on Poisson sampling Mitzenmacher and Upfal [60]. we can convert the bounds to the setting where each user samples exactly m tokens.
We will use 1 as shorthand for 1 1[[ent.] t) and similarly use 1 shorthand for 1 (cnntu, t). The following quantities define the global and local frequencies of average and local counts respectively:
Φ Σ1 and Qu,t PEV
W.O -EV
Σ 1
(16)
For F; defined using the genie-aided estimator, to bound the competetive KL estimation first need to bound the KL distance between ground truth M, and estimate Fe under high probability over the draw of the meta-dataset S.
KT(MI[ E Mr log
Me
S log
Σ
Fe
F.
Mr
F.2
(M,
F.)2
logg
1+
(17)
Fi
Ft
Using Lemma E.6 we lower bound the denominator in the KL upper bound in Eq. 17.
Lemma E.6 (Lower bound on Ft). With probability atleast nmn -2 F. 5
Ide
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 o some token in the vocabulary that has an average count of t, i.e., 1 1. From Lemma 15, 16 and Claim 20 in Acharya al. [2], with probability least m 3, Qun <<cat...=/m). where Qun is the unnormalized local Good-Turing Turing estimate. Additionally, from Lemma 17 in Acharya et al. [2], the normalization factor N,, for user u satisfies with probability at least
10m N2, 1 + O(!m²/*). Since m. is Poisson distribution from Lemma D. the number of data points m. for the user satisfies [m, m| < / /m, with constant probability. Putting these together, and applying a union bound over all users we conclude:
1
1 cetty,v
ent de
tt
F = c
á
=
m
m
m
PEV
with probability at least nm
19

---

Next we try to upper bound (M, - Ft)² for some t 2 1. But before that we introduce some more helpful results.
Lemma E.7 (User's Qu concentrates around center P). With probability least 1 - 6, Ve V:
P. (0 Po log 4/6
1 Σ Qu. Po = n
(n( (v
Proof. Since Qu, E [0, 1] a bounded random variable with variance (1))-)))(1+aa), we can apply the Bernstein inequality in Lemma D.2. This would give us confidence interval for single token which we can then union bound over all tokens v to get the result.
Lemma E.8 (variance of linear estimators; Claim 21 in Acharya et al. [2]) For every distribution
Var ΣΣIr(a ΣΣ ((
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 ΣΣ (
Σ Var
Σ 1((())))
27
Σ E
Σ (x,())
Σ E
Σ (1)f(xxtt)
(a)
Σ Σ ((1))
(b)
For L + *1', (1 = 0 and hence (a). (b) uses the fact that 1 is an indicator random variable.
Lemma E.9 (expected sensitivity of the local frequencies; Claim 20 in [2]). For every user distribu- tion over V. from which we draw n i.i.d. samples, let be the local count and de = Evev 1 is the frequency of the local count,
E [bbb) B[o++1]
O Bl(ba] max
log m t+1'
log t t+
We consider the two cases + 1 2 logn and + < log n separately. Consider the case when log n. We first show that
Proof.
E[111 - E[1 = np. )²
nps
mp (npa)² < 5e
log n
(4)
t!
t
t!
++
n²
The first equality follows by substituting E[11]] e en (npa) /t.. For the inequality, note that if Inp. 1 s 25(t + 1) log n, then the inequality follows. If not, then by the Chernoff bound
E[1`] Pr(t, t) S n and hence [[1
E[
+1 [[ + E[1 < 2/n
By definition, E| [[ Σ.
E E
E E
'j. Substituting,
E +
E
B -
88)
nyz
(npx))
np.
t!
t
20

---

E
mpr (np²)
np
+
npa (mpx)
np
+
t!!
t +
xmmp>>1
(b)
Tpp
Σ
5e
-np:
mp=
(mpx))
log
11
++1
n
ampp,
xmppx>>
log n
2n
log n
+
+
S
O
B[[d]]
12
t ++1
n
where (a) follows from the fact that E[1
/ /tl. (b) follows from the fact that np, S 1 in the first summation and Equation (4). The proof for the case + < log n is similar and hence omitted.
upx (np=)1
Lemma E.10 (Upper bound on M FiD. For any t 2 1. with high probability at least l-npoly(m).
M Fel = =
( (1_++1))
11 m )
Proof.
I Fil s
1 - Quvi
+E> DEV
Σ
19 u.
Qu (.
(18)
NEV
Ai:iiichlet concentration
B=:Locu Good-Turing
Our proof mainly relies on the concentration properties of Dirichlet distribution and some properties of the local Good-The Good-Turing estimate. While the first term is fairly streaightforward to bound, the second term requires us handle the fact that 1 is a random variable that depends on the global statistics for the token v, while Qst. Qu,e 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:
Var
(Po cmtu/m))
E 1 (Po mm)
(19)
P.P
)
1.
))
o
Σ
E
(20)
(1 (1+a)
m 1 aa)
LVEV
LVEV
E[
0
(21)
m(1 +
This completes our derivation for the upper bound on the first term (A) in Eq. 18. For the second term (B), the definition of the genie-aided local Good-Turing estimator and bound its variance in the folliowing way:
Var
Σ
1 12,r - Qct
Σ E
Σ Qu, - Qut
BEV
VEV
1 Σ Qu,
ent + Elob.ontu..ttll
Σ 1 "EV
my,
(0,centu.
2 Σ
B[1
Σ
Qu,
cnt
+
(22)
Ma
VEV
2
(enta,v + 1) Du
Qu Pu.cunt....)
+ Σ
[
m. E El
DEV
21

---

The first inequality above uses Lemma E.8, and the second equality uses the definition of the genie- aided estimator. 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
ΣEELJ
Σ
Qu,
ent +
1 11
Q 2.Y
ent
m
Ma
Mh
NEV
cntu,v logo (4/6) mi
B[
VEV
cntp log (4/6) m2
NEV
m²
The final inequality above the definition of t, and the concentration of mu, L.e, with high probabibility: [m, [[mm]l |m, ml (m)).
Next we will bound the second term in Eq. 22, for which use some properties of the expected sensitivity of the local frequency counts (Lemma E.9).
(entu.e v + (EElou.cntu... m. E Blu,,
Pu.cotw..ll)
Σ VEV
72
(c v + +1)
Σ
Y
my cutu,v
NEV
Applying Jensen's inequality it is easu to see that we can bound the above term with:
(entu,o + 1) (u, Vu.entu,
[E
ΣELL BEV
myBloou.nt....
=
m
(23)
Combining 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 setup in Section 4 in two ways: (1) Q. = i.e., each user's true distribution matches exactly the cluster center (this matches the +
in our setup); and (2) each user independently samples m. Poi(m), which is the number samples their local dataset, that is used to estimate the empirical distribution Q.. (note that = 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 with some collection of users Ze. Then, we know that cnt, is distributed Poisson random variable with mean Σmez. Qu(v]. Thus, with probability at least. 1 8:
Jentu/n EmEz, mQulol/m) S log(3/6).
We mainly rely on Lemma 13 for the Genie aided estimator from [2]. The key difference being that E[øx], which is words with total count A would differ from B[oentu.... For, this we first bound [//] enly, with high probability. We do this, by:
ent wv E(cntu,,) =s Ö(y B8 E(cntu,))
Plu a
=x (00mmmm.u. + mõl
22

---

since the counts follow Poisson distribution under indpendent sampling. Thus, when a = [mm), the first term dominates, and we see that each this gives us:
k E [entu..
V cnt ont cnt
with high probability. Then from Claim 20 in [2], qwe know that [B[øx] E el O(k)
The rest of our analysis relies with the genie-aided estimator mainly relies on Theorem 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:
KL(MI)F) SO =0( (V a + m
(24)
+
Next, we use Lemma 15 from Acharya et al. [2] to prove the same claim for the estimator that is not genie-aided. And finally, reuse the conversion of the high probability result to in expectation using arguments similar to our proof of Lemma 4.1. This completes our analysis the upper bound on competetive regret for the average of local Good-Turing (our estimator).
Finally, for the lower bound in Theorem 4.3, we invoke known lower bounds on the concentration of the Dirichlet distribution ((//++)), and minimax statistical lower bounds on the risk of the empirical estimator O d ), to conclude that the competetive risk of FedAvg at least ś (4/m + 1/o+1). This m completes our proof of Theorem 4.3.
Lemma E.11 (Stirling approximation [46]). The I' function is sandwiched as follows. There is positive constant C, such that for all 0:
Cx2-1/2 S T(x) S Cx=
This implies log (x) log C + (x log + 0(1=). For large enough x. we substitute [(x) = log C + (a 1/2) log x.
Theorem E.12 (Clustering approximates MIEE objective; restated). The joint likelihood for Pa and zu can be upperbounded with the clustering objective in Eq. 4, and this approximation is right upto an additive term of O1/a).
Let V be fixed set of words, d =: | is the size of the set, and A(V) is the set of all possible discrete probability distributions over Each user u the federated setup has an unknown distribution Qu E A(V). Further, each Qu belongs to one of k clusters, and the cluster membership is denoted by 2n, where E R is 1-hot k-dimensional vector. Each cluster E [k] is associated with a center P. E A(V), and the matrix with column vectors as cluster centers is denoted as P E BRdx i.c., P (P,,,??.....??
Independently for each user u, given the cluster membership Aus and some a > 0, the user's distribution Qu is sampled from Dirichlet around the cluster center Pz,, i.e., Qu Dir(aPz.). Further, each user samples an i.i.d. dataset S, of size nu. When all users sample dataset of fixed size n, n, we refer this setting heterogeneous. Given the set of datasets (S,, hu and the distribution assumptions above, we can write down the maaximun-likellhoog estimate (MLE) for the cluster centers P and cluster memberships (zu)x:
hizz-homogeneous, and the more general case size-
arg max bM
Σ log Pr(S, P. (eutu)
(25)
Σ log DirMuuln, P. (z(hu
arg max P, h
D log I(a) + Inn. + 1) log I'nn + +
arg max P. P,
Σ log [(Qu[a] + [PPu[>]
log [(aPz.(<)) log [(Q.(z)) + +11)
23

---

where Qu is the empirical distribution obtained from the dataset Sy and y(x)= t-- e e- dr is the Gamma function. In the above derivation, the second inequality follows from n. Qu being sufficient statistic for the class of distributions parameterized by (Quhu.
For the hizzeeooogesuus case,
simplify the MLE objective further to be:
arg max P
ΣΣ log IQg[[Qu[2] + a [x]) log ((aPzulx))
(26)
For large enough a, we can use the Stirling's approximation of the Gamma function to approximate the above MLE objective (see Lemma E.11):
Σ ΣC + P.Pulla [x] 1/2) log (o..(....lee [x]] + <<Pu(>))) - (Qu(z) + a (x])
arg max Pa
1
- (aP [x]] - 1/2) log(a (x)) [z
(27)
Qu(xt + al nm
Σ Qn] a (PP 1/2) log
arg max (a
aP 2x
Qu( [xc]
arg max Pa
Σ Σ Qu(x] log log P + (Q,(z) + aP2u(a) 1/2)
1
a 2m [2]]
Q [x]
Σ KL( || P2.) + [(Qu[=] aPPul [-|2 1/2) log
arg max Pu
a Zu [r]
Qu(z)
arg min P. Pitsuha
Σ KL(Qu P2.) + Σ 1/2log
1+
aP nm
Qu(a]
Σ (Q((x)) aPzulz)) log
1
ar zu [ar]
2u(r
Σ kI, P2.) + Σ 1/2
S arg min Pa
aP [2]]
Σ KL(Â, || P2.,) +
aar min Pa
when cluster center Plu] 2(1).
F Missing proofs from Section 3
Lemma E.1 Qu variance). For user u from cluster c, the variance of the estimate Qu(vo) is:
Polu] (1- 1 [1|u][1-A] P.(ul)
a a+1 ma
ma
Proof. By the Law of Total Variation, the variance of Qulu] is:
Var(Qu(()) = EQq. [Varq...) (Q.[e) Q.(+))) + Vare
EE - Qu(v)[ Q.(v)) + Var(Qu
Ma
1
Po(c)]3 + [c(
(Po[u]
maa
a+1
1
-PPpo[ - Po(e))) + (1 - P.
1
+
24

---

F1 Proof of Theorem 3.2
Theorem F.2 (Alg. 3 convergence; restated). For the model in Section 4, if Vc, U, ne ), = 2(11,, centers are sufficiently separated. i.e., A =: mine (L((P))P))
+022+150))) and x =: maxizj KT/P)))))//. Given initialization with assignment error 0 ((11/)) for any cluster, after t = log((([) iterations, w.h.p. the assignment error is ((xxp(-A(c + 1))).
Lemma F.3 (Pinsker [68]). For P,Q E AlV-2, the KL(P || Q) 2 =1l QQ - Q).
Lemma F.4 (KL upper bound). For P.Q. E A the KL divergence KL(P Q) S (P
Proof Overview. The main technique we follow is to use Pinsker's inequality to lower bound KL
divergence with TV (P, P, Q) and upper bound KL with x². This allows us to then treat the clustering
in the l3 metric, and pay the d-dimensional penalty term z. 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 need to upper and lower bound KL, we present the full proof for completeness.
Let us begin by introducing some notaion. For any S S [m]] define Ws EiEs WA Recall that
T {i [n], 2N g} and T
{i E [m], =/ = let define
s gh
{i E [n], 2, g, 2(()) = h T nTt
(5) Then we have 1/ The mis-clustering rate at iteration can be written
E y " (5)
and n = Σ gE|N R'ng
(4)
22
A,
19
+ 2 n g
n (a)
gh
71
We define a cluster-wise mis-clustering rate at iteration 8 as
Eythein
nih
(8) Eythe [g] T'g
(5)
G. max kR[R]
12
The first term in the maximum operator of definition of G, can 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, I Gg.
Let A mingg4he[k) 00 - 0₂| be the signal strength. For E [k). let 0(.)) be the estimated center of cluster h at iteration Define our error rate of estimating centers at iteration as
A, max KE
10 02
Lemma F.5. Based on our definition of M above, Go = O ((+ vie) 2).
Proof. First we show À = Ovd. Since Pc[m] = a we can invoke the KL upper bound above to show that X (2), and furthermore, the KL for each pair is O(M). Thus, this tells us:
kd
1) 12
(28)
n
In other words the signal to noise ratio is A/O 2 / + kdd
Based on the above result, we also get:
A
1/K
OVK.
+ kd/n
25

---

Recall for us is given by v1/a. 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:
1
4
Go
Or
Ao
(29)
rk
2
Vrk
Lemma F.6.
|| S V 3(n + d)
for all S S [m].
(30)
with probability greater than - Lemma F.7.
Xmaz Σ mmmi S <60 + d).
(31)
with probability greater than - -cxpp( 1-cxp--0.5m) 0.5m).
Lemma F.8. For any fixed E [m)) S [n]. 0 and 5
1 IS
30 30 (1
|S| + log(1/0))
[ 6
{
Pr
Wi:
w
S
S exp
(-min
4dd
+ 8.
JES
Lemma F9.
|| S 30 / (d + log [T) for all h E [k]
(32)
with probability greater than - Lemma F.10. For any fixed 01. ... 0b E Rd and > we have
Σ Ealalo - 1/ s (wi, 110 8g 0g11b} S m exp (- aaa22
+
V
5ng log
(33)
20
RET;
for all g # h [K]² with probability greater than n
The following two lemmas give the iterative relationship between the error of estimating centers and the error of estimating labels. Let & 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*} < 3n - + v. In the rest part of the proof, if not otherwise stated, we all condition the event E and the following analysis are deterministic.
Lemma F.11. On event Gf S 1, then we have
3 + min {- -VIC. 2G,AA--1,
3
A,
(34)
Lemma F.12. On event E, A, I 1-c and 2 236 then
2
(
28 An)
5k log ann
G.41 Ap2
+
(2
+
(35)
Proof of Lemma F.I.. For any B C [n], define YB step s can be written as
[B] EiEB Vi. The error of estimated centers
0(0)
0n
Σ (yi 0x)
(yi 0%)
nh VESha
nh a SES.A
Mah (Ys.. 0n)
wi +
Th
Thh
VESWA
ath
According to our label update step, we have 11y/
0(5-1))1
< y: 06--1) for any i e Sah. This
means for any E Sah, Yi is closer to 0fs-1 than 00 11) so is the average of {ui, E S.z). Thus, we have
ys.. 06.-11) - ||s.., 06.-1)11.
26

---

Consequently, triangle inequality gives us
||S.S 8a1 YS... 0a 110(--1) Ball + 110/ 0nll
which, combined with Lemma F.6 and the definition of As-1, yields
Ys.. - 0nll S 3(n + d)/nar + -A.-1
Taking weighted over # [k:]. we get
Σ
Bah
Vsa. - 0n
n d)
v Mah + 2A
Σ axh
aah
nh
aghh
apph
3
V - - 1)Gs + 2G
Thh
where the Last inequality is due to Cauchy-Schwartz and the fact that Eathh nah S Gynk. Note that Wss. = WTn Layh Wsn. Triangle inequality and Lemma F.10 imply
[W lI
30 / n /
3(n d) n/h--
Since Gg 5.2 we have
nh 2 nah 2 12 (1- - G.) 2 ni 2 an.
(36)
2
Combining the pieces, we obtain
|0% 0nl|
d log an
k(n + d) + /
an G. 2G
S 30
3 (1 / /G.) + 22GAA-1)
(37)
Therefore, we get the first term in (34). To prove the second term, we decompose 8/.²) differently.
0/ =
Σ
(0z. + ws) 1 (2((0) = hn}
Th
i=1
0al 2, a, 2(0) = 1n) + Σ w;
The
Thh
TETh
Dah
0a +
WT.-
(38)
Th
nh
Then, the error of () can be upper bounded as
[000) 0n|
Nah (0.
0n.)
1 WT.
S
Σ
Nah (0a 0n)
Wr
ni
nn
nhb
n4
By triangle inequality,
Bah (0.
0n)
nab
110. - O S AA
nah
< AAG..
(39)
Thh
agfh
h
azth
This, together with Lemma F.6 and (36), implies
0 0al|= AAG, + top
3(n + d)
S (AG, + 3
A
(40)
for all à E [k). The proof is complete.
27

---

Proof of Lemma F.12. For any # h E [k] x [k]
I 1( 9, ( =n] - 11 1 + wi 0/ 1 S 10, + wi 03
1 1(100, 11 110 1
12 s wi 0%.) - -06 0/ (41)
Triangle inequality implies
110 8/ | 2 [10 8n1l - |10, 8/x²) 1) 2 2 (1 A., 110 -0,11
Using the fact that (1 - x²² - y² 2 (1 ---1 y)² when y(1 - y) 2 0. we obtain
10g / 1| - 110g 0/ (1 (1 2A,) 10g lal 2 e 2110g 1x1
(42)
Denote by An 0 0% for [kt] Then,
I 2 = g, { 1( ) = =h
< 1 0nl1² S 2
(w;, On 0g + An-
-
e² 10g 11 < 2 (w;; On 0g)
< I
+1
Aa
2 (wi; An An)
2
Taking a sum over E Tÿ and using Markov's inequality on the second term, we obtain
(a+1) " ph
2 10g 0nl1 S (woi, On 10-0 -0g)
Σ
16
(w/(An Ag))²
(43)
A
TET
IET
Note that [ 1 11,00, 1a1 S (ws, Gn - 0gg) } are independent Bernoulli random variables. By
Lemma F. 10, the first term in RHS of (43) can be upper bounded by
eA22)
ny exp
+ V 5ng log n.
(44)
320²
By Lemma F.7, the second term in RHS of (43) can be upper bounded by
Σ
16
(w - ,,)) 96(n + d)o²
114, An112.
(45)
E
EAA
:ET
Combining (43), (44) and (45) and using the fact that ||
Dn 12 S 4A7A2, we get
dA²
384(n + g
+ +d)
(x+1) ah
< ng
exp
v 5ng log +
A3
320²
eA2
Consequently,
n1 gh
(5-1)
AA2
5 5log n
384
max k
< l
exp
+ ./
c4,22 A3
(46)
n
320
on
Since A, < 1/2 and r 2 20c 2 the RHS of (46) is smaller that 1/2 when an 2 32k log n. Thus,
(o+1)
(++1)
n
2
1hh
>
14
> an
2
2
for all h E [k] and have
(e+1)
2
A1
5k log n
768
KEIR]
max
(++1)
<
exp
32o²
+
ann
+
c4p2
A3
(47)
9yth
which, together with (46), implics
AA
5k logn n
768 e1,2 A
G S exp
+ log((/ai)
+
320²
ann
Under the assumptions that e' 2/2 > 22 36, we have the desired result (35).
28

---

Proof. From Lemma F.11, necessary condition for Ao 2 4 is Go S ( ) Setting
vr
c= / in Lemma F.12, we have G1 I 0.35. Plugging into Lemma F.11 gives us A - 0.4,
under the assumption that r 2 16v Then it can be easily proved by induction that G, S 0.35 and As < 0.4 for all 1. Consequently, Lemma F.1l yields
3
A,
VKG. + S G
which, combined with (35), implies
C C G5+1 S p2
+ 72 (1 +G, + C? + 5 log
20 p2
3C G,
5k log
a²n
a
for some constant C. Here have chosen
1/5 in Lemma 35 to get the first inequality.
L
F1111 Combining Analysis in lg Norm and Mapping it Back KL divergence.
Proof. From the proof of Lemma F.11, the error of estimating 0, at iteration s can be written as ((a) On
1 Wr; + uhs with
nn
3 (2VKG, + G. AS /
Junll S
(48)
In addition, by Lemma F.1l and Lemma F.12, there is a constant C₁ such that
Ag S 3 VG. + 2G C1 C1 A--1 0.7A,-1+ for all s 2 1. Therefore, when r is large enough, we have
Ciklogg n
1/4
r
a²n
A: llog 1/4
$ C2r + +C2
a²n
for all s 2 logn. n. Then by (42), we have
1 2 = g, 2 2(x+1) h/ s 1113:10,, 0n11 12 < (wi, On 0g + An - -
1/4
where (1 - 2A, 2 2 B1 = 40pr 4C₂ log n an
In order to prove that A, attains convergence rates, first upper bound the expectation of A, 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; An Ag) more carefully. Note that {T [b]] are deterministic sets, we could use concentration equalities to upper bound Wr; and un parts separately,
Let un
1 Wr; for h [R:] and we decompose na
{
9,
2/ 2(0+1) i into three terms.
I I(25 2. 2(0+1) = =n s [ - 12 s (wi, On - 0g)
+11 (32 2 (wi, un ugy
+11 (3, < (w,, un vg) }
where 3g and B4 will be specified later and | 81 82 31. Taking over h [K] and i [m]. we obtain
EAs+1 I E.J1 R.2 FEJa
with
J1
Σ
Σ 1 [31102. - 0,1 < 2 (w;, On 0z.))
(49)
he A
i=1
n2
J2
Σ
22 $ uz }
(50)
d=1
J3 = Σ Σ
mm
(51)
NE 211
29

---

Let us first consider the expectation of J1. Using Chernoff's ' bound, we have
0gll 802
S exp (- 1822)
Pr ( 0nl1 < 2 (w; 8r - =8g) S
exp
802
Thus.
(_8AA2
_AA2
EJ1 S kexp
exp
80²
80²
with 7 _2
S0 log k
32 12---811 8/72.
We use Markov Inequality to upper bound J2. Markov's inequality and Lemma F.7 give us
Σ
132 $ (w,, un -uz)}}
Σ
(u,(un - ug))²
n. A
gelk]
TE
24a
(mg + d)|lun - u2
n A
gE(k]
(48) implies
96o2G. G
960 2 k(nn + kd)
12v Ag.
J2
S
Σ
Σ
nng + d) 1
Ag
n5 A
an,33 A²
helk] sE(k]
Here the second inequality is due to the fact that G, A,/a. And we choose B2 V8k/r in the last equality.
Finally, we upper bound the expectation of J3. Given 2, = 9, we have
Pr
A < 2
2W1 Eig
81
{
31 4
< Pr
A²
(wi 11))
Pr
A
(wiy Up)
4
34
1
34
1
}
< Pr
A
Wis
WT;
+ Pr
2
Wis
WTi
8
ny
8
nip
Choosing t= = m
VdA
02
},
8
exp
A2 4r72
in Lemma F.8, and
64
8
max V do A, A
3d
B1
r 2
Van
an
we obtain Pr (,,,22 < (wi, Ug
mn)} S eexp( -A2/(40²))
where we have used the assumption
that n 2 an and on 2 36. Thus,
(- 1
EJ < 2k
0z
Combining the pieces, we have
EAst1 I
E [1]
+
E []][[E]]
+ E [Ja] + P(E)}
__42
12
1
exp
E.A, + 2k exp
802
with y= = (B1 8k/r - 64/r) - 8/72 = -
1 o(1). Here only prove the case that r - co. For
the finite case, all the o(1) in the following proof can be substituted by a small constant.
1
2 exp (-(1 - 7)
2
2 n
< 2 exp (-(1-)) 7)
A
E.,, <
+
2- r
802
802
n
30

---

when s 2 4logn. By Markov's inequality, for any t > 0,
Pr (A, 2 - BA, - exp (-11- - 77)
A
3
(52)
8
n
If (1 1)) 80 I 2oo choose t = exp
=xxp
(-(1-1)
17
80
A
Ba?
Aa
and we have
Pr A. (
exp (111 -1 A 80
80
A
exp
Otherwise, since A₂ only takes discrete values of (0, n ... 1). choosing t t= 1 in (52) leads to
3
4
Pr (A, > 0) Pr A,
2n expl (eexppp2lloppp 2oog n)
n² n
Now, from Lemma F.3 and Lemma F.4, for Pi[e] (11). vi, we conclude that (L(P.) O - Pylle). Applying the result from above we then get:
Tk 2 CVk
A
2 cv (1
Since n. = !(k2) we have n 2,(1.),, which implies:
na? 2 Ck log
which is the first condition in the EM analysis. The second condition requires the KL to be sandwiched by le norm. But this means that we we can replace À with vÀ and A with /A the result in Eq. 52, and complete the proof of Theorem 3.2.
F22
Privacy Analysis from Section
F.2.1 Technical Lemmas
We begin by recalling the definition differential privacy, and the variant concentrated differential privacy that we use in this work.
Definition F.13 (Differential Privacy (DP) [24]). A randomized algorithm M 2 - y satisfies (e, diffefrerettiall privacy ((e, 8)-DP) if for every pair of neighboring datasets X, X" E Xm (i.e., datasets that differ in exactly one entry),
VY y E Se P 6
When = O, we say that M satisfies e-differential privacy pure differential privacy.
Definition F.14 (Concentrated Differential Privacy (zCDP) [15]). A randomized algorithm M x satisfies p-cCDP if for every pair of neighboring datasets X, X' E X
Va E (1, Da (M()))(M(X') pa,
where Da (XX))) is the a-Rényi divergence between M(X) and M(').
Note that zCDP and DP are on different scales, but otherwise can be ordered from most-to-least restrictive, Specifically, (e, DD- implies --ZCDP, which implies roughly (e (e log(1/)) 6D- for every d [15].
Both these definitions are closed under post-processing and can be composed with graceful degrada- tion of the privacy parameters.
Lemma F.15 (Post Processing [15]). If M : X" + y is (),6)-DP. and P y + Z is any randomized function, then the algorithm P M is (c,6)-DP. Similarty if M is p-zCDP then the algorithm P o M is p-rCDP.
Given probability distributions over Q. (all)) - log [, (6))
31

---

Lemma F.16 (Composition of CDP [15]). If M is an adaptive composition of differentially private algorithms M...... Mr, then
(eT,,rr) (rr- then M is E. and 1. if M1, Mr are (er.
2. if Mi,.... Mr are P P.r....P-CCCPP then M is (EP:t)
We can achieve differential privacy via noise addition proportional to sensitivity [24].
Definition F.17 (Sensitivity). Let f X" + F be function, its fz-sensitivity is defined to be A, maxx~x'ex" (X f(x')l2, Here, X ~ X denotes that X and X are neighboring datasets (i.e., those that differ exactly one entry).
For functions with bounded & --sensitivi,,, we can achieve c-DP by adding noise from Laplace distribution proportional to -ssesiittivy For functions taking values in R for large more useful to add noise from Gaussian distribution proportional to the fz-sensitivity, to get (c, 0DD and p-ZCDP.
Lemma F.18 (Gaussian Mechanism). Let f xm - IR be a function with ez-seiiitiv A 1.2. Then the Gaussian mechanism
- +N (0.( (Az2)
-axd
22p
satisfies p-zCDP.
F.2.2 Proof of Theorem 3.3
From the Lemmas above for Gaussian and Laplace mechanism, is easy see that the adaptive clipping step in Algorithm 4 satisfies 2p/3-7CDP, where as the Laplace mechanism is p/3-zCDP. Then, we can use composition lemma to complete show p-7CDP of Algorithm 4.
For Algorithm 5, each step uses satisfies p/ K-ZCDP from the guarantees for exponential mechanism, since we clip the KL difference with c. Here too, composition gives us the result.
For 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 p/3T D/3P-ZCDP. Finally, composing over the T rounds clustering and with initial private center and private init get the final result. We then use Lemma 3.6 from [15] to convert zCDP guarantees to (E, 6)-DP.
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Claims from the introduction and abstract are reflected in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Appendix A
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
33

---

Theorem statements are in the main paper and proofs in Appendix E and Appendix
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Experimental details are in Section 5 and Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
34

---

Answer: [No]
Justification: While the datasets we use are all open-source with publicly available links, code not yet ready for anonymous open sourcing. We plan open-source the code with appropriate licensing for the updated version of the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.cm cc/
public/gue Coleee/bddSSSubmsssc..
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results, See the NeurIPS code and data submission guidelines (https:
guggguieee deuu/de///e Codeeubmiis asison caisoooniccyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details are in Section 5 and Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All our results report error bars with 95% confidence intervals.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
35

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
In Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https htt////uuriip.pp/ cc/ppblic/
Answer: [Yes]
Justification: Our work is studies private and personalized frequency estimation and does not violate any Code of Ethics
Guidelines:
The answer NA means that the authors have 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., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
discuss some limitations of 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.
Justification: In Appendix
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 why the paper does not address societal impact.
36

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work is not focused on releasing any pretrained models datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We make sure 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 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 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.
37

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 not available online, the authors are encouraged to reach out to the asset's s creators.
13. 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/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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.
14. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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.
15. 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 equivalent approval/review based the requirements of your country 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 research with human subjects.
Depending 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.
38

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
39


Reviews:
[{"review_id": "O0fO4Omr9R", "reviewer": "Reviewer_ymNE", "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.", "questions": "What is the relationship of the privacy parameter rho compared to standard parameters epsilon and delta in differential privacy?\n\nPlease see the main points raised in the weaknesses section.", "limitations": "The authors have a specific section discussing the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is well-written and the setting considered in the paper is clearly introduced. The challenges related to frequency estimation in the federated learning setting as well as the related work are also reviewed in a thorough manner.\n\nOne of the novelty of the proposed approach is the idea of clustering users first before performing the frequency estimation. The authors that the proposed approach offers strong theoretical guarantees and does not suffer from the same limitations as other methods from the state-of-the-art, such as the sensitivity to small datasets.\n\nBased on the proposed approach, the authors have developed several variants of the personalized frequency estimation histograms, including for the private and no-private settings as well as for the size-heterogenous one. The private version of the algorithm combines several clever privacy-preserving subroutines in a clever manner.\n\nThe experiments conducted are quite extensive as they have been done on three datasets and have compared the approach to three different baselines.", "weaknesses": "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.\n\nThe 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. \n\nA 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. \n\nIn 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.\n\nA few typos :\n-« Baysian prior » -> « Bayesian prior »\n-« privacy preserving analysis » -> « privacy-preserving analysis »\n-« estimatation in KL divergence » -> « estimation in KL divergence »\n-« For competetive estimator » -> « For competitive estimator »\n-« LLoyd’s K-means » -> « Llloyd’s K-means »\n-« differentially private » -> « differentially-private »\n-« Privacy preserving noisy sum » -«  « Privacy-preserving noisy sum »", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and a positive assessment of our work. To address concerns, we refer to plots in the paper that ablate $K$, clarify guarantees in Thm. 5.2 are in terms of the typical $(\\varepsilon, \\delta)$ approx. DP, and discuss the choice of the stricter user-level DP definition. We also explain the motivation to study next-word prediction, and the choice of the stylized model: mixture of Dirichlets. Thank you for pointing us to typos, we will correct them, add discussion on the above points and a conclusion section in the extra page we get for the final. **Please let us know if your concerns are addressed, and if so, we would be grateful if you might raise your score.** \n\n>> **Choosing the number of clusters K**\n\nIn Fig 3a, 5 in our paper, we plot the test NLL when varying the number of clusters $K$ parameter for our clustering Alg 3. **Please see the global response for more discussion on this**.\n\n>> **Converting cJDP to user-level DP**\n\nOur final guarantees in Thm 5.2 are in terms of user-level $(\\varepsilon, \\delta)$-DP which is equivalent to JDP in our setting. We only use the $\\rho$-cJDP (stricter than RDP) in Def 3.2 for easier and tighter composition (than approx. DP), and convert them to user-level approx. DP finally.\n\nOur setting is akin to the billboard model (Hsu et al. [41]) where any $(\\varepsilon, \\delta)$ user-level DP (Dwork [25]) algorithm is also $(\\varepsilon, \\delta)$-JDP. Similarly, if the algorithm satisfies user-level $\\rho$-zCDP (Bun and Steinke [16]) then, it is also ρ-zCJDP (Def 3.2). Thus, we first use zero-concentrated DP (zCDP) for privacy accounting and composition, and then use Lemma 3.6 from [16] to convert the zCDP guarantees to approx. DP in Thm. 5.2. \n\n>> **User vs. record level DP**\n\nWe defined our problem in Sec. 3 with the goal of achieving personalized predictions for each user’s frequency counts, while satisfying user-level privacy. As you correctly point out, this is indeed stricter than item-level privacy, but it is also a common choice for privacy guarantees in cross-user federated learning (e.g., Cummings et al. [22], Levy et al. [52], Wang et al. [71]). In cross-user FL, e.g. learning to predict the next word on users’ mobile devices, it is natural to protect the participation of every user; as a user may type many different words that are correlated, hiding a single word may not be sufficient privacy protection. Indeed if we had the weaker item level DP protection with $\\varepsilon = 2$ (say), the overall user-level privacy loss for a user that contributed $m=200$ words would theoretically be $\\varepsilon \\cdot m = 400$. Even if 10 of the words are correlated with sensitive information (e.g. related to their medical diagnosis), the $\\varepsilon$ parameter increases by that factor. Our user-level guarantee ensures that even if the user were to change all 400 words they contributed, the models will not be significantly different.\n\n>> **Using the mixture of Dirichlets model for our analysis in Sec 4**\n\nIn Sec. 4, we introduce the mixture of Dirichlets model with the goal of guiding algorithm design when the underlying population is heterogeneous, yet is structured. While each user has a different token distribution it is likely that there exists clusters of ``similar’’ users. This model motivates a setup where adapting a single global model for each user would be less effective than identifying clusters, learning a single model for the cluster and then adapting the cluster-level model for the user. We also observe this empirically in our experiments in Sec 6, where FedAvg+FT performs worse than our approach. \n\nNote that the specific choice of Dirichlet for analysis is only so that we can easily derive the optimal Bayes estimator (Thm 4.2), which gives us our finetuning algorithm in Eq 2. For all other theoretical results in Sec. 4, 5 we only need each user’s true distribution to be concentrated along the cluster center, i.e., w.h.p., for all users $u$ in cluster $c$, and any token $v$, $|P_c[v] - Q_u[v]| = O(1)$. The $\\alpha$ parameter of Dirichlet further allows us to analyze our setting with varying degrees of cluster-level concentration, since under the Dirichlet assumption $|P_c[v] - Q_u[v]| = O(1/\\sqrt{\\alpha})$.  \n\n>> **Why focus only on the next-word prediction problem in the FL setting? Extending to other datasets and losses.**\n\nLarge language models are next-token prediction models that can be used to solve tasks in natural language like math reasoning, coding etc. In FL, the performance of these next-token prediction models can be improved fairly across users when they are personalized (Hard et al. [38], Salemi et al. [66]). At the same time, since each user has few data points for purely local learning, global collaboration is needed to bias the local learning algorithm. Thus, we study frequency estimation (estimating the marginal distribution) as a first step towards private and personalized language models. We highlight how measuring the error in KL divergence and reducing privacy-utility trade offs requires algorithmic innovations (e.g., clustering with Good-Turing, data-dependent private initialization) that are specific to our problem and much required. \n\nWe provide an extended analysis on three real-world datasets with varying levels of statistical heterogeneity, each of which are a popular choice for analyzing next-token predictors in FL (Wu et al. [76]). Please let us know if there is a specific dataset that is characteristically different from these, and we would be happy to include experiments on it.\n\nAs we mention in Sec 1, we focus on error measured in KL as is common in language modeling (Hoffmann et al. [40]), and is equivalent to minimizing the negative log-likelihood of the sample. Distribution estimation in KL distance is also studied more generally (Drukh et al. [24], Acharya et al. [2]). We believe that extending KL to other losses is an interesting direction of future work, as opposed to a weakness.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PfxXDutOsk", "reviewer": "Reviewer_zabr", "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.", "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'?\n- Since the authors defined their own DP relaxation definition, it will be helpful if the authors could clarify:\n    - 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?\n    - 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": "see above", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The paper empirically demonstrates that the proposed method have better performance than the baselines. \n- The authors also conducted thorough ablation experiments to explain the importance of various sub-parts of the algorithm.", "weaknesses": "- The largest concern to me is the presentation clarity of the paper. \n    - 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. \n    - 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.\n    - 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback! To address concerns we provide an algorithm’s box for the end-to-end algorithm, clarify that assumptions in Sec 4 are only for theoretical motivation, and results in Sec 6 are on real datasets where they needn't hold. We clarify our DP definition is not relaxed at all, in fact it is more stringent than RDP and the final guarantees in Thm. 5.2 are for the approx DP definition. We will use the 1 extra page in the final to incorporate them. **Please let us know if your concerns are addressed, and if so, we would be grateful if you might raise your score.** \n\n>> **Final alg in Sec 5**\n\nIn the 1 page PDF (Fig 1), we show an algorithm box for the end-to-end algorithm that uses Alg 5 for private initialization of cluster centers, and Alg 3 for clustering, where the re-centering step uses Alg 4. Each user finetunes their corresponding cluster center using Eq. 2.  \n\n>> **Assumptions in stylized model and concentration of user histograms**\n\n- In Sec 4, to guide algorithm design we introduce a model where the user distributions are distributed as a mixture of Dirichlets. This model has two latent variables: cluster memberships and cluster centers. For simplicity, we analyze estimators under increasing levels of knowledge in this model: cluster center known -> cluster center unknown but cluster membership known -> both cluster center and membership unknown. \n- The simpler analysis in a clean model leads us to three key findings (end of Sec 4) using which we propose the clustering algorithm (Alg. 3), with the recentering step using Good-Turing estimates (Eq. 4), and finetuning given by Eq. 2. \n- In Sec 6, we evaluate our proposed algorithm in the most general setting, where cluster centers/memberships are unknown. Here, we also do not assume anything about the distribution of $Q_u$’s. \n- Regarding the **well-concentrated** assumption: We only need that w.h.p., for all users $u$ in cluster $c$, and any token $v$, $|P_c[v] - Q_u[v]| = O(1)$. This is naturally satisfied by our Dirichlet assumption since the $\\alpha$ parameter of Dirichlet allows us to analyze our setting with varying degrees of cluster-level concentration where $|P_c[v] - Q_u[v]| = O(1/\\sqrt{\\alpha})$.  \n\n\n>> **Novelty**\n\nThe novelty of our Algorithm lies beyond the proposed K-means style clustering method. Following is a set of novel contributions we make, where the final four claims are verified empirically in Sec 6 and defined, analyzed formally in Sec 5.:\n- Introduce the personalized and private frequency estimation problem in KL divergence.\n- Contrary to popular beliefs (Cheng et al. [20], Wu et al. [76]) in the FL community, we show that there exists practical FL problems where finetuning a single global model (FedAvg+FT) is not sufficient for personalization. \n- Specifically to improve clustering error in terms of KL divergence (as opposed to squared loss in k-means) we motivate and propose a Good-Turing based center estimator (Sec 4).\n- Propose a private version of the clustering algorithm with approximate differential privacy guarantees (Thm. 5.2).\n- Propose a novel data-dependent and yet private initialization for clustering (Alg. 5), which is crucial to improve clustering performance (Fig. 3b). Prior works on private-clustering (e.g., Chang et al.: Locally private k-means) do not analyze expectation maximization (EM) algorithms given their dependence on initialization which is harder to privatize (most initializers for EM directly output a set of points from the dataset (Arthur et al. [5])).\n- Propose a practical 2-stage algorithm for the size-heterogeneous setting (Section 5). Naive algorithms in this setting typically hurt the privacy-utility trade off since hiding the participation of data-rich users now becomes very hard when other users have less data. We show that our proposed algorithm is able to improve estimation error from including both data-poor and data-rich users in the collaboration, without hurting the privacy utility trade off (Fig 4).\n\n\n>> **Our final guarantees are in terms of approx DP and Def 3.2 is stricter than RDP**\n\nOur final guarantees in Thm 5.2 are in terms of user-level $(\\varepsilon, \\delta)$-DP which is equivalent to JDP in our setting. We only use the $\\rho$-cJDP (stricter than RDP) in Def 3.2 for easier and tighter composition (than approx. DP), and provide final guarantees in terms of user-level approx. DP.\n\n**In billboard model, JDP=user-level DP**: Note that our setting is akin to the billboard model (Hsu et al. [41]) where any $(\\varepsilon, \\delta)$ user-level DP (Dwork [25]) algorithm is also $(\\varepsilon, \\delta)$-JDP. Similarly, if the algorithm satisfies user-level $\\rho$-zCDP (Bun et al. [16]) then, it is also ρ-zCJDP (Def 3.2). \n\n**Final guarantees are for user-level DP**: We use zCDP for composition at the user-level, and then use Lemma 3.6 from [16] to convert zCDP to approx. DP in Thm. 5.2. \n\n**Why use zCDP?**: Yes, zDP bounds the $\\alpha$ Renyi divergence between the mechanism outputs on neighboring datasets by $\\rho\\alpha$. Thus, zCDP is a stricter notion than Renyi DP (RDP) since zCDP requires Renyi divergence of multiple orders to be simultaneously bounded. Further, zCDP enjoys the same benefits as RDP for composition, since zCDP’s sequential composition is also asymptotically tight for repeated applications of the Gaussian mechanism [16].\n\n>> **[New Experiment] Evaluating performance at different privacy budgets** \n\nIn Fig 2 (1 page PDF) we evaluate our algorithm and other baselines under different privacy budgets. Each algorithm needs to satisfy ($\\varepsilon, \\delta$)-JDP (which is ($\\varepsilon, \\delta$) user-level DP). On the x-axis we vary $\\varepsilon$ and fix the same $\\delta$ reported in the paper. We find that our approach continues to outperform baselines. While the clustering baseline IFCA and MAML suffer from severe performance degradation under low privacy budgets, our approach observes a much better privacy-utility trade off.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lw8wBlUFUw", "reviewer": "Reviewer_DFyV", "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.", "questions": "See weaknesses.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-Private and personalized histogram estimation is an important problem.\n\n-The paper is easy to understand, presentation and justification of the methods are well organized.\n\n-The experiments are conducted on multiple datasets, the method consistently outperforms others.", "weaknesses": "- Lack of comparison to local only training and other personalized estimation methods.\n\n- 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).\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and a positive assessment of our work. To address your concerns, we present new ablations of our algorithms under different values of per-client training data points $m$, local only training baselines, highlight existing ablations in our paper on the number of clusters, and address the point of analyzing a simpler setting. We will add the new results and discussion on the above points in the extra page we get for the final version. **Please let us know if your concerns are addressed, and if so, we would be grateful if you might raise your score.** \n\n\n>> **[New experiment] Lack of comparison to local only training and other personalized estimation methods.**\n\nIn Fig 5 (1 page PDF), we find that the local only training performance where the clients never collaborate is worse than our approach (for a privacy budget of $\\epsilon=15$), even as we increase the number of local data points $m$. Further, comparing this result to Fig 2 (1 page PDF), we find this trend to hold even under a very low privacy budget of $\\varepsilon=2$.\n\n\n>> **The main results are rather in simpler settings, the analysis mostly builds on top of previous work.**\n\nWhile our analysis technique adapts results on clustering and Good-Turing estimation to our setting, its main goal is to serve as a guide that formally validates our algorithmic choices, for e.g. the local finetuning algorithm (in Eq. 2), or using the average of local Good-Turing estimates for estimating the cluster center, as opposed to the typical average of empirical estimates that is used in the K-means algorithm (see our comparison in Thm. 4.3). In the FL and privacy literature, it is also common for algorithms to be analyzed under such models of data heterogeneity (e.g., Cummings et al. [22], Cheng et al. [20]) that mimic real world distributions. \n\n>> **[New experiment] From the current experiments, performance dependence on variables such as number of clusters, clients, data samples per clients, heterogeneity etc. is not clear.**\n\n**Number of clusters**: In Fig. 3a (Sec 6) and Fig. 5 (App B) we show how the choice of $K$ (number of clusters) affects the final performance of our algorithm. As expected, the performance improves as $K$ roughly matches the natural heterogeneity/clusters present in the data distribution. In practice, we can cross-validate to find $K$ and use $K=10$ (close to the optimal choice for all datasets), for all our runs. \n\n**Number of clients and privacy budget**:In the attached PDF, we show performance evaluations under varying values of $m$ and fixed privacy budget of $\\epsilon=15$ (Fig 3), and varying $\\epsilon$ for the values of $m$ in our paper (Fig 2). We note that while the performance of each method drops with reducing $\\epsilon$ and $m$, our algorithm has better privacy-utility tradeoffs and is more sample efficient \n \n**Heterogeneity**: Given that the natural datasets we consider already have varying levels of heterogeneity (see discussion in global response on a different choice of $K$ being optimal on each dataset) and allow us to investigate the key phenomenons, we see no compelling reason to empirically study synthetic data. We are also not aware of a sufficiently faithful model of data for our problem (and even for our own data model there are many parameters, for which it is unclear how to choose synthetic values). At the same time we welcome and will certainly consider specific suggestions about synthetic data experiments that would make our work more informative.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2osVGGCF9f", "reviewer": "Reviewer_a5Ce", "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.", "questions": "1. How does the choice of $K$ affect the performance? Do the authors have a suggestion on how to choose $K$ in practice?\n2. 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?\n3. 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": "The authors addressed the limitations and negative impact of their work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The problem of data heterogeneity is very important in federated learning.\n2. The proposed algorithms are practical and can be applied to large-scale datasets.\n3. Empirical results demonstrate improvement over FedAvg and MAML baselines.", "weaknesses": "1. The cluster assumption requires choosing the cluster number $K$.\n2. 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.\n3. 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$.\n4. 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.\n5. In the experiments, the authors only show results for $\\varepsilon=15$.\n6. Relevant papers to cite and discuss:\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n=====Update=====\n\nIncreased my score to 5 after the rebuttal", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback. To address concerns, we add new experiments with different values of $\\varepsilon$, number of local points $m$, and plots showing the sparsity of the estimated cluster centers in practice. We also point to our ablations on $K$ in our paper, and clarify concerns with Thm. 5.1, 5.2. We thank you for the relevant citations, and we will use the 1 extra page in the final version to incorporate the new experiments, clarifications, and citations. **Please let us know if your concerns are addressed, and if so, we would be grateful if you might raise your score.** \n\n>> **Different values of $\\varepsilon$ and local data points $m$**\n\nWe show performance evaluations under varying values of $m$, and fixed privacy budget of $\\varepsilon=15$ (Fig 3 in PDF), and varying $\\varepsilon$ for the values of $m$ in our paper (Fig 2 in PDF). We note that while the performance of each method drops with reducing $\\epsilon$ and $m$, our algorithm has better privacy-utility tradeoffs and is more sample efficient. We also note that these privacy parameters are for a relatively small dataset (10k to 50k users). If the population size was, say, 10x larger, the privacy parameters would be 10x smaller for the same kind of results. In practical FL settings, we would often expect the population sizes to be much larger and thus the epsilons to be proportionately smaller.\n\n\n>> **Choice of number of clusters $K$**\n\nIn Fig. 3a (Sec 6) and Fig. 5 (App B) we show how the choice of $K$ (number of clusters) affects the final performance of our algorithm. As expected, the performance improves as $K$ roughly matches the natural heterogeneity/clusters present in the data distribution. In practice, we can use hold out to validate $K$. See our global response for more discussion on this.\n\n>> **Privacy accounting via RDP**\n\nWe do not use composition results for approximate DP. All our composition results rely on zero concentrated DP (zCDP) (Bun and Steinke [16]) for privacy accounting of user-level DP. We then convert the zCDP guarantees to approximate DP in Thm. 5.2. Note that zCDP is a stricter notion of privacy than RDP since it requires Renyi divergence of multiple orders to be simultaneously bounded. Further, it enjoys the same benefits as RDP for composition, since zCDP’s sequential composition is also asymptotically tight for repeated applications of the Gaussian mechanism [16]. \n\n>> **[New Experiment] Sparsity of $P_c[v]$ in practice**\n\nIn practice, the estimated cluster centers are sparse, as we see on Reddit and StackOverflow, where about 10% of the tokens carry 90% of the probability mass (Fig 4 in 1-page PDF). The assumption we make in Thm. 5.1 on P_c[v] being lower bounded by a constant is needed solely for the theoretical convergence guarantees for clustering. This is needed in theory because KL divergence does not obey triangle inequality, but $l_2^2$ distance does, and for KL between estimated to be roughly tracked by $l_2^2$, we need $P_c[v] = \\Omega(1)$. \n\n>> **Dependence on $d$ in the KL separation condition in Thm. 5.1**\n\nThere is a typo in the condition on the minimum cluster separation $\\Delta$. The correct condition is $\\Delta = \\Omega(k^2 + \\frac{k^3 d}{n})$, as derived in the proof of Thm. 5.1 in  App. D.1 (below L831). We apologize for this typo and will correct it in the final version. Note that the number of clients $n$ appears in the separation condition on $\\Delta$, and the convergence is determined by the condition number. Thus, the condition on KL divergence only scales as $poly(k)$ when $n=\\Theta(d)$. For the datasets we use, the number of users is sufficiently large ($d/n < 2$ for all). \n\n>> **$\\log(1/\\delta)$ term in Thm. 5.2**\n\nYes, for the privacy guarantee to be meaningful $\\delta$ needs to be $o(1/n)$. But, note that the guarantee in Thm. 5.2 is $\\sqrt{\\log(1/\\delta)}$, which for the datasets of our size (around $10^5$), translates to $\\approx \\sqrt{5}$. \n\nTo elaborate more on the term, the result in Thm. 5.2 is obtained by first composing the privacy mechanisms in each clustering iteration with zero-concentrated DP (zCDP) and then converting the zCDP guarantees into approx. DP using Lemma 3.6 in Bun & Steinke [16]. The $\\log (1/\\delta)$ term is an artifact of this conversion. On the other hand, if we were to instead compute approx. DP guarantees (instead of zCDP) at each clustering iteration and then use the advanced composition theorem to compose them, we would still end with a $\\sqrt{\\log(1/\\delta)}$ term from the composition. So, this term is unavoidable either way. \n\n\n>> **L256 Without any assumptions on the relationship … the best one can do is locally estimate**\n\nConsider the following worst-case setting where there is no mutual information between the training datapoints sampled by any pair of users. Let each user’s true token distribution $Q_u$ be supported over subsets of tokens that are disjoint for any pair of users. Further, for any user $u$ and token $v$ with $Q_u[v]>0$, the value of $Q_u[v]>0$ is independent of the token distribution of other users. In such a case, no estimator can improve the average estimation error through collaboration/sharing of data. The optimal estimator is purely local. We will add this example to the paper.\n\n>> **Missing citations**\n\nThanks for pointing out these references, we will cite & discuss them. We believe our contributions go beyond these: for e.g., while [2, 3, 4] look at private histogram estimation, they concern with algorithms that estimate a single global quantity (e.g., mean histogram) whereas ours deals with the nuances of privatizing personalization algorithms that estimate a separate distribution for each user, and learn multiple global centers. Further, [1, 4] conduct their analysis in the $l_2$ metric, whereas we look into the more appropriate KL distance for distribution estimation. KL is harder to analyze and develop algorithms for, given that it does not even obey the triangle inequality.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their feedback and list some of the new experiments we add to the 1-page PDF which we will incorporate using the extra page in the final version. We also address a common concern on the choice of number of clusters $K$ for our Algorithm 3.

___
## **List of new experiments in 1-page PDF.**
- Figure 1: Outline of the end-to-end algorithm that first privately initializes cluster centers and then runs private clustering. Finally, the learned cluster centers are finetuned by each user locally. 
- Figure 2: We evaluate our method and baselines under different privacy budgets (varying $\varepsilon$), where each algorithm needs to satisfy $(\varepsilon,10^{-10})$-JDP. We find that our approach continues to improve over baselines even at low privacy budgets and has better privacy-utility tradeoff.
- Figure 3: We evaluate our method and baselines under different values of the number of local datapoints $m$ and find that our approach is more sample efficient than baselines.
- Figure 4: We plot the CDF plot for the distribution over tokens given by the cluster centers estimated on Reddit and StackOverflow. The plots suggest that 90% probability mass is present in 10% of the tokens indicating that the underlying cluster centers are indeed sparse. But note that the same is not true for the FedAvg estimate indicating that the global center is denser. This means that the discovered underlying clusters are sparse and diverse.
- Figure 5: We evaluate our method at $\varepsilon=15, \delta=10^{-10}$ privacy budget, compared with perfectly private local only training when local dataset size is varied. Even at higher values of $m$ our algorithm is able to learn a more accurate estimate by first collaborating with users to privately learn cluster centers and then adapting them locally for each user using their local datasets only.

___
## **Choosing the number of clusters $K$ in practice.**

In Fig 3a and Fig 5 in our original submission we plot the performance of our end-to-end algorithm as we vary the choice of number of clusters $K$ used in the private initialization (Algorithm 5) and histogram clustering (Algorithm 3). We use a hold out validation set on the Reddit dataset to identify the number of clusters $K=10$ to be the optimal choice (Fig 5). We use the same choice for the other two datasets: Amazon Reviews and StackOverflow and a post hoc reveals that $K=10$ is close to optimal for those as well (Fig 3a). In fact we note that the choice of number of clusters K is not very sensitive close to the optimal. Furthermore, we also note the varying levels of heterogeneity in the different datasets we evaluate as $K=8$ and $K=16$ does better for StackOverflow and Amazon Reviews respectively. In practice, we show that one can use a hold-out validation set to more accurately judge the choice of this hyperparameter (Fig 5).

Meta-review: 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.
Decision: Accept (poster)