Abstract:Modern data-driven and distributed learning frameworks deal with diverse massive data generated by clients spread across heterogeneous environments. Indeed, data heterogeneity is a major bottleneck in scaling up many distributed learning paradigms. In many settings however, heterogeneous data may be generated in clusters with shared structures, as is the case in several applications such as federated learning where a common latent variable governs the distribution of all the samples generated by a client. It is therefore natural to ask how the underlying clustered structures in distributed data can be exploited to improve learning schemes. In this paper, we tackle this question in the special case of estimating $d$-dimensional parameters of a two-component mixture of linear regressions problem where each of $m$ nodes generates $n$ samples with a shared latent variable. We employ the well-known Expectation-Maximization (EM) method to estimate the maximum likelihood parameters from $m$ batches of dependent samples each containing $n$ measurements. Discarding the clustered structure in the mixture model, EM is known to require $O(\log(mn/d))$ iterations to reach the statistical accuracy of $O(\sqrt{d/(mn)})$. In contrast, we show that if initialized properly, EM on the structured data requires only $O(1)$ iterations to reach the same statistical accuracy, as long as $m$ grows up as $e^{o(n)}$. Our analysis establishes and combines novel asymptotic optimization and generalization guarantees for population and empirical EM with dependent samples, which may be of independent interest.

A unified framework for hard and soft clustering with regularized optimal transport

Regularised optimal self-transport is approximate Gaussian mixture maximum likelihood

Regularized EM algorithm

Convex Clustering via Optimal Mass Transport

EM for Mixture of Linear Regression with Clustered Data

Convergence Rates of the Regularized Optimal Transport : Disentangling Suboptimality and Entropy

Regularized EM Algorithms: A Unified Framework and Statistical Guarantees

Data-driven regularization of Wasserstein barycenters with an application to multivariate density registration

Transfer Operators from Batches of Unpaired Points via Entropic Transport Kernels

High-dimensional logistic entropy clustering

Jointly Modeling and Clustering Tensors in High Dimensions

Fast Entropic Regularized Optimal Transport Using Semidiscrete Cost Approximation

An Information-Theoretic Analysis of Hard and Soft Assignment Methods for Clustering

Universal Lower Bounds and Optimal Rates: Achieving Minimax Clustering Error in Sub-Exponential Mixture Models

Flexible Clustering with a Sparse Mixture of Generalized Hyperbolic Distributions

Adaptative clustering by minimization of the mixing entropy criterion

Optimal Transport for Kernel Gaussian Mixture Models

SP$^2$OT: Semantic-Regularized Progressive Partial Optimal Transport for Imbalanced Clustering

Semi-Discrete Optimal Transport: Nearly Minimax Estimation With Stochastic Gradient Descent and Adaptive Entropic Regularization

Optimal transport for Gaussian mixture models

A Consistent Extension of Discrete Optimal Transport Maps for Machine Learning Applications