Collaborative Learning of Discrete Distributions under Heterogeneity and Communication Constraints
Anonymous Author(s)
AffiliationAddress email
Abstract
In modern machine learning, users often have to collaborate to learn distributionsthat generate the data.
Communication can be a significant bottleneck.
Priorwork has studied homogeneous users—i.e., whose data follow the same discretedistribution—and has provided optimal communication-efficient methods.
However, these methods rely heavily on homogeneity, and are less applicable in thecommon case when users’ discrete distributions are heterogeneous.
Here weconsider a natural and tractable model of heterogeneity, where users’ discretedistributions only vary sparsely, on a small number of entries.
We propose a noveltwo-stage method named SHIFT:
First, the users collaborate by communicatingwith the server to learn a central distribution; relying on methods from robuststatistics.
Then, the learned central distribution is fine-tuned to estimate the individual distributions of users.
We show that our method is minimax optimal in ourmodel of heterogeneity and under communication constraints.
Further, we provideexperimental results using both synthetic data and n -gram frequency estimation inthe text domain, which corroborate its efficiency.
Introduction
Research on learning from distributed data has grown in recent years, as data is commonly generatedby multiple users, such as smart devices and wireless sensors.
Communication costs and bandwidthare often bottlenecks on the performance of learning algorithms in these settings [23, 6, 18].
Thebottlenecks become even more severe in federated learning and analytics [31] where many userscoordinate with a server to learn a central model while the wireless links are typically expensive andoperate at low rates.
This paper considers learning high-dimensional discrete distributions from user data in the distributedsetting.
Many communication-efficient methods have been proposed and their optimality undercommunication constraints has been established under various models [27, 7, 28, 2, 3, 15, 16].
However, the key challenge of heterogeneity , i.e., that users’ distributions can differ, is rarely considered.
Heterogeneity is common, as users inevitably have unique characteristics [14].
Meanwhile, heterogeneity can cause a significant performance drop for learning algorithms designed only for i.i.ddata [37, 34, 21].
To use all the data, one needs to learn some central structure, transferable to allindividual users.
Then one may locally learn the unique components for each user [22, 51, 17, 54].
To study this paradigm, we first need to introduce a suitable model of heterogeneity.
We consider,as an example, the heterogeneous frequencies of words across different texts, e.g., news articles,books, plays (tragedies and comedies), viewed as users.
Most words appear with almost the same
[Table]
∗ This method [16] requires interactive communication protocols, while other methods are non-interactive.
probabilities in different texts, however, a few can have very different probabilities, such as “sorrowful”being common in tragedies and “convivial” being common in comedies.
Motivated by this, weformulate a model of sparse heterogeneity.
Specifically, suppose that the discrete distributions of allusers differ from an underlying central distribution in at most s entries, where s is much smaller thanthe dimension d .
Sparse heterogeneity is relevant to applications such as recommendation systems[29, 41, 35, 8] and medical risk scoring [48, 42, 40].
However, given data generated by multiple distributions with sparse heterogeneity, previous works[27, 2, 3, 16] either do not use all the data, or suffer from bias due to heterogeneity that does not vanishas the sample size increases.
Here we propose a novel s parse h eterogeneity i nspired collaborationand f ine- t uning method (SHIFT) where we first collaboratively learn the central distribution, and thenfine-tune the central estimate to individual distributions.
Our method makes full use of heterogeneousdata, leading to a significant improvement in error rates compared to prior methods.
See Table 1 foran overview, explained in detail later.
1.1 Contributions
We consider the problem of learning d -dimensional distributions with s -sparse heterogeneity.
Weassume there are T clusters of user datapoints, and allow each datapoint to be transmitted with bbits of information to the server.
Our setting embraces heterogeneous data and thus is a significantgeneralization of the models from [27, 7, 28, 2].
Our technical contributions are as follows:
• We propose the SHIFT method to learn heterogeneous distributions with collaboration
and tuning, in a sample-efficient manner.
Our method can in principle be used with an
arbitrary robust estimate for the probability of each entry/coordinate.
When entry-wise
median and trimmed mean are used, we provide upper bounds on the estimation error of
individual distributions in the ℓ 2 and ℓ 1 norms.
We show a factor of min { T, d/ max { s, 2 b }}
improvement in sample complexity compared to previous works; showing the benefit
of collaboration (large T ) and sparsity (small s )
, despite communication constraints and
heterogeneity.
• To justify the optimality of our method, we prove minimax lower bounds on the estimation
errors for individual distributions in the ℓ 2 and ℓ 1 norms, holding for all, possibly interactive,
methods.
These lower bounds, combined with our upper bounds, imply that our median-
based method is minimax optimal.
• We support our method with experiments on both synthetic and empirical datasets, showing
a significant improvement over previous methods.
1.2 Related Works
Learning with Heterogeneity.
Learning with heterogeneity is commonly found in the broadercontext of multi-task learning [13, 50, 9] and federated learning [5, 43, 25, 45], where a central modelor representation is learned from multiple heterogeneous datasets.
These central representations canbe useful for few-shot learning tasks [49, 24] due to their ability to adapt to new tasks efficiently.
In heterogeneous linear regression, [51, 22] show improved sample complexities by assuming alow dimensional central representation, compared to the i.i.d.
setting [25, 36].
Related results areproved by [17] for personalized federated learning.
[54] study a bandit problem where the unknownparameter in each dataset equals a global parameter plus a sparse instance-specific term.
We study adifferent setting of learning distributions with sparse heterogeneity under communication constraints.
Estimating Distributions under Communication Constraints.
Estimating discrete distributionshas a rich literature [11, 19, 44, 20].
Under communication constraints, [27, 7, 28, 2] consider thenon-interactive scenario and establish the minimax optimal rates, in terms of data dimension andcommunication budget, via potentially shared randomness, when all users’ data is homogeneous.
The optimality for the general interactive (blackboard) methods is developed by [1].
A few worksstudy the estimation of sparse distributions.
[3, 15] consider s -sparse distributions and establishesminimax optimal rates under communication and privacy constraints with various minimal samplesizes needed.
Complementary to minimax rates, [16] provides pointwise rates, governed by thehalf-norm of the distribution instead of its dimension.
Our setting embraces heterogeneous data, andthus is a generalization of the one studied in above works.
Robust Estimation & Learning.
Robust statistics and learning study algorithms resilient tounknown data corruption [30, 26].
The median-of-means method [32, 38, 39], partitions the datainto subsets, computes an estimate from each, and takes their median.
Similarly, some works studyrobustness from the optimization perspective, proposing to robustly aggregate gradients of the lossfunctions [46, 47, 10, 55].
We adapt some analysis techniques from [39, 55] to the significantlydifferent setting of estimation with heterogeneity and communication constraints.
1.3 Notations
Throughout the paper, for an integer d ≥ 1 , we write [ d ] for both { 1 , . . .
, d } and { e 1 , . . .
, e d } ⊆ R d ,where e k is the k -th canonical basis vector of R d .
For a vector v ∈ R d , we refer to the entries of vby both [ v ] 1 , . . .
, [ v ] d and v 1 , . . .
, v d .
We denote ∥ v ∥ p
= ( (cid:80) k ∈
[ d ] | v k | p ) 1 p for all p > 0 with ∥ v ∥ 0defined additionally as the number of non-zero entries.
We let P d := { p = ( p 1 , . . .
, p d ) ∈ [0 , 1] d :p 1 + · · · + p d = 1 } be the simplex of all d -dimensional discrete probability distributions.
For p ∈ P d ,we denote by B s ( p )
the s -distinct neighborhood { p
′ ∈ P d : ∥ p ′ − p ∥ 0 ≤ s } .
For a random variable
X , we denote n i.i.d.
copies of X by X [ n ] .
Given any index set I , we write |I| for its cardinality anddenote by [ v ] I the sub-vector
([ v ] k ) k ∈I indexed by I .
We use the Bachmann-Landau asymptoticnotations Ω( · ) , Θ( · ) , O ( · ) to hide constant factors, and use ˜Ω( · ) , ˜ O ( · ) to also hide logarithmic factors.
We denote the categorical distribution with class probability vector p ∈ P
d by Cat( p ) .
We use q · and(cid:98) · to indicate the intermediate estimate and the final estimate, respectively.
Problem Setup
We consider the problem of collaboratively learning distributions defined according to the followingmodel of heterogeneity (see Figure 1 for an illustration).
There are T ≥ 1 clusters {C t ≜ ( X t,j ) j ∈ [ n ] :t ∈ [ T ] } of user datapoints, each of which contains n i.i.d.
local datapoints.
Each datapoint X t,j isin a one-hot format, i.e., X t,j ∈ { e 1 , e 2 , . . .
, e d } , and follows the categorical distribution Cat( p t )
where p t ∈ P
d is unknown.
Thus, user datapoints in the same cluster have an identical distributionp t , while the distribution p t can vary, i.e., be heterogeneous, across
clusters t ∈
[ T ] .
The datapoint
X t,j is encoded by its user into a message Y t,j , and then transmitted to a central server.
We assumethe server knows which cluster t ∈
[ T ] each Y t,j belongs to, as well as the number of clusters T .
Weassume that the communication with the server for each datapoint takes no more than b bits.
Thispaper mainly addresses the communication bottleneck and does not involve privacy concerns.
The goal here is to collaboratively learn the distributions p t from the collection of messages { Y t ′ , [ n ] ≜( Y t ′ ,j ) j ∈ [ n ] : t ′ ∈
[ T ] } despite heterogeneity.
More precisely, we aim to design per-cluster estimators
(cid:98) p t : { Y t ′ , [ n ] : t ′ ∈
[ T ] } → P d to minimize the ℓ 2 errors
[Equation]
for all t ∈ [ T ] .
We also study the widely-used ℓ 1 error metric (in addition to the ℓ 2 metric).
When T = 1 , i.e., alluser datapoints are homogeneous and there is a single distribution to learn, the problem reduces tothe one studied by [27, 7, 28, 2, 1].
Model of heterogeneity.
In heterogeneous settings, collaboration among the users is most beneficial if the local distributions are related.
Wemodel this by assuming that the local distributions are sparse perturbations of an unknowncentral distribution p ⋆ ∈ P d .
The distributionp t of each cluster t differs from p ⋆ in at mosts ≥ 0 entries:
[Equation]
The central distribution p ⋆ can be viewed asthe central structure across heterogeneous clusters of datapoints.
The level of heterogeneity iscontrolled by the parameter s .
When s is muchsmaller than d , the local distributions differ fromthe center in a small number of entries.
While motivated by word frequencies of different texts, our model of sparse heterogeneity is also relevant for recommendation systems, where thehigh-dimensional item-preference vectors of users can vary sparsely [29, 41, 35, 8]; and medicalrisk scoring, where hospitals enjoy similar predicted results but with a few systematic differences indiagnosis behavior, healthcare utilization, etc [48, 42, 40].
Algorithm
We now introduce our method for leveraging heterogeneous data to improve per-cluster sampleefficiency.
We first discuss a hashing-based method to handle the communication constraint.
Sincecommunication between each user datapoint and the server is restricted to at most b > 0 bits(where we may have 2 b ≪ d ), the datapoint X t,j needs to be encoded by an encoding function
W t,j : X ≜
[ d ] → Y .
The b -bit constraint enforces |Y| ≤ 2 b .
Then the encoded message
Y t,j :=
W t,j ( X t,j ) is sent to the server, where it is decoded and used.
Under relatively sophisticated protocols, the design of encoding functions can be interactive [16, 15,1], i.e., depend on previously sent messages.
Here we adopt a non-interactive encoding-decodingscheme, based on uniform hashing [4, 15] where W t,j depends only on X t,j and is independentof other messages.
Specifically, each datapoint X t,j is encoded via an independent random hashfunction h t,j : [ d ] → [2 b ] .
Upon receiving all messages, the server counts the empirical frequenciesof all symbols, leading to hashed estimates q b t .
The communication scheme based on uniform hashingis summarized below.
(Encoding) :
Send the message Y t,j = h t,j ( X t,j ) encoded by a hash function h t,j : [ d ] → [2 b ]; (Decoding) :
Count N tk ( Y t, [ n ] ) =
|{ j ∈ [ n ] : h t,j ( k ) =
Y t,j }| and return [ q b t ] k
= N tk /n. It is easy to verify that E [ q b t ] =
[(2 b − 1) p t + 1] / 2 b ≜ b t ; and thus the hashed estimate q b t is biasedfor p t .
We also write b ⋆ =
[(2 b − 1) p ⋆ + 1] / 2 b for the mean of a hashed datapoint sampled fromthe central distribution.
More details on the hashed estimator q b t are given in Appendix A. 3.1 The SHIFT Method We now introduce the SHIFT method, which consists of two stages: collaborative learning andfine-tuning . The first stage estimates the central hashed distribution b ⋆ using all hashed estimates
Algorithm input: individual hashed estimators q b 1 , , . . .
, q b T , threshold parameter α ▷ Stage
I: Collaborative Learning
Estimate b ⋆ via robust statistical methods: q b ⋆ ← robust _ estimate( { q b t : t ∈ [ T ] } )
▷ Stage II: Fine-Tuning for k = 1 , . . .
, d do for t = 1 , . . .
, T do [ (cid:98) b t ] k ← [ q b ⋆ ] k if | [ q b ⋆ ] k − [ q b t ] k | ≤ (cid:113) α
[ q b t ] k n else [ q b t ] k
[ (cid:98) p t ] k ← Proj [0 , 1] ( 2 b [ (cid:98) b t ] k − 1 2 b − 1 ) end for end for output: estimates (cid:98) p 1 , . .
, (cid:98) p T { q b t : t ∈ [ T ] } .
This is achieved via tools from robust statistics such as the median or trimmed mean.
The key insight here is that, since the heterogeneity is sparse, for each entry at which the individualdistributions less mismatch with the central one, most datapoints (used to estimate that entry) aresampled from the probability of the central distribution.
Hence, to estimate those entries of the centraldistribution, we can treat the datapoints generated by heterogeneous users as corrupted, and leveragerobust statistical methods to mitigate their influences.
In the second stage—fine-tuning—we detect mismatched entries between individual hashed estimatesq b t and the central estimate q b ⋆ .
Recall that the central and individual distributions differ in only afew entries.
For entries k ∈ [ d ] such that | [ q b ⋆ ] k − [ q b t ] k | is below ( α [ q b ⋆ ] k /n )
1 / 2 for some thresholdparameter
α , we may expect that p tk = p ⋆k .
As a result, the estimate [ q b ⋆ ] k of b ⋆k can be more accurateas it is learned collaboratively.
Thus, we assign [ q b ⋆ ] k as the final estimate [ (cid:98) b t ] k of b tk .
On the other hand, for the entries where the central and individual distributions differ, i.e., p tk ̸ = p ⋆k ,the threshold is more likely to be exceeded.
In this case, we keep the individual estimate [ q b t ] k as[ (cid:98) b t ] k .
Finally, since the hashed distributions b t are biased, we debias them in the final estimates ofp t where Proj [0 , 1] ( · ) indicates truncating the input with upper/lower bound 1 and 0 .
Our methoddoes not require sample splitting, despite using two stages, leading to increased sample-efficiency.
We remark that the SHIFT does not require the knowledge of the sparse heterogeneity s .
Knowledge transfer to new clusters.
The (collaboratively) learned central distribution from Algorithm 1 is adaptable to new clusters, with possibly few datapoints.
This can be particularly beneficialfor sample efficiency because most entries of the target distribution are well-estimated throughcollaborative learning.
One can transfer those entries, and it suffices to estimate the few remainingentries, instead of the whole distribution.
See Theorem 2 for the details.
The knowledge transferutility further motivates the importance of collaborative learning.
3.2 The Median-Based SHIFT
In this section, we provide statistical guarantees, in the form of upper bounds on the error, forthe median-based SHIFT method, where robust _ estimate( { q b t : t ∈ [ T ] } ) in Algorithm 1 is theentry-wise median.
Specifically, we let
[Equation]
for each k ∈ [ d ] .
When there is no ambiguity, we write q b ⋆ = median (cid:0) { q b t } t ∈
[ T ] (cid:1) .
We also provide results for thetrimmed-mean-based SHIFT method, see Appendix D.
By setting the threshold parameter α in Algorithm 1 as α = Θ(ln( n )) , we prove the upper bounds onthe final individual ℓ 2 estimation errors as follows.
The results for the ℓ 1 error are in Appendix C.
Theorem 1.
Suppose n ≥ 2 b +6 ln( n ) and α = Θ(ln( n )) 1 .
Then, for the median-based SHIFTmethod, for any t ∈ [ T ] ,
[Equation]
When n ≥ 2 b +6 ln( n ) and n = Ω(2 b min { T, d/ max { 2 b , s }} ) , the rate further becomes
[Equation]
The upper bound in (2) consists of two terms.
The first term max { 2 b , s } / (2 b n ) is independent ofthe dimension d , and is a factor d/ max { 2 b , s } smaller compared to the rate d/ (2 b n ) obtained bythe minimax optimal method using only homogeneous datapoints [27].
Thus, it brings a significantbenefit under sparse heterogeneity, i.e., when s ≪ d .
Meanwhile, the second, dimension-dependent,term d/ (2 b Tn ) is T times smaller than d/ (2 b n ) , since it depends on the total sample-size Tn used collaboratively, despite heterogeneity.
Therefore, our method shows a factor of min { T, d/ max { 2 b , s }}improvement in sample efficiency, compared to previous work designed for homogeneous datapoints.
For completeness, we also consider a heuristic version of estimators of prior works [27, 16] inwhich all datapoints are pooled to learn a global distribution T − 1 (cid:80) t
∈ [ T ] p t , which is then usedby each cluster.
While this uses all datapoints, it inevitably introduces a non-vanishing bias δ t =p t − T − 1 (cid:80) t ′ ∈
[ T ] p t ′ in estimating individual distributions, and can behave poorly when the bias islarge.
See Table 1 for more details.
Finally, we discuss our results on knowledge transfer.
The central estimator q b ⋆ is adaptable to anew cluster C T +1 in the following way.
We adjust the fine-tuning procedure in Algorithm 1 to[ (cid:98) b
T
+1 ] k ← [ q b ⋆ ] k if | [ q b ⋆ ] k − [ q b T +1 ]
k | ≤ (cid:113) α
[ q b T +1 ] k / ˜ n , and [ (cid:98) b T +1 ]
k ← [ q b T +1 ] k otherwise,where ˜ n is the size of C T +1 .
We then show the following result.
Theorem 2.
Let q b T
+1 be the hashed estimate of any new cluster C T +1 with ˜ n datapoints suchthat n ≥ ˜ n ≥ 2 b
+6 ln(˜ n ) and ˜ n = Ω(2 b min { T, d/ max { 2 b , s }} ) .
Let the threshold parameter beα = Θ(ln(˜ n )) .
Then, the median-based SHIFT method has error bounded by
[Equation]
Similarly, one can see that adaptation to new clusters with the median-based SHIFT method achieves afactor of min { Tn/ ˜ n, d/ max { 2 b , s }} improvement in sample-efficiency compared to the case wherethe distribution of the new cluster is estimated from scratch (i.e. without any knowledge transfer).
3.2.1 Highlights of Theoretical Analysis
In this section, we introduce the key analytical ideas behind the proof of Theorems 1 and 2.
Ouranalysis is novel compared to previous analyses for methods with homogeneous datapoints.
Thefinal individual estimation errors relate to the error of estimating the central hashed distribution b ⋆ .
However, we only expect high accuracy at the center for entries at which the heterogeneity is mild.
To quantify the influence of heterogeneity, for any 0 < η ≤ 1 , we define the set of η -well-alignedentries as I η := { k ∈ [ d ] : |B k | < ηT } , where B k ≜ { t ∈ [ T ] : b tk
̸ = b ⋆k , i.e., p tk ̸ = p ⋆k } is the set ofclusters whose distribution differs from p ⋆ in the k -th entry.
We aim to estimate the η -well-alignedentries accurately by using robust statistical methods.
Further, we argue that there are few poorly-aligned entries, and they affect the final per-cluster erroronly mildly.
By the pigeonhole principle, the number of entries that are not η -well-aligned is upperbounded by |I cη | ≜ | [ d ] \I η | ≤ sT ηT = s η .
Therefore, given an estimator q b ⋆ that is accurate for theη -well-aligned entries, the entries of b ⋆ can be estimated accurately except for at most s/η entries.
The following technical lemma bounds the error for each entry k ∈ I η .
Lemma 1.
Suppose q b ⋆ = median (cid:0) { q b t } t ∈ [ T ] (cid:1) .
Then for any 0 < η ≤ 1 / 5 and k ∈ I η , it holds that
[Equation]
Tn [Equation] [Equation] Lemma 1 provides an upper bound relating to the frequency |B k | /T of misalignment (smaller thanη ), and a variance term b ⋆k (1 − b ⋆k ) .
This result cannot be obtained by directly applying the standard
Chernoff or Hoeffding bounds to random variables distributed in [0 , 1] as in previous works [16]for two reasons: 1) the datapoints are heterogeneous, 2) the variance b
⋆k (1 − b ⋆k ) here can be small,compared to general random variables in [0 , 1] , implying more concentration than follows from
Hoeffding’s inequality.
To address these issues, we analyze the concentration of the empirical(1 / 2 ± |B k | /T )
-quantiles to mitigate the influence of heterogeneity, and we also use Bernstein’sinequality, which is variance-dependent [52], to obtain bounds relying on both the sample size Tnand the variance b ⋆k (1 − b ⋆k ) .
Also, the constant 1 / 5 , controlling the heterogeneity, is not essential (we choose 1 / 5 for clarity).
Itcan be replaced with any number below 1 / 2 so that estimating the central probability distributionbecomes possible, as the information conveyed by homogeneous datapoints dominates.
Lemma 1 reveals that well-aligned entries of the central distribution are accurately estimated.
Thusone can use the central estimate for the entries where the central distribution p t aligns with the targetdistribution p t .
The remaining entries, that are neither well-aligned nor satisfy p ⋆k =
p tk , can beestimated by the individual estimator.
We argue that a properly chosen threshold parameter α filters out the desired entries to be estimated individually with high probability, leading to Theorems 1 and2.
While estimating p ⋆ is not our main goal, one can readily obtain from Lemma 1 the following boundon estimating p ⋆ by summing up the errors for all entries k ∈ [ d ] = I η with η = max k ∈ [ d ] |B k | .
Corollary 1 reveals that the central distribution can be accurately estimated if the mismatching ofdistributions happens uniformly across all entries, i.e., each entry differs in O ( sT/d ) of clusters.
Corollary 1.
Let (cid:98) p
⋆ = Proj [0 , 1] ( 2 b q b ⋆ − 1 2 b − 1 ) be obtained by the debiasing operation from Algorithm1.
Suppose |B k | =
O ( sT/d ) for any k ∈ [ d ] , with η = max k ∈ [ d ] |B k | /T , the median-based SHIFTmethod enjoys
[Equation]
Lower Bounds
To complement our upper bounds, we now provide minimax lower bounds for estimating distributionsunder heterogeneity.
Since our setting contains T heterogeneous clusters of datapoints, our minimaxerror metric is slightly different from the one studied in [27, 7, 28, 2].
Using the ℓ 2 error as the loss,the lower bound metric is defined as
[Equation]
[Equation] where the supremum is taken over all possible central distributions p ⋆ ∈ P d and individual distributions { p t : t ∈ [ T ] } in B s ( p ⋆ ) ≜ { p ∈ P d : ∥ p − p ⋆ ∥ 0 ≤ s } , and the infimum is taken over all estimation methods ˆ p t that use all heterogeneous messages { Y t ′ ,j ≜ W t ′ ,j ( X t ′ ,j ) : j ∈ [ n ] , t ′ ∈
[ T ] }encoded (possibly interactively) by any encoding functions { W t ′ ,j : j ∈ [ n ] ,
t ′ ∈
[ T ] } with outputin [2 b ] , e.g., the random hashing maps.
The measure (3) characterizes the best possible worst-caseperformance of estimating distributions under our model of heterogeneity.
Since the supremum is taken over all distributions p ⋆ , p 1 , . .
. , p T in P d such that ∥ p t − p ⋆ ∥ 0 ≤ sfor all t ∈ [ T ] , we consider two representative cases therein: 1)
The homogeneous case wherep 1 = · · · = p T = p ⋆ ∈ P d .
Then the setting essentially reduces to the single-cluster problem butwith nT datapoints, and the goal here is to estimate p ⋆ , leading to the lower bound Ω( d/ (2 b Tn )) .
The s/ 2 -sparse case where ∥ p ⋆ ∥ 0 ≤ s/ 2 and ∥ p t ∥ 0 ≤ s/ 2 for all t ∈
[ T ] .
Then it naturally holdsthat { p t : t ∈
[ T ] } ⊆
B s ( p ⋆ ) .
By constructing independent priors for { p t : t ∈ [ T ] } and p ⋆ , onecan show that only datapoints generated by p t itself are informative for estimating p t .
In this case,we show the lower bound Ω(max { 2 b , s } / (2 b n )) .
Combining the two cases, we find the followinglower bound.
The formal argument is provided in Appendix E.
Theorem 3.
For any—possibly interactive—estimation method, and for any t ∈ [ T ] , we have
[Equation]
By a similar argument but with an additional ( T + 1) -st cluster of ˜ n users, we obtain a lower boundfor adapting to a new cluster.
Theorem 4.
For any—possibly interactive—estimation method, and a new cluster C T +1 , we have
[Equation]
Theorem 3 and 4, combined with the upper bounds in Section 3, imply that our method is minimaxoptimal up to logarithmic terms.
We provide similar lower bounds for the ℓ 1 error in Appendix E.
Experiments
We test SHIFT on synthetic data as well as the Shakespeare dataset [12].
As a baseline method,we use the estimator in [27] that is minimax optimal under homogeneity.
We apply the baselinemethod both locally and globally.
In the local case, the estimator (cid:98) p t for each cluster is computedwithout datapoints from other clusters.
In the global case, we pool data from all clusters, and computeestimators (cid:98) p
= (cid:98) p 1 = · · · = (cid:98) p T .
The performance measure for estimating (cid:98) p t , t ∈ [ T ] is taken as
T − 1 (cid:80) Tt =1 ∥
p t − (cid:98) p t ∥ 22 .
5.1 Synthetic Data
We set the uniform distribution, p ⋆ = (1 /d, . . .
, 1 /d )
as the central distribution.
In Appendix F, wealso experiment on the truncated geometric distribution and compare our method with the localizationrefinement method [16].
Among the d entries of p ⋆ , we draw s entries uniformly at random andassign new values for them uniformly at random over [0 , 1] , with re-normalization to preserve theirsum.
We repeat this procedure T times to obtain sparsely perturbed distributions p 1 , . .
, p T ∈ P d .
Then, n i.i.d.
datapoints X t, 1 , . . .
, X t,n ∼ Cat( p t ) are generated for each cluster t ∈ [ T ] .
We setthe dimension to d = 300 and run the simulation by varying n, T, s .
As we see from (2), the error ofour method depends on s only when 2 b < s .
For this reason, we let b = 2 in our experiments.
We run SHIFT with the entry-wise median and entry-wise trimmed mean as the robust estimate.
Weset the threshold parameter α
=
ln( n ) and the trimming proportion ω
= 0 .
In Appendix F, weprovide results for other choices of the hyperparameters α, ω and discuss a heuristic for choosingα .
Figure 2 illustrates that our method outperforms the baseline method for most choices of n, T, s .
Specifically, as Theorem 1 predicts, the ℓ 2 error of our method decreases as the number of clusters Tincreases.
On the other hand, when the baseline methods are applied globally, without consideringheterogeneity, they show a bias that does not disappear as the sample size n or the number ofclusters T increases.
This shows that the fine-tuning step in SHIFT is crucial for the estimation ofheterogeneous distributions.
Finally, the right panel of Figure 2 shows that our method is effectiveonly when s is small compared to the dimension d ; which highlights the crucial role of sparseheterogeneity.
When s is close to d , the distributions p 1 , . . .
p T could be considerably differentwithout any meaningful central structure, making collaboration less useful than local estimation.
5.2 Shakespeare Dataset
The Shakespeare dataset was proposed as a benchmark for federated learning in [12].
The datasetconsists of dialogues of 1,129 speaking roles in Shakespeare’s 35 different plays.
In our experiment, [Figure] we study the distribution of k -grams ( k -tuples of consecutive letters from the 26-letter Englishalphabet, see Chapter 3 of [33]) appearing in the dialogues.
We consider each play as a cluster C tand estimate
the distribution p t ∈ P d , d = 26 k of k -grams.
Since the ground-truth distribution p t isunknown, we regard the empirical frequency as p t .
To verify the heterogeneity, we run the chi-squared goodness-of-fit test for each pair of distributionsfrom distinct clusters p u and p v .
Resulting p-values were essentially zero within machine precision,which suggests that the distributions of k -grams are strongly heterogeneous.
We also perform entrywise tests comparing [ p u ] i and [ p v ] i for all u ̸ = v ∈
[ T ] , i ∈
[ d ] .
In total, 25.8% of the tests wererejected at the significance level of 5%.
This again supports the heterogeneity.
We draw n = 20 , 000 datapoints with replacement from each cluster and test SHIFT with communication budgets b ∈ { 2 , 4 , 6 , 8 } .
We set the fine-tuning threshold α = ln( n ) and the trimmingproportion ω
= 0 .
1 , which we choose following the heuristic discussed in Appendix F. We repeatthe experiment ten times by taking different datapoints and report the average ℓ 2 error of estimationin Table 2.
The standard deviations are small even over ten repetitions.
SHIFT shows competitiveperformance on the empirical dataset even without the sparse heterogeneity (1) rigorously verified.
[Table]
Conclusion and Future Directions
We formulate the problem of learning distributions under sparse heterogeneity and communicationconstraints. We propose the SHIFT method, which first learns a central distribution, and then finetunes the estimate to individual distributions. We provide both theoretical and experimental results toshow its sample-efficiency improvement compared to classical methods that target only homogeneousdata. We provide both theoretical and experimental results to corroborate its efficacy in leveragingheterogeneous data.
Many interesting directions to be explored, including investigating if there is apoint-wise optimal method with rate depending on { p t : t ∈ [ T ] } and p ⋆ ; and designing methods forother information constraints, such