TLDR: Proto-EVFL is a novel Vertical Federated Learning (VFL) framework that addresses the challenges of class imbalance and unaligned data. It introduces a probabilistic dual prototype learning scheme for unbiased feature representation, a mixed prior guided module to mitigate local model bias, and an adaptive gated feature aggregation strategy for consistent feature contributions. Experiments show Proto-EVFL outperforms existing methods in accuracy, especially in zero-shot scenarios, and offers improved communication efficiency and privacy protection by avoiding gradient transmission.
In the evolving landscape of artificial intelligence, Vertical Federated Learning (VFL) stands out as a crucial paradigm for collaborative model training while preserving data privacy. VFL allows multiple organizations to build a shared model using their distinct feature sets for the same samples, without directly sharing raw data. However, a significant challenge in VFL is the scarcity of aligned samples – data points that are common across all participating parties. To overcome this, VFL often relies on massive amounts of locally unaligned and unlabeled data, which introduces its own set of problems, particularly class imbalance.
A new research paper, Proto-EVFL: Enhanced Vertical Federated Learning via Dual Prototype with Extremely Unaligned Data, introduces a novel framework designed to tackle these complex class imbalance issues. The authors, Wei Guo, Yiyang Duan, Zhaojun Hu, Yiqi Tong, Fuzhen Zhuang, Xiao Zhang, Jin Dong, Ruofan Wu, Tengfei Liu, and Yifan Sun, highlight that class imbalance in VFL manifests in two ways: intra-party class imbalance (imbalance within each party’s local dataset) and inter-party class imbalance (inconsistent class distributions across different parties).
These imbalances can lead to significant drawbacks, such as biased local models, inconsistent feature contributions when combining data from different parties, and a limited ability for the model to predict rare or previously unseen classes. Traditional VFL methods often overlook these issues, resulting in suboptimal performance, especially in scenarios where data is extremely unaligned or certain classes are barely represented.
Introducing Proto-EVFL’s Innovative Approach
Proto-EVFL, short for Prototype-Enhanced Vertical Federated Learning, addresses these challenges through a three-pronged approach:
First, it introduces a **Probabilistic Dual Prototype Learning Scheme**. This scheme uses ‘class prototypes’ for each party, which are essentially representative points for each class in the model’s internal feature space. By learning the relationships between these prototypes and the unaligned data, the system can dynamically select the most relevant unaligned samples for training. This process is guided by a novel dual optimal transport cost, which measures the distance between samples and prototypes in both directions, ensuring that no class patterns are overlooked and improving the confidence of pseudo-labels for unlabeled data.
Second, a **Mixed Prior Guided Module** is designed to mitigate the bias caused by intra-party class imbalance. This module intelligently combines local and global class prior probabilities. This means that while each party considers its own data distribution, it also incorporates a broader, global understanding of class frequencies, preventing local models from overfitting to their majority classes and improving their ability to learn from less frequent ones.
Third, Proto-EVFL employs an **Adaptive Gated Feature Aggregation Strategy**. In VFL, features from different parties are combined by the active party (the one with labels) to make predictions. However, due to varying class imbalances, the contribution of features from each party can be inconsistent. This strategy uses a ‘gating network’ to dynamically weigh and aggregate the features from each party, ensuring that the combined features are optimally balanced for the final prediction task, thereby improving the overall generalization of the model.
Also Read:
- Enhancing Federated Learning with DAG-AFL: A New Approach for Asynchronous and Heterogeneous Environments
- H2Tune: Adapting Foundation Models in Diverse Federated Learning Environments
Performance and Privacy Advantages
The researchers conducted extensive experiments on various datasets, including image datasets like ModelNet-10 and Fashion-MNIST, and tabular datasets such as Credit and Adult. Proto-EVFL consistently outperformed existing VFL frameworks, especially in challenging scenarios with very few or even zero examples of certain classes (few-shot and zero-shot learning). For instance, in zero-shot scenarios, Proto-EVFL showed at least a 6.97% improvement over baselines.
Beyond accuracy, Proto-EVFL also demonstrates superior communication efficiency, requiring fewer communication rounds and lower overall communication costs to achieve optimal performance compared to other methods. Its training time is also significantly reduced, making it a more practical solution for real-world applications.
A critical aspect of federated learning is privacy. Proto-EVFL enhances privacy by avoiding the direct transmission of raw gradient information, a common vulnerability in other VFL methods. Instead, it relies on the exchange of class prototypes and intermediate representations. Experiments showed that Proto-EVFL significantly reduces the risk of label inference attacks, with an attack accuracy of only 8.89% compared to over 56% for other methods. While adding differential privacy noise can further protect intermediate representations, the study notes a trade-off with model performance, suggesting careful calibration of noise levels.
In conclusion, Proto-EVFL represents a significant advancement in Vertical Federated Learning, offering a robust and efficient solution to the pervasive problem of class imbalance in unaligned data. Its innovative use of dual prototypes, mixed prior guidance, and adaptive feature aggregation paves the way for more accurate, private, and scalable collaborative AI models.
