Local Performance vs. Out-of-Distribution Generalization: An Empirical Analysis of Personalized Federated Learning in Heterogeneous Data Environments

TL;DR

This paper empirically analyzes how personalized federated learning approaches perform locally and in out-of-distribution scenarios, proposing a modified FedAvg method with individualization to improve robustness in heterogeneous data environments.

Contribution

It introduces Federated Learning with Individualized Updates (FLIU), a simple extension to FedAvg that enhances personalization and robustness in heterogeneous federated learning settings.

Findings

FLIU improves local performance over standard FedAvg.

Evaluation on MNIST and CIFAR-10 shows enhanced generalization in non-IID conditions.

Modified approach demonstrates robustness in complex data heterogeneity environments.

Abstract

In the context of Federated Learning with heterogeneous data environments, local models tend to converge to their own local model optima during local training steps, deviating from the overall data distributions. Aggregation of these local updates, e.g., with FedAvg, often does not align with the global model optimum (client drift), resulting in an update that is suboptimal for most clients. Personalized Federated Learning approaches address this challenge by exclusively focusing on the average local performances of clients' models on their own data distribution. Generalization to out-of-distribution samples, which is a substantial benefit of FedAvg and represents a significant component of robustness, appears to be inadequately incorporated into the assessment and evaluation processes. This study involves a thorough evaluation of Federated Learning approaches, encompassing both their local performance and their generalization capabilities. Therefore, we examine different stages within a single communication round to enable a more nuanced understanding of the considered metrics. Furthermore, we propose and incorporate a modified approach of FedAvg, designated as Federated Learning with Individualized Updates (FLIU), extending the algorithm by a straightforward individualization step with an adaptive personalization factor. We evaluate and compare the approaches empirically using MNIST and CIFAR-10 under various distributional conditions, including benchmark IID and pathological non-IID, as well as additional novel test environments with Dirichlet distribution specifically developed to stress the algorithms on complex data heterogeneity.