A federated learning framework for Non-IID data : data feature optimization and personalization strategy

Federated learning is an emerging distributed machine learning paradigm that enables multiple edge devices to collaboratively train task-specific models while preserving data privacy. However, the non-independent and identically distributed nature of client data often leads to client drift, slow convergence, and degraded model performance. To address these challenges, this paper proposes a multi-stage collaborative optimization framework for federated learning under Non-IID settings. Specifically, the proposed framework first employs feature distillation to construct noise-perturbed proxy data for cross-client knowledge interaction. Then, to reduce proxy data redundancy and accelerate model convergence, an entropy-based dynamic selection mechanism is designed to prioritize informative proxy samples for local training. At the communication level, a partial parameter selection strategy is proposed to dynamically distinguish shared parameters from personalized parameters according to parameter update magnitude. In this way, only shared parameters are uploaded to the server for aggregation, while personalized parameters are retained locally, thereby reducing unnecessary communication transmission. In addition, a late-stage personalized fine-tuning strategy is further introduced to enhance the model’s adaptability to local client data distributions. Experimental results on multiple benchmark Non-IID datasets demonstrate that the proposed method outperforms existing baseline methods, achieving up to 6.87% higher classification accuracy while effectively reducing the number of communication rounds required to reach the target performance. The results verify the effectiveness of the proposed framework in improving model performance, accelerating convergence, and reducing communication overhead.