Leveraging Foundation Models for Efficient Federated Learning in Resource-Restricted Edge Networks

Generally speaking, KD refers to the method of transferring knowledge from one or multiple teacher models to a student model [22]. Specifically, the softened output of the teacher model on a public dataset, along with the ground-truth output, is used to train the student model. In this context, a KD loss function is employed in addition to the supervised learning loss to minimize the discrepancy between the soft labels of the teacher model and the predictions made by the student model. In KD assisted FL scenarios, depending on the method, both clients and the server can assume the roles of either the teacher or the student, i.e., mutual KD [23]. In such scenarios, knowledge from clients is shared with the server to create a global knowledge base. This knowledge is then distilled back to the clients to enhance their performance, allowing for the transfer of knowledge from other clients to each individual client. Typically, KD methods [24] utilize a public dataset shared among all entities (also referred to as the transfer set) to align the extracted knowledge of local models and the server’s one. This assumption of a publicly available dataset, however, is unrealistic in practice, as it can be accessed by third parties and raise privacy concerns. In this paper, KD is employed in a data-free fashion (without reliance of a transfer set), and to establish a framework for knowledge exchange between local models and the FM of the server.

Fine-tuning FMs for downstream tasks has shifted the dominant paradigm of machine learning from “training from scratch” to the “pretrain-then-finetune” framework. Fully fine tuning of a large FM, however, involves updating all its parameters, which increases the risk of overfitting. This challenge has driven the development of partial fine-tuning methods [10]. With the advent of LLMs, a novel capability known as prompting [25] has been introduced for such models. This technique involves prepending learnable parameters to the embedding space of the input—whether it be the embedding space of tokens in LLMs or the embedding space of image patches in Transformer-based vision models. This approach provides pre-trained FMs with hints about downstream tasks while keeping their parameters frozen. For further details please refer to Reference [26].

We consider an edge environment consisting of $N$ IoT devices, denoted by $\mathbb{U}=\{u_{1},\cdot\cdot\cdot,u_{N}\}$, which are coordinated by an edge server. Each IoT device $u_{n}$ for ($1\leq n\leq N$), aims to perform a $C$-class classification task with the assistance of an FM located at the edge server. Each IoT device $u_{n}\in\mathbb{U}$ possesses a local dataset, represented as $\mathbb{D}^{n}$. Local datasets are distributed heterogeneously among IoT devices and collectively form the entire dataset $\mathbb{D}=\{\mathbb{D}^{1},\cdot\cdot\cdot,\mathbb{D}^{N}\}$. Each IoT device, based on its available storage and computation resources, employs a lightweight local model, denoted by $f^{n}(\cdot;\bm{\theta}^{n})$ parameterized by $\bm{\theta}^{n}$. This enables the framework to handle model heterogeneity. To develop a FL framework that trains local models with the support of an FM at the server, we solve the following personalized FL problem

$$\begin{array}[]{rrclcl}\operatorname*{argmin}_{\{\bm{\theta}^{1},\cdot\cdot\cdot,\bm{\theta}^{n}\}}\mathbb{E}_{\mathbb{D}^{n}\in\mathbb{D}}\{\mathcal{J}^{n}(\bm{\theta}^{n},\mathcal{D}^{n})\},\end{array}$$

(1)

where $\mathcal{J}^{n}(\cdot)$ is the loss function of local model $f^{n}(\cdot,\bm{\theta}^{n})$ over local dataset $\mathbb{D}^{n}$.

The proposed FedD2P framework includes an initialization stage followed by iterating four steps, summarized as follows:

To construct the LA prompt generator and without loss of generality, we assume that the linguistic description of classes within the downstream task are available at the server and are represented by $\{s_{1},\ldots,s_{C}\}$. Their corresponding semantic representation vectors are extracted by the text encoder of the vision-language FM as $\bm{e}_{c}=F_{text}(s_{c})\in\mathbb{R}^{d}$, where $F_{text}(\cdot)$ denotes the text encoder of the FM, and $d$ represents its embedding dimension. The prompt generator $G(\cdot;\bm{\omega})$, then uses these semantic representation of linguistic description of classes to generate class-specific prompts. To account for the correlation among the semantic representations of classes, we employ a multi-head self-attention mechanism to construct the LA prompt generator. The prompt for class $c$ is, therefore, calculated as

