Title: Private and Personalized Frequency Estimation in a Federated Setting
Venue: NeurIPS 2024
Accepted: yes
Keywords: federated learning

personalization

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

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 datadependent 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.

1 Introduction

Federated learning algorithms jointly learn from decentralized user data, addressing statistical challenges in local learning [49, 58]. At the same time it presents two key challenges, among few others [44, 52]. Firstly, data distributions can vary across users (often called statistical heterogeneity [27, 76]), which reduces the effectiveness of a single, shared model for all users. Secondly, the data may be sensitive and mutually beneficial collaboration may compromise user privacy [31, 71]. An important practical problem where user data distributions are diverse, and user-privacy is paramount is next-word prediction for the keyboard input on user devices. Users often have diverse vocabulary, writing styles, and topics that leads to varied data distributions [32, 36, 65]. Naturally, frequency estimation is also one of the most basic statistical tasks with numerous other applications [14, 47].

Motivated by personalized next-word prediction, we introduce and study the problem of personalized frequency estimation in a federated setting . In this problem, each user has very few ( O ( d ) ) samples from an unknown, user-specific distribution over a large, finite domain of size d . The users interact with a server that needs to provide the user with a personalized estimate of their distribution. We focus on error measured in KL divergence that is common in language modeling [3, 38], and is equivalent to minimizing the negative log-likelihood of samples from the user distribution. For more general loss minimization problems, numerous personalized federated learning

(PFL) algorithms [67, 76], tackle statistical heterogeneity by first learning a single global model across all users (for e.g. , FedAvg [57]), and then finetuning (FT) the global model locally for each user (FedAvg+FT [18, 20]), thus balancing collaboration and personalization. The downside of this approach is that it is agnostic to the structure of the user population which is highly heterogeneous, but often consists of a number of concentrated subpopulations of similar users.

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-Turing estimators, we propose to estimate global cluster centers by averaging user-level Good-Turing estimates for the current users in the cluster. Furthermore, since the performance of clustering is significantly affected by the initialization [5] of cluster centers, we also propose a data-dependent initialization approach, also specific to estimation in KL. We improve the practicality of our algorithm by separately handling data-poor users in cases where users vary significantly in their local data sizes (size-heteogeneity).

Privacy: To address the challenge of user-level privacy we give a joint differentially private (DP) version of our iterative algorithm which requires that the estimator for user i is differentially private with respect to the data of all the other users (but may depend arbitrarility on user i 's data) [48]. This definition is necessary due to the final personalization step which happens locally on user device and thus does not present a privacy risk. Specifically, we make each iteration of our clustering algorithm provably private by relying on adaptive clipping, secure aggregation, and common noise addition mechanisms [23]. Initialization of clustering algorithms typically uses data points themselves as initial centers and hence presents a significant challenge for privacy preserving analysis. Our provably private data-dependent initialization algorithm for cluster centers runs the exponential mechanism over candidates randomly sampled around the estimated population mean.

Empirical evaluation: We validate both non-private and private versions of our algorithm on the real world data datasets: Reddit [16], StackOverflow [6], and Amazon Reviews [62], where Reddit has a token vocabulary of size 10k and the others have 32k tokens. We find that our method furnishes significant gains over standard and clustering baselines, reducing the error by over 26% and 42% (averaged over datasets) in the non-private and private settings respectively. In the private, sizeheterogeneous setting, we improve over the clustering baseline IFCA [33] by over 30% . Finally, we show that our data model inspired adjustments are pivotal in yielding the performance improvements in practice. Particularly, we justify the use of our Good-turing based estimator, and our private initialization scheme with a favorable privacy/utility tradeoff.

Formal guarantees. While our focus is on the empirical performance, to guide intuition and design of algorithms, we introduce a relatively simple generative model where user distributions are sampled from a mixture of Dirichlets, where Dirichlets are sufficiently separated globally and concentrated locally. We derive and analyze different estimators ( e.g. , Bayes optimal, FedAvg [57], Good-Turing [34]). Even when the cluster identities are known, we show that FedAvg (average of user empirical distributions) has poor guarantees compared to our proposed estimator (average of user Good-Turing estimates), where the minimax error rates for the latter do not scale with dimension, and is optimal for some regimes of problem parameters. Thus, using this model we demonstrate how additional structure in the user data can be exploited in the context of our more concrete frequency estimation problem while preserving the privacy of user data.

2 Problem definition

We use P to denote a distribution, P [ v ] for the probability of event v , ˜ P for an estimate of P (but ˆ P if the estimate is computed without user collaboration), P for sets. P for matrices ( P : ,i / P i, : index into the i th column/row respectively), and { P i } for a collection of elements P i indexed by i .

Setup. Let V be a finite vocabulary (set) of tokens, with vocabulary size d =: |V| , and ∆( V ) is the set of all probability distributions over V . In a federated setup we have a collection of users U , and each user u ∈ U has an unknown distribution Q u ∈ ∆( V ) over the tokens. The data available to each user u is a dataset S u comprising of m u i.i.d. samples from Q u . We will use ̂ Q u to refer to the empirical distribution estimated from S u . The combined collection of datasets is referred to as meta-dataset S =: {S u } , and similarly ˆ Q =: { ̂ Q u } u is the collection of all empirical distributions.

Unless specified otherwise, we assume that we are in the size-homogeneous setting where all users have a dataset of the same size m = m u . We refer to the more general case as size-heterogeneous .

Goal. Given the meta-dataset S , the goal is to compute estimates { ˜ Q u } u such that:

utility : ∑ u E [KL( Q u ∥ ˜ Q u )] is minimized, where the expectation is taken over sampling of meta-dataset S and any randomness in computing { ˜ Q u } u .

privacy : ˜ Q u is computed privately with respect to all users except u , i.e. the estimate is joint differentially private (see Definition 2.2).

Definition 2.1 ( u -neighboring meta-datasets) . Meta-datasets S and S ′ are u -neighbors, i.e. , S ∼ u S ′ if they differ only in inclusion or deletion of user u 's private data. For any algorithm M , we denote M -u as the output of the algorithm to all other users except u .

Definition 2.2 (Joint Differential Privacy [48]) . Given ε ≥ 0 , δ ∈ (0 , 1] , and neighbouring relation ∼ u , a randomized mechanism M : D → Y from the set of meta-datasets D to an output space Y is ( ϵ, δ ) -joint differentially private if for all u -neighboring S ∼ u S ′ , and all events E ⊆ Y ,

where probabilities are taken over the randomness of M . We say that M is ρ -zero-concentrated joint DP ( ρ -zCJDP) if for all u -neighboring meta-datasets S ∼ u S ′ :

where D α ( M -u ( S ) ∥ M -u ( S ′ )) is the α -Rényi Divergence between the outputs of all users except u .

3 Private algorithm for learning personalized histograms

In this section, we present our algorithm for learning private and personalized histograms in both size-homogeneous and size-heterogeneous settings. Without assumptions on the distribution of users, in the worst case, the best estimator would simply estimate each user's histogram locally, i.e. , estimate each user's Q u using their local dataset S u . Avoiding such worst-cases, we assume that the real-world problem is more structured. Specifically, we assume that each user distribution Q u is well clustered around K unknown cluster centers in the population of users. Note, we assume no knowledge of the cluster centers or number of clusters K . In Section 5, we empirically validate the presence of clusters in real-world datasets by evaluating the performance of our clustering based algorithm on the same datasets, and in Section 4 we theoretically analyze some of the algorithmic design choices we make, in a stylized Bayesian setting that simulates a user-distribution with clusters.

As a warmup, we first discuss our local finetuning algorithm, where the setup is simpler. Say an algorithm runs a private collaborative protocol, and hands a distribution P to user u , with the guarantee that KL( Q u ∥ P ) is 'reasonably' small, then how should the user u finetune P using the local dataset S u available to them? For this question, we now present our local finetuning algorithm.

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 a specific cluster. In the next section, we discuss how to estimate P . Here, we discuss how to personalize P for each user. For a user u , with empirical distribution ̂ Q u from m u samples, we use FT( P, ̂ Q u ) to denote the personalized, locally finetuned estimate of user distribution Q u (see Eq. 3). Note that P is moved closer to ̂ Q u in KL divergence (depending on choice of α > 0 ). When the population of users is well concentrated, FT lowers variance by relying more on the global estimate P , while adding some bias ( ̂ Q u is an unbiased estimate of Q u , but P is biased). In Section 4, we show that FT is the Bayes optimal estimator in KL divergence, under certain assumptions on the user distribution.

Next, we discuss how to collaboratively estimate P . For a user-distribution with multiple clusters, a natural approach is to first identify the groups of clustered users. Then, we can use the distribution that is close in KL distance, and in expectation over the distribution of users in the cluser. We refer to these as the cluster centers. In the next Section, we present our algorithm to learn cluster centers, which we personalize for each user in the cluster using our local finetuning update (Equation 3).

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 { ˜ z u } u , and the matrix of K cluster centers ˜ P ∈ R d × K . To learn cluster centers ˜ P , akin to the popular clustering algorithms like Lloyds [53], the objective we optimize is:

with the main difference being that we are concerned with the KL distance, instead of ℓ 2 2 . Given ˜ P , user u 's assigned cluster ˜ z u is simply given by arg min k KL( Q u ∥ P : ,k ) . Conversely, applying Lemma 3.1, given cluster memberships the cluster center is determined by the cluster average. Lemma 3.1. For users with assigned cluster i , i.e., ˜ Z i =: { u : ˜ z u = i } , the following is true:

Since, we do not have access to Q u , we replace Q u in Lemma 3.1 and the cluster assignment step with the local Good-Turing [34] estimate ̂ Q gt u . Typically, one would use the empirical distribution ̂ Q u , which is an unbiased estimate of Q u . But, as we will see in Section 5, most real-world distributions are long-tailed and the Good-Turing estimate's approximation of the user distribution does not suffer as heavily from the high-dimensional nature of the problem. We discuss the Good-Turing estimate and analyze its error more formally in Section 4. This simplification means that the cluster center estimator averages the local Good-Turing estimates from each user in the cluster. In theory, we can also estimate the cluster center P in a better way, e.g. , by directly running the Good-Turing correction on the cluster level token counts. This is a possible direction of improvement in future works.

To iteratively optimize the objective in Equation 4, we present our algorithm HistogramCluster (see Alg. 3). We start with initial center estimates ˜ P (0) , and run T iterations to refine them. In each iteration we peform two steps: (i) collect user outputs for each user via UserOutput (see Alg. 1); and (ii) re-center clusters using an algorithm, which for now is the NonPrivateCenter (see Alg. 2). Later, we see how we need only update the re-centering algorithm to satisfy user-level privacy contraints. From Lemma 3.1 and Alg. 1, it is not hard to see that NonPrivateCenter which computes the mean of the local estimates of users in a cluster, greedily reduces the objective value in Eq. 4 after each iteration. When the misclustering rate is sufficiently low to begin with, averaging local Good-Turing estimates brings the estimated center closer (in KL divergence) to the true centers that minimize Eq. 4, thereby reducing the number of iterations for convergence.

An extension of the convergence analysis in Balakrishnan et al. [7], gives us an upper bound on the misclustering error of Alg. 3 which goes down exponentially with each iteration of clustering (Theorem 3.2). Similar to most clustering analyses [8, 17, 33], we make suitable conditions on the initialization of cluster centers ˜ P (0) and the separation of true clusters in the user population.

Algorithm 1 UserOutput

Require:

Centers

˜

P

, cluster

k

, user

u

.

Ensure:

User output

(

b

u

, Q

u

)

.

1:

b

u

✶

(arg min

j

KL(

̂

Q

gt

u

∥

˜

P

:

,j

) =

k

)

2:

̂

Q

u

b

u

·

ˆ

Q

gt

u

+(1

-

b

u

)

·

0

d

3:

Return

(

b

u

,

̂

Q

u

)

.

Algorithm 2 NonPrivateCenter

Require:

User data A =: { ( b u , ̂ Q u ) } .

Ensure:

Center ˜ P .

1:

Return

˜

P

∑

u

̂

Q

u

/

∑

u

b

u

.

Algorithm 3 HistogramCluster

Require:

Set of users

U

, initial centers

˜

P

(0)

, iterations

T

.

Ensure:

Estimates of cluster memberships

{

˜

z

u

}

u

and centers

˜

P

.

1:

Initialize

t

1

.

2:

while

t

T

do

3:

for

j

= 1

to

K

do

4:

Collect:

A

(

t

)

k

UserOutput( ˜

P

(

t

-

1)

, k, u

)

}

