[Paper Review] Personalized Federated Learning with Moreau Envelopes

Canh T. Dinh, Nguyen H. Tran+1 more|arXiv (Cornell University)|Jun 16, 2020

Privacy-Preserving Technologies in Data56 references214 citations

TL;DR

Introduces pFedMe, a personalized federated learning algorithm using Moreau envelopes to decouple personalized model optimization from global model learning, achieving state-of-the-art convergence rates and improved local accuracy over FedAvg and Per-FedAvg.

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.

Motivation & Objective

• Address non-i.i.d. data across clients in federated learning by enabling personalization.
• Formulate a bi-level optimization problem using Moreau envelopes to separate global and personalized model updates.
• Prove convergence rates for strongly convex and nonconvex objectives.
• Empirically validate pFedMe against FedAvg and Per-FedAvg on real and synthetic datasets.

Proposed method

• Formulate a bi-level optimization where F(w)=1/N sum Fi(w) and Fi(w)=min_thetai { fi(theta_i) + (lambda/2) ||theta_i - w||^2 }.
• Use Moreau envelopes to derive the proximal-like personalized updates and a gradient-based outer update for the global model.
• At each round, perform R local steps to optimize theta_i(w) via delta-approximation tilde_theta_i, then update w via gradient of Fi with respect to w.
• Provide convergence analysis for strongly convex and nonconvex smooth settings, establishing quadratic and 2/3 sublinear speedups under respective assumptions.
• Compare pFedMe with FedAvg and Per-FedAvg on MNIST and synthetic data, examining hyperparameters R, K, |D|, lambda, and beta.

Experimental results

Research questions

• RQ1 How can Moreau envelopes be leveraged to decouple personalized model optimization from global model learning in FL?
• RQ2 What convergence rates can be achieved for pFedMe under strongly convex and nonconvex objectives?
• RQ3 Does pFedMe provide better local (personalized) performance and/or faster convergence than FedAvg and Per-FedAvg in non-i.i.d. settings?
• RQ4 How do hyperparameters (R, K, |D|, lambda, beta) affect performance and convergence?

Key findings

• pFedMe achieves state-of-the-art convergence speeds: quadratic speedup for strongly convex and 2/3 sublinear speedup for smooth nonconvex objectives.
• Empirically, pFedMe’s personalized models outperform FedAvg and Per-FedAvg in local accuracy and convergence rate on MNIST and synthetic data.
• The Moreau envelope-based formulation effectively decouples personalized optimization from global learning, enabling parallelized updates.
• Delta-approximation of the inner minimizer and first-order gradients suffice, avoiding Hessian computations as in other meta-learning approaches.
• Hyperparameter tuning shows larger R can improve convergence but trades off computation; appropriate lambda balances personalization and global alignment.