Federated Domain Generalization via Prompt Learning and Aggregation

TL;DR

This paper proposes a novel FedDG framework through Prompt Learning and AggregatioN (PLAN), which comprises two training stages to collaboratively generate local prompts and global prompts at each federated round, and introduces prompt learning to adapt pre-trained vision-language models in the FedDG scenario.

Abstract

Federated domain generalization (FedDG) aims to improve the global model generalization in unseen domains by addressing data heterogeneity under privacy-preserving constraints. A common strategy in existing FedDG studies involves sharing domain-specific knowledge among clients, such as spectrum information, class prototypes, and data styles. However, this knowledge is extracted directly from local client samples, and sharing such sensitive information poses a potential risk of data leakage, which might not fully meet the requirements of FedDG. In this paper, we introduce prompt learning to adapt pre-trained vision-language models (VLMs) in the FedDG scenario, and leverage locally learned prompts as a more secure bridge to facilitate knowledge transfer among clients. Specifically, we propose a novel FedDG framework through Prompt Learning and AggregatioN (PLAN), which comprises two training stages to collaboratively generate local prompts and global prompts at each federated round. First, each client performs both text and visual prompt learning using their own data, with local prompts indirectly synchronized by regarding the global prompts as a common reference. Second, all domain-specific local prompts are exchanged among clients and selectively aggregated into the global prompts using lightweight attention-based aggregators. The global prompts are finally applied to adapt VLMs to unseen target domains. As our PLAN framework requires training only a limited number of prompts and lightweight aggregators, it offers notable advantages in computational and communication efficiency for FedDG. Extensive experiments demonstrate the superior generalization ability of PLAN across four benchmark datasets.

Figures (7)

• Figure 1: The novel problem setting of FedDG aims to learn a global model from multiple decentralized source domains, enabling it to directly generalize to completely unseen target domains.
• Figure 2: Overall framework of PLAN for FedDG. PLAN consists of two training stages at each federated round. In stage (a), the server sends the global text and visual prompts, denoted as $\bm{T}^g$ and $\bm{V}^g$, to each client. Client $k$ then learns its own text and visual prompts, $\bm{T}^k$ and $\bm{V}^k$, using local data and aligns these learned prompts with the global prompts. The learned prompts are subsequently uploaded to the server. In stage (b), the server distributes the global aggregators, $\mathcal{A}_t$ and $\mathcal{A}_v$, along with all clients' local prompts to each client. Each client then learns to selectively aggregate local prompts from all clients into the global prompts, and the learned aggregators, $\mathcal{A}_t^k$ and $\mathcal{A}_v^k$, are uploaded back to the server.
• Figure 3: Effects of prompt depth and prompt length.
• Figure 4: Comparison of computation and communication costs of PLAN and FedDG methods. Each round includes two communications between clients and the server.
• Figure 5: Performance comparison in few-shot settings on PACS and OfficeHome. The values on the x-axis represents the number of shots.