Layer-Wise Adaptive Weighting for Faster Convergence in Federated Learning
Vedant S. Lanjewar, Hai Anh Tran, Truong X. Tran · 2024
Federated learning (FL), a decentralized approach utilizing multiple devices for training machine learning algorithms, presents an innovative solution to the challenges associated with handling sensitive personal data in neural network training. This research explores the application of FL for training neural network models, emphasizing its benefits in safeguarding privacy, optimizing efficiency, and addressing communication costs. While FL offers advantages such as enhanced data efficiency, promotion of heterogeneity, and scalability, challenges arise from the presence of non-independent and non-identically distributed data in the real world. This research introduces the FedLayerWise algorithm, a novel approach that dynamically assigns weights to each layer of clients’ models based on their contributions to the global model. The algorithm leverages the relationship between gradients of the client model and the global model to adaptively modify weights post-training, expediting convergence towards optimal loss. Experimental results on the MNIST dataset with non-identical distributions demonstrate the algorithm’s effectiveness. FedLayerWise outperforms existing algorithms (FedAdp and FedAvg) under highly skewed non-IID data distribution, achieving faster convergence. The proposed modifications showcase improvements and converge faster to non-IID data distributions, reducing computation rounds by about $\mathbf{4 0 \%}$ compared to FedAdp and about $\mathbf{5 0 \%}$ compared to FedAvg in such scenarios.