u

5:

Re-center:

˜

P

(

t

)

:

,j

A

(

t

)

k

)

6:

end for

7:

t

t

+1

8:

end while

9:

Return:

{

˜

z

(

T

)

u

}

u

∈U

,

˜

P

(

T

)

.

̸

Theorem 3.2 (Alg. 3 convergence) . For the model in Section 4, if ∀ c, v , n c = Θ( K 2 ) , P c [ v ] = Ω(1) , centers are sufficiently separated, i.e., ∆ =: min i = j KL( P i ∥ P j ) = Ω( k 2 + k 3 d / n ) , and λ =: max i = j KL( P i ∥ P j ) / ∆ . Given an initialization with assignment error O ( √ 1 / λ ) for any cluster, after t = log( |U| ) iterations, w.h.p. the assignment error is O (exp( -∆( α +1))) .

̸

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. 2 applied to all the current clusters is differentially private in each iteration we can then bound the total privacy leak across iterations via advanced composition for DP algorithms. In our algorithm PrivateCenter (Alg. 4) we first use the Laplace mechanism in step 1 to privately compute the number of users in the cluster. Then, we use the Gaussian mechanism in step 2 to get an initial estimate of the center since all local Good-Turing estimates are distributions in ∆( V ) and thus bounded in ℓ 2 . This initial estimate B is used to identify the quantile B ± c √ B/b , where the mean lies with high probability [69], so that we can then get a refined estimate in step 3 when we add noise (with std. deviation scaling with quantile interval c ). This adaptive clipping procedure is fairly common for practical and private estimation [4, 12]. Finally, to ensure the noised output is still a distribution contained in ∆( V ) , we use ℓ 2 projection, Π ∆( V ) ( P ) =: arg min Q ∈ ∆( V ) ∥ Q -P ∥ 2 2 . In practice, we find that projection with ℓ 2 is comparable with a KL projection.

Privacy preserving noisy sum. As can be seen from the description of PrivateCenter , it only requires being able to compute a noisy sum over a subset of users. In the federated setting implementing this summation as part of algorithm HistogramCluster may seem to require the server to know the (private) cluster membership of each user. However, this can be avoided by computing each such sum as a sum over all the users, where a user submits value 0 if they do not belong to the cluster, i.e., b u = 0 , and b u = 1 otherwise (Alg. 1). Similarly, Alg. 1 also emits the true local estimate if it is part of the cluster and 0 d (vector of d 0 s) otherwise. Privacy-preserving computation of a noisy sum over all the users is one of key primitives in FL with a number of known implementations, e.g., Bonawitz et al. [13]. Crucially for the privacy analysis, while each user now participates in the sum computations for every cluster, the inclusion/deletion of a single user can only affect the sums for a single cluster.

Algorithm 4 PrivateCenter

Require:

̂

u

Data collection

,

,

Q

)

}

privacy parameter

u