$$\bm{h}_{c}=\text{Softmax}(\frac{\bm{q}_{c}\bm{\mathcal{K}}}{\sqrt{d}})\bm{\mathcal{V}}\bm{\mathcal{W}}_{h},$$

(4)

where $\bm{q}_{c}=\bm{e}_{c}\bm{\mathcal{W}}_{q},\bm{\mathcal{K}}=\bm{E}\bm{\mathcal{W}}_{k}$, and $\bm{\mathcal{V}}=\bm{E}\bm{\mathcal{W}}_{v}$ represent the query vector, and the key and value weight matrices, respectively. Here, $\bm{E}=[\bm{e}_{c}]_{c=1}^{C}$, and $\mathcal{W}_{h}$ denotes the parameters of the head layer. Accordingly, the class-specific prompts is defined as $\bm{P}_{c}=\{s_{c},\bm{h}_{1},\cdot\cdot\cdot,\bm{h}_{C}\}$, which is forwarded to the vison-language FM to generate the global knowledge as follows

$$\bm{g}_{c}=\frac{\text{exp}\big{(}\text{cos}(F_{image}(\bm{h}_{c}),F_{text}(s_{c}))\big{)}}{\sum_{c=1}^{C}\text{exp}\big{(}\text{cos}(F_{image}(\bm{h}_{c}),F_{text}(s_{c}))\big{)}},$$

(5)

where $F_{image}(\cdot)$ denotes the image encoder of the FM.

To distill the per-class aggregated knowledge $\bm{a}_{c}$ into the LA prompt generator, the prompt-tuning is performed as follows

$$\operatorname*{argmin}_{\bm{\omega}}\sum_{c=1}^{C}\mathcal{L}_{c}(\bm{g}_{c},y_{c})+\mathcal{L}_{k}(\bm{g}_{c},\bm{a}_{c}),$$

(6)

where $y_{c}$ represents the ground-truth corresponding to the class $c$. This completes presentation of the proposed FedD2P framework, next, evaluation experiments are presented.

Table 1 presents the performance comparisons of FedD2P with the baselines across two statistical distribution settings. The results indicate that leveraging an FM at the server is more effective than relying solely on the aggregated knowledge from edge devices. We attribute this consistently superior performance to the robust representations provided by the CLIP model and the effective distillation of global knowledge to IoT devices. It can also be inferred that distilling the distributed knowledge of IoT devices to the LA prompt generator fine-tunes the FM for the downstream task more effectively. The most significant performance gain compared with other baselines is observed for the OxfordPets dataset. This is likely due to the dataset’s challenging nature for training a local CNN model from scratch, whereas according to  [12] the few-shot accuracy of the CLIP model on this dataset is superior compared with other datasets.

The effectiveness of the mutual KD framework is significantly determined by the entropy of soft labels. Low entropy for the local models and high entropy for the CLIP model necessitate adjusting the temperature parameters when computing local and global soft labels. Fig. 2(a) illustrates the performance of the FedD2S framework over CIFAR10 dataset and across different temperature parameters. As shown, the most effective performance is achieved by decreasing the entropy of global soft labels and increasing the entropy of local ones. This can be attributed to the high generalization ability of the robust CLIP model and the low generalization capability of simple local models.

For comparison, we employ a Multi-Layer Perceptron (MLP)-based prompt generator, wherein each semantic representation vector, $\bm{e}_{c}$, is directly mapped to $\bm{h}_{c}$. Fig. 2(b) illustrates the impact of the self-attention mechanism within the LA prompt generator module on the CIFAR10 and EuroSat datasets. The superior performance of the multi-head self-attention mechanism is attributed to its ability to consider the relationships among different semantic representations of classes. This capability enables the generation of effective prompts that extract soft labels with sufficient generalization content, thereby assisting local models.