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

Divyansh Jhunjhunwala, Pranay Sharma, Zheng Xu, Gauri Joshi

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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

* Added discussion on extending the setup to multi-class classification (after Definition 1).

* Included proof sketches following Theorem 1 and Theorem 2.

* Expanded the discussion on extending insights to deeper CNNs and Transformer architectures (Conclusion section).

* Introduced an additional experiment on two-layer CNNs with different initialization strategies (Appendix, Section: Additional Experimental Details and Results).

* Corrected typos, including the definition of the loss function, a missing citation on Page 3, and a reference to Condition 2 (Page 3), among others.

All new changes have been highlighted in blue for ease of reference.

Assigned Action Editor: ~Anastasios_Kyrillidis2

Submission Number: 5085