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
Program (AFOSR-YIP) award #FA9550-20-1-
A Properties of Uniform Hashing
Algorithm input: cluster C t with n ≥ 1 users having data X t,j , j = 1 , . . . , n for j = 1 , . . . , n do
Generate a uniformly random hash function h t,j : [ d ] → [2 b ] using shared randomness
Encode message Y t,j = h t,j ( X t,j ) and send it to the server ▷ Encoding end for for k = 1 , . . .
, d do
Count N tk
( Y t, [ n ] ) ← |{ j ∈ [ n ] : h t,j ( k ) =
Y t,j }| ▷ Decoding
Estimate q b t k ← N t k /n end for output: (cid:98)
b t
Recall that for all t ∈ [ T ] and k ∈ [ d ] , b tk = p tk (2 b − 1)+1
b ∈ (cid:2) 12 b , 1 (cid:3) .
Proposition 1 (Properties of Hashed Estimates) .
For each t ∈ [ T ] , suppose q b t is computed in cluster
C t as in Algorithm 2 with i.i.d datapoints X t,j ∼ Cat( p t ) , ∀ j ∈ [ n ] .
Then, it holds that
q b 1 , . .
, q b T ∈ [0 , 1] are independent; 2.
for any t ∈ [ T ] and k ∈ [ d ] , N tk ∼ Binom( n, b tk ) ; 3. supp( p t − p ⋆ )
= supp( b t − b ⋆ ) and p ⋆k = 1 (or 0 ) is equivalent to b ⋆k = 1 (or 12 b , respectively).
Proof.
Property 1 holds because (cid:98) b 1 , . . .
, (cid:98) b T are obtained by cluster-wise encoding-decoding of independent datapoints.
To see property 2, we have for any j ∈ [ n ] and k ∈
[ d ] that
[Equation]
Thus, I ( h t,j ( k ) =
Y t,j ) is a Bernoulli variable with success probability b tk .
Since each datapoint is encoded with an independent hash function, N t k has a binomial distribution with n trials and parameter b tk .
Property 3 directly follows from b t − b ⋆ =
( p t − p ⋆ )(2 b − 1) / 2 b and as b > 0 .
Proposition 2 (Property of Debiasing) .
For any y , y ⋆ ∈ R d , let x = Proj [0 , 1] ( 2 b y − 1 2 b − 1 ) and x ⋆ = Proj [0 , 1] ( 2 b y ⋆ − 1 2 b − 1 ) .
Then it holds that for q = 1 , 2 , E [ ∥ x − x ⋆ ∥ qq ]
=
O ( E [ ∥ y − y ⋆ ∥ qq ]) .
In particular, we have for q = 1 , 2 and any t ∈ [ T ] , E [ ∥ (cid:98) p t − p t ∥ qq ]
=
O ( E [ ∥ (cid:98) b t − b t ∥ qq ]) , where (cid:98) p t =
Proj [0 , 1] ( 2 b (cid:98) b t − 1 2 b − 1 ) is the final per-cluster estimate obtained in Algorithm 1.
Proof.
Using the inequality that | Proj [0 , 1] ( x ) − Proj [0 , 1] ( y ) | ≤ | x − y | for any x, y ∈ R , we have
E [ ∥ x − x ⋆ ∥ qq ] = (cid:88)
[Equation]
In the last step, we used that 2 b / (2 b − 1) ≤ 2 for all b ≥ 1 , and thus the O ( · ) only depends on universal constants.
B General Lemmas
In this section, we state some general lemmas that will be used in the analysis.
Lemma 2 (Berry-Esseen Theorem; [55]) .
Assume that Z 1 , . . .
, Z n are i.i.d.
copies of a random variable Z with mean µ , variance σ 2 > 0 , and such that E (cid:2) | Z − µ | 3 (cid:3) < ∞ .
Then,
[Equation]
where ¯ Z = 1 n (cid:80) ni
=1 Z i and γ ( Z ) =
E [ | Z − µ | 3 ] /σ 3 is the absolute skewness of Z .
Lemma 3 (Hoeffding’s Inequality; [55]) .
Let Z 1 , . . .
Z n ∈ [ l, r ] , l < r , be independent random variables and let ¯ Z = 1 n (cid:80) nj =1 Z j .
Then for any δ ≥ 0 ,
[Equation]
Lemma 4 (Bernstein’s Inequality; [54]) .
Let Z 1 , . . .
, Z n be i.i.d.
copies of a random variable Z with | Z − E
[ Z ] | ≤ M , M > 0 and Var( Z 1 )
=
σ 2 > 0 , and let ¯ Z = 1 n (cid:80) nj =1 Z j .
Then for any δ ≥ 0 ,
[Equation]
The second inequality above directly follows from 1 a + b ≥ 12 min { 1 a , 1 b } for any a, b > 0 .
Note that (6) also allows σ = 0 because P ( | ¯ Z − E [ ¯ Z ] | > δ )
=
and min (cid:110)
δ 2 /σ 2 ≜ + ∞ , δ/M (cid:111)
=
δM .
Therefore, we use this lemma for all σ ≥ 0 below.
B.1 Analysis Framework
For each t ∈ [ T ] , we denote by
[Equation]
the set of entries in which the central estimate [ (cid:98) b ⋆ ] k is adapted to cluster C t .
In this language, the final estimates can be expressed as (cid:98) b tk = q b ⋆k
[Equation]
+ q
[Equation]
{ k / ∈ K tα } for t ∈ [ T ] .
Therefore, it holds that, for q = 1 , 2 ,
[Equation]
Let I t ≜ { k ∈ [ d ] : b tk = b ⋆k , i.e. , p tk = p ⋆k } be set of entries at which the t -th cluster’s distribution p t aligns with the central distribution p ⋆ .
We next bound the two terms from (8) in Lemmas 5 and 6.
These do not need the independence of q b ⋆ and q b t , and hence do not require sample splitting despite the division between stages.
Lemma 5.
For any t ∈ [ T ] , α ≥ 1 and η ∈ (0 , 1] , with K tα from (7) , we have, for q = 1 ,(cid:88)
[Equation]
Proof.
We first take q = 1 .
For any k ∈ [ d ] , clearly
[Equation]
We use this bound for k ∈ I η ∩ I t .
For k / ∈ I η ∩ I t , we instead bound
[Equation]
If k ∈ K tα , it holds by definition that | q b ⋆k − q b tk | ≤ (cid:113) α q b tk /n , thus we further have
[Equation]
By Jensen’s inequality and since n q b tk ∼ Binom( n, b tk ) , we have
[Equation]
and
[Equation] [Equation]
Plugging (11) and (12) into (10), we find
[Equation]
Summing up (9) over all entries in I η ∩ I t and summing up (13) over all entries not in I η ∩ I t leads to the claim for q = 1 in Lemma 5.
The case q = 2 follows by a similar argument.
Lemma 6.
For any t ∈ [ T ] , α ≥ 1 and η ∈ (0 , 1] , with K tα from (7) , we have, for q = 1 ,(cid:88) k ∈ [
[Equation]
Proof.
For q = 1 , note that
[Equation]
and
[Equation]
Combining (14), (15) with the first inequality in (12) for k ∈ I η ∩ I t , and using the last inequality in (12) for k / ∈ I η
∩
I t leads to the claim with q = 1 .
We can similarly obtain the bound with q = 2 .
Combing Lemma 5 and 6 with (8), we find the following proposition:
Proposition 3.
For any α ≥ 1 , and q = 1 , 2 , it holds that
[Equation]
Proposition 3 does not rely on how q b ⋆ is obtained.
The next part is devoted to proving that when q b ⋆ is obtained via a certain robust estimate, the bounds in Proposition 3 are small for certain values of α and η .
C Median-Based Method
In this section, we provide the proofs for the median-based SHIFT method.
We first re-state the detailed version of some key results that apply to both the ℓ 2 and ℓ 1 errors.
Below, we use σ k = (cid:112) b
⋆k (1 − b ⋆k ) to denote the standard deviation of the Bernoulli variable with success probability
b ⋆k = p ⋆k
+(1 − p ⋆k ) / 2 b .
We also recall that B k is defined as the set of clusters with distributions mismatched with the central distribution at the k -th entry, i.e. , B k = { t ∈
[ T ] : p tk
̸ = p ⋆k } , and I η is defined as the η -well-aligned entries, i.e. , I η = { k ∈ [ d ] : |B k | < ηT } .
Lemma 7 (Detailed statement of Lemma 1) .
Suppose q b ⋆ = median (cid:0) { q b t } t ∈
[ T ] (cid:1) .
Then for any 0 < η ≤ 15 , k ∈ I η , and q = 1 , 2 , it holds that
[Equation]
Let us define, for q = 1 , 2 ,
[Equation]
Proposition 4.
Suppose q b ⋆ = median (cid:0) { q b t } t ∈ [ T ] (cid:1) .
Then for any 0 < η ≤ 15 and q = 1 , 2 , it holds that
[Equation]
 .
We omit the proofs of Proposition 4 and Theorem 6 (below), as Proposition 4 is a direct corollary of Lemma 7 by using (cid:80) k ∈
[ d ] σ qk = O ( d/ 2 bq/ 2 ) for q = 1 , 2 , and Theorem 6 follows from the same analysis as Theorem 5.
Theorem 5 (Detailed statement of Theorem 1) .
Suppose n ≥ 2 b +6 ln( n ) and α ≥ 2(8 + (cid:112) 8 ln( n ))with α = O (ln( n )) .
Then for the median-based SHIFT method, for any 0 < η ≤ 15 , q = 1 , 2 , and t ∈ [ T ] ,
[Equation]
Furthermore, by setting η = Θ(1) with η ≤ 15 , we have
[Equation]
Theorem 6 (Detailed statement of Theorem 2) .
Suppose n ≥ ˜ n ≥ 2 b +6 ln(˜ n ) and α ≥ 2(8 + (cid:112) 8 ln(˜ n )) 2 with α
= O (ln(˜ n )) .
Then the median-based SHIFT method for predicting the distribution of the new cluster with ˜ n users achieves, for q = 1 , 2 ,
[Equation]
C.1 Proof of Lemma
We first consider T ≤ 20 ln( n ) .
In this case, by Bernstein’s inequality (Lemma 4) with M = 1 , we have for any t ∈
[ T ] \B k that for any δ ≥ 0 ,
[Equation]
Taking δ = max { σ k (cid:112) 8 ln( n ) /n, 8 ln( n ) /n } in (16), we find
[Equation]
Since | [ T ] \B k | > T 2 for any k ∈ I η with η ≤ 15 , we have, since q b ⋆k = median (cid:0) { q b tk } t ∈
[ T ] (cid:1) , that there are t − , t + ∈
[ T ] \B k with q b t ′ k ≤ q b ⋆k ≤ q b t ′ k .
Hence, | q b ⋆k − b ⋆k | ≤ max t ∈
[ T ] \B k | q b tk − b ⋆k | .
Recall that for any random variable 0 ≤ X ≤ 1 and any δ ≥ 0 , E [ X ] ≤ δ + P ( X ≥ δ ) .
Therefore, by taking the union bound of (17) over k ∈ [ T ] \B k , and by the assumption that T ≤ 20 ln( n ) , we have
[Equation]
Similarly, we have
[Equation]
For each k ∈ [ d ] with b ⋆k ̸ = 1 (recall that b ⋆k ≥ 1 / 2 b by definition), let γ k =
(1 − 2 b ⋆k (1 − b ⋆k )) / (cid:112) b
⋆k (1 − b ⋆k ) , and let ˜ F k ( x ) := 1 T −|B k | (cid:80) t ∈
[ T ] \B k 1 ( q b tk ≤ x ) be the empirical distribution function of { q b tk : b tk = b ⋆k } .
Let ε ∈
(0 , 1 / 2) and C ε = √ 2 π exp((Φ − 1 (1 − ε )) 2 / 2) .
For δ ≥ 0 , define, recalling ηT > |B k | for all k ∈ I η ,
[Equation]
where the term 10 − 8 Tn is used to overcome some challenges due to the discreteness of empirical distributions, and can be replaced with other suitably small terms (see the proof of Lemma 9).
Further, define
[Equation]
To prove Lemma 1 for T > 20 ln( n ) , we need the following additional lemmas:
Lemma 8.
For any δ ≥ 0 such that
[Equation]
it holds with probability at least 1 − 4 e − 2 δ that
[Equation]
and
[Equation]
Proof.
The proof essentially follows Lemma 1 of [58].
We provide the proof for the sake of being self-contained.
Let Z tk =
( q b tk − b tk ) / (cid:113) Var( q b tk ) be a standardized version of q b tk for each t ∈
[ T ] and k ∈ [ d ] , with b ⋆k ̸ = 1 .
Let ˜Φ k ( z ) = 1 T −|B k | (cid:80) t ∈
[ T ] \B k 1 ( Z tk ≤ z ) be the empirical distribution of { Z tk : t ∈ [ T ] \B k } .
The distribution of Z tk is identical t ∈
[ T ] \B k , and we denote by Φ k their common cdf.
By definition, E [˜Φ k ( z )] =
Φ k ( z ) for any z ∈ R .
Let z 1 > 0 >
z 2 be such that Φ( z 1 )
=
12 + G ′ k,T,δ and Φ( z 2 )
=
12 − G ′ k,T,δ , which exist due to (20).
Then, by Lemma 2, we have
[Equation]
Further, by the Hoeffding’s inequality, we have for any δ ≥ 0 and z ∈ R ,
[Equation]
with probability at least 1 − 2 e − 2 δ .
Then, by a union bound of (22) for z = z 1 , z 2 , and by (21), it holds with probability at least 1 − 4 e − 2 δ that
[Equation]
[Equation]
Finally, we bound the values of z 1 and z 2 .
By the mean value theorem, there exists ξ ∈ [0 , z 1 ] such that
[Equation]
[Equation]
By (20) and the definition of z 1 , we have z 1 ≤ Φ − 1 (1 − ε ) , and thus, by (24), we have
[Equation]
[Equation]
Similarly, we have
[Equation] [Equation]
Since for all z , ˜Φ k ( z ) =
˜ F k ( σ k z/ √ n + b ⋆k ) , plugging (25) and (26) into (23), we find the conclusion of this lemma.
This leads to our next result.
Lemma 9.
For any k ∈ [ d ] such that condition (20) holds, we have with probability at least 1 − 4 e − 2 δ that
[Equation]
Proof.
Let ˆ F k be the empirical distribution function of { q b tk : t ∈ [ T ] } , such that for all x ∈ R , ˆ F k ( x ) := 1 T (cid:80) t ∈
[ T ] 1 ( q b t k ≤ x ) .
We have
[Equation]
≤ max (cid:26)
[Equation] [Equation]
Define ˜ F − k ( x ) := 1 T −|B k | (cid:80) t ∈
[ T ] \B k 1 ( q b tk < x ) ≤ ˜ F k ( x ) .
Then by (28) and Lemma 8, we have, with probability at least 1 − 4 e − 2 δ that ˆ F k
[Equation]
and
[Equation]
Let q b ( j ) k , ∀ j ∈
[ T ] be the j -th smallest element in { q b tk : t ∈
[ T ] } .
Recalling the definition of the median, if T is odd
, then q b ⋆k = q b (( T +1) / 2) k .
Therefore, 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) ≤ 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) ≥
+ 12 T , leading to a contradiction with (29).
On the other hand, if T is even, q b ⋆k =
( q b ( T/ 2) k
+ q b ( T/ 2+1) k ) / 2 .
Therefore, 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) ≤ 12 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) ≥ 12 , which is also contradictory to (29).
To summarize, it holds that
[Equation]
with probability at least 1 − 4 e − 2 δ .
If T ≤ 20 ln( n ) , Lemma 7 follows directly from (18) and (19).
Now, given Lemma 8 and Lemma 9, we turn to prove Lemma 7 with T ≥ 20 ln( n ) .
We first check condition (20).
Since |B k | ≤ ηT for any k ∈ I η , η ≤ 15 , and γ k σ k ≤ 1 , we have for each k ∈ I η that
[Equation]
When T ≥ 20 ln( n ) , for any k ∈ [ d ] such that σ k ≥ 20 √ n (
[Equation]
− 2 η ) , taking δ = ln( n ) above, we have
[Equation]
Therefore, condition (20) in Lemma 9 is satisfied with ε = 0 .
035755 , for which we can check that
C ε ≤ 13 .
Thus, for any δ ≤ ln( n ) ,
[Equation]
Therefore, by (30), we have, using that for any random variable 0 ≤ X ≤ 1 and any 0 ≤ r ≤ 1 , E [ X ] ≤ r + P ( X ≥ r ) , and that for δ = (ln n ) / 2 , one has 4 e − 2 δ = 4 /n , we find
[Equation]
Similarly, by the Cauchy-Schwarz inequality, we also have
[Equation]
On the other hand, for any k ∈ [ d ] \B k such that σ k < 20 √ n (1 − 2 η ) , by Bernstein’s inequality and a union bound, we have
[Equation]
Since | [ T ] \B k | > T 2 , we have as before that | q b ⋆k − b ⋆k | ≤ max t ∈
[ T ] \B k | q b tk − b ⋆k | .
Taking δ = 4 max { ln( Tn 2 ) , 10 (cid:112) ln( Tn 2 ) } /n in (33), with the same steps as above, we find
[Equation]
and
[Equation]
To summarize, combining (31), (32) with (34),(35), we complete the proof when T > 20 ln( n ) .
Furthermore, by using (cid:80) k ∈ [ d ] σ qk = O ( d/ 2 bq/ 2 ) for q = 1 , 2 , we directly reach Proposition 4.
C.2 Proof of Theorem
We first consider the case where T ≤ 20 ln( n ) .
By definition, (cid:98) b tk is either equal to q b tk or q b ⋆k , and the latter happens only when k ∈ K tα ,
i.e. , | q b ⋆k − q b tk | ≤ (cid:113) α q b tk /n .
In this case, we have
[Equation]
Therefore, we have | (cid:98) b
tk − b tk | ≤ | q b tk − b tk | + (cid:113) α q b tk /n for all k ∈
[ d ] .
This leads to
[Equation]
where (36) holds by Jensen’s inequality.
By further using the Cauchy–Schwarz inequality, we have
[Equation]
and (cid:88)
[Equation]
Plugging (37) and (38) into (36), we find
[Equation]
We can similarly prove
[Equation]
Next we prove the case where T ≥ 20 ln( n ) =
Ω(ln( n )) .
We first consider the estimation errors over k ∈ I η ∩ I t such that σ k ≥ 20 √ n (1 − 2 η ) .
Let E tk :=
{ q b tk ≥ 12 b tk and | q b ⋆k − b ⋆k | ≤ 8 (cid:112)
b ⋆k /n } .
If n ≥ 2 b +6 ln( n ) and 0 < η ≤ 1 / 5 , then since b ⋆k ≥ 12 b for any k ∈ [ d ] , we have
[Equation]
Hence, by (30), it holds that
[Equation]
By Bernstein’s inequality and as b ⋆k ≥ 12 b , we have
[Equation]
where the last inequality holds because n ≥ 2 b +6 ln( n ) .
Combining (40) with (39), we find
P (( E tk ) c ) ≤ 6 n 2 .
By definition, k / ∈ K tα implies | q b ⋆k − q b tk | > (cid:113) α q b tk /n .
On the event E tk , this further implies | q b
⋆k − q b tk | > (cid:112) αb tk / 2 n .
Combined with (39) and that b ⋆k =
b tk for any k ∈ I t , we have on the event E tk
[Equation]
Let ζ ≜ (cid:112)
[Equation]
2 − 8 ≥
[Equation]
ln
[Equation]
n
[Equation]
and F tk :=
[Equation]
b tk /n (cid:111) .
By Bernstein’s inequality, and using n ≥ 2 b +6 ln( n ) , we have
[Equation]
Combining (41) with (42), we find that for any k ∈ I η ∩ I t with σ k ≥ 20 √ n (1 − 2 η ) , it holds that
[Equation]
On the other hand for any k ∈ I η ∩ I t with σ k < 20 √ n (1 − 2 η ) , we have
[Equation]
Therefore, we have for all k ∈ I η ∩ I t , and q = 1 ,min
[Equation]
Since α = O (ln( n )) , by (43) and Proposition
[Equation]
we obtain
[Equation]
Combining (44) with Proposition 4 and using that σ k ≤ (cid:112) b
⋆k = (cid:112) b tk for any k ∈ I t , we have
[Equation]
Since | ( I t ) c | =
∥ p t − p ⋆ ∥ 0 ≤ s
[Equation]
by the Cauchy-Schwarz inequality
[Equation]
we have (cid:88)
[Equation]
Plugging (46) into (44), we further obtain
[Equation]
) 
 .
