# A Wasserstein Minimax Framework for Mixed Linear Regression

arXiv: https://arxiv.org/abs/2106.07537 · HTML: https://mlsanity.com/paper/2106.07537/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 40 references that are in the map.
- First posted on arXiv 14 Jun 2021.

## Builds on (most cited first, 20 of 20)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Three Approaches for Personalization with Applications to Federated Learning](https://mlsanity.com/paper/2002.10619/) (2020)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [NIPS 2016 Tutorial: Generative Adversarial Networks](https://mlsanity.com/paper/1701.00160/) (2016)
- [On Gradient Descent Ascent for Nonconvex-Concave Minimax Problems](https://mlsanity.com/paper/1906.00331/) (2019)
- [Generalization and Equilibrium in Generative Adversarial Nets (GANs)](https://mlsanity.com/paper/1703.00573/) (2017)
- [Statistical guarantees for the EM algorithm: From population to sample-based analysis](https://mlsanity.com/paper/1408.2156/) (2014)
- [Distributionally Robust Federated Averaging](https://mlsanity.com/paper/2102.12660/) (2021)
- [Robust Federated Learning: The Case of Affine Distribution Shifts](https://mlsanity.com/paper/2006.08907/) (2020)
- [On the Convergence and Robustness of Training GANs with Regularized Optimal Transport](https://mlsanity.com/paper/1802.08249/) (2018)
- [EM Converges for a Mixture of Many Linear Regressions](https://mlsanity.com/paper/1905.12106/) (2019)
- [Global Convergence of EM Algorithm for Mixtures of Two Component Linear Regression](https://mlsanity.com/paper/1810.05752/) (2018)
- [Spectral Experts for Estimating Mixtures of Linear Regressions](https://mlsanity.com/paper/1306.3729/) (2013)
- [A Convex Duality Framework for GANs](https://mlsanity.com/paper/1810.11740/) (2018)
- [Sinkhorn EM: An Expectation-Maximization algorithm based on entropic optimal transport](https://mlsanity.com/paper/2006.16548/) (2020)
- [GAT-GMM: Generative Adversarial Training for Gaussian Mixture Models](https://mlsanity.com/paper/2006.10293/) (2020)

## Built on by (most cited first, 7 of 7)

- [Collaboration Equilibrium in Federated Learning](https://mlsanity.com/paper/2108.07926/) (2021)
- [Learning Mixtures of Linear Dynamical Systems](https://mlsanity.com/paper/2201.11211/) (2022)
- [An Optimal Transport Approach to Personalized Federated Learning](https://mlsanity.com/paper/2206.02468/) (2022)
- [Tensor Decompositions Meet Control Theory: Learning General Mixtures of Linear Dynamical Systems](https://mlsanity.com/paper/2307.06538/) (2023)
- [Imbalanced Mixed Linear Regression](https://mlsanity.com/paper/2301.12559/) (2023)
- [EM for Mixture of Linear Regression with Clustered Data](https://mlsanity.com/paper/2308.11518/) (2023)
- [Certifiably Robust Model Evaluation in Federated Learning under Meta-Distributional Shifts](https://mlsanity.com/paper/2410.20250/) (2024)

## Abstract (arXiv)

Multi-modal distributions are commonly used to model clustered data in statistical learning tasks. In this paper, we consider the Mixed Linear Regression (MLR) problem. We propose an optimal transport-based framework for MLR problems, Wasserstein Mixed Linear Regression (WMLR), which minimizes the Wasserstein distance between the learned and target mixture regression models. Through a model-based duality analysis, WMLR reduces the underlying MLR task to a nonconvex-concave minimax optimization problem, which can be provably solved to find a minimax stationary point by the Gradient Descent Ascent (GDA) algorithm. In the special case of mixtures of two linear regression models, we show that WMLR enjoys global convergence and generalization guarantees. We prove that WMLR's sample complexity grows linearly with the dimension of data. Finally, we discuss the application of WMLR to the federated learning task where the training samples are collected by multiple agents in a network. Unlike the Expectation Maximization algorithm, WMLR directly extends to the distributed, federated learning setting. We support our theoretical results through several numerical experiments, which highlight our framework's ability to handle the federated learning setting with mixture models.