Generative AI-Powered Plugin for Robust Federated Learning in Heterogeneous IoT Networks

TL;DR

This paper introduces a generative AI-powered plugin for federated learning that enhances model convergence and robustness by balancing data distributions across heterogeneous IoT devices.

Contribution

It proposes a novel plugin that uses generative AI for data augmentation and balanced sampling to address Non-IID data issues in federated learning.

Findings

Improved convergence speed in federated learning.

Enhanced robustness against data imbalance.

Effective data augmentation with generative AI.

Abstract

Federated learning enables edge devices to collaboratively train a global model while maintaining data privacy by keeping data localized. However, the Non-IID nature of data distribution across devices often hinders model convergence and reduces performance. In this paper, we propose a novel plugin for federated optimization methods that approximates Non-IID data distributions to IID through generative AI-enhanced data augmentation and balanced sampling strategy. The key idea is to synthesize additional data for underrepresented classes on each edge device, leveraging generative AI to create a more balanced dataset across the FL network. Additionally, a balanced sampling approach at the central server selectively includes only the most IID-like devices, accelerating convergence while maximizing the global model's performance. Experimental results validate that our approach significantly improves convergence speed and robustness against data imbalance, establishing a flexible, privacy-preserving FL plugin that is applicable even in data-scarce environments.