🤖 machine learning

Prompt Estimation from Prototypes for Federated Prompt Tuning of Vision Transformers

The paper proposes PEP-FedPT, a unified federated learning framework for Vision Transformers that achieves both generalization and personalization by introducing a Class-Contextualized Mixed Prompt (CCMP) which adaptively combines class-specific prompts using global prototypes and client priors without storing client-dependent trainable parameters.

Original authors: M Yashwanth, Sharannya Ghosh, Aditay Tripathi, Anirban Chakraborty

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: M Yashwanth, Sharannya Ghosh, Aditay Tripathi, Anirban Chakraborty

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, super-smart library of books (a pre-trained Vision Transformer) that knows a little bit about everything but isn't an expert in any one specific topic. You want to teach this library to become an expert in a specific subject, like identifying rare birds or spotting diseases in X-rays, but you can't move all the books to one place. Instead, you have hundreds of small, local branches (clients) scattered around the world, each holding a different, unique collection of books.

This is the challenge of Federated Learning: training a smart global model without ever moving the private data from the local branches.

The Problem: The "One-Size-Fits-All" vs. "Too Personal" Dilemma

The paper identifies a frustrating tug-of-war in trying to teach this library:

The Global Approach (Too Rigid): If you give every branch the exact same set of instructions (a "global prompt"), the library works well for the average branch but fails miserably for the ones with weird or unique data. It's like giving a generic "How to Cook" manual to a branch that only has seafood; the instructions don't fit their specific ingredients.
The Personal Approach (Too Narrow): If you let every branch write its own custom instructions, they become experts at their specific local data but forget how to talk to the rest of the network. They become so specialized they can't help anyone else, and if a new branch joins, they have no idea what to do.

The Solution: PEP-FedPT (The "Smart Mixologist")

The authors propose a new method called PEP-FedPT. Think of this as a "Smart Mixologist" who creates a custom drink for every branch, but without needing to store a secret recipe book at every branch.

Here is how it works, using a simple analogy:

1. The Shared Ingredients (Global Prompts)

The central server (the library head) sends out a set of Shared Prompts. These are like basic, universal ingredients (salt, pepper, water) that everyone agrees on. They help the library understand the basics.

2. The Special Flavors (Class-Specific Prompts)

The server also maintains a set of Class-Specific Prompts. Imagine these as distinct flavor concentrates: one for "Birds," one for "Cars," one for "Trees." These are shared globally, so everyone has access to the same flavor bottles.

3. The Magic Mixing (CCMP)

This is the paper's big innovation. Instead of just handing a branch a single flavor or a generic mix, the system creates a Class-Contextualized Mixed Prompt (CCMP) for every single image the branch sees.

How does it decide the mix?

The "Smell" Test (Prototypes): The system looks at the image and asks, "Does this look more like a bird or a car?" It uses a "global map" (class prototypes) to estimate the likelihood.
The "Local Menu" (Class Priors): It also checks the branch's local menu. If a branch mostly sees birds, the system knows to lean heavier on the "Bird" flavor.

The system mixes the "Bird" and "Car" flavors together in a precise ratio based on these two factors. It's like a mixologist saying, "This image looks 80% like a bird and 20% like a car, so I'll mix 80% bird flavor and 20% car flavor for this specific drink."

Why This is a Game-Changer

No Secret Recipe Books: The branches don't need to store their own unique, heavy instruction manuals. They just use the shared global flavors and mix them on the fly. This saves massive amounts of memory and communication space.
Best of Both Worlds: The resulting "drink" is personalized enough to handle the branch's weird data (generalization) but still connected to the global network (personalization).
Privacy Friendly: The branches only send back tiny summaries (prototypes) of what they've learned, not the actual images. It's like sending a description of the flavors used, not the actual food.

The Results

The paper tested this on several "hard" datasets (like CIFAR-100 and TinyImageNet) where data is messy and unevenly distributed.

The Winner: PEP-FedPT consistently beat all other methods.
The Proof: It didn't just get a higher average score; it also helped the worst-performing branches improve significantly. It managed to be both a great local expert and a helpful global neighbor simultaneously.

In short, the paper presents a way to train a giant AI model across many different, messy data sources by letting the model "mix and match" its own instructions on the fly, using a shared set of tools and a little bit of local context. It achieves the perfect balance between being a generalist and a specialist.

Technical Summary: PEP-FedPT

Problem Statement

Federated Learning (FL) faces significant challenges when adapting large pre-trained Vision Transformers (ViTs) to downstream tasks under data heterogeneity. While Visual Prompt Tuning (VPT) offers a parameter-efficient solution for fine-tuning, existing approaches struggle to balance generalization and personalization:

Global Prompt Tuning: Methods like FedPR use shared prompts across all clients. While they generalize well, they fail to adapt to highly heterogeneous client data distributions (non-iid), leading to "client drift" and suboptimal performance.
Personalized Prompt Tuning: Methods like pFedPG generate client-specific prompts. While they adapt well to local data, they often overfit, resulting in poor generalization to unseen or non-participating clients. Furthermore, many personalized approaches assume full client participation in every round, which is impractical in real-world FL systems.
Existing Hybrid Approaches: Methods like SGPT attempt to bridge this gap using group-specific prompts but introduce optimization complexity through non-differentiable selection mechanisms and two-stage training processes.