(

{

ρ

b

, clipping param- eter

, projection operator

c

Ensure:

1:

b

2:

3:

4:

5:

6:

Algorithm 5 PrivateInit

∆(

)

V

.

Π

Private center

.

P

˜

∑

u

B

Σ

˜

P

u

b

+Lap

1

b

6

c

1

b

¯

Q

u

˜

P

Π

Return

( ∑

u

2

bρ

/

( √

Q

̂

u

+

·

diag(

(∑

u

/

3

N

B

¯

u

Q

+

2

(0

)

N

Q

V

∆(

)

˜

P

˜

P

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 PrivateInit (Alg. 5), we provide a private data-dependent initialization technique, that takes as input a private center Q 0 (private FedAvg), and samples K 2 points around it as an initial candidate set. Then, it chooses K points iteratively from this set by sampling via the exponential mechanism. The utility function for each candidate Q is determined by the average reduction in the clustering objective when adding Q into the current set of initializations I , vs. not. The privacy budget ρ controls the temperature, with lower budget enforcing higher smoothing. In the non-private case, we replace the candidate set with the set of empirical user estimates, and select a high value of ρ . This iterative procedure of selecting candidates with highest regret of omission at each step is meant to ensure that each ̂ Q u is reasonably close to some initial center (akin to [5]). We remark that the execution of the sampling step does not lend itself easily for implementation in the distributed setting. In practice, one would need to use a small dataset collected centrally for this.

(0

]

)

ρ

,

,

B

+

B

c

-

c

2

ρ

/

3

I

)

Σ)

√

B

b

B

b

,

u

√

d

)

)

3.4 Putting it all together: Algorithm for private and personalized frequency estimation

In Alg. 6 we present our complete algorithm. We first compute a private center of all user histograms using Alg. 4 and privacy parameter ρ / 3 . We pass this to Alg. 5 with privacy parameter ρ / 3 to get initial centers for Alg. 3, where the re-centering step applies Alg. 4 with privacy parameter ρ / 3 T . The privately learned cluster centers from Alg. 6 are then finetuned for each user locally, using Eq. 3.

Theorem 3.3 (End-to-end privacy guarantee) . Algorithms 4, 5 are both user-level ρ -zCJDP. Our end-to-end private and personalized estimation algorithm is ( ρ +2 √ ρ log( 1 / δ ) , δ ) -JDP, ∀ δ > 0 .

Algorithm 6 EndToEnd Algorithm

Require: Set of users U , dataset S , number of clusters K , number of clustering iterations T , finetuning parameter α , sampling parameter ω , privacy parameter ρ , clipping parameter c , projection operator Π ∆( V ) .

Ensure: Private and personalized estimates ̂ Q u for each u ∈ U . 1: Q 0 ← PrivateCenter( U , ρ, c, Π ∆( V ) ) # Estimate global center 2: ˜ P (0) ← Algorithm PrivateInit( Q 0 , S , K, ρ, ω, c ) # Estimate initial cluster centers 3: { ˜ z u } u and centers ˜ P ← Algorithm HistogramCluster( U , ˜ P (0) , T ) # Estimate cluster memberships 4: ∀ u : ˜ Q u ← FT( ˜ P ˜ z u , α ) # Local finetuning using Eq. 3 5: Return { ˜ Q u : u ∈ 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 u =: ( 1 / σ 2 u ) / ∑ u ∈Z c ( 1 / σ 2 u ) . Here σ 2 u is the variance of the user's local estimate, i.e. , Var[ ̂ Q u ] (from Lemma 3.4). When P z u [ v ] is known this weighted estimate is optimal under ℓ 2 [37]. Since P z u [ v ] is unknown we replace P z u [ v ] in Lemma 3.4 with a uniform average of ̂ Q u from heavier users, following Cummings et al. [21].

In the private setting though this approach cannot be used directly since the sensitivity of the weighted mean estimate is too high for users with a lot of data, and more so for small sized clusters. Cummings et al. [21] propose an algorithm to estimate the private heterogeneous mean of user data from different Bernoulli distributions, but their approach does not directly transfer to our setting due to the highdimensional nature of our user means and the relatively poor sensitivity of Alg. 3 to errors in the estimation of cluster centers at each iteration. Thus we consider a two-stage approach, split across data-rich and data-poor users. In the first stage we privately cluster users with sufficiently large datasets applying Alg. 3, and treating them as size-homogeneous data-rich users. Once cluster centers are learned privately, using Alg. 1, each data-poor user assigns itself to the closest center in KL divergence. Finally for each cluster, we apply Algorithm 2 from Cummings et al. [21] that re-centers each cluster privately, based on the newly added users.

4 Formal analysis in a stylized data model

In Section 3, we presented a clustering based iterative algorithm (Alg. 3) for our frequency estimation problem introduced in Section 2, wherein we made several key design choices. In particular, we used average of Good-Turing estimators to estimate the cluster center, and used the update in Equation 3 to locally finetune the learnt cluster centers. Now, we analyze these algorithms in a stylized model.

Bayesian Model. Each user u belongs to a cluster z u ∈ [ K ] with cluster center P z u . The user histogram Q u ∼ Dir( αP z u ) , and the user dataset S u ∼ Q m u u is sampled i.i.d. from Q u . We use Z c =: { u : z u = c } to denote the set of users from cluster c , and n c =: |Z c | is the number of users in c . Higher value of α implies more concentrated clusters since Var( Q u [ v ]) = O ( P zu [ v ] / 1+ α ) , and as α →∞ , Q u → P z u in weak topology [28]. Please note that the Dirichlet assumption is mainly for simplicity and our results ( e.g., Theorem 4.3) only require each cluster's user distribution to be exponentially concentrated along each token, i.e. , Pr( | Q u [ v ] -P z u [ v ] | ≥ t ) = O ( e -t ) .

Purely local learning. We present two local estimators that estimate Q u : empirical ̂ Q u . and GoodTuring [30] ̂ Q gt u . The naïve estimate ̂ Q u is the average of user data in S u . Next, we define the

Good-Turing estimator. Let the frequency of token v in user's dataset S u (of size m u ) be cnt u,v . We denote the local frequency of the count j as ϕ u,j =: ∑ v ∈V ✶ (cnt u,v = j ) , i.e. , the number of token in S u with count j . Following Orlitsky and Suresh [63], for a token v that appears j = cnt u,v times, the Good-Turing estimator for Q u [ v ] is given by:

where N u is the normalization factor so that ∑ v ̂ Q gt u [ v ]=1 . Let E u be the class of natural estimators that assign the same probability to tokens with equal counts in S u . Then, following Lemma 6 in Orlitsky and Suresh [63] and Theorem 2 in Acharya et al. [2] we conclude that with respect to E u , the Good-Turing estimate ̂ Q gt u has a worst case suboptimality gap of ˜ O (min( √ 1 / m , d / m ) (see Lemma 4.1 for the full statement). The key point to note here is that when the user dataset is small, we do not suffer from the vocabulary size d , unlike the empirical estimate ̂ Q u . This is mainly because Good-Turing more accurately estimates the probability of unseen words [56].

Lemma 4.1. ( ̂ Q gt u suboptimality gap) For any Q u , the suboptimality gap of ̂ Q gt u with respect to E u is E [KL( Q u ∥ ̂ Q gt u )] -min ˜ A ∈E u E [KL( Q u ∥ ˜ A )] = ˜ O (min( √ 1 / m , d / m )) , that matches minimax rates.

Bayes optimal finetuning when the cluster center is given. When P z u is known, FT in Eq. 3 is Bayes optimal in KL and when only an estimate of the center is known, the error of the 'plug-in' estimate scales linearly with the error in the center's estimate (Theorem 4.2).

Theorem 4.2 (Bayes optimal local learning) . Given P z u , α , the estimator FT( P z u , ̂ Q u ) is Bayes optimal in KL divergence, for the Dirichlet prior Dir( αP z u ) . When ˜ P z u is the estimated cluster center, and ˜ Q opt u is the Bayes optimal estimate of Q u , then KL( ˜ Q opt u ∥ FT( ˜ P z u , ̂ Q u )) ≤ α / α + m u KL( P z u ∥ ˜ P z u )) .

Estimating the cluster center given cluster members. Given all users in a cluster c , one estimate of the cluster center P c is the solution to the maximum log-likelihood objective in Eq. 6, which is typically solved with the FedAvg [49] algorithm. We denote this estimate as ˜ P fa c and is a simple weighted average of empirical estimates (see proof in Appendix E). Similary, for user u in cluster c , the personalized output from the PFL algorithm FedAvg+FT is given by FT( ˜ P fa c , ̂ Q u ) .

Our algorithms (Alg. 6, Alg. 2) use a different estimator. In the size-homogenous case, to estimate the cluster center we average the local Good-Turing estimates ̂ Q gt u for all users in the cluster:

To analyze the accuracy guarantees of our estimate ˜ P gt c , we first define the class of competing estimators for it. Let ¯ cnt v =: ⌊ ∑ u ∈Z cnt u,v / ∑ u m u ⌋ be the average count of token v across all users in a cluster. We define ¯ E to be the class of estimators that assign the same mass to tokens v 1 , v 2 if ¯ cnt v 1 = ¯ cnt v 2 . Such a class is natural for estimators that rely on aggregated statisitics, and includes the Good-Turing correction applied to cluster level counts. In Theorem 4.3, for any cluster center P c , we compare the suboptimality gap in KL divergence for our estimate ˜ P gt c in Eq. 7 with FedAvg.

Theorem 4.3 ( ˜ P gt c in Eq. 7 has lower suboptimality gap than FedAvg) . For the competetive class of estimators with the same average count, i.e., ¯ E (defined above), and suboptimality gap from Lemma 4.1, the suboptimality gap for ˜ P gt in Eq. 7, w.r.t. ¯ E is ˜ O ( α + m / α +1 · min( √ 1 / m , d / m )) . The suboptimality gap over ¯ E for FedAvg ˜ P fa is ˜ Ω( d ( α + m ) / m ( α +1) ) . ˜ O, ˜ Ω hides polylog factors in m,n .

When m is small and comparable to the intra-cluster std. deviation, i.e., m / α +1 = O (1) , then the suboptimality gap for our estimator does not suffer from large dimension d . On the other hand, FedAvg guarantees are much weaker and degrade with dimension for small user datasets (common in practice). From, Theorem 4.2, we also conclude that our final estimate for Q u , given by the finetuned center: FT( ˜ P, ̂ Q u ) would also have stronger error guarantees than FedAvg+FT.

Partitioning users into clusters. In practice, neither the cluster membership z u nor the center P c is known. The Bayes optimal estimate for P c , z u is the mean of the posterior Pr( Q u | S , α )

FedAvg

MAML

IFCA

Ours

6.0

6.2

6.4

6.6

6.8

Test Negative Log Likelihood

Reddit (w/o finetuning)

H

(

˜ P

f a

)

Non-private

Private

(a) Reddit

FedAvg

MAML

IFCA

Ours

7.5

8.0

8.5

9.0

9.5

Test Negative Log Likelihood

StackOverflow (w/o finetuning)

H

(

˜ P

f a

)

Non-private

Private

(b) StackOverflow

FedAvg

MAML

IFCA

Ours

7.5

8.0

8.5

9.0

9.5

Test Negative Log Likelihood

Amazon Reviews (w/o finetuning)

H

(

˜ P

f a

)

Non-private

Private

(c) Amazon Reviews

(see Theorem 4.2 proof), but since the posterior is intractable, we instead compute the maximum likelihood estimate (MLE). Asymptotically, MLE's accuracy is vindicated by Doob's theorem [59], and its variance matches the Cramér-Rao bound [69]. Still, the joint MLE for P c , z u is a non-concave maximization problem. Akin to solving the MLE for mixture distributions [61], we maximize an evidence lower bound with Expectation-Maximization (EM) [7]. Since each Q u belongs to only one cluster, it must be that z u ∈ arg min c ∈ [ k ] KL( Q u ∥ P c ) . Thus, the MLE for the cluster centers is given by objective in Eq. 4 from Section 3, which is optimized iteratively by our Algorithm 6.

Takeaways. We note: (i) when cluster memberships are known, average of local Good-Turings has a lower suboptimality gap than FedAvg; and (ii) solutions to the clustering objective in Eq. 4 estimates cluster centers for each user; and (iii) if the center estimates are accurate then finetuning them with FT yields estimates that are equally close to the Bayes optimal solution. These findings validate our algorithmic design choices for the clustering and finetuning algorithms in Section 3.

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-based baselines, and the gap amplifies in the size-heterogeneous case; and (iii) we show our method's performance improvements can be largely attributed to algorithmic design choices influenced by our data model and analysis in Section 4.

Datasets. We evaluate methods on three real-world datasets: Reddit [16], StackOverflow [6], and Amazon Reviews [62]. For Reddit, we use the NLTK tokenizer [11] with a vocabulary of size 10 k tokens, and for the other two datasets, we use the Huggingface (bert-case-uncased) tokenizer [74] with a vocabulary size of 32 k tokens. For more details on datasets and hyperparameters see Appendix C.

Baselines. We evaluate two baselines that learn a single global model: FedAvg [57] and MAML [29]; and a clustering-based approach which learns multiple models, one for each cluster: IFCA [33]/HypCluster [54]. Given a single global model (or a cluster level model), each user can finetune this locally in different ways: full batch gradient descent (GD) initialized at global model, RTFA [19], or with our problem specific method in Eq. 3. We compare these on FedAvg model

FedAvg

MAML

IFCA

Ours

6.0

6.2

6.4

6.6

Test Negative Log Likelihood

Reddit (with finetuning)

H

(

˜ P

f a

)

Non-private

Private

(a) Reddit

FedAvg

MAML

IFCA

Ours

7.5

8.0

8.5

9.0

Test Negative Log Likelihood

StackOverflow (with finetuning)

H

(

˜ P

f a

)

Non-private

Private

(b) StackOverflow

FedAvg

MAML

IFCA

Ours

7.25

7.50

7.75

8.00

8.25

8.50

Test Negative Log Likelihood

Amazon Reviews (with finetuning)

H

(

˜ P

f a

)

Non-private

Private

(c) Amazon Reviews

4

6

8

10

12

14

16

K

8.0

8.2

8.4

8.6

8.8

Test Negative Log Likelihood

Stackoverflow

Amazon Reviews

(a)

Reddit

StackOverflow

Amazon

6.1

6.5

6.9

7.3

7.7

8.1

8.5

8.9

9.3

Test Negative Log Likelihood

Importance of Alg 5 (PrivateInit)

Alg 5 (PrivateInit)

(b)

Reddit

StackOverflow

Amazon

6.0

6.5

7.0

7.5

8.0

8.5

9.0

Test Negative Log Likelihood

Importance of Good Turing Average

Empirical average

Good-Turing average

(c)

Reddit

StackOverflow

Amazon

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Test Negative Log Likelihood

Comparing Local Personalization on FedAvg

RTFA

Ours

(d)

(Figure 3(d)), and fix the best one (ours) as the local finetuning approach for all algorithms. In the private setting, we adapt the baselines by using standard techniques [1, 25] that privatize l 2 bounded gradients for MAML/IFCA, and l 1 bounded probability distributions for FedAvg.

Evaluation metric. For distribution Q u , estimate ˜ Q u , KL( Q u , ˜ Q u ) = NLL( Q u , ˜ Q u ) -H( Q u ) , where NLL( Q u , ˜ Q u ) =: -∑ w ∈V Q u,w log Q u,w is the negative log-likelihood loss, and H( Q u ) is the entropy of Q u . Following language modeling [3, 38], we use the test set of the user to get an unbiased estimate of the NLL loss, by replacing Q u with the empirical distribution on the test set. Since the test set is insufficient to get an unbiased estimate of H( Q u ) , as a reference point, we instead report the entropy of the global center or FedAvg estimate H( ˜ P fa ) , which is a rough estimate of the 'hardness' of estimating the user distribution. We refer to this difference in test NLL and H( ˜ P fa ) as NLL sub-optimality gap . In the private case, we ensure the algorithms satisfy (15 , 10 -10 ) -JDP.

Our approach significantly reduces test NLL in private and non-private settings. In Figure 1 we report a significant reduction in test NLL sub-optimality gap even before local finetuning, by atleast 25 -45% in the non-private setting and upto 50% in the private case. Non-privately MAML/IFCA perform better than FedAvg on StackOverflow and Amazon, but their private versions perform similar or worse than the private FedAvg. This is because the gradient-based optimization in MAML and IFCA can incur very high per-iteration privacy overheads which only accumulates more for the latter that additionally suffers from poor convergence due to imperfect cluster initialization. On the other hand, our method achieves a privacy-utility tradeoff that is comparable to the gradient-free FedAvg, and converges in fewer iterations when clustering is initialized with centers from Alg. 5. In Figure 2 we compare the performance after we finetune the global/cluster-level estimates locally for each user (using Eq. 3), and verify that the relative gains remain consistent with pre-finetuning. In particular we note that finetuning global FedAvg/MAML models does not do better than finetuning cluster-level models, vindicating the presence of concentrated subpopulations in an otherwise highly heterogeneous real-world user distribution.

Algorithmic design decisions influenced by our model in Section 4. In Section 4 we identified the Good-Turing based estimator (Eq. 7) to suffer less from vocabulary size, compared to FedAvg (Theorem 4.3), and in practice too we observe a big improvement in the test NLL loss of Alg. 3 when it uses Good-Turing as the local estimate, vs. empirical (Figure 3(c)). The relative gap is particularly wider on the large vocabulary datasets Amazon and StackOverflow. In Figure 3(b), we show that our data-dependent private initialization (Alg. 5), that is also based on our mixture of Dirichlet model, plays a crucial role in lowering the final test loss, and we attribute this gain to a significant reduction in clustering iterations (from 200 to 20 ), thereby reducing privacy overhead. In Figure 3(d) we note that the Bayes optimal finetuning algorithm in our data model (Eq. 3) does better than typical gradient-based (GD) or proximal term based (RTFA) approaches. Finally, in Figure 3(a), we plot test NLL as we vary the number of clusters assumed by Alg. 3. We note that for runs on all datasets we choose K =10 which was found to be optimal on the Reddit validation set (see Appendix C), even though it is clearly (slightly) suboptimal on the other two, suggesting that Alg. 3's performance is not too sensitive to the choice of K practice.

Size-heterogeneous case. In Figure 4 plot test NLL in the size-heterogeneous case where our algorithm's gains ( > 40% error reduction) over single global model baselines FedAvg and MAML are even more pronounced than the size-homogeneous results in Figures 1, 2. For the better suited clustering baseline IFCA, the privacy utility tradeoff worsens (compared to size-homogenous) due to poorer higher sensitivity of gradient estimates when dataset sizes vary. On the other hand, our two-stage approach reduces privacy noise in the clustering stage by partitioning users into different stages and still improves performance for data scarce users in the second stage by incorporating them

FedAvg

MAML

IFCA

Ours

7.0

7.5

8.0

8.5

9.0

Test Negative Log Likelihood

Reddit (size-heterogeneous)

H

(

˜ P

f a

)

Non-private

Private

(a) Reddit

FedAvg

MAML

IFCA

Ours

8.0

8.5

9.0

9.5

10.0

10.5

Test Negative Log Likelihood

StackOverflow (size-heterogeneous)

H

(

˜ P

f a

)

Non-private

Private

(b) StackOverflow

FedAvg

MAML

IFCA

Ours

8

9

10

11

Test Negative Log Likelihood

Amazon Reviews (size-heterogeneous)

H

(

˜ P

f a

)

Non-private

Private

(c) Amazon Reviews

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 a global model and then uses it to partition users based on their losses, most works learn diverse models from scratch using gradient based algorithms, as done by FedAvg [57]. The key difference being that they first partition clients based on their loss [33] or gradients [72] and then use the same gradients to update different models, one for each cluster. Additionally, their analysis also holds mainly for smooth/strongly convex loss functions, for e.g. , least squares [33, 54]. This is a natural and promising approach but we are not aware of practical results showing that it can improve on the more direct combination of FedAvg and FT [19, 76]. In fact, Wu et al. [76] raise concerns about mode collapse with clustering iterations. Contrary to the works above, we focus on a non-gradient based approach specifically for frequency estimation in KL divergence, which can be ill-conditioned in practice, and also provide an algorithm for the more challenging size-heterogeneous setting. As we show in our work, this problem requires solutions that are different from the ones proposed for general loss families in machine learning. Moreover, unlike the above, we provide privacy guarantees for our clustering based personalization. For more discussion on related works please see Appendix B

Distribution estimation in KL divergence. A multitude of works on mixture of Gaussians, and mixed linear regression propose and analyze distribution estimation algorithms [9, 22, 45, 61], but their guarantees are mainly for parametric estimation errors in ℓ 1 /ℓ 2 metric. In contrast, we are concerned with histogram estimation in KL divergence which presents interesting challenges since this is not a proper distance metric ( e.g. , does not satisfy triangle inequality). On the other hand, [30, 34, 64] study Good-Turing estimators and give estimation error guarantees in KL divergence for categorical distributions over a fixed alphabet. Our work extends estimators of Acharya et al. [2], Orlitsky and Suresh [63] to the federated setting where the goal is to estimate a full population of distributions that share a latent structure. Relevant to our objective and metric is [17] that analyzes guarantees for information theoretic clustering. Their analysis shows that one can adapt analysis for other metrics ( e.g. , Hellinger) to obtain worst-case approximation guarantees. In contrast, we investigate practical and private algorithms for the federated setting.

7 Conclusion

We introduce the problem of private and personalized frequency estimation in KL divergence. For this, we propose an iterative algorithm that privately learns clusters in the population of all user frequencies. Each user in a cluster locally finetunes their corresponding cluster center to produce personalized and private frequency estimates with formal joint DP guarantees. We improve the privacy-utility tradeoff of our algorithm by proposing a novel data-dependent private initialization for clustering that empirically reduces number of clustering iterations. We also present a two-stage version of our approach to separately handle the harder size-heterogenous setting. In a Bayesian model where the user distributions are distributed as Dirichlets around well-separated centers, we reason about different collaborative and local estimators, and provide formal guarantees for some of our algorithmic design choices, like Good-Turing estimators, and the choice of the local finetuning algorithm. Empirically, we test our algorithm on three real-world datasets and show a significant reduction in test NLL, by 25 -45% in the non-private setting and upto 50% in the private case.

References

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: 10.1145/2976749.2978318. URL http://dx.doi.org/10.1145/2976749.2978318 .

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.

Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems , 35:2230022312, 2022.

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.

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.

The TensorFlow Federated Authors. TensorFlow Federated Stack Overflow dataset, 2019. URL https://www.tensorflow.org/federated/api_docs/python/tff/simulation/ datasets/stackoverflow/load_data .

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.1214/16-AOS1435. URL https://doi.org/10.1214/16-AOS1435 .

Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, Joydeep Ghosh, and John Lafferty. Clustering with bregman divergences. Journal of machine learning research , 6(10), 2005.

Mikhail Belkin and Kaushik Sinha. Polynomial learning of distribution families. In 2010 IEEE 51st Annual Symposium on Foundations of Computer Science , pages 103-112. IEEE, 2010.

Alberto Bietti, Chen-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.

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. URL https://www.nltk.org/book/ .

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.

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://arxiv.org/abs/1611.04482 .

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.

Mark Bun and Thomas Steinke. Concentrated differential privacy: Simplifications, extensions, and lower bounds. In Theory of Cryptography Conference , pages 635-658. Springer, 2016.

Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Koneˇ cn` y, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097 , 2018.

Kamalika Chaudhuri and Andrew McGregor. Finding metric structure in information theoretic clustering. In COLT , volume 8, page 10, 2008.

Gary Cheng, Karan Chadha, and John Duchi. Fine-tuning is fine in federated learning. arXiv preprint arXiv:2108.07313 , 3, 2021.

Gary Cheng, Karan Chadha, and John Duchi. Federated asymptotics: a model to compare federated learning algorithms. In International Conference on Artificial Intelligence and Statistics , pages 10650-10689. PMLR, 2023.

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.

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:29139-29151, 2022.

Sanjoy Dasgupta. Learning mixtures of gaussians. In 40th Annual Symposium on Foundations of Computer Science (Cat. No. 99CB37039) , pages 634-644. IEEE, 1999.

Cynthia Dwork. Differential privacy in new settings. In Proceedings of the twenty-first annual ACM-SIAM symposium on Discrete Algorithms , pages 174-183. SIAM, 2010.

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.

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.

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.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07948 , 2020.

Thomas S Ferguson. A bayesian analysis of some nonparametric problems. The annals of statistics , pages 209-230, 1973.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

William A Gale and Geoffrey Sampson. Good-turing frequency estimation without tears. Journal of quantitative linguistics , 2(3):217-237, 1995.

Jonas Geiping, Hartmut Bauermeister, Hannah Dröge, and Michael Moeller. Inverting gradientshow easy is it to break privacy in federated learning? Advances in Neural Information Processing Systems , 33:16937-16947, 2020.

Robin C Geyer, Tassilo Klein, and Moin Nabi. Differentially private federated learning: A client level perspective. arXiv preprint arXiv:1712.07557 , 2017.

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.

Irving J Good. The population frequencies of species and the estimation of population parameters. Biometrika , 40(3-4):237-264, 1953.

Dennis Grinwald, Philipp Wiesner, and Shinichi Nakajima. Solution simplex clustering for heterogeneous federated learning. arXiv preprint arXiv:2403.03333 , 2024.

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:1811.03604 , 2018.

Joachim Hartung, Guido Knapp, and Bimal K Sinha. Statistical meta-analysis with applications . John Wiley & Sons, 2011.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Justin Hsu, Zhiyi Huang, Aaron Roth, and Zhiwei Steven Wu. Jointly private convex programming. In Proceedings of the twenty-seventh annual ACM-SIAM symposium on Discrete algorithms , pages 580-599. SIAM, 2016.

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.

Shengyuan Hu, Zhiwei Steven Wu, and Virginia Smith. Private multi-task learning: Formulation and applications to federated learning. arXiv preprint arXiv:2108.12978 , 2021.

Prateek Jain, Om Dipakbhai Thakkar, and Abhradeep Thakurta. Differentially private matrix completion revisited. In International Conference on Machine Learning , pages 2215-2224. PMLR, 2018.

Prateek Jain, John Rush, Adam Smith, Shuang Song, and Abhradeep Guha Thakurta. Differentially private model personalization. Advances in Neural Information Processing Systems , 34: 29723-29735, 2021.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun 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):1-210, 2021.

Adam Tauman Kalai, Ankur Moitra, and Gregory Valiant. Efficiently learning mixtures of two gaussians. In Proceedings of the forty-second ACM symposium on Theory of computing , pages 553-562, 2010.

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.

Vishesh Karwa and Aleksandra B Slavkovi´ c. Differentially private graphical degree sequences and synthetic graphs. In Privacy in Statistical Databases: UNESCO Chair in Data Privacy, International Conference, PSD 2012, Palermo, Italy, September 26-28, 2012. Proceedings , pages 273-285. Springer, 2012.

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.

Jakub Koneˇ cn` y, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492 , 2016.

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:12466-12479, 2021.

Jeffrey Li, Mikhail Khodak, Sebastian Caldas, and Ameet Talwalkar. Differentially private meta-learning. arXiv preprint arXiv:1909.05830 , 2019.

Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE signal processing magazine , 37(3):50-60, 2020.

Stuart Lloyd. Least squares quantization in pcm. IEEE transactions on information theory , 28 (2):129-137, 1982.

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.

Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. Advances in Neural Information Processing Systems , 34:15434-15447, 2021.

David A McAllester and Robert E Schapire. On the convergence rate of good-turing estimators. In COLT , pages 1-6, 2000.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private language models without losing accuracy. arXiv preprint arXiv:1710.06963 , page 84, 2017.

Jeffrey W Miller. A detailed treatment of doob's theorem. arXiv preprint arXiv:1801.03122 , 2018.

Michael Mitzenmacher and Eli Upfal. Probability and computing: Randomization and probabilistic techniques in algorithms and data analysis . Cambridge university press, 2017.

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.

Jianmo Ni, Jiacheng Li, and Julian McAuley. Justifying recommendations using distantlylabeled reviews and fine-grained aspects. In Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP) , pages 188-197, 2019.

