A Distributed Generative AI Approach for Heterogeneous Multi-Domain Environments under Data Sharing constraints

TL;DR

HuSCF-GAN presents a comprehensive framework for distributed generative AI in heterogeneous, multi-domain environments under strict data-sharing constraints by combining four-cut heterogeneous U-Shaped Split Learning with Clustered Federated Learning. A genetic algorithm optimizes per-client cut points, while KL-based intra-cluster weighting and discriminator activation clustering address non-IID data, enabling scalable federation with privacy-preserving exchanges of activations/gradients. Empirical evaluations across MNIST family, higher-resolution images, medical imaging, and audio demonstrate consistent gains in image generation quality and classifier performance, along with substantial latency reductions relative to existing decentralized GAN approaches. The approach is extensible to other generative architectures and modalities, offering a practical pathway toward privacy-preserving, scalable distributed training on edge and IoT devices, with noted limitations and directions for future work including fully decentralized cut-point strategies and stronger privacy guarantees.

Abstract

Federated Learning has gained increasing attention for its ability to enable multiple nodes to collaboratively train machine learning models without sharing their raw data. At the same time, Generative AI -- particularly Generative Adversarial Networks (GANs) -- have achieved remarkable success across a wide range of domains, such as healthcare, security, and Image Generation. However, training generative models typically requires large datasets and significant computational resources, which are often unavailable in real-world settings. Acquiring such resources can be costly and inefficient, especially when many underutilized devices -- such as IoT devices and edge devices -- with varying capabilities remain idle. Moreover, obtaining large datasets is challenging due to privacy concerns and copyright restrictions, as most devices are unwilling to share their data. To address these challenges, we propose a novel approach for decentralized GAN training that enables the utilization of distributed data and underutilized, low-capability devices while not sharing data in its raw form. Our approach is designed to tackle key challenges in decentralized environments, combining KLD-weighted Clustered Federated Learning to address the issues of data heterogeneity and multi-domain datasets, with Heterogeneous U-Shaped split learning to tackle the challenge of device heterogeneity under strict data sharing constraints -- ensuring that no labels or raw data, whether real or synthetic, are ever shared between nodes. Experiments show that our approach demonstrates significant improvements across key metrics, where it achieves an average 10% boost in classification metrics (up to 60% in multi-domain non-IID settings), 1.1x -- 3x higher image generation scores for the MNIST family datasets, and 2x -- 70x lower FID scores for higher resolution datasets. Find our code at https://github.com/youssefga28/HuSCF-GAN.

Figures (20)

• Figure 1: Comparison between Federated Learning, traditional Split Learning, and U-shaped Split Learning. In Federated Learning, each client trains a local model for several epochs and then sends its model weights to a central server. The server aggregates these weights—typically by averaging—and sends the updated global model back to the clients. This process is repeated for multiple rounds. In traditional Split Learning, the model is divided into two parts: the client holds the initial segment, and the server holds the remaining part. In U-shaped Split Learning, the model is split into three segments: the client retains both the initial and final segments, while the server manages the middle segment.
• Figure 2: SFL architecture: In the first step (split learning), clients train their local (client-side) model segments and send the resulting activations to the server, which continues training on the server-side segments. During backpropagation, the server returns the gradients to the clients. In the second step, after several local epochs, clients send their updated client-side models to the server for federated aggregation and redistribution.
• Figure 3: Architecture of a Conditional GAN (cGAN). The generator and discriminator are both conditioned on auxiliary information, such as class labels, allowing the model to generate data that adheres to specific conditions.
• Figure 4: HuSCF-GAN Overview: Clients first send device capabilities to the server, which uses a Genetic Algorithm to assign optimal cut points. Clients then perform U-shaped split learning, exchanging intermediate activations/gradients with the server. Every $E$ epochs, the server clusters discriminator activations and computes intra-cluster KLD scores and perform an intra-cluster federated learning round whose aggregation weights consider both data size and KLD.
• Figure 5: Overview of the proposed HuSCF-GAN procedure.