FedPDG: Prediction Discrepancy–Guided Data Generation for Heterogeneous Federated Learning

One emerging approach to mitigating data heterogeneity in Federated Learning (FL) is to employ diffusion models to generate synthetic data for clients, thereby aligning local data distributions with the global distribution. Prior work has primarily focused on balance-oriented augmentation, which assumes a balanced global class distribution and thus generates samples of rare classes to rebalance each client's local dataset. However, in practice, global data distributions are often inherently imbalanced. Moreover, privacy constraints in FL hinder the server’s ability to accurately estimate the global distribution, rendering balance-oriented augmentation suboptimal. This raises a key, underexplored challenge: How can synthetic data be generated and selected to align local distributions with the true, yet unknown, global distribution? Our key insight is that a model’s performance implicitly reflects the data distribution it has been trained on. Based on this observation, we use the performance discrepancy between local and global models to identify the regions where each client’s local dataset is lacking, and generate corresponding samples for clients. Furthermore, we adapt the diffusion model via preference optimization, enabling it to generate data that better aligns with the true global distribution. Extensive experiments on multiple benchmarks demonstrate that FedPDG outperforms state-of-the-art methods, achieving up to 3.82\% improvement.

Federated learning allows many users or organizations to train a shared model without sending their private data to a central server. However, each participant may own data with very different patterns. For example, one client may have many images of cats but few images of ships, while another may have the opposite. A common way to reduce this mismatch is to generate extra synthetic data for each client, but existing methods often assume that the overall data should be balanced across classes. This assumption can be wrong when the real-world data distribution is naturally imbalanced. We propose FedPDG, a method that decides what synthetic data each client needs by comparing how a client model and the global model perform on generated samples. If the global model handles a sample well but a client model struggles with it, this suggests that the client lacks similar data. FedPDG then sends such useful synthetic samples to the client and further adapts the generator to produce more relevant data. This helps local models better match the true global data distribution without exposing private data. Experiments on several image classification benchmarks show that FedPDG improves federated learning under challenging non-identical and imbalanced data settings.