Alon Orlitsky and Ananda Theertha Suresh. Competitive distribution estimation: Why is good-turing good. Advances in Neural Information Processing Systems , 28, 2015.

Liam Paninski. Variational minimax estimation of discrete distributions under kl loss. Advances in Neural Information Processing Systems , 17, 2004.

Alireza Salemi, Sheshera Mysore, Michael Bendersky, and Hamed Zamani. Lamp: When large language models meet personalization. arXiv preprint arXiv:2304.11406 , 2023.

Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning: Modelagnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural networks and learning systems , 32(8):3710-3722, 2020.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Ramon Van Handel. Probability in high dimension. Lecture Notes (Princeton University) , 2014.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

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 IEEE INFOCOM 2019-IEEE conference on computer communications , pages 2512-2520. IEEE, 2019.

Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Towards provably personalized federated learning via threshold-clustering of similar clients. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022) , 2022.

Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Provably personalized and robust federated learning. arXiv preprint arXiv:2306.08393 , 2023.

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 transformers: State-of-the-art natural language processing, 2020.

Huiwen Wu, Cen Chen, and Li Wang. A theoretical perspective on differentially private federated multi-task learning. arXiv preprint arXiv:2011.07179 , 2020.

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.

Dun Zeng, Xiangjing Hu, Shiyu Liu, Yue Yu, Qifan Wang, and Zenglin Xu. Stochastic clustered federated learning. arXiv preprint arXiv:2303.00897 , 2023.

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 an interesting future direction would be to extend this to conditional next-token distributions. Nevertheless, we find that existing algorithms in the federated learning literature when applied to our problem setting perform suboptimally, and require problem specific interventions we introduce, like Algorithms 3, 5, and Good-Turing based estimators. Our current work also does not handle a federated version of the private initialization algorithm we propose. We suggest handling the initialization at the central node that maintains a small dataset of clients. This is a known limitation for sampling from the exponential mechanism in a federated setting. At the same time we would like to point out, that all our other algorithms, including private ones can be implemented in the federated setting.

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 a single model. Cummings et al. [21] also study the problem of private heterogeneous mean estimation but do not consider the clustered setting, or the subtleties of estimatation in KL divergence. More recent works ( e.g. , [42, 43]) propose personalization algorithms that are user-level private, but are mainly tailored for settings where the different models for each user share representations or lie in a hidden low rank subspace. In contrast, we explore another latent structure more suited for histogram estimation, i.e. , one where histograms need only be concentrated around the vertices of a low-dimensional polytope. Other works [10, 41, 75] propose gradient-based algorithms for private multi-task learning where they learn different models for each user (task), regularized using inter-task prior relationship matrices. In contrast, our work does not assume any such prior knowledge.

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 a vocabulary of size 10 k tokens, and use the Huggingface (bert-case-uncased) tokenizer [74] with a size of 32 k tokens for the other two. We have ≈ 10 k users in Reddit, each with at least 1 k tokens. For StackOverflow and Amazon Reviews we have ≈ 19 k and 50 k clients respectively, each with 500 tokens For all datasets, we partition the data for each client into 60 : 40 train/test splits. Additionally, we set aside 5% of users in each, as a validation, to tune cluster count K , privacy clip bound c , etc. In the private setting, we adapt the baselines to the private setting by using standard techniques [1, 25] to privatize l 2 bounded gradients for MAML/IFCA, and l 1 bounded probability distributions for FedAvg, so as to ensure all algorithms achieve JDP(15 , 10 -10 ) . For more details see Appendix C.

