2025/02/11 by Divyansh Jhunjhunwala, Pranay Sharma, Jhunjhunwala, Divyansh +5
Computer Science · #Advanced Statistical Modeling Techniques #Distributed #FOS: Computer and information sciences #Machine Learning (cs.LG) #Parallel #Privacy-Preserving Technologies in Data #and Cluster Computing (cs.DC)
paper · pdf · doi:10.48550/arxiv.2502.08024
openalex publication_date 2025/02/11 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
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.