DiVerFed: Distribution-Aware Vertical Federated Learning for Missing Information

Vertical Federated Learning (VFL) is a prevalent paradigm designed to facilitate collaboration between multiple entities possessing distinct feature sets yet sharing a common user base for model training. However, challenges emerge when generating predictions for users with partial data, such as the scenario where a new user registers with only one participating company and submits a limited number of features. The typical method of addressing missing data involves padding the absent values with zeros or mean figures, which can cause an out-of-distribution issue, thereby precipitating a notable deterioration in model accuracy. To address this issue, we introduce DiVerFed, a distribution-aware VFL framework for missing information. DiVerFed’s primary objective is to maintain robust model performance even when confronted with incomplete user information. Within this framework, we treat the VFL’s top model as an encoder responsible for capturing the underlying data distribution. To make this distribution effectively, we incorporate a reconstruction component. Furthermore, we simulate the scenario of missing information by integrating a Feature-wise Dropout module during the training phase. To enhance the framework’s efficacy in classification tasks, we also incorporate label information within the loss function to leverage a class-aware distribution to bolster the model’s accuracy. Our experimental analyses confirm that DiVerFed significantly outperforms conventional approaches in classification tasks when information from only one party is accessible.