One-Shot Heterogeneous Federated Learning with Local Model-Guided Diffusion Models

Mingzhao Yang, Shangchao Su, Bin Li, Xiangyang Xue

Abstract: In recent years, One-shot Federated Learning (OSFL) methods based on Diffusion Models (DMs) have garnered increasing attention due to their remarkable performance. However, most of these methods require the deployment of foundation models on client devices, which significantly raises the computational requirements and reduces their adaptability to heterogeneous client models. In this paper, we propose FedLMG, a heterogeneous one-shot Federated learning method with Local Model-Guided diffusion models. In our method, clients do not need access to any foundation models but only train and upload their local models, which is consistent with traditional FL methods. On the clients, we employ classification loss and batch normalization loss to capture the broad category features and detailed contextual features of the client distributions. On the server, based on the uploaded client models, we utilize backpropagation to guide the server’s DM in generating synthetic datasets that comply with the client distributions, which are then used to train the aggregated model. By using the local models as a medium to transfer client knowledge, our method significantly reduces the computational requirements on client devices and effectively adapts to scenarios with heterogeneous clients. Extensive quantitation and visualization experiments on three large-scale real-world datasets, along with theoretical analysis, demonstrate that the synthetic datasets generated by FedLMG exhibit comparable quality and diversity to the client datasets, which leads to an aggregated model that outperforms all compared methods and even the performance ceiling, further elucidating the significant potential of utilizing DMs in FL.

Lay Summary: We present a new one-shot federated learning method that enables devices to collaboratively train a model without sharing their private data. Unlike previous methods that require large, complex models on each device, the clients in our method train a lightweight model locally, which is then used by the server to guide a diffusion model that generates a synthetic dataset. This synthetic dataset compy the data distributions on the clients without leaking the privacy-sensitive information, enabling the server to build a strong aggregated model. Our method achieves strong results on real-world datasets and highlights the great potential of utilizing diffusion models in federated learning.

Link To Code: https://github.com/MingzhaoYang/FedLMG

Primary Area: Deep Learning->Everything Else

Keywords: Federated Learning, Diffusion Model

Submission Number: 6707