[Paper Review] Improving Federated Learning Personalization via Model Agnostic Meta Learning

Yihan Jiang, Jakub Konečný+2 more|arXiv (Cornell University)|Sep 27, 2019

Privacy-Preserving Technologies in Data22 references368 citations

TL;DR

The paper shows Federated Averaging (FedAvg) can be interpreted as a meta-learning algorithm and introduces Personalized FedAvg by combining FedAvg with a fine-tuning stage (Reptile/Adam) to improve personalization, initial model quality, and convergence speed in non-i.i.d. FL settings.

ABSTRACT

Federated Learning (FL) refers to learning a high quality global model based on decentralized data storage, without ever copying the raw data. A natural scenario arises with data created on mobile phones by the activity of their users. Given the typical data heterogeneity in such situations, it is natural to ask how can the global model be personalized for every such device, individually. In this work, we point out that the setting of Model Agnostic Meta Learning (MAML), where one optimizes for a fast, gradient-based, few-shot adaptation to a heterogeneous distribution of tasks, has a number of similarities with the objective of personalization for FL. We present FL as a natural source of practical applications for MAML algorithms, and make the following observations. 1) The popular FL algorithm, Federated Averaging, can be interpreted as a meta learning algorithm. 2) Careful fine-tuning can yield a global model with higher accuracy, which is at the same time easier to personalize. However, solely optimizing for the global model accuracy yields a weaker personalization result. 3) A model trained using a standard datacenter optimization method is much harder to personalize, compared to one trained using Federated Averaging, supporting the first claim. These results raise new questions for FL, MAML, and broader ML research.

Motivation & Objective

• Motivate and formalize the need for personalization in Federated Learning amidst data heterogeneity.
• Reveal the connection between FedAvg and MAML-like meta-learning methods.
• Propose a two-stage FedAvg-based approach (Personalized FedAvg) to optimize initial model quality, personalization, and fast convergence.
• Demonstrate empirically that FedAvg inherently optimizes personalized performance and that a fine-tuning stage improves personalization stability and effectiveness.
• Highlight how global model accuracy does not guarantee strong personalization and discuss implications for FL and MAML research.

Proposed method

• Interpret FedAvg as a meta-learning process; relate FedAvg to Reptile and first/second-order MAML components.
• Show that FedAvg equals a linear combination of FedSGD and FOMAML updates under certain conditions.
• Propose Personalized FedAvg: (i) run FedAvg with multiple local epochs and momentum on the server, (ii) switch to Reptile-like fine-tuning with small K and Adam on the server to improve the initial model, (iii) continue personalization using the same client optimizer as training.

Experimental results

Research questions

• RQ1 Can FedAvg be interpreted as a meta-learning algorithm oriented toward personalization in FL?
• RQ2 How does the choice of local epochs (K) and optimizer affect personalization vs. initial model quality?
• RQ3 Does a two-stage training (FedAvg training followed by fine-tuning with Reptile/Adam) improve both initial and personalized performance in FL?
• RQ4 How does personalization performance vary across different data regimes (i.e., client data availability and non-i.i.d. settings)?

Key findings

• FedAvg can be viewed as a meta-learning algorithm that inherently optimizes personalized performance, not just global accuracy.
• Increasing local epochs (E) improves personalization up to a point, but may destabilize the initial model; a two-stage approach helps stabilize and improve personalization.
• Fine-tuning the initial model with Reptile(K) and Adam yields better and more stable personalized performance across a broader hyperparameter range.
• Center-placed initial models trained in a centralized way are harder to personalize than FedAvg-trained models, especially under practical personalization scenarios.
• Different global models with the same accuracy can have very different personalization capacity, highlighting the need to focus on personalization metrics in FL evaluation.