A Method for Data Augmentation in Vertical Federated Learning Addressing Data Heterogeneity

Area of Science:

• Artificial Intelligence
• Machine Learning
• Data Science

Background:

• Vertical Federated Learning (VFL) enables data collaboration but faces challenges with data heterogeneity.
• Imbalanced local data and shifting parameter distributions degrade VFL model performance.

Purpose of the Study:

• To propose a framework addressing data heterogeneity in VFL.
• To improve the performance, generalization, and stability of VFL models in multi-data-source environments.

Main Methods:

• Conditional Generative Adversarial Networks (CGANs) for targeted data augmentation (FeCWGAN-GP).
• A parameter aggregation algorithm (WFedDA) using sample size and Wasserstein distance.
• A smoothed Stochastic Gradient Descent (SGD-MA) with exponential moving average (EMA).

Main Results:

• FeCWGAN-GP alleviates performance degradation from imbalanced local data.
• WFedDA optimizes the global model by accounting for parameter distribution shifts.
• SGD-MA enhances training stability by reducing gradient and parameter fluctuations.
• Experiments on MNIST, CIFAR-10, Fashion-MNIST, and MIMIC-III demonstrate significant improvements.

Conclusions:

• The proposed framework effectively tackles data heterogeneity challenges in VFL.
• The methods significantly enhance the generalization capability and stability of the global VFL model.