Optimization of Federated Learning Algorithm for Non-lID Data: Improvements to the FedDyn Algorithm
Wenrui Bai · 2024
This paper investigates the optimization of the FedDyn algorithm in Federated Learning (FL). Federated Learning is a distributed machine learning framework that enables model training on local data across multiple devices while ensuring data privacy. However, the existing FedDyn algorithm, while performing well in handling non-independent and identically distributed (non-IID) data and partial client participation, still faces challenges in large-scale scenarios, such as slow convergence, low training efficiency, and high memory overhead. To address these challenges, this paper systematically optimizes the FedDyn algorithm and proposes several improvements. First, a pre-trained model initialization is introduced, where the global model is pre-trained on a public dataset and then distributed to clients, accelerating early-stage convergence. Second, parallel processing and multi-threading techniques are employed to enhance local training efficiency, and mixed-precision training is used to reduce memory usage and accelerate computation. Additionally, a data volume-weighted hierarchical aggregation method is proposed, ensuring that clients with larger datasets have a greater influence on global model updates, thus improving the model's generalization capability and stability. The effectiveness of the optimized algorithm is validated through experiments on the CIFAR-10 and CIFAR-100 datasets. Compared to FedAvg, SCAFFOLD, and the original FedDyn, the optimized FedDyn shows significant improvements in test accuracy, convergence speed, and computational efficiency, particularly in handling highly heterogeneous non-IID data. Future work could explore more advanced quantization methods, such as low-bit quantization, to further reduce computational resource consumption, and incorporate techniques like differential privacy to enhance the application in privacy-sensitive domains.