Scaling Law Analysis in Federated Learning: How to Select the Optimal Model Size?

TL;DR

This paper analyzes how to select the optimal model size in federated learning by deriving a theoretical bound on generalization error and validating it through extensive experiments, addressing the challenges of scaling large models in decentralized data settings.

Contribution

It introduces a PAC-Bayes bound for federated learning, providing a theoretical framework to determine optimal model size based on client number and compute, filling a key research gap.

Findings

Optimal model size decreases with more clients if total compute is fixed.

Switching to federated learning reduces the maximum achievable generalization performance.

Estimating optimal model size should consider average training compute across clients.

Abstract

The recent success of large language models (LLMs) has sparked a growing interest in training large-scale models. As the model size continues to scale, concerns are growing about the depletion of high-quality, well-curated training data. This has led practitioners to explore training approaches like Federated Learning (FL), which can leverage the abundant data on edge devices while maintaining privacy. However, the decentralization of training datasets in FL introduces challenges to scaling large models, a topic that remains under-explored. This paper fills this gap and provides qualitative insights on generalizing the previous model scaling experience to federated learning scenarios. Specifically, we derive a PAC-Bayes (Probably Approximately Correct Bayesian) upper bound for the generalization error of models trained with stochastic algorithms in federated settings and quantify the impact of distributed training data on the optimal model size by finding the analytic solution of model size that minimizes this bound. Our theoretical results demonstrate that the optimal model size has a negative power law relationship with the number of clients if the total training compute is unchanged. Besides, we also find that switching to FL with the same training compute will inevitably reduce the upper bound of generalization performance that the model can achieve through training, and that estimating the optimal model size in federated scenarios should depend on the average training compute across clients. Furthermore, we also empirically validate the correctness of our results with extensive training runs on different models, network settings, and datasets.