Hyperparameters. We validate hyperparameters for our algorithms and baselines using a hold out validation set of users. For clustering we use K = 10 , as validated by Figure 5, and fix this for all datasets and clustering baselines. We run clustering for T = 50 iterations non-privately and T = 20 iterations privately on all datasets. We tune the finetuning parameter λ by sweeping across { 0 . 05 , 0 . 1 , 0 . 15 , 0 . 2 , 0 . 25 , 0 . 3 , 0 . 35 , 0 . 4 , 0 . 45 , 0 . 5 } and find λ = 0 . 25 , λ = 0 . 15 and λ = 0 . 1 to be optimal on Reddit, StackOverflow and Amazon respectively. We train MAML and IFCA using SGD with momentum, with learning rate 0 . 01 and momentum parameter 0 . 9 . For RTFA, we find the proximal regularization parameter of 0 . 2 to be optimal on Reddit and StackOverflow, 0 . 1 to perform better on Amazon Reviews. For private training, we use a clipping threshold of c = 0 . 1 in Alg. 4, and use c = 4 . 0 for Alg. 4.

Computational resources. None of our experiments require very high computational requirements and can be run with one 3090Ti card. All runs can be reproduced in approximately 500 GPU hours.

and for x < λ ,

4

6

8

10

12

14

16

K

0.925

0.950

0.975

1.000

1.025

Test Negative Log Likelihood

Reddit

D Useful Lemmas

Lemma D.1 (Chernoff bound) . If X ∼ Poi ( λ ) , then for x ≥ λ ,

