[Paper Review] Federated Multi-Task Learning under a Mixture of Distributions

Othmane Marfoq, Giovanni Neglia+3 more|arXiv (Cornell University)|Aug 23, 2021

Privacy-Preserving Technologies in Data67 references79 citations

TL;DR

The paper introduces FedEM and D-FedEM, federated EM-like algorithms for personalized multi-task learning under a mixture-of-distributions assumption, with convergence guarantees and evidence of improved accuracy and fairness over state-of-the-art methods.

ABSTRACT

The increasing size of data generated by smartphones and IoT devices motivated the development of Federated Learning (FL), a framework for on-device collaborative training of machine learning models. First efforts in FL focused on learning a single global model with good average performance across clients, but the global model may be arbitrarily bad for a given client, due to the inherent heterogeneity of local data distributions. Federated multi-task learning (MTL) approaches can learn personalized models by formulating an opportune penalized optimization problem. The penalization term can capture complex relations among personalized models, but eschews clear statistical assumptions about local data distributions. In this work, we propose to study federated MTL under the flexible assumption that each local data distribution is a mixture of unknown underlying distributions. This assumption encompasses most of the existing personalized FL approaches and leads to federated EM-like algorithms for both client-server and fully decentralized settings. Moreover, it provides a principled way to serve personalized models to clients not seen at training time. The algorithms' convergence is analyzed through a novel federated surrogate optimization framework, which can be of general interest. Experimental results on FL benchmarks show that our approach provides models with higher accuracy and fairness than state-of-the-art methods.

Motivation & Objective

• Motivate federated learning when a single global model is insufficient due to client data heterogeneity.
• Propose a flexible mixture model assumption where each client distribution is a mixture of M underlying distributions.
• Develop federated EM-like algorithms (FedEM and D-FedEM) to learn shared components and personalized mixtures.
• Provide convergence guarantees within a federated surrogate optimization framework.
• Demonstrate empirical advantages in accuracy and fairness over existing personalized and non-personalized FL methods.

Proposed method

• Adopt a generative mixture model: each client distribution is a mixture over M underlying distributions with shared component models.
• Represent each client’s predictor as a weighted sum of M component models, with learnable mixture weights per client.
• Formulate learning as minimizing negative log-likelihood; apply federated EM with E-step computing latent assignments and M-step updating mixture weights and component parameters.
• Provide centralized (client-server) and decentralized (fully peer-to-peer) implementations, FedEM and D-FedEM, with convergence guarantees to stationary points.
• Embed the approach in a federated surrogate optimization framework that generalizes across settings and solvers.

Experimental results

Research questions

• RQ1 Under what conditions can federated learning exploit cross-client data to improve personalized models when client data are heterogeneous?
• RQ2 Can a mixture-of-distributions assumption enable principled, scalable personalization in FL while preserving data locality?
• RQ3 Do EM-like federated algorithms converge to stationary points and provide practical generalization to unseen clients?
• RQ4 How do centralized and decentralized FL architectures compare under the proposed framework in terms of convergence and performance?
• RQ5 Do the proposed methods offer improvements in accuracy and fairness over existing personalized and non-personalized FL baselines?

Key findings

• FedEM and D-FedEM converge to stationary points of the federated objective under standard assumptions, with quantified rates for the gradient and mixture updates.
• The mixture-of-distributions framework yields personalized models as weighted combinations of component models, enabling knowledge transfer across clients.
• Empirical results on FL benchmarks show higher accuracy and improved fairness compared to state-of-the-art personalized and non-personalized FL methods.
• The approach facilitates inference for unseen clients by fixing components and learning new mixture weights at test time.
• Federated surrogate optimization provides a unifying convergence framework for the proposed EM-like algorithms across client-server and decentralized settings.