The central question addressed is: Can we achieve effective personalization while relying solely on globally shared prompts?

Methodology: PEP-FedPT

The authors propose PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning), a unified framework that achieves personalization without storing client-dependent trainable parameters. The core innovation is the Class-Contextualized Mixed Prompt (CCMP).

1. Prompt Architecture

The framework utilizes two sets of trainable parameters shared globally across all clients:

Shared Prompts ( $P_S$ ): Inserted at the very first layer of the ViT. These capture low-level, class-agnostic representations common across the manifold.
Class-Specific Prompts ( $P_C$ ): A set of prompts corresponding to each class, shared globally.

2. Class-Contextualized Mixed Prompt (CCMP)

Instead of inserting a single static prompt, PEP-FedPT dynamically generates a mixed prompt ( $m$ ) for each input sample at intermediate layers (specifically layers 5, 6, and 7 in the experiments). This mixed prompt is a soft combination of the global class-specific prompts, weighted by scores derived from:

Global Class Prototypes ( $\mu$ ): Aggregated centroids of CLS-token representations for each class, computed by averaging CLS tokens of data points within a class across clients. These are updated periodically via federated averaging with momentum.
Client Class Priors ( $\delta$ ): The empirical label distribution (frequency of classes) specific to each client's local data.

3. Prompt Mixing Mechanism

For a given input $x$ at client $k$ , the system computes a soft weight vector $s$ for the class-specific prompts:

Similarity Score: The cosine similarity between the input's CLS token representation ( $cls_{l-1}$ ) at layer $l$ and the global class prototypes ( $\mu$ ) is calculated.
Weighted Scoring: These similarities are scaled by the client's local class priors ( $\delta$ ) and a temperature parameter ( $\tau$ ).
Normalization: The scores are normalized to form a probability distribution.
Mixing: The final CCMP is computed as the weighted sum of the global class-specific prompts: $m = P_C \times s$ .

This allows each client to dynamically personalize the prompt based on their local data distribution and the current input, using only global information (prompts and prototypes) and local statistics (priors), without transmitting raw data or client-specific model weights.

4. Theoretical Foundations

The paper provides two theoretical justifications for the CCMP design:

Quadratic Upper Bound Minimization: The authors prove that the CCMP minimizes a quadratic upper bound on the global loss function, assuming the loss is Lipschitz smooth and the prompt space is compact. The optimal mixing strategy aligns with the weighted average of class prompts based on class priors.
MMSE Optimality: The CCMP is shown to be the Minimum Mean Squared Error (MMSE) estimator of the true class prompt given the observed CLS token representation.

Key Contributions

Unified Framework: Introduction of PEP-FedPT, which jointly optimizes shared and class-specific prompts to address data heterogeneity in ViT federated learning, striking a balance between generalizability and personalization.
Novel Mixing Strategy: Design of the CCMP, which generates per-sample prompts as a function of global class-specific prompts, modulated by local class priors and global prototypes. This eliminates the need for local prompt storage or server-side generation of client-specific prompts.
Theoretical Insights: Demonstration that CCMP minimizes a quadratic upper bound on the loss and is optimal in the MMSE sense.
Empirical Superiority: Extensive validation showing the method outperforms state-of-the-art baselines across diverse heterogeneity scenarios.

Experimental Results

The method was evaluated on CIFAR-100, TinyImageNet, DomainNet, and iNaturalist under various non-iid settings (pathological splits and Dirichlet distributions).

Performance: PEP-FedPT consistently outperformed baselines (including FedVPT, FedPR, SGPT, and pFedPG) in both Mean Accuracy and Worst Local Accuracy.
- On CIFAR-100 (Pathological), it achieved 95.46% mean accuracy, surpassing pFedPG (92.96%) and FedVPT-D (85.15%).
- On TinyImageNet, it improved mean accuracy by 8.59% over pFedPG.
Generalization: In held-out evaluations (testing on clients not participating in training), PEP-FedPT maintained high accuracy (93.71% on CIFAR-100), whereas personalized methods like pFedPG failed to generalize (marked as NA or significantly lower).
Efficiency: The method required significantly fewer communication rounds (12 rounds to reach 83% accuracy on CIFAR-100) and lower communication overhead (4.6M parameters communicated) compared to SGPT (13M) and FedVPT-D.
Ablation Studies: Experiments confirmed that combining Shared Prompts with CCMP yields significant gains over using Shared Prompts alone. The inclusion of Class Priors further improved performance, particularly in pathological settings.

Significance and Claims

The paper claims that PEP-FedPT establishes a strong foundation for efficient and generalizable federated prompt tuning of Vision Transformers. Its primary significance lies in resolving the trade-off between personalization and generalization in FL:

It achieves personalization by leveraging local data distributions (via priors) to dynamically mix prompts.
It maintains generalization by relying exclusively on globally shared prompts and prototypes, avoiding the overfitting issues of fully personalized models.
It offers scalability by not requiring client-specific parameter storage or full client participation, making it suitable for resource-constrained edge devices.

The authors note limitations, specifically that the method relies on empirical estimates of class priors and centroids, which may be unreliable in semi-supervised or unsupervised settings, and that the advantage diminishes in strictly IID data distributions where the method effectively reduces to standard Fed-VPT. However, within the context of non-IID federated learning, the proposed approach offers a robust, theoretically grounded, and empirically superior solution.