Lemma D.2 (Bernstein's Inequality) . Let X 1 , X 2 , . . . , X n be n independent zero mean random variables such that with probability ≥ 1 -ϵ i , | X i | < M . Then,

If t = √ 2 ( ∑ i E [ X 2 i ]) log 1 δ + 2 3 M log 1 δ , then

Lemma D.3 (Banerjee et al. [8]) . The solution for the optimization problem:

is given by (1 -λ ) P 0 + λP 1 .

E Missing proofs from Section 4

Lemma E.1. ( ̂ Q gt u achieves minimax suboptimality gap; restated) For any Q u , the suboptimality gap of ̂ Q gt u with respect to E u is E [KL( Q u ∥ ̂ Q gt u )] -min ˜ A ∈E E [KL( Q u ∥ ˜ A )] = ˜ O (min( √ 1 / m , d / m )) , that matches minimax rates.

Proof. From Lemma 8 in [63] we know that for the class of natural estimators E u , the suboptimality gap defined above is given by E [KL( M k ∥ ˜ M k )] , where

We use the definition of ϕ u,t from [63], where:

Now, from Theorem 2 in [63] we also know that KL( M k ∥ ˜ M k ) = ˜ O ( √ 1 /m ) with probability at least 1 -1 /m over the randomness in draw of the user dataset S u . Since their estimate is lower bounded by 1 / 2 m , we can convert their high probability guarantees to guarantees in expectation. Further, we note that in their proof of Theorem 2, we can easily upper bound the following

in the worst case. Thus, in the worst case, E [KL( M k ∥ ˜ M k )] is upper bounded by the minimum of O ( d / m ) and ˜ O ( √ 1 / m ) . This completes the proof of Lemma 4.1.

Theorem E.2 (Finetuning is Bayes optimal given center; restated) . If cluster center P z u , and Dirichlet parameter α is known then FT( P z u , ̂ Q u ) is Bayes optimal in KL divergence.

Proof. Wenote that the distribution of ̂ Q u given the center P c u is a Dirichlet-Multinomial distribution. Further, the Dirichlet distribution is a conjugate prior for the Multinomial.

Hence, the posterior distribution for Q u | ̂ Q u , P c u is a Dirichlet Multinomial with mean: ( α α + m ) P u + ( m α + m ) ̂ Q u . To see why, mean is the Bayes optimal estimator, we invoke Lemma D.3, which implies:

Now the Bayes risk of algorithm M , given the center estimate P is :

Recursively applying Lemma D.3 on every measurable subspace (under measure defined by the posterior Pr( Q u | ̂ Q u , P ) ) of the set of token distributions ∆( V ) , we conclude that the optimal solution is E [ Q u | ̂ Q u , P ] when Q u ∼ Pr( Q u | ̂ Q u , P ) . This concludes our proof. Note that the same proof also applies for Lemma 3.1, but applied iteratively on each element in the summation.

LemmaE.3 (Accuracy of plugging center estimate into finetuning; restated) . Let ˜ P z u be the estimated cluster center for user u , and ˜ Q opt u be the Bayes optimal estimate of Q u . Then,

Proof. This result is a simple application of Jensen's inequality for KL divergence.

Lemma E.4 (FedAvg estimate) . The FedAvg model is given by Q fa = 1 n ∑ ̂ Q u ∈S ̂ Q u .

Proof. The log likelihood objective log Pr( S u ∥ Q ) is equivalent upto additive constant with the objective KL( ̂ Q u ∥ Q u ) . We can once again use the result from Lemma D.3 to conclude that the optimal is simply the average of local means.

Theorem E.5 ( ˜ P gt c in Eq. 7 has lower suboptimality gap than FedAvg; restated) . For the competetive class of estimators with the same average count, i.e., ¯ E (defined above), and suboptimality gap from Lemma 4.1, the suboptimality gap for ˜ P gt in Eq. 7, w.r.t. ¯ E is ˜ O ( α + m / α +1 · min( √ 1 / m , d / m )) . The suboptimality gap over ¯ E for FedAvg ˜ P fa is ˜ Ω( d ( α + m ) / m ( α +1) ) . ˜ O, ˜ Ω hides polylog factors in m,n .

Proof. For notational convenience, we denote the cluster center with P (instead of P c ), and the number of users in the cluster as n (instead of n c ). To derive bounds on competetive estimation, where the competetive class can only use the average count for each token, we need to first ensure that our Good-Turing based estimator can actively estimate the following quantity:

Similarly, the local equivalent of M t for each user is defined as:

As a starting point, we will first derive our results with the genie-aided estimator, where we have oracle access to the ration of expected frequencies for every user, and then switch the analysis to the more general version when we do not. The genie-aided local estimator ˆ Q gen u,v is:

and the corresponding average of genie-aided Good-Turing estimates:

Next, we derive results for the setting where m u ∼ Poi( m ) so that the local frequency counts of words become independent. Then, using the results on Poisson sampling in Mitzenmacher and Upfal [60], we can convert the bounds to the setting where each user samples exactly m tokens.

We will use ✶ t v as a shorthand for ✶ ( ⌊ cnt v ⌋ = t ) and similarly use ✶ t u,v as a shorthand for ✶ (cnt u,v = t ) . The following quantities define the global and local frequencies of average and local counts respectively:

For F t defined using the genie-aided estimator, to bound the competetive KL estimation error we first need to bound the KL distance between ground truth M t and estimate F t under high probability over the draw of the meta-dataset S .

Using Lemma E.6 we lower bound the denominator in the KL upper bound in Eq. 17.

Lemma E.6 (Lower bound on F t ) . With probability atleast 1 -nm -2 , F t = ˜ Ω ( tϕ t m ) .

Proof. The proof relies on three results for local Good-Turing (per-user) estimates. We will combine the three results, and then apply a union bound over all the users. Let v be some token in the vocabulary that has an average count of t , i.e. , ✶ t v = 1 . From Lemma 15, 16 and Claim 20 in Acharya et al. [2], with probability at least 1 -m -3 , ̂ Q un u = ˜ Ω( cnt u,v / m ) , where ̂ Q un u is the unnormalized local Good-Turing estimate. Additionally, from Lemma 17 in Acharya et al. [2], the normalization factor N u for user u satisfies with probability at least 1 -10 m -2 , N u = 1 + ˜ O ( 1 / m 1 / 4 ) . Since m u is a Poisson distribution from Lemma D.1, the number of data points m u for the user satisfies | m u -m | ≤ √ m , with constant probability. Putting these together, and applying a union bound over all n users we conclude:

with probability at least 1 -nm -2 .

Next we try to upper bound ( M t -F t ) 2 for some t ≥ 1 . But before that we introduce some more helpful results.

Lemma E.7 (User's Q u concentrates around center P ) . With probability at least 1 -δ , ∀ v ∈ V :

Proof. Since Q u,v ∈ [0 , 1] is a bounded random variable with variance P v (1 -P v ) / (1+ α ) , we can apply the Bernstein inequality in Lemma D.2. This would give us a confidence interval for a single token v , which we can then union bound over all tokens v ∈ V to get the result.

Lemma E.8 (variance of linear estimators; Claim 21 in Acharya et al. [2]) . For every distribution p ,

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,

̸

For t = t ′ , E [ t v 1 t ′ v ] = 0 and hence (a). (b) uses the fact that t v is an indicator random variable.

✶ ✶ Lemma E.9 (expected sensitivity of the local frequencies; Claim 20 in [2]) . For every user distribution p over V , from which we draw n i.i.d. samples, let t be the local count and ϕ t = ∑ v ∈V ✶ t v is the frequency of the local count,

Proof. We consider the two cases t +1 ≥ log n and t +1 < log n separately. Consider the case when t +1 ≥ log n . We first show that

The first equality follows by substituting E [1 t x ] = e -np x ( np x ) t /t ! . For the inequality, note that if | np x -t -1 | ≤ 25( t +1)log n , then the inequality follows. If not, then by the Chernoff bound E [1 t x ] = Pr( t x = t ) ≤ n -3 and hence ∣ ∣ E [1 t x ] -E [1 t +1 x ] ∣ ∣ ≤ E [1 t x ] + E [1 t +1 x ] ≤ 2 /n 3 .

By definition, E [Φ t ] -E [Φ t +1 ] = ∑ x E [1 t x ] -E [1 t +1 x ] . Substituting,

where (a) follows from the fact that E [1 t x ] = e -np x ( np x ) t /t ! . (b) follows from the fact that np x ≤ 1 in the first summation and Equation (4). The proof for the case t +1 < log n is similar and hence omitted.

LemmaE.10 (Upper bound on | M t -F t | ) . For any t ≥ 1 , with high probability at least 1 -n poly( m ) .

Proof.

Our proof mainly relies on the concentration properties of Dirichlet distribution and some properties of the local Good-Turing estimate. While the first term is fairly streaightforward to bound, the second term requires us to handle the fact that ✶ t v is a random variable that depends on the global statistics for the token v , while ˆ Q gt u,v -Q u,v depends only on the local statistics. This requires us to analyze this differently from Acharya et al. [2].

Using Lemma E.8, followed by Lemma E.7 we can bound term A in Eq. 18 in the following way:

This completes our derivation for the upper bound on the first term ( A ) in Eq. 18. For the second term ( B ), we use the definition of the genie-aided local Good-Turing estimator and bound its variance in the folliowing way:

The first inequality above uses Lemma E.8, and the second equality uses the definition of the genieaided estimator. In Eq. 22 we have two terms, and we bound the first term using Jensen inequality followed by the DKW inequality [70] for empirical estimators of the cumulative distribution. With high probability 1 -δ ,

The final inequality above uses the definition of t , and the concentration of m u , i.e. , with high probabibility: | m u -E [ m u ] | = | m u -m | = ˜ O ( √ m ) .

Next we will bound the second term in Eq. 22, for which we use some properties of the expected sensitivity of the local frequency counts (Lemma E.9).

Applying Jensen's inequality it is easu to see that we can bound the above term with:

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 our setup in Section 4 in two ways: (1) Q u = P , i.e. , each user's true distribution matches exactly the cluster center (this matches the α →∞ case in our setup); and (2) each user independently samples m u ∼ Poi( m ) , which is the number of samples in their local dataset, that is used to estimate the empirical distribution ̂ Q u (note that m u = m for our results in Section 5). We shall now derive results in this setting, and then consider the more general case.

Without loss of generality, let us fix the center to be c , with some collection of users Z c . Then, we know that cnt v is distributed as a Poisson random variable with mean m ∑ u ∈Z c Q u [ v ] . Thus, with probability at least 1 -δ :

We mainly rely on Lemma 13 for the Genie aided estimator from [2]. The key difference being that E [ ϕ k ] , which is words with total count k would differ from E [ ϕ cnt u,v ] . For, this we first bound ⌊ k / n ⌋ -cnt u,v with high probability. We do this, by:

since the counts follow Poisson distribution under indpendent sampling. Thus, when α = Ω( m ) , the first term dominates, and we see that each this gives us:

with high probability. Then from Claim 20 in [2], qwe know that | E [ ϕ k ] -E [ ϕ cnt u,v ] | = O ( k )

The rest of our analysis relies with the genie-aided estimator mainly relies on Theorem 2 in Acharya et al. [2]. Using the result in Lemma E.10, we conclude that with probability 1 -n poly( m ) over the draw of the meta-dataset S , the KL distance:

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, we reuse the conversion of the high probability result to one in expectation using arguments similar to our proof of Lemma 4.1. This completes our analysis of the upper bound on competetive regret for the average of local Good-Turings (our estimator).

Finally, for the lower bound in Theorem 4.3, we invoke known lower bounds on the concentration of the Dirichlet distribution ˜ Ω( 1 / α +1 ) , and minimax statistical lower bounds on the risk of the empirical estimator ˜ Ω( d m ) , to conclude that the competetive risk of FedAvg is at least ˜ Ω( d / m + 1 / α +1 ) . This completes our proof of Theorem 4.3.

Lemma E.11 (Stirling approximation [46]) . The Γ function is sandwiched as follows. There is a positive constant C , such that for all x > 0 :

This implies log Γ( x ) = log C +( x -1 / 2 ) log x -x + O ( 1 / x ) . For large enough x , we substitute Γ( x ) ≈ log C +( x -1 / 2 ) log x -x .

Theorem E.12 (Clustering approximates MLE objective; restated) . The joint likelihood for P c and z u can be upperbounded with the clustering objective in Eq. 4, and this approximation is tight upto an additive term of O ( 1 / α ) .

Let V be a fixed set of words, d =: |V| is the size of the set, and ∆( V ) is the set of all possible discrete probability distributions over V . Each user u in the federated setup has an unknown distribution Q u ∈ ∆( V ) . Further, each Q u belongs to one of k clusters, and the cluster membership is denoted by z u , where z u ∈ R k is a 1 -hot k -dimensional vector. Each cluster c ∈ [ k ] is associated with a center P c ∈ ∆( V ) , and the matrix with column vectors as cluster centers is denoted as P ∈ R d × K , i.e., P =: [ P 1 , P 2 , . . . , P k ] .

Independently for each user u , given the cluster membership z u , and some α > 0 , the user's distribution Q u is sampled from a Dirichlet around the cluster center P z u , i.e., Q u ∼ Dir( α P z u ) . Further, each user samples an i.i.d. dataset S u of size n u . When all users sample dataset of fixed size n u = n , we refer to this setting as size-homogeneous, and the more general case as sizeheterogeneous. Given the set of datasets {S u } u and the distribution assumptions above, we can write down the maximum-likelihood estimate (MLE) for the cluster centers P and cluster memberships { z u } u :

where ̂ Q u is the empirical distribution obtained from the dataset S u and γ ( x ) =: ∫ t x -1 e -x d x is the Gamma function. In the above derivation, the second inequality follows from n u ̂ Q u being a sufficient statistic for the class of distributions parameterized by { Q u } u .

For the size-homogenous case, we can simplify the MLE objective further to be:

For large enough α , we can use the Stirling's approximation of the Gamma function to approximate the above MLE objective (see Lemma E.11):

when cluster center P [ v ] = Ω(1) .

F Missing proofs from Section 3

Lemma F.1 ( ̂ Q u variance) . For user u from cluster c , the variance of the estimate ̂ Q u [ v ] is:

Proof. By the Law of Total Variation, the variance of ̂ Q u [ v ] is:

F.1 Proof of Theorem 3.2

Theorem F.2 (Alg. 3 convergence; restated) . For the model in Section 4, if ∀ c, v , n c = Θ( K 2 ) , P c [ v ] = Ω(1) , centers are sufficiently separated, i.e., ∆ =: min i = j KL( P i ∥ P j ) = Ω( k 2 + k 3 d / n ) , and λ =: max i = j KL( P i ∥ P j ) / ∆ . Given an initialization with assignment error O ( √ 1 / λ ) for any cluster, after t = log( |U| ) iterations, w.h.p. the assignment error is O (exp( -∆( α +1))) .

̸

Lemma F.4 (KL upper bound) . For P, Q ∈ ∆ |V|1 , the KL divergence KL( P || Q ) ≤ χ 2 ( P Q )

Proof Overview. The main technique we follow is to use Pinsker's inequality to lower bound KL divergence with TV 2 ( P, Q ) and upper bound KL with χ 2 . This allows us to then treat the clustering in the l 2 2 metric, and pay the d -dimensional penalty term 1 d . Based on this general principle, the following section presents the EM convergence analysis of Algorithm 3, extending results from [7] to our setting. While the steps below closely mirrors their analysis upto the final terms where we need to upper and lower bound KL, we present the full proof for completeness.

Let us begin by introducing some notaion. For any S ⊆ [ n ] , define W = ∑ w . Recall that

S i ∈ S i T ∗ g = { i ∈ [ n ] , z i = g } and T ( s ) g = { i ∈ [ n ] , ˆ z ( s ) i = g } , let us define

Then we have n ( s ) h = ∑ g ∈ [ k ] n ( s ) gh and n ∗ h = ∑ g ∈ [ k ] n ( s ) hg .

The mis-clustering rate at iteration s can be written as

̸

̸

We define a cluster-wise mis-clustering rate at iteration s as

̸

̸

The first term in the maximum operator of definition of G s can be understood as the false positive rate of cluster h and the second term is the true negative rate of cluster h . It is easy to see the relationship that A s ≤ G s .

̸

Let ∆ = min g = h ∈ [ k ] ∥ θ g -θ h ∥ be the signal strength. For h ∈ [ k ] , let ˆ θ ( s ) h be the estimated center of cluster h at iteration s . Define our error rate of estimating centers at iteration s as

Lemma F.5. Based on our definition of M above, G 0 = O (( 1 2 -6 √ r k ) 1 λ ) .

Proof. First we show λ = O √ d . Since P c [ v ] = 1 d , we can invoke the KL upper bound above to show that ∆ λ = O ( 1 d ) , and furthermore, the KL for each pair is Θ( M ) . Thus, this tells us:

In other words the signal to noise ratio is ∆ /σ ≥ √ k 2 + k 3 d n .

Based on the above result, we also get:

̸

Recall σ for us is given by √ 1 /α . 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:

Lemma F.6.

with probability greater than 1 -exp( -0 . 3 n ) .

Lemma F.7.

with probability greater than 1 -exp( -0 . 5 n ) .

Lemma F.8. For any fixed i ∈ [ n ] , S ⊆ [ n ] , t > 0 and δ > 0 ,

with probability greater than 1 -n -3

Lemma F.10. For any fixed θ 1 , · · · , θ k ∈ R d and a > 0 , we have

̸

for all g = h ∈ [ k ] 2 with probability greater than 1 -n -3 .

The following two lemmas give the iterative relationship between the error of estimating centers and the error of estimating labels. Let E be the intersection of high probability events in Lemma F.6, Lemma F.7 Lemma F.9, Lemma F.10 and the initialization condition (29). Then we have P {E c } ≤ 3 n -3 + ν . In the rest part of the proof, if not otherwise stated, we all condition on the event E and the following analysis are deterministic.

Lemma F.11. On event E , if G s ≤ 1 2 , then we have

Lemma F.12. On event E , if Λ s ≤ 1 -ϵ 2 and r ≥ 36 ϵ -2 , then

Proof of Lemma F.11. For any B ⊆ [ n ] , define ¯ Y B = 1 | B | ∑ i ∈ B y i . The error of estimated centers at step s can be written as

̸

̸

According to our label update step, we have ∥ y i -ˆ θ ( s -1) h ∥ ≤ ∥ y i -ˆ θ ( s -1) a ∥ for any i ∈ S ah . This means for any i ∈ S ah , y i is closer to ˆ θ ( s -1) h than ˆ θ ( s -1) a , so is the average of { y i , i ∈ S ah } . Thus, we have

Consequently, triangle inequality gives us

which, combined with Lemma F.6 and the definition of Λ s -1 , yields

̸

Taking a weighted sum over a = h ∈ [ k ] , we get

̸

̸

̸

̸

where the Last inequality is due to Cauchy-Schwartz and the fact that ∑ a = h n ah ≤ G s n h . Note that W S hh = W T ∗ h -∑ a = h W S ha . Triangle inequality and Lemma F.10 imply

∥

̸

W

S

hh

Since G s ≤ 1 2 , we have

∥

3

σ

d

+log

n

n

∗

h

+

σ

3(

n

+

d

)

n

∗

h

-

n

hh

.

Combining the pieces, we obtain

Therefore, we get the first term in (34). To prove the second term, we decompose ˆ θ ( s ) h differently.

Then, the error of ˆ θ ( s ) h can be upper bounded as

By triangle inequality,

̸

̸

This, together with Lemma F.6 and (36), implies

for all h ∈ [ k ] . The proof is complete.

√

√

√

√

̸

̸

̸

Proof of Lemma F.12. For any g = h ∈ [ k ] × [ k ] ,

̸

Triangle inequality implies

Using the fact that (1 -x ) 2 -y 2 ≥ (1 -x -y ) 2 when y (1 -x -y ) ≥ 0 , we obtain

Denote by ∆ h = ˆ θ ( s ) h -θ h for h ∈ [ k ] . Then,

Taking a sum over i ∈ T ∗ g and using Markov's inequality on the second term, we obtain

Note that I { ϵ 2 4 ∥ θ g -θ h ∥ 2 ≤ ⟨ w i , θ h -θ g ⟩ } are independent Bernoulli random variables. By Lemma F.10, the first term in RHS of (43) can be upper bounded by

By Lemma F.7, the second term in RHS of (43) can be upper bounded by

Combining (43), (44) and (45) and using the fact that ∥ ∆ g -∆ h ∥ 2 ≤ 4Λ 2 s ∆ 2 , we get

Consequently,

Since Λ s ≤ 1 / 2 and r ≥ 20 ϵ -2 , the RHS of (46) is smaller that 1 / 2 when αn ≥ 32 k 2 log n . Thus,

for all h ∈ [ k ] and we have

̸

which, together with (46), implies

Under the assumptions that ϵ 4 α ∆ 2 /σ 2 ≥ r 2 ϵ 4 ≥ 36 , we have the desired result (35).

Proof. From Lemma F.11, a necessary condition for Λ 0 ≤ 1 2 -4 √ r is G 0 ≤ ( 1 2 -6 √ r ) 1 λ . Setting ϵ = 7 √ r in Lemma F.12, we have G 1 ≤ 0 . 35 . Plugging it into Lemma F.11 gives us Λ 1 ≤ 0 . 4 , under the assumption that r ≥ 16 √ k . Then it can be easily proved by induction that G s ≤ 0 . 35 and Λ s ≤ 0 . 4 for all s ≥ 1 . Consequently, Lemma F.11 yields

which, combined with (35), implies

for some constant C . Here we have chosen ϵ = 1 / 5 in Lemma 35 to get the first inequality.

F.1.1 Combining Analysis in ℓ 2 Norm and Mapping it Back KL divergence.

Proof. From the proof of Lemma F.11, the error of estimating θ h at iteration s can be written as ˆ θ ( s ) h -θ h = 1 n h W T ∗ h + u h , with

In addition, by Lemma F.11 and Lemma F.12, there is a constant C 1 such that

for all s ≥ 1 . Therefore, when r is large enough, we have

for all s ≥ log n . Then by (42), we have

In order to prove that A s attains convergence rates, we first upper bound the expectation of A s and then derive the high probability bound using Markov's inequality. Similar to the two-mixture case, we need to upper bound the inner product ⟨ w i , ∆ h -∆ g ⟩ more carefully. Note that { T ∗ h , h ∈ [ k ] } are deterministic sets, we could use concentration equalities to upper bound W T ∗ h and u h parts separately.

where β 2 and β 4 will be specified later and β = β 1 -β 2 -β 4 . Taking a sum over h ∈ [ k ] and i ∈ [ n ] , we obtain

with

Let us first consider the expectation of J 1 . Using Chernoff's bound, we have

Thus,

with γ = β 2 -8 σ 2 log k ∆ 2 ≥ β 2 -8 /r 2 .

We use Markov Inequality to upper bound J 2 . Markov's inequality and Lemma F.7 give us

(48) implies

Here the second inequality is due to the fact that G s ≤ A s /α . And we choose β 2 = √ 8 k/r in the last equality.

Finally, we upper bound the expectation of J 3 . Given z i = g , we have

Choosing t = max { √ d ∆ σ , ∆ 2 σ 2 } , δ = exp ( -∆ 2 4 σ 2 ) in Lemma F.8, and

we obtain Pr { β 4 ∆ 2 ≤ 2 ⟨ w i , v g -v h ⟩ } ≤ 2 exp( -∆ 2 / (4 σ 2 )) , where we have used the assumption that n ∗ g ≥ αn and αn ≥ 36 r 2 . Thus,

Combining the pieces, we have

with γ = ( β 1 -√ 8 k/r -64 /r ) 2 -8 /r 2 = 1 -o (1) . Here only prove the case that r →∞ . For the finite case, all the o (1) in the following proof can be substituted by a small constant.

when s ≥ 4 log n . By Markov's inequality, for any t > 0 ,

If (1 -η ) ∆ 2 8 σ 2 ≤ 2 log n , choose t = exp ( -(1 -η -8 σ ∆ ) ∆ 2 8 σ 2 ) and we have

Otherwise, since A s only takes discrete values of { 0 , 1 n , · · · , 1 } , choosing t = 1 n in (52) leads to

Now, from Lemma F.3 and Lemma F.4, for P i [ v ] = Ω(1) , ∀ i , we conclude that KL( P i ∥ P j ) = Θ( ∥ P i -P j ∥ 2 ) . Applying the result from above we then get:

Since n c = Ω( k 2 ) we have n = Ω( k 3 ) , which implies:

which is the first condition in the EM analysis. The second condition requires the KL to be sandwiched by ℓ 2 norm. But this means that we we can replace λ with √ λ and ∆ with √ ∆ in the result in Eq. 52, and complete the proof of Theorem 3.2.

F.2 Privacy Analysis from Section 5

F.2.1 Technical Lemmas

We begin by recalling the definition of differential privacy, and the variant of concentrated differential privacy that we use in this work.

Definition F.13 (Differential Privacy (DP) [24]) . A randomized algorithm M : X n →Y satisfies ( ϵ, δ ) -differential privacy ( ( ϵ, δ ) -DP) if for every pair of neighboring datasets X,X ′ ∈ X n (i.e., datasets that differ in exactly one entry),

When δ = 0 , we say that M satisfies ϵ -differential privacy or pure differential privacy.

Definition F.14 (Concentrated Differential Privacy (zCDP) [15]) . A randomized algorithm M : X n →Y satisfies ρ -zCDP if for every pair of neighboring datasets X,X ′ ∈ X n ,

where D α ( M ( X ) || M ( X ′ )) is the α -Rényi divergence between M ( X ) and M ( X ′ ) . 2

Note that zCDP and DP are on different scales, but otherwise can be ordered from most-to-least restrictive. Specifically, ( ϵ, 0) -DP implies ϵ 2 2 -zCDP, which implies roughly ( ϵ √ 2 log(1 /δ ) , δ ) -DP for every δ > 0 [15].

Both these definitions are closed under post-processing and can be composed with graceful degradation of the privacy parameters.

Lemma F.15 (Post Processing [15]) . If M : X n → Y is ( ϵ, δ ) -DP, and P : Y → Z is any randomized function, then the algorithm P ◦ M is ( ϵ, δ ) -DP. Similarly if M is ρ -zCDP then the algorithm P ◦ M is ρ -zCDP.

Lemma F.16 (Composition of CDP [15]) . If M is an adaptive composition of differentially private algorithms M 1 , . . . , M T , then

We can achieve differential privacy via noise addition proportional to sensitivity [24].

Definition F.17 (Sensitivity) . Let f : X n → R d be a function, its ℓ 2 -sensitivity is defined to be ∆ f, 2 = max X ∼ X ′ ∈X n ∥ f ( X ) -f ( X ′ ) ∥ 2 , Here, X ∼ X ′ denotes that X and X ′ are neighboring datasets (i.e., those that differ in exactly one entry).

For functions with bounded ℓ 1 -sensitivity, we can achieve ϵ -DP by adding noise from a Laplace distribution proportional to ℓ 1 -sensitivity. For functions taking values in R d for large d it is more useful to add noise from a Gaussian distribution proportional to the ℓ 2 -sensitivity, to get ( ϵ, δ ) -DP and ρ -zCDP.

Lemma F.18 (Gaussian Mechanism) . Let f : X n → R d be a function with ℓ 2 -sensitivity ∆ f, 2 . Then the Gaussian mechanism

satisfies ρ -zCDP.

F.2.2 Proof of Theorem 3.3

From the Lemmas above for Gaussian and Laplace mechanism, it is easy to see that the adaptive clipping step in Algorithm 4 satisfies 2 ρ/ 3 -zCDP, where as the Laplace mechanism is ρ/ 3 -zCDP. Then, we can use composition lemma to complete show ρ -zCDP of Algorithm 4.

For Algorithm 5, each step uses satisfies ρ/K -zCDP from the guarantees for exponential mechanism, since we clip the KL difference with c . Here too, composition gives us the result.

For the final end-to-end algorithm, we use the private aggregation scheme, we describe, and changing the user data would only affect two cluster estimates, but since we use add-one, leave-one definition of privacy, the entire step is ρ/ 3 T -zCDP. Finally, composing over the T rounds of clustering and with initial private center and private init we get the final result. We then use Lemma 3.6 from [15] to convert zCDP guarantees to ( ε, δ ) -DP.

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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 are not discussed in the paper.

The authors are encouraged to create a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

Justification: Theorem statements are in the main paper and proofs in Appendix E and Appendix F.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs 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 to 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 be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experimental details are in Section 5 and Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: While the datasets we use are all open-source with publicly available links, code us not yet ready for anonymous open sourcing. We plan to 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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence 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 are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: In Appendix A we discuss some limitations of our work. This paper presents work whose goal is to advance the field of Private Federated Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work is not focused on releasing any pretrained models or datasets.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We make sure to cite relevant work including references to the three publicly available federated learning datasets we use in our work.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We do not do any crowdsourcing experiments.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We do not conduct any experiments with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors addressed the limitations and negative impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
-Private and personalized histogram estimation is an important problem.

-The paper is easy to understand, presentation and justification of the methods are well organized.

-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.

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

- From the current experiments, performance dependence on variables such as number of clusters, clients, data samples per clients, heterogeneity etc. is not clear. I would suggest doing more ablations studies possibly using e.g. synthetic datasets to understand these.

QUESTIONS:
See weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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$?\n\nLIMITATIONS:\nsee above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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 »\n\nQUESTIONS:\nWhat 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.\n\nLIMITATIONS:\nThe authors have a specific section discussing the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-08-02"}