Optimal Model Averaging: Towards Personalized Collaborative Learning

TL;DR

This paper provides a theoretical analysis of weighted model averaging in federated learning, demonstrating how it can improve personalized models by reducing expected error under minimal assumptions.

Contribution

It introduces a formal framework for understanding the benefits of weighted model averaging for personalization in federated learning, with quantification of its advantages.

Findings

Weighted averaging reduces expected squared error when local variance is non-zero.

The benefit of averaging depends on the choice of weight and the optimal weight.

The work offers a foundation for testing personalization benefits in more complex settings.

Abstract

In federated learning, differences in the data or objectives between the participating nodes motivate approaches to train a personalized machine learning model for each node. One such approach is weighted averaging between a locally trained model and the global model. In this theoretical work, we study weighted model averaging for arbitrary scalar mean estimation problems under minimal assumptions on the distributions. In a variant of the bias-variance trade-off, we find that there is always some positive amount of model averaging that reduces the expected squared error compared to the local model, provided only that the local model has a non-zero variance. Further, we quantify the (possibly negative) benefit of weighted model averaging as a function of the weight used and the optimal weight. Taken together, this work formalizes an approach to quantify the value of personalization in collaborative learning and provides a framework for future research to test the findings in multivariate parameter estimation and under a range of assumptions.