Similarly, we can reach the following ℓ 2 counterpart:
[Equation]
Note that (cid:80)
k ∈
[ d ] |B k | /T ≤ s and for any set I with |I| =
⌈ sη ⌉ , (cid:88) k ∈I
[Equation]
Now, recalling the definition of I η , we apply Lemma 10 in (47) with ( r k , x k ) =
( (cid:112) b tk /n, |B k | /T )
for all k ∈ [ d ] , to find
[Equation]
Therefore, for any η = Θ(1) with η ≤ 15 , we finally have
E
[Equation]
O
[Equation]
[Equation] [Equation]
[Equation] [Equation]
(cid:33) .
Similarly, by combining (48) with Lemma 10, we have for any η = Θ(1) with η ≤ 15 ,
E [ ∥ (cid:98) b
[Equation]
− b t ∥ 22 ] =
˜ O (cid:18) max { 2 b , s }
[Equation]
n
[Equation] [Equation]
(cid:19) .
The result directly follows Proposition 2.
Lemma 10.
Given η ∈ (0 , 1] , r k ≥ 0 for all k ∈ [ d ] , and for q = 1 , 2 , consider the functions
f q : { x ∈ R d : 0 ≤ x k ≤ 1 , ∀ k ∈ [ d ] and (cid:80) k ∈ [ d ] x k ≤ s } → R , f q ( x 1 , . . .
, x d ) := (cid:80)
k ∈ [ d ] r qk ( 1 { x k ≥ η } + x qk 1 { x k < η } ) .
Then it holds that
[Equation]
where r (1) ≥ · · · ≥ r ( d ) is the non-decreasing rearrangement of { r 1 , . .
. , r d } .
Proof.
We only prove the result for f 1 , and the result for function f 2 follows similarly.
Note that r k ( 1 { x k ≥ η } + x k 1 { r k ≥ η } ) is increasing with respect to r k and x k .
To consider the maximum of the sum in f , by the rearrangement inequality, without loss of generality, we can assume r 1 ≥ r 2 ≥ · · · ≥ r d ≥ 0 and 1 ≥ x 1 ≥ x 2 ≥ · · · ≥ x d ≥ 0 .
In this case, we claim that the maximum is attained at x 1 = · · · = x
⌊ s/η ⌋ =
η , x
⌊ s/η ⌋ +
=
s − η
⌊ s/η ⌋ , and x k =
for all k >
⌊ s/η ⌋ + 1 .
Further, the maximum is (cid:80)
⌊ s/η ⌋
k =1 r k + r ⌊ s/η ⌋
+
( s − η
⌊ s/η ⌋ ) 2 , which is upper bounded by the right-hand side of (49).
We now use the exchange argument to prove the claim.
Step 1: If there is some k such that x k > η ≥ x k +1 , then defining x ′ by letting ( x ′ k , x ′ k +1 )
=
( η, x k + x k +1 − η ) while for other j , x ′ j = x j , increases the value of f .
Therefore, the maximum is attained by x such that for some j , x 1 = · · · = x j = η > x j +1 ≥ · · · ≥ x d .
Step 2: If there is some k such that η > x k ≥ x k +1 > 0 , then defining x ′ by letting ( x ′ k , x ′ k +1 )
= (min { η, x k + x k +1 } , max { 0 , x k + x k +1 − η } ) while for other j , x ′ j = x j , increases the value of f .
Therefore, combined with Step 1, the maximum is attained by x such that for some j , x 1 = · · · = x j = η > x j +1 ≥ 0 and x k = 0 for all k > j + 1 .
Thus most one element lies in (0 , η ) .
Combining Step 1 and Step 2 above, we complete the proof of the claim, which further leads to (49).
D Trimmed-Mean-Based Method
In this section, we study the trimmed-mean-based estimator.
Fix ω ∈ (0 , 1 / 2) .
Specifically, for each k ∈ [ d ] , let U k be the subset of { [ q p t ] t ∈
[ T ] } obtained by removing the largest and smallest ωT elements 3 .
Then, the trimmed-mean-based method can be expressed as
[Equation]
We also write q b ⋆ = trmean (cid:0) { q b t } t ∈ [ T ] , ω (cid:1) .
For any chosen trimming proportion 0 ≤ η ≤ ω ≤ 15 , we control the estimation error of each η -well aligned entry.
Intuitively, this is small because there are at most a fraction of η elements from heterogeneous distributions.
These are trimmed if they behave as outliers, and otherwise kept in U k .
The error control for a single entry k ∈ I η is in Lemma 11.
Lemma 11.
Suppose q b ⋆ = trmean (cid:0) { q b t } t ∈
[ T ] , ω (cid:1) such that 0 ≤ ω ≤ 15 .
Then for each k ∈ I η with 0 < η ≤ ω and any q = 1 , 2 , it holds that
[Equation]
Proof.
To prove Lemma 11, we need the following lemma.
Lemma 12.
For each k ∈ I η with 0 < η ≤ ω ≤ 15 , and any ε k , δ k ≥ 0 , it holds with probability at least
[Equation]
Proof of Lemma 12.
By Bernstein’s inequality and the union bound, we have for any ε k , δ k > 0 that
[Equation]
and
[Equation]
By the definition of q b ⋆k , we have
[Equation]
It is clear that
[Equation]
Then we claim that (cid:12)(cid:12)(cid:12)(cid:80) t ∈B k ∩U k q b tk − b ⋆k (cid:12)(cid:12)(cid:12) ≤ |B k | max t ∈
[ T ] \B k | q b tk − b ⋆k | .
Let Q k, l and Q k, r be the indices of the trimmed elements on the left side and right side, respectively, i.e. , the smallest and largest ωT elements among { q b tk } t ∈ [ T ] .
If B k ∩ U k ̸ = ∅ , then |U k \B k | <
T (1 − 2 ω ) .
Furthermore, we have |Q k, l ∪ ( U k \B k ) | =
|Q k, r ∪ ( U k \B k ) | = ωT + |U k \B k | <
T (1 − ω ) ≤ | [ T ] \B k | , which leads to ([ T ] \B k ) ∩ Q k, l ̸ = ∅ and ( T \B k )
∩ Q k, r ̸ = ∅ .
In conclusion, we have max t ∈U k | q b tk − b ⋆k | ≤ max t ∈
[ T ] \B k | q b tk − b ⋆k | , which completes the proof of the claim.
Therefore, we have
[Equation]
with probability at least 1 − 2 e −
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
{ δ 2 kσ 2 k ,δ k } .
Given Lemma 12, by setting
[Equation]
and δ
[Equation]
using that 1 / (1 − 2 ω ) ≤ 53 , and recalling σ k ≤ (cid:112)
b
⋆k , we have with probability at least 1 − 4 T 2 n 2 that | q
[Equation]
≤
[Equation]
ω
[Equation]
= ˜
O
[Equation]
which implies
[Equation]
Similarly, we can obtain
[Equation]
Given these results, we readily establish the following bound on the total error over all η -well-aligned entries.
Proposition 5.
Suppose q b ⋆ = trmean (cid:0) { q b t } t ∈
[ T ] , ω (cid:1) such that 0 ≤ ω ≤ 1 / 5 .
Then for each k ∈ I η with 0 < η ≤ ω and any q = 1 , 2 , it holds that
[Equation]
By setting α = Θ(ln( Tn )) , we find the following result.
Theorem 7.
Suppose n ≥ 2 b +5 ln
( Tn ) and α ≥ 2(8 + (cid:112) 8 ln( Tn ))
2 with α
= O (ln( Tn )) .
Then for the trimmed-mean-based SHIFT method, for any 0 < ω ≤ 15 , t ∈ [ T ] and q = 1 , 2 ,
[Equation]
Proof.
To apply Proposition 3, we need to bound (cid:80) k ∈I η ∩I t min { P ( k / ∈ K tα ) , (cid:112) b tk (
1 − b tk ) /n } and (cid:80) k ∈I
η ∩I t min { P ( k / ∈ K tα ) , b tk (1 − b tk ) /n } .
Let E tk :=
{ q b tk ≥ 12 b tk and | q b ⋆k − b ⋆k | ≤ 8 (cid:112)
b
⋆k ln
( T 3 n 2 ) /n } .
For each entry k ∈ I η ∩ I t , since n ≥ 2 b
ln
( T 3 n 2 ) and b ⋆k ≤ 12 b , we have 1 n ≤ (cid:113) b
⋆k n ln( T 3 n 2 ) .
By (52), we have with probability at least 1 − 4 T 2 n 2 that
[Equation]
By Bernstein’s inequality and as b ⋆k ≥ 12 b , we have
[Equation]
where the last inequality is because n ≥ 2 b +5 ln
( Tn ) .
Combining (53) with (54), we find P (( E tk ) c ) ≤
T 2 n 2 .
Now following the argument from (41)-(43), we can obtain that for all k ∈ I η ∩ I t ,
[Equation]
Since α = O (ln( Tn )) , by applying (43) to Proposition 3 with η = ω and using Proposition 5 with n = Ω(2 b ) , we find
[Equation]
55) and
E [ ∥
[Equation]
d
[Equation]
 .
Note that | ( I ω ∩ I t ) c | ≤ |I cω | + | ( I t ) c | ≤ s/ω
+ s =
O ( s/ω ) and (cid:88)
[Equation]
Plugging (56) into (55) and using E [ ∥ (cid:98) p
t − p t ∥ 1 ]
=
O ( E [ ∥ (cid:98) b t − b t ∥ 1 ]) , we find the conclusion in terms of the ℓ 1 error.
The results in terms of the ℓ 2 error can be obtained similarly.
E Lower Bounds
In this section, we provide the proofs for the minimax lower bounds for estimating distributions under our heterogeneity model.
We first re-state the detailed version the lower bounds that apply to both the ℓ 2 and ℓ 1 errors.
Theorem 8 (Detailed statement of Theorem 3) .
For any—possibly interactive—estimation method, and for any t ∈ [ T ] and q = 1 , 2 , we have
[Equation]
Theorem 9 (Detailed statement of Theorem. 4) .
For any—possibly interactive—estimation method, and a new cluster C T +1 , we have
[Equation]
We omit the proof of Theorem 9 since it follows from the same analysis as Theorem
[Equation]
E.1 Proof of Theorem
As discussed in Section 4, we will prove (57) by considering two special cases of our sparse heterogeneity model:
The homogeneous case where p 1 = · · · = p T = p ⋆ ∈ P d . 2.
The s/ 2 -sparse case where
∥ p ⋆ ∥ 0 ≤ s/ 2 and ∥ p t ∥ 0 ≤ s/ 2 for all t ∈ [ T ] .
Therefore, it naturally holds that
[Equation]
and
[Equation]
For the first case, combining (58) with the existing lower bound result [6, Cor 7] and [26, Thm 2] for the homogeneous setup, where all datapoints are generated by a single distribution, that for any estimation method (possibly based on interactive encoding),
[Equation]
we prove that the lower bound is at least of the order of the second term in (57).
For the second case, without loss of generality, we assume s is even.
This can be achieved by considering s − 1 instead of s , if necessary.
Recall that supp( · ) denotes the indices of non-zero entries of a vector.
Fixing any t ∈ [ T ] , we further consider the scenario where
[Equation]
One example where (60) holds is when supp( p t ) ⊆
[ s/ 2] and supp( p t ′ ) ⊆ { s/ 2 + 1 , . . .
, d } for all t ′ ̸
= t .
If (60) holds, then the support of the datapoints generated by { p t ′ :
t ′ ̸
= t } does not
[Figure]
overlap with the support of those generated by p t , and hence former are not informative for estimating p t .
Therefore, by further combining (59) with the existing lower bound result [14, Thm 2] for the s/ 2 -sparse homogeneous setup, where all datapoints are generated by a single s/ 2 -sparse distribution, that for any estimation method (possibly based on interactive encoding),
[Equation]
Thus, we have
[Equation]
This proves that the lower bound is at least of the order of the first term in (57).
Overall, we conclude the desired result.
F Supplementary Experiments
Truncated geometric distribution We consider the truncated geometric distribution with parameter β ∈ (0 , 1) , p ⋆ = 1 − β 1 − β d (1 , β, . .
, β d − 1 ) , as the central distribution and repeat the experiment in
Section 5.1.
We use d = 300 , β = 0 .
95 , b = 2 and vary n, T, s .
Figure 3 summarizes the results.
As in Section 5.1, we observe that our methods outperform the baseline methods in most cases, especially when s is small.
Also, we see the benefit of collaboration, i.e. , decreasing trend of the error as T increases, only in our methods.
[Figure]
Hyperparmeter selection.
We provide additional experiments using different hyperparameters α and ω from discussed in Section 5.1.
All other settings are identical to Section 5.1.
We test the hyperparameters ( α, ω ) =
(2 r ln( n ) , 0 . 1) for r ∈ {− 5 , − 4 , . .
. , 4 } and ( α, ω ) =
(ln( n ) , ω ) for ω ∈ { 0 . 05 , 0 . 1 , . .
Figure 4 summarizes the results.
We find that setting the threshold α too small leads to replacing almost all coordinates of the central estimate (cid:98) p
⋆ with local ones.
In the extreme case of α ≈ 0 , our method is essentially returns the local minimax estimates.
On the other hand, we observe that the performance of our method is less sensitive to the trimming proportion ω .
While the choice of α is crucial to the performance of our method, we argue that it is possible to select a reasonably good α by checking the number of fine-tuned entries, i.e. ,
[Equation]
n  
[Equation]
In Figure 5, we observe that more than half ( d/ 2 = 150 ) of the entries are fine-tuned when r ∈ {− 5 , − 4 , − 3 } .
These correspond to the three curves in the top left of Figure 4 that perform no better than the baseline methods.
In conclusion, by selecting α such that the number of fine-tuned entries are small enough compared to d , it is possible to reproduce the results in Section 5.
[Figure]