Fractal Augmented Pre-training and Gaussian Virtual Feature Calibration for Tackling Data Heterogeneity in Federated Learning

Federated learning (FL) enables collaborative model training across multiple clients while preserving privacy. In practical situations, the heterogeneous and unbalanced distribution of the data has a significant impact on the performance of the model. Although some work has been carried out to address this issue, such as adding regularization terms, employing specific server aggregation strategies, and utilizing deep generative models to augment the training data, there is still a lack of efficient approaches to derive intrinsic representation of the local data to improve the global model without compromising client privacy. Through our careful observations and analysis, we found that incorporating pre-training and calibrating of the global model using virtual data and virtual features that are generated based on the client data distribution can improve model generalization. In this work, we propose Virtual Data Augmented Federated Learning (FedVDA) to resolve this problem. Specifically, FedVDA combines unsupervised pre-training with Augmented Fractal (AF) virtual images and Gaussian Mixture Model (GMM) virtual feature calibration. By integrating color tone transformations into the virtual data generated by fractals, we bridge the gap between virtual and client data distributions. Multi-modal feature modeling using variances on each client allows the server to efficiently calibrate the classifier with balanced sampled virtual features, reducing both computational and communication overhead. Compared to other data augmentation methods, our method directly calibrates model features, significantly improving model performance in scenarios with data heterogeneity and imbalance, while minimizing additional computational and communication costs. Our experiments demonstrate that FedVDA outperforms existing federated learning methods and can seamlessly integrate with other algorithms.