Collaborative Learning of Discrete Distributions under Heterogeneity and Communication Constraints
Abstract
In modern machine learning, users often have to collaborate to learn distributions that generate the data.
Communication can be a significant bottleneck.
Prior work has studied homogeneous users— i.e. , whose data follow the same discrete distribution—and has provided optimal communication-efficient methods.
However, these methods rely heavily on homogeneity, and are less applicable in the common case when users’ discrete distributions are heterogeneous.
Here we consider a natural and tractable model of heterogeneity, where users’ discrete distributions only vary sparsely, on a small number of entries.
We propose a novel two-stage method named SHIFT:
First, the users collaborate by communicating with the server to learn a central distribution; relying on methods from robust statistics.
Then, the learned central distribution is fine-tuned to estimate the individual distributions of users.
We show that our method is minimax optimal in our model of heterogeneity and under communication constraints.
Further, we provide experimental results using both synthetic data and n -gram frequency estimation in the text domain, which corroborate its efficiency.
Introduction
Research on learning from data distributed over multiple computational units (machines, users, devices) has grown in recent years, as data is commonly generated by multiple users, such as smart devices and wireless sensors.
While many works focus on learning predictive models with distributed data, learning the data distribution itself is also increasingly popular [56, 3, 14, 15, 60].
In various applications, it is often required to reconstruct the data distribution from scattered measurements.
Examples include sensor networks and P2P (Peer2Peer) systems, load balancing and query processing (see, e.g. , [41, 61, 46] and references therein).
In these scenarios, communication costs and bandwidth are often bottlenecks on the performance of learning algorithms [22, 5, 29, 17].
The bottlenecks become even more severe in federated analytics [31], where many users coordinate with a server to learn central models, while communication via the wireless links is typically expensive and operates at low rates.
[Table]
This paper considers learning high-dimensional discrete distributions from user data in the distributed setting.
Many communication-efficient methods have been proposed
[Table]
and their optimality under communication constraints has been established under various models [26, 6, 27, 2, 3, 14, 15].
However, the key challenge of heterogeneity ,
i
[Table]
e. , that users’ distributions can differ, is rarely considered.
Heterogeneity is common, as users inevitably have unique characteristics [13].
Meanwhile, heterogeneity can cause a significant performance drop for learning algorithms designed only for i.i.d data [37, 34, 20, 59].
To use all the data, one needs to learn some central structure, transferable to all individual users.
Then one may locally learn the unique components for each user [21, 53, 16, 57].
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 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, we formulate a model of sparse heterogeneity.
Specifically, suppose that the discrete distributions of all users differ from an underlying central distribution in at most s entries, where s is much smaller than the dimension d .
Sparse heterogeneity is relevant to applications such as recommendation systems [28, 42, 35, 7] and medical risk scoring [50, 43, 40].
However, given data generated by multiple distributions with sparse heterogeneity, previous works [26, 2, 3, 15] either do not use all the data, or suffer from bias due to heterogeneity that does not vanish as the sample size increases.
Here we propose a novel s parse h eterogeneity i nspired collaboration and f ine- t uning method (SHIFT) where we first collaboratively learn the central distribution, and then fine-tune the central estimate to individual distributions.
Our method makes full use of heterogeneous data, leading to a significant improvement in error rates compared to prior methods.
See Table 1 for an overview, explained in detail later.
1.1 Contributions
We consider the problem of learning d -dimensional distributions with s -sparse heterogeneity.
We assume there are T clusters of user datapoints, and allow each datapoint to be transmitted with b bits of information to the server.
Our setting embraces heterogeneous data and thus is a significant generalization of the models from [26, 6, 27, 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 medianbased 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 broader context of multi-task learning [12, 52, 8] and federated learning [4, 44, 24, 47], where a central model or representation is learned from multiple heterogeneous datasets.
These central representations can be useful for few-shot learning tasks [51, 23] due to their ability to adapt to new tasks efficiently.
In heterogeneous linear regression, [53, 21] show improved sample complexities by assuming a low dimensional central representation, compared to the i.i.d.
setting [24, 36].
Related results are proved by [16] for personalized federated learning.
[57] study a bandit problem where the unknown parameter in each dataset equals a global parameter plus a sparse instance-specific term.
We study a different setting: learning distributions with sparse heterogeneity under communication constraints.
Estimating distributions under communication constraints.
Estimating discrete distributions has a rich literature [10, 18, 45, 19].
Under communication constraints, [26, 6, 27, 2] consider the non-interactive scenario and establish the minimax optimal rates, in terms of data dimension and communication 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 works study the estimation of sparse distributions.
In particular, [3] consider s -sparse distributions and establish minimax optimal rates under communication and privacy constraints, which are further improved by localization strategies in [14].
Complementary to minimax rates, [15] provides pointwise rates, governed by the half-norm of the distribution instead of its dimension.
Our setting embraces heterogeneous data, and thus is a generalization of the one studied in above works.
Robust estimation & learning.
Robust statistics and learning study algorithms resilient to unknown data corruption [30, 25].
The median-of-means method [32, 38, 39], partitions the data into subsets, computes an estimate from each, and takes their median.
Similarly, some works study robustness from the optimization perspective, proposing to robustly aggregate gradients of the loss functions [48, 49, 9, 58].
We adapt some analysis techniques from [39, 58] to the significantly different 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 v by 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 ∥defined 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 and denote by [ v ] I the sub-vector
([ v ] k ) k ∈I indexed by I .
We use the Bachmann-Landau asymptotic notations Ω( · ) , Θ( · ) , 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 following model 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 is in 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 distribution p 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 assume that the message sent by each datapoint is encoded into no more than b bits and b can be significantly smaller than log 2 d so that the communication is efficient.
We also assume the server knows which cluster t ∈
[ T ] each Y t,j belongs to, as well as the number of clusters T .
This paper 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. , all user datapoints are homogeneous and there is a single distribution to learn, the problem reduces to the one studied by [26, 6, 27, 2, 1].
Model of heterogeneity.
In heterogeneous settings, collaboration among the users is most beneficial if the local distributions are related.
We model this by assuming that the local distribu
- tions are sparse perturbations of an unknown central distribution p ⋆ ∈ P d .
The distribution p t of each cluster t differs from p ⋆ in at most s ≥ 0 entries:
[Equation]
The central distribution p ⋆ can be viewed as the central structure across heterogeneous clusters of datapoints.
The level of heterogeneity is controlled by the parameter s .
When s is much smaller than d , the local distributions differ from the 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 the high-dimensional item-preference vectors of users can vary sparsely [28, 42, 35, 7]; and medical risk scoring, where hospitals enjoy similar predicted results but with a few systematic differences in diagnosis behavior, healthcare utilization, etc [50, 43, 40].
Algorithm
We now introduce our method for leveraging heterogeneous data to improve per-cluster sample efficiency.
We first discuss a hashing-based method to handle the communication constraint.
Since communication 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 [15, 14, 1], i.e. , depend on previously sent messages.
Here we adopt a non-interactive encoding-decoding scheme, based on uniform hashing [3, 14] where W t,j depends only on X t,j and is independent of other messages.
Specifically, each datapoint X t,j is encoded via an independent random hash function h t,j : [ d ] → [2 b ] .
Upon receiving all messages, the server counts the empirical frequencies of all symbols,
[Table]
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 is summarized below.
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 biased for p t .
We also write b ⋆ =
[(2 b − 1) p ⋆ + 1] / 2 b for the mean of a hashed datapoint sampled from the 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 and fine-tuning .
The first stage estimates the central hashed distribution b ⋆ using all hashed estimates { 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 individual distributions less mismatch with the central one, most datapoints (used to estimate that entry) are sampled from the probability of the central distribution.
Hence, to estimate those entries of the central distribution, we can treat the datapoints generated by heterogeneous users as corrupted, and leverage robust statistical methods to mitigate their influences.
In the second stage—fine-tuning—we detect mismatched entries between individual hashed estimates q b t and the central estimate q b ⋆ .
Recall that the central and individual distributions differ in only a few entries.
For entries k ∈ [ d ] such that | [ q b ⋆ ] k − [ q b t ] k | is below ( α [ q b ⋆ ] k /n )
1 / 2 for some threshold parameter α , we may expect that p tk = p ⋆k .
As a result, the estimate [ q b ⋆ ] k of b ⋆k can be more accurate as 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 of p t where Proj [0 , 1] ( · ) indicates truncating the input with upper/lower bound 1 and 0 .
Our method does 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 beneficial for sample efficiency because most entries of the target distribution are well-estimated through collaborative learning.
One can transfer those entries, and it suffices to estimate the few remaining entries, instead of the whole distribution.
See Theorem 2 for the details.
The knowledge transfer utility 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, for the median-based SHIFT method, where robust _ estimate( { q b t : t ∈ [ T ] } ) in Algorithm 1 is the entry-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 the trimmed-mean-based SHIFT method, see Appendix D.
By setting the threshold parameter α in Algorithm 1 as α = Θ(ln( n )) , we prove the upper bounds on the 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 )) 2 .
Then, for the median-based SHIFT method, 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 of the dimension d , and is a factor d/ max { 2 b , s } smaller compared to the rate d/ (2 b n ) obtained by the minimax optimal method using only homogeneous datapoints [26].
Thus, it brings a significant benefit 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 [26, 15] in which all datapoints are pooled to learn a global distribution T − 1 (cid:80) t ∈
[ T ] p t , which is then used by 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 is large.
See Table 1 for more details.
Finally, we discuss our results on knowledge transfer.
The central estimator q b ⋆ is adaptable to a new 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
such that 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 a factor of min { Tn/ ˜ n, d/ max { 2 b , s }} improvement in sample-efficiency compared to the case where the 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.
Our analysis is novel compared to previous analyses for methods with homogeneous datapoints.
The
The final 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-aligned entries as
[Equation]
where B k ≜ { t ∈ [ T ] : b tk ̸ =
b ⋆k , i.e. , p tk ̸ = p ⋆k } is the set of clusters whose distribution differs from p ⋆ in the k -th entry.
We aim to estimate the η -well-aligned entries accurately by using robust statistical methods.
Further, we argue that there are few poorly-aligned entries, and they affect the final per-cluster error only mildly.
By the pigeonhole principle, the number of entries that are not η -well-aligned is upper bounded 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]
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 [15] 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’s inequality, which is variance-dependent [54], to obtain bounds relying on both the sample size Tn and the variance b ⋆k (1 − b ⋆k ) .
Also, the constant 1 / 5 , controlling the heterogeneity, is not essential (we choose 1 / 5 for clarity).
It can be replaced with any number below 1 / 2 so that estimating the central probability distribution becomes possible, as the information conveyed by homogeneous datapoints dominates.
Lemma 1 reveals that well-aligned entries of the central distribution are accurately estimated.
Thus one can use the central estimate for the entries where the central distribution p t aligns with the target distribution p t .
The remaining entries, that are neither well-aligned nor satisfy p ⋆k =
p tk , can be estimated 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 and 2.
While estimating p ⋆ is not our main goal, one can readily obtain from Lemma 1 the following bound on 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 of distributions 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 Algorithm 1.
Suppose |B k
|
=
O ( sT/d ) for any k ∈ [ d ] , with η = max k ∈ [ d ] |B k | /T , the median-based SHIFT method enjoys
[Equation]
Lower Bounds
To complement our upper bounds
[Equation]
we now provide minimax lower bounds for estimating distributions under heterogeneity.
Since our setting contains T heterogeneous clusters of datapoints
[Equation]
our minimax error metric is slightly different from the one studied in [26, 6, 27, 2].
Using the ℓ 2 error as the loss
[Equation]
the lower bound metric is defined as
[Equation]
[Figure]
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 output in [2 b ] , e.g. , the random hashing maps.
The measure (3) characterizes the best possible worst-case performance 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 ≤ s for all t ∈
[ T ] , we consider two representative cases therein: 1) The homogeneous case where p 1 = · · · = p T = p ⋆ ∈ P d .
Then the setting essentially reduces to the single-cluster problem but with 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 holds that { p t : t ∈ [ T ] } ⊆
B s ( p ⋆ ) .
By constructing independent priors for { p t : t ∈ [ T ] } and p ⋆ , one can 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 following lower 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 bound for 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 minimax optimal 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 [11].
As a baseline method, we use the estimator based on uniform grouping in [26] that is minimax optimal under homogeneity, i.e. , in the single-task regime.
We apply the baseline method both locally and globally.
In the local case, the estimator (cid:98) p t for each cluster is computed without datapoints from other clusters.
In the global case, we pool data from all clusters, and compute estimators (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 .
[Table]
5.1 Synthetic Data
We set the uniform distribution, p ⋆ = (1 /d, . . .
, 1 /d )
as the central distribution.
In Appendix F, we also experiment on the truncated geometric distribution and compare our method with the localizationrefinement method [15].
Among the d entries of p ⋆ , we draw s entries uniformly at random and assign new values for them uniformly at random over [0 , 1] , with re-normalization to preserve their sum.
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 set the dimension to d = 300 and run the simulation by varying n, T, s .
As we see from (2), the error of our 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.
We set the threshold parameter α = ln( n ) and the trimming proportion ω = 0 .
In Appendix
F, we provide results for other choices of the hyperparameters α, ω and discuss a heuristic for choosing α .
Figure 2 illustrates that our method outperforms the baselines for most choices of n, T, s .
Specifically, as Theorem 1 predicts, the ℓ 2 error of our method decreases as the number of clusters T increases.
On the other hand, when the baseline methods are applied globally, without considering heterogeneity, they show a bias that does not disappear as the sample size n or the number of clusters T increases.
This shows that the fine-tuning step in SHIFT is crucial for the estimation of heterogeneous distributions.
Finally, the right panel of Figure 2 shows that our method is effective only when s is small compared to the dimension d ; which highlights the crucial role of sparse heterogeneity.
When s is close to d , the distributions p 1 , . . .
p T could be considerably different without 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 [11].
The dataset consists of dialogues of 1,129 speaking roles in Shakespeare’s 35 different plays.
In our experiment, we study the distribution of k -grams ( k -tuples of consecutive letters from the 26-letter English alphabet, see Chapter 3 of [33]) appearing in the dialogues.
We consider each play as a cluster C t and estimate the distribution p t ∈ P d , d = 26 k of k -grams.
Since the ground-truth distribution p t is unknown, 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 distributions from 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 were rejected 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 trimming proportion ω = 0 .
1 , which we choose following the heuristic discussed in Appendix F.
We repeat the experiment ten times by taking different datapoints and report the average ℓ 2 error of estimation in Table 2.
The standard deviations are small even over ten repetitions.
SHIFT shows competitive performance on the empirical dataset, even though we do not rigorously know if the sparse heterogeneity model (1) applies.
Conclusion and Future Directions
We formulate the problem of learning distributions under sparse heterogeneity and communication constraints.
We propose the SHIFT method, which first learns a central distribution, and then finetunes the estimate to adapt to individual distributions.
We provide both theoretical and experimental results to show its sample-efficiency improvement compared to classical methods that target only homogeneous data.
Many interesting directions remain to be explored, including investigating if there is a point-wise optimal method with rate depending on { p t : t ∈ [ T ] } and p ⋆ ; and designing methods for other information constraints, such as local differential privacy constraints.
Acknowledgements
During this work, Xinmeng Huang was supported in part by the NSF TRIPODS 1934960, NSF DMS 2046874 (CAREER), NSF CAREER award CIF-1943064; Donghwan Lee was supported in part by ARO W911NF-20-1-0080, DCIST, Air Force Office of Scientific Research Young Investigator
[-Recall that for all t ∈ [ T ] and k ∈ [ d ] , b tk = p tk (2 b − 1)+1 
-]
[-= 
-]
[-q b ⋆k implies ˆ F k (cid:16) b ⋆k + C ε σ k √ n G ′ k,T,δ (cid:17) ≤ 12 − 12 T and b ⋆k − C ε σ k √ n G ′ k,T,δ > q b ⋆k implies ˆ F − k (cid:16) b ⋆k − C ε σ k √ n G ′ k,T,δ (cid:17) ≥ 
-]
Program (AFOSR-YIP) award #FA9550-20-1-
[-⌊ s/η ⌋ +
-]
[-⌊ s/η ⌋ , and x k = 
-]
[-+
-]