Prompt Estimation from Prototypes for Federated Prompt Tuning of Vision Transformers

TL;DR

This paper introduces PEP-FedPT, a federated prompt tuning framework for Vision Transformers that balances personalization and generalization using class-specific prompts and prototypes, improving performance across heterogeneous clients.

Contribution

It proposes a novel class-contextualized prompt method and a federated training approach that enhances both personalization and generalization in federated vision transformer tuning.

Findings

PEP-FedPT outperforms state-of-the-art methods on multiple datasets.

The framework effectively handles data heterogeneity in federated settings.

Class-specific prompts improve personalization without increasing client-side parameters.

Abstract

Visual Prompt Tuning (VPT) of pre-trained Vision Transformers (ViTs) has proven highly effective as a parameter-efficient fine-tuning technique for adapting large models to downstream tasks with limited data. Its parameter efficiency makes it particularly suitable for Federated Learning (FL), where both communication and computation budgets are often constrained. However, global prompt tuning struggles to generalize across heterogeneous clients, while personalized tuning overfits to local data and lacks generalization. We propose PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning), a unified framework designed to achieve both generalization and personalization in federated prompt tuning of ViTs. Within this framework, we introduce the novel Class-Contextualized Mixed Prompt (CCMP) - based on class-specific prompts maintained alongside a globally shared prompt. For each input, CCMP adaptively combines class-specific prompts using weights derived from global class prototypes and client class priors. This approach enables per-sample prompt personalization without storing client-dependent trainable parameters. The prompts are collaboratively optimized via traditional federated averaging technique on the same. Comprehensive evaluations on CIFAR-100, TinyImageNet, DomainNet, and iNaturalist datasets demonstrate that PEP-FedPT consistently surpasses the state-of-the-art baselines under diverse data heterogeneity scenarios, establishing a strong foundation for efficient and generalizable federated prompt tuning of Vision Transformers.