A new paper from Vector Faculty Member Xiaoxiao Li presents a new approach combining generalized and personalized learning into an efficient system capable of handling data heterogeneity. Called shared and group prompt tuning (SGPT), the method improves performance and enhances safety, and interpretability.

“Unlocking the Potential of Prompt-Tuning in Bridging Generalized and Personalized Federated Learning,” co-authored by Wenlong Deng and Christos Thrampoulidis, showcases how this innovative approach combines the strengths of generalized learning (where an AI learns from various sources) and personalized learning (where an AI is tailored to specific users). The design allows the algorithm to capture both common and specialized features, facilitating better alignment with diverse local data distributions without requiring local fine-tuning.

Federated learning aims to train machine learning models across multiple clients without sharing their data, making it crucial in domains like computer vision. However, data heterogeneity, characterized by domain discrepancies or imbalanced class distributions, presents a significant hurdle. Traditional generalized federated learning methods, which learn a single global model, often struggle with significant data heterogeneity. Personalized federated learning methods, which tailor models to individual clients, can lead to overfitting.

Background and Motivation

Traditional FL approaches can be broadly categorized into generalized FL (GFL) and personalized FL (PFL). GFL aims to learn a single global model that generalizes well across all clients, while PFL focuses on tailoring models to individual clients or client groups. Both approaches have limitations: GFL struggles with significant data heterogeneity, while PFL may overfit to local data and fail to generalize to out-of-federation clients.

To tackle these challenges, the authors introduce SGPT, a novel algorithm that blends the advantages of both GFL and PFL. SGPT harnesses the power of vision transformers (ViTs), which, while traditionally seen as computationally intensive, have recently benefited from parameter-efficient tuning methods like prompt tuning that greatly improve their efficiency, making them well-suited for FL. By applying prompt-tuning techniques, SGPT establishes a flexible and efficient FL framework optimized for model tuning in distributed environments.

SGPT Methodology

The core idea behind SGPT is to learn both shared prompts and group-specific prompts, allowing the model to capture common features across all clients while also adapting to group-specific characteristics. Here’s a breakdown of the key components:

1. Shared prompts: These are designed to capture common representations across all clients. They are attached to the early layers of the ViT model, where features tend to be more uniform across different classes.
2. Group prompts: These prompts are designed to extract specialized information for different data groups. They are inserted into higher layers of the ViT, where features become more diverse and specialized.
3. Prompt selection module: This module uses a similarity-based clustering approach to assign data points to specific groups. It learns a set of keys for each group and selects the appropriate group prompt based on the similarity between the input features and the learned keys.
4. Block coordinate descent (BCD) optimization: To effectively train the prompts, SGPT employs a BCD approach. It first optimizes the shared prompts to learn common information, then optimizes the group prompts to extract more specialized knowledge

The authors introduce several techniques to improve the stability and effectiveness of their approach:

• Calibration of the selection function using accumulated selection probability to avoid collapse into few groups.
• Momentum parameter aggregation for both keys and group prompts to ensure selection consistency and knowledge consistency.

Theoretical Analysis

The paper provides a theoretical analysis of the gap between the global and local performance of the SGPT model. The authors identify two key factors affecting this gap:

1. Generalization: related to the number of samples in each group.
2. Distribution discrepancy: the difference between the global group distribution and the local group distribution of each client.

SGPT addresses these factors by using shared prompts in early layers to maximize the sample size for common features, and group prompts in higher layers to minimize distribution discrepancy for diverse features.

Experimental Setup and Results

The authors conducted extensive experiments on various datasets to evaluate SGPT’s performance under both label heterogeneity and feature heterogeneity conditions:

Label Heterogeneity:

• CIFAR-100: 100 clients, with each client assigned data from a specific number of classes (s).
• Five-dataset: a sequence of 5 datasets (SVHN, CIFAR10, not-MNIST, Fashion-MNIST, and MNIST) distributed across 20 clients.

Feature Heterogeneity:

• Office-Caltech10: four data domains with 10 classes each.
• DomainNet: six domains with the top ten most frequent classes.

The experiments compared SGPT against several baseline methods, including FedVPT, FedMix, pFedPG, FedEM, and FedPR. The results demonstrated that SGPT consistently outperformed these baselines across different heterogeneity levels and datasets.

Key findings include

1. SGPT achieved higher global accuracy and worst-local accuracy compared to other methods, indicating better performance on both global and local data distributions.
2. SGPT showed robustness to increasing levels of data heterogeneity, with smaller performance drops compared to other methods as heterogeneity increased.
3. In feature heterogeneity experiments, SGPT achieved the highest average accuracies on both Office-Caltech10 and DomainNet datasets.

The authors also conducted ablation studies to analyze the impact of different components of SGPT:

• The combination of shared and group prompts led to significant improvements in both global and worst-local accuracy.
• The proposed Block Coordinate Descent optimization strategy proved crucial for effective training of the prompts.
• The prompt selection module with momentum updating improved clustering performance and stability.

Conclusion and Implications

The SGPT algorithm represents a significant advancement in federated learning, effectively bridging the gap between generalized and personalized approaches. By leveraging prompt-tuning techniques and the power of vision transformers, SGPT demonstrates superior performance in handling data heterogeneity across clients.

The key innovations of SGPT – shared and group prompts, the prompt selection module, and the BCD optimization strategy – provide a flexible framework that can adapt to both global and local data distributions without requiring local fine-tuning. This approach not only improves performance but also maintains efficiency, with significantly fewer trainable parameters compared to traditional FL methods.

As federated learning continues to gain importance in privacy-preserving machine learning applications, methods like SGPT that can effectively handle heterogeneous data distributions will be crucial for real-world deployments. Future research could explore the application of similar prompt-tuning techniques to other types of models beyond vision transformers, as well as investigating the scalability and communication efficiency of such approaches in large-scale federated learning systems.