Token-Level Prompt Mixture with Parameter-Free Routing for Federated Domain Generalization

TL;DR

TRIP is proposed, a Token-level pRompt mIxture with Parameter-free routing framework for FedDG that treats prompts as multiple experts, and assigns individual tokens within an image to distinct experts, facilitating the capture of fine-grained visual patterns.

Abstract

Federated domain generalization (FedDG) aims to learn a globally generalizable model from decentralized clients with heterogeneous data while preserving privacy. Recent studies have introduced prompt learning to adapt vision-language models (VLMs) in FedDG by learning a single global prompt. However, such a one-prompt-fits-all learning paradigm typically leads to performance degradation on personalized samples. Although the mixture of experts (MoE) offers a promising solution for specialization, existing MoE-based methods suffer from coarse image-level expert assignment and high communication costs from parameterized routers. To address these limitations, we propose TRIP, a Token-level prompt mixture with parameter-free routing framework for FedDG, which treats multiple prompts as distinct experts. Unlike existing image-level routing designs, TRIP assigns different tokens within an image to specific experts. To ensure communication efficiency, TRIP incorporates a parameter-free routing mechanism based on token clustering and optimal transport. The instance-specific prompt is then synthesized by aggregating experts, weighted by the number of tokens assigned to each. Additionally, TRIP develops an unbiased learning strategy for prompt experts, leveraging the VLM's zero-shot generalization capability. Extensive experiments across four benchmarks demonstrate that TRIP achieves optimal generalization results, with communication of only 1K parameters per round. Our code is available at https://github.com/GongShuai8210/TRIP.

Figures (8)

• Figure 1: Comparison of our method, TRIP, with prior prompt-based methods. (a) PromptFL guo2023promptfl aims to learn a shared global prompt applied to all samples, which limits its generalization ability. (b) pFedMoAP luo2024mixture trains a parameterized routing network to assign personalized prompts for each sample, but this introduces additional parameters that must be communicated between clients and the server. (c) Our method constructs a parameter-free router to assign image tokens to prompt experts, significantly reducing communication overhead while improving generalization performance.
• Figure 2: Conceptual framework of TRIP for FedDG. TRIP consists of two essential components: (a) a parameter-free routing mechanism and (b) unbiased prompt expert learning. In (a), an image is first divided into patches and encoded into token embeddings. These embeddings are then grouped into clusters based on capacity-aware clustering. Next, clusters are assigned to experts using OT by minimizing the total cost in a cost matrix constructed from the distances between cluster centroids and static keys. In (b), the local prediction distribution generated by local prompt experts are aligned with CLIP’s zero-shot inference distribution to mitigate potential bias toward localized data patterns.
• Figure 3: Comparison of learnable (L) and our parameter-free (P) image-level and token-level router designs across different datasets
• Figure 4: Visualization of model inference stability under different random seeds on Office-Home and VLCS.
• Figure 5: Effect of capacity factor $\alpha$ during training and inference.