Personalized Federated Learning with Moreau Envelopes

TL;DR

This paper introduces pFedMe, a personalized federated learning algorithm using Moreau envelopes, achieving state-of-the-art convergence rates and superior empirical performance over existing methods.

Contribution

The paper proposes a novel personalized FL algorithm with theoretical convergence guarantees and improved empirical results, addressing client data heterogeneity.

Findings

pFedMe achieves quadratic convergence for strongly convex objectives.

pFedMe outperforms FedAvg and Per-FedAvg in empirical tests.

Theoretical analysis shows state-of-the-art convergence rates.

Abstract

Federated learning (FL) is a decentralized and privacy-preserving machine learning technique in which a group of clients collaborate with a server to learn a global model without sharing clients' data. One challenge associated with FL is statistical diversity among clients, which restricts the global model from delivering good performance on each client's task. To address this, we propose an algorithm for personalized FL (pFedMe) using Moreau envelopes as clients' regularized loss functions, which help decouple personalized model optimization from the global model learning in a bi-level problem stylized for personalized FL. Theoretically, we show that pFedMe's convergence rate is state-of-the-art: achieving quadratic speedup for strongly convex and sublinear speedup of order 2/3 for smooth nonconvex objectives. Experimentally, we verify that pFedMe excels at empirical performance compared with the vanilla FedAvg and Per-FedAvg, a meta-learning based personalized FL algorithm.