Global Convergence of Federated Learning for Mixed Regression

TL;DR

This paper introduces a federated learning algorithm for mixed regression that guarantees global convergence even with highly unbalanced local data, by combining moment descent and clustering refinement.

Contribution

It proposes a novel federated learning method that handles cluster-structured clients with limited data, achieving convergence from arbitrary initialization.

Findings

Algorithm converges globally from any initialization.

Effective in scenarios with highly unbalanced local data.

Provides theoretical guarantees using VC dimension bounds.

Abstract

This paper studies the problem of model training under Federated Learning when clients exhibit cluster structure. We contextualize this problem in mixed regression, where each client has limited local data generated from one of $k$ unknown regression models. We design an algorithm that achieves global convergence from any initialization, and works even when local data volume is highly unbalanced -- there could exist clients that contain $O(1)$ data points only. Our algorithm first runs moment descent on a few anchor clients (each with $\tilde{\Omega}(k)$ data points) to obtain coarse model estimates. Then each client alternately estimates its cluster labels and refines the model estimates based on FedAvg or FedProx. A key innovation in our analysis is a uniform estimate on the clustering errors, which we prove by bounding the VC dimension of general polynomial concept classes based on the theory of algebraic geometry.