Collaborative Learning of Discrete Distributions under Heterogeneity and Communication Constraints
This work addresses the challenge of communication-efficient distribution estimation for heterogeneous users in machine learning, representing an incremental improvement over prior methods that assumed homogeneity.
The paper tackles the problem of collaborative learning of discrete distributions under heterogeneity and communication constraints by proposing a two-stage method called SHIFT, which first learns a central distribution using robust statistics and then fine-tunes it for individual distributions, achieving minimax optimality and demonstrating efficiency in synthetic and text data experiments.
In modern machine learning, users often have to collaborate to learn the distribution of 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 for estimating that distribution. 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 their respective individual distribution. We show that SHIFT 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.