Source-linked AI summary
Improving Federated Learning Personalization via Model Agnostic Meta Learning
Yihan Jiang, Jakub Konečný, Keith Rush, Sreeram Kannan
TL;DR
The paper asks how decentralized FL models can personalize across heterogeneous devices while retaining a strong global model and fast convergence. It connects FedAvg with MAML, proposes training and fine-tuning around this connection, and reports that personalization should be a primary FL objective because training choices strongly affect it.
Problem
Heterogeneous, non-i.i.d. FL clients make device-level performance difficult, while existing approaches disconnect global-model training from local personalization.
Method
The paper interprets FedAvg through the MAML framework and studies a two-stage training and fine-tuning modification for joint FL objectives.
Results
FedAvg can be interpreted as meta learning that optimizes personalized performance, while global-model accuracy alone tends to reduce personalization capacity.
Takeaways & Limitations
FL objectives should include personalized performance alongside initial-model quality and fast convergence, and MAML methods should consider these practical constraints.
Takeaways & Limitations
Some MAML datasets do not admit a good notion of initial accuracy.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Federated Learning addresses decentralized, heterogeneous device data, but personalization remains difficult when training and adaptation are disconnected. The paper connects FL with MAML and argues that personalization, initial-model quality, and fast convergence should be optimized together.
- FL trains a shared model without sending raw device data to the cloud, but non-i.i.d. and unbalanced client data makes performance across devices challenging.
- Existing personalization evaluates gradient-based adaptation from a converged initial model, leaving training and personalization disconnected and potentially suboptimal.
- MAML optimizes post-adaptation performance on heterogeneous tasks, while FL clients similarly perform local gradient updates before sending aggregate updates to a coordinator.
- The paper proposes addressing improved personalized models, solid initial models, and fast convergence simultaneously.
- It connects FedAvg with MAML-style meta learning and proposes a two-stage FedAvg modification involving training and fine-tuning.
2 INTERPRETING FEDAVG AS A META LEARNING ALGORITHM
FedAvg can be interpreted through the meta-learning framework because both FL clients and MAML tasks perform local adaptation before a coordinator aggregates updates. Under equal client weights, FedAvg and Reptile become identical, while FedAvg can also be viewed as combining gradient-based methods with different adaptation lengths.
- FedAvg–Reptile equivalence: FedAvg and Reptile become the same algorithm when clients have equal data amounts and therefore identical weights.Both sample clients, perform local optimization, and aggregate updates into a new global model.
- Shared structure: The conceptual algorithm connects MAML meta-training, Reptile, and FL training through shared inner-loop adaptation and outer-loop aggregation.MAML aggregates task gradients, while FL aggregates locally optimized client updates.
- Linear-combination view: FedAvg can be rearranged as a linear combination of Federated SGD and First Order MAML with different numbers of local steps.This interpretation follows from expressing FedAvg client updates as sums of local gradient updates.
- Adaptation length: FedSGD corresponds to FOMAML(0), optimizing performance after zero local updates, whereas FOMAML uses the gradient after K local steps.MAML's full update directly optimizes performance after local adaptation, while FOMAML omits the second-order term.
- Adaptation length: Increasing the number of local adaptation steps substantially improves personalization until the initial model becomes unstable beyond a certain point.The paper reports that K = 1 produces a model hard to personalize and that larger K eventually destabilizes initial-model performance.
3 PERSONALIZED FEDAVG
Personalized FedAvg combines FedAvg training with a Reptile-based fine-tuning stage to improve the initial model while preserving personalized performance. Experiments find that personalization is more stable with suitable local training, whereas centralized optimization is harder to personalize.
- Algorithm: Personalized FedAvg first runs FedAvg with momentum SGD, then fine-tunes the initial model using Reptile(K) with Adam.Personalization uses the same client optimizer as training.
- FedAvg behavior: FedAvg with several local epochs converges quickly and already optimizes personalized performance, but its initial model is relatively unstable.The authors motivate fine-tuning to improve the initial model without sacrificing personalization.
- Convergence: Increasing local epochs initially improves results; E in the range of 5−10 appears best in the EMNIST-62 convergence experiment.The experiment compares initial and personalized accuracy after 500 rounds and communication rounds to 80% accuracy.
- Dataset scope: On Shakespeare, personalization produces only a small positive improvement, which the authors associate with next-character prediction being similar across users.The paper does not study this problem further and suggests next-word prediction might show a larger personalization difference.
- Fine-tuning: Reptile(10) improves initial accuracy, slightly improves personalized accuracy, and maintains roughly the same performance across a wide range of local personalization epochs.Reptile(1) improves initial accuracy but reduces personalized accuracy.
- Centralized comparison: Centralized initial models are significantly harder to personalize than federated initial models, requiring smaller learning rates and yielding smaller personalization improvements.With centralized models, Adam performs better than SGD at the tested learning rate but remains below the federated fine-tuning performance.
4 DISCUSSION AND FUTURE WORK
The paper argues that federated learning should prioritize adaptation to client heterogeneity alongside global-model quality, connecting FL with MAML and raising broader optimization questions.
- Federated Averaging improves personalized performance while also improving the global model, challenging the conventional shared-model objective.
- Challenges for Model Agnostic Meta Learning: MAML methods should consider initial-model performance and fast communication-round convergence because many FL clients lack personalization data or connectivity.
- Challenges for broader Machine Learning: Models with similar initial accuracy can differ substantially in personalization capacity, motivating measures or optimization methods that predict adaptability.
- FL and MAML algorithms may provide base models that improve domain adaptation after replacing the final output layer.
- Analyzing inner-outer optimization structures could clarify connections between optimization and generalization.
A APPENDIX
The appendix provides additional details referenced from the paper’s main body.
- The appendix contains further details referenced from the main body of the paper.
A.1 FINE TUNING OPTIMIZERS
Fine-tuning experiments compare server optimizers and report that Adam is consistently better and more stable for initial accuracy under matched client optimizers.
- Adam consistently provides better and more stable initial accuracy than other server optimizers when comparing the same client optimizers.
A.2 PER-CLIENT PERSONALIZATION RESULTS
Per-client results visualize initial and personalized accuracies for sampled EMNIST-62 and Shakespeare test clients, emphasizing client-level variation and deployment risk.
- Figure 4 shows the distribution of initial and personalized accuracies separately for sampled EMNIST-62 and Shakespeare test clients.
- Each dot represents a random sample of test clients used in the personalization experiments.
- Client-level analysis matters because degrading one user’s experience may impose disproportionate costs relative to comparable improvements.