Initialization Matters: Unraveling the Impact of Pre-Training on Federated Learning

TL;DR

Pre-trained models in federated learning reduce data heterogeneity effects by initializing with aligned filters, leading to better test performance, as shown through theoretical analysis and experiments.

Contribution

This work provides a theoretical analysis of how pre-training influences filter alignment and improves federated learning performance under data heterogeneity.

Findings

Pre-trained models lead to fewer misaligned filters at initialization.

Data heterogeneity mainly affects misaligned filters during training.

Experiments confirm theoretical predictions in synthetic and real FL settings.

Abstract

Initializing with pre-trained models when learning on downstream tasks is becoming standard practice in machine learning. Several recent works explore the benefits of pre-trained initialization in a federated learning (FL) setting, where the downstream training is performed at the edge clients with heterogeneous data distribution. These works show that starting from a pre-trained model can substantially reduce the adverse impact of data heterogeneity on the test performance of a model trained in a federated setting, with no changes to the standard FedAvg training algorithm. In this work, we provide a deeper theoretical understanding of this phenomenon. To do so, we study the class of two-layer convolutional neural networks (CNNs) and provide bounds on the training error convergence and test error of such a network trained with FedAvg. We introduce the notion of aligned and misaligned filters at initialization and show that the data heterogeneity only affects learning on misaligned filters. Starting with a pre-trained model typically results in fewer misaligned filters at initialization, thus producing a lower test error even when the model is trained in a federated setting with data heterogeneity. Experiments in synthetic settings and practical FL training on CNNs verify our theoretical findings.