The Effect of Personalization in FedProx: A Fine-grained Analysis on Statistical Accuracy and Communication Efficiency

FedProx is a simple yet effective federated learning method that enables model personalization via regularization. Despite remarkable success in practice, a rigorous analysis of how such a regularization provably improves the statistical accuracy of each client's local model hasn't been fully established. Setting the regularization strength heuristically presents a risk, as an inappropriate choice may even degrade accuracy. This work fills in the gap by analyzing the effect of regularization on statistical accuracy, thereby providing a theoretical guideline for setting the regularization strength for achieving personalization. We prove that by adaptively choosing the regularization strength under different statistical heterogeneity, FedProx can consistently outperform pure local training and achieve a minimaxoptimal statistical rate. In addition, to shed light on resource allocation, we design an algorithm, provably showing that stronger personalization reduces communication complexity without increasing the computation cost overhead. Finally, our theory is validated on both synthetic and real-world datasets and its generalizability is verified in a non-convex setting. Federated Learning (FL) has emerged as an attractive framework for aggregating distributed data, enabling clients to collaboratively train a shared global model while preserving data privacy. In the currently prevalent paradigm (McMahan et al., 2017), FL is formulated as a finite sum minimization problem focusing on a single shared model. Nevertheless, it has been well recognized that one of the key challenges in FL is the statistical heterogeneity of the client datasets. As the participants collect their own local data, it often reflects client-specific characteristics and is not identically distributed. With high statistical heterogeneity, training a single model for all clients by minimizing their average in-sample loss becomes questionable. To address this challenge, one solution is to relax the common model constraint and solve alternatively the following objective in FedProx (Li et al., 2020a): ( min p The smaller λ is, the weaker the coupling of the local models the formulation would enforce thus the higher personalization is.