Global Convergence of Federated Learning for Mixed Regression

This paper studies the problem of model training under Federated Learning when clients exhibit cluster structures. 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 arbitrary initialization, and works even when local data volume is highly unbalanced – there could exist clients that contain <inline-formula> <tex-math notation="LaTeX">$O(1)$ </tex-math></inline-formula> data points only. Our algorithm is intended for the scenario where the parameter server can recruit one client per cluster referred to as “anchor clients”, and each anchor client possesses <inline-formula> <tex-math notation="LaTeX">$\tilde {\Omega }(k)$ </tex-math></inline-formula> data points. Our algorithm first runs moment descent on this set of anchor clients to obtain coarse model estimates. Subsequently, every 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 of the clustering errors, which we prove by bounding the Vapnik-Chervonenkis dimension of general polynomial concept classes based on the theory of algebraic geometry.