An Efficient Diffusion Model-Based Federated Learning with Non-IID Data for 6G Networks

Non-independent and identically distributed (non-IID) data significantly impact the performance of federated learning (FL) for 6G networks. Existing external solutions cannot address the underlying distribution imbalances or data scarcity. In this paper, we propose, for the first time, a novel FL training framework based on diffusion models (DMs). Within this framework, the DM and FL models are trained simultaneously. The DM training process is specifically designed to adapt to the non-IID characteristics of the data distribution in FL, effectively achieving a more balanced data distribution. To account for the resource limitations of client devices, DM training is assigned to clients with data scarcity, while other clients focus on FL model training. Additionally, a subset of data-rich clients participates in DM training to guide the generation of higher-quality synthetic data. Experimental results demonstrate that the proposed DM-based FL framework outperforms existing FL approaches by improving performance and achieving a more balanced data distribution.