Token-Level Prompt Mixture with Parameter-Free Routing for Federated Domain Generalization
Gong, Shuai Cui, Chaoran Dong, Xiaolin Nie, Xiushan Zhu, Lei Chang, Xiaojun
Gong, Shuai
Cui, Chaoran
Dong, Xiaolin
Nie, Xiushan
Zhu, Lei
Chang, Xiaojun
Supervisor
Department
Computer Vision
Embargo End Date
Type
Journal article
Date
License
Language
English
Collections
Research Projects
Organizational Units
Journal Issue
Abstract
Domain Generalization (FedDG) aims to train a globally generalizable model on data from decentralized, heterogeneous clients. While recent work has adapted vision-language models for FedDG using prompt learning, the prevailing "one-prompt-fits-all" paradigm struggles with sample diversity, causing a marked performance decline on personalized samples. The Mixture of Experts (MoE) architecture offers a promising solution for specialization. However, existing MoE-based prompt learning methods suffer from two key limitations: 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. TRIP treats prompts as multiple experts, and assigns individual tokens within an image to distinct experts, facilitating the capture of fine-grained visual patterns. To ensure communication efficiency, TRIP introduces a parameter-free routing mechanism based on capacity-aware clustering and Optimal Transport (OT). First, tokens are grouped into capacity-aware clusters to ensure balanced workloads. These clusters are then assigned to experts via OT, stabilized by mapping cluster centroids to static, non-learnable keys. The final instance-specific prompt is synthesized by aggregating experts, weighted by the number of tokens assigned to each. Extensive experiments across four benchmarks demonstrate that TRIP achieves optimal generalization results, with communicating as few as 1K parameters. Our code is available at https://github.com/GongShuai8210/TRIP.
Citation
S. Gong, C. Cui, X. Dong, X. Nie, L. Zhu, X. Chang, "Token-Level Prompt Mixture with Parameter-Free Routing for Federated Domain Generalization," IEEE Transactions on Image Processing, vol. 35, no. 99, pp. 656-669, 2026, https://doi.org/10.1109/tip.2026.3652431.
Source
IEEE Transactions on Image Processing
Conference
Keywords
46 Information and Computing Sciences, 4611 Machine Learning
Subjects
Source
Publisher
IEEE