PromptFL: Let Federated Participants Cooperatively Learn Prompts Instead
of Models — Federated Learning in Age of Foundation Model

Contrastive Language-Image Pre-Training (CLIP) is a neural network trained on hundreds of millions of (image, caption) pairs Radford et al. (2021). CLIP encodes images and captions separately as vectors, enabling users with visual modality samples to retrieve, score, or classify samples from textual modalities. Models are often very fragile and only know very specific things you trained them to do. CLIP extends the knowledge of classification models to a wider range of things by leveraging semantic information in text. Standard classification models completely discard the semantic meaning of class labels and simply enumerate numeric classes behind the scenes; CLIP works by understanding the meaning of the classes. ALIGN is another CLIP-like vision-language pre-training Jia et al. (2021).

CLIP pre-trains an image encoder and a text encoder to predict which images are paired with which texts. We can use this behavior to convert the CLIP to an image classifier. We may convert all [class] to captions such as “picture of [class]” and predict the caption class CLIP estimates the best pairing with the given image. In many previous works, this has involved prompt template engineering, in which human engineers or algorithms search for the best template for each class Fürst et al. (2021); Li et al. (2021); Singh et al. (2022); Yuan et al. (2021).

Given the unstable client availability, for the round $t$ of federated learning, the federated server samples a small subset of $m$ clients meeting eligibility requirements out of all $n$ clients to participate in the learning.

Upon notification of being selected at the round $t$, each selected client downloads the current parameters $\theta$ of global model and a training program from the federated server. Each selected client locally computes an update to the global model on its local training data by executing the training program. More specifically, the gradients updated at one client (denoted as $G$), are computed by $\frac{\partial \ell (X,y,\theta )}{\partial \theta }$, where $X$, $y$ denote the batches of training data and corresponding labels, and $\ell (\cdot )$ refers to the loss function.

The gradients $G$ in typical federated learning settings are the minimum that must be shared to the server, corresponding to FedSGD method. In FedAVG McMahan et al. (2017), models are consecutively updated on more batches of local data, which can be several epochs of training, and then shared. We note that a common way is to share the updated model $\theta +G$, but this practically amounts to sharing $G$ since all participants know $\theta$.

Upon having received local updates from $m$ clients, the federated server aggregates these updates and update its global model, and initiates next round learning. In addition to the federated learning framework that relies on the centralized server node, there are also some federated learning implementations based on the decentralized framework Roy et al. (2019); Lalitha et al. (2018); Hu et al. (2019). This means that the aggregation of gradients does not necessarily occur in a fixed federation server, but may also occur in some clients.

Given a pre-trained CLIP backbone, the input to the text encoder is designed in the form of [prompt vectors][class]. Inspired by the simple and straightforward prompt design in Zhou et al. (2021), we introduce a set of $p$ continuous embeddings of dimension $d$ in the [prompt vectors]. $d$ is same as the dimension of word embeddings in the text encoder, thus 512 by default. $p$ is a hyperparameter specifying the number of embeddings. In a word, [prompt vectors] are $p$ learnable $d$-dimensional vectors.

Given a batch of image-text pairs, CLIP will maximize the cosine similarity for matched pairs while minimize the cosine similarity for all other unmatched pairs. Since CLIP is pre-trained to predict whether an image matches a textual description, it can compute the classification loss and logits by aligning the two embedding spaces for images and texts (i.e., [prompt vectors][class]) respectively. Formally, let $g(\cdot )$ and $h(\cdot )$ be the feature extraction function of the image and text encoder. Let $w_i=h(\text{P},{\text{K}}_i)$ be the weight vector generated by the text encoder, where $i\in [1,k]$. $k$ denotes the number of classes and each $(\text{P},{\text{K}}_i)$ is derived from the prompt in the form of [prompt vectors][class]${}_i$, where [class]${}_i$ is replaced by the word embedding vector of specific class label name. Let $\cos [\cdot |\cdot ]$ denote the cosine similarity used by CLIP. By forwarding a $(\text{P},{\text{K}}_i)$ and an image x, the classification prediction probability and logits are computed as

where P is the only part that is updated in local back propagation and aggregated in the federated server.

Prompting are particularly useful in the FL case, as using prompts to push the model in the correct direction is particularly effective. This feature enables prompting to converge quickly in FL, requires less data per user, and is less affected by adverse factors in the process, e.g., non-IID and unbalanced data, limited communication bandwidth, and unreliable and limited device availability. In this paper, the prompt learner employed in PromptFL though simple and straightforward as a bridge to our core idea is easy to follow. We also envision that more complex and effective bridges would be there to replace the role and should be a valuable direction.

The average download speed within the globe for mobile internet was 54 Mbps, and the average upload speed for mobile internet was 12 Mbps that reported by 2021 O’Dea (2021). PromptFL requires locally downloading while FL requires communicating the model repeatedly between users and the federated server. Thus, the communication cost in terms of file transfer volume is that it takes only 1.4 minutes to transfer 600MB for PromptFL, and 9 hours for FL to transfer 40GB.

FL requires FLOPs computed by (2$\times$3$\times$model parameters$\times$local training epoch$\times$local training batch size$\times$input sequence length) for training, while the training FLOPs of PromptFL is much smaller and negligible compared to FL. For both PromptFL and FL, inference requires FLOPs computed by (2$\times$model parameters$\times$input sequence length), in the setting where the key and value vectors for attention computation are cached. Compared to the acceptable computational and storage costs, the RAM on the modern cell phones is a key bottleneck. We believe that this bottleneck will no longer be a problem in the near future as the techniques evolve: (1) Out-of-the-box offloading inference Rajbhandari et al. (2021). (2) Trends for more RAM Patterson (2022) and tiny CLIPs Sisodia (2021). (3) Inference with quantization methods Gholami et al. (2021).

Apart from image classification, many different vision tasks are compatible with PromptFL, such as object detection Gu et al. (2021), video understanding Xu et al. (2021) and visual question answering Shen et al. (2021). This means that the system cost of PromptFL is shared by many tasks. The prompt learner incurs these costs per personal task specific user subset requires. PromptFL is thus competitive in terms of economics.

While pre-trained CLIPs are available for download at the time of writing this paper, high-performance models in these domains are often costly to train. For example, the CLIP model trained on 400 million labeled images. The training process took 30 days across 592 V100 GPUs Radford et al. (2021). This would have cost million dollars to train on AWS on-demand instances. The value of these models and their exposure over publicly-accessible APIs make us rethink the framework of PromptFL. As illustrated in Figure 1, we hypothesize that the model APIs typically return low-dimensional outputs like confidence scores or logits, so information leakage is significantly reduced Dziedzic et al. (2022). In such a case, the prompt learner can still be trained normally, because the CLIP backbone is kept frozen during the training process. The difference is that users need to make queries to the model APIs with their private images. Some lightweight secure inference techniques like Liu et al. (2020) can be used in the framework to protect privacy.

We select a representative collection of recognition datasets used in CLIP as our benchmarks. General Objects: Caltech101 Fei-Fei et al. (2004) for general object detection. Fine-grained Categories: Flowers102 Nilsback and Zisserman (2008), OxfordPets Parkhi et al. (2012) and Food101 Bossard et al. (2014) for fine-grained classification from diversified categories.

As compared to our proposed PromptFL, we choose current representative framework in FL, FedAVG, by updating and averaging the model weights collaboratively among server and clients. We compare both training from the scratch and fine-tuning with pretrained models as our baseline method. We select the most prevailing models, Vit b16 and Retnet50, as our backbone in both our image encoder of PromptFL and the corresponding backbone in the baseline method.

How does the prompting differ from the existing adaptation method in FL? Currently in vision, the standard adaptation method is fine-tuning. Therefore we consider fine-tuning as the de-facto way of adapting visual pre-trained models in FL. Fine-tuning is highly flexible in its usage: it can adapt the pre-trained models to new input domains or new tasks with different output semantics. Yet it also requires some level of access to the pre-trained models: often entire parameters. Unlike fine-tuning, prompting adapts the inputs to a pre-trained model by modifying the model’s inputs. This opens up unique applications: the input-space adaptation puts control in the hands of the FL user; FL users only need to find the prompts, they don’t need to control the pre-trained model itself while training and testing. In this way, FL users can provide adapted images and prompts to an online API that can only operate on their inputs. On the other hand, fine-tuning is typically conditioned on inputs. Its update also directly contains some embeddings of visual feature information. In contrast, the prompts we explore in this paper are input-agnostic across the training data. So the prompting can prevent leaking of user’s private information from FL update to a certain extent.

For CLIP, an image-language model, PromptFL organizes users to collaboratively learn prompts as the CLIP’s output transformation function. Given a frozen pre-trained CLIP $F$ and a task dataset $D\left\{\right({\text{x}}_m,{\text{y}}_m\left)\right\}$ across clients, the target of PromptFL is to learn a single, static, task-specific prompting $f_{\text{prompt}}$ on class space parameterized by [prompt vectors]. Image classes are represented by labels (e.g., ‘panda’) which are then prompted (i.e., ‘[prompt vectors][panda]’) to specify the context of the user’s task. We follow CLIP’s protocol and compute the cosine similarity of the embeddings for each class, normalized to a probability distribution via softmax. The class with the highest probability is selected as the model output. The prompting is added to the class space to form a prompted output $\text{y}+v_f$. During training, PromptFL will maximize the likelihood of the correct label y,

while the gradient updates are applied only to the [prompt vectors] $v_f$ and the CLIP parameters $F$ remain frozen. During validation, the optimized prompt is added to all test-time classes, $D_{\text{test}}\left\{\right({\text{x}}_m,{\text{y}}_m+v_f\left)\right\}$, which will be then processed through the frozen $F$.

To validate the effectiveness of our method, we compare the performance of PromptFL with existing framework by 1)training the collaborative model from the scratch and 2)fine-tuning the full model with pretrained weights. We evaluate the performance across four representative dataset used in CLIP for both general objects and fine-grained classification. We report the performance with two representative and influential backbone, Resnet50(38.3M parameters) and Vit b16(86.6M parameters). For the evaluation metrics, we select three aspects to assess the performance of each method, 1)representative Top-1 accuracy on the test set, 2)F1 score to measure the weighted and unified average of precision and recall, which is more useful especially on unbalanced class distribution, 3)as well as the computational and communication cost reported in Fig. 4. We presuming that higher result on accuracy and F-1 score as well as lower result on computation latency will lead to better a framework, detailed comparison results show the superior if PromptFL in Tab. 2.

All experiments are conducted with Pytorch on GeForce RTX 3090 GPU. Training is performed with SGD with 0.001 learning rate. Tab. 2 measures the overall performance of PromptFL against existing framework from the perspective of two data distribution settings. For the iid setting, each client shares the same classes, while for the extreme non-iid setting, each client owns the independent and non-overlapping classes. We can see that from Tab. 2, PromptFL obtains superior results with similar or better accuracy and F1 value, but with only 0.01% $\sim$ 0.1% learnable parameters with the iid setting. Further more, with the non-iid setting PromptFL achieves competitive performance on both accuracy and efficiency against existing framework. Superior outcome on both settings manifest the advantage of our proposed PromptFL.

After obtaining the decent performance in both extreme iid and non-iid setting, we hope to further testify the stability of PromptFL and figure out the impact of different data distribution on clients to the performance of PromptFL. To observe the intermediate status, we select $p$% overlapped ratio of classes from 0% to 10%, 20% and 50%, which means that $p$% of classes will appear on more than one client and the remaining $1-p$% classes only shows on single client. Fig.  2 reports the accuracy and F1 with corresponding distribution. From the result, we observe that given the circumstance that the class on each client is sufficient, the distribution of class has not much impact on the performance of PromptFL, only a tiny improvement when the overlapping of classes reaches 50%. On the contrary, existing framework shows miserable stability when encountering shifted class distribution other than unified mode by observing the Tab.  2

Following the few-shot evaluation setting adopted in CLIP, we use 2, 4, 8, 16 shots in training PromptFL and validate the performance with corresponding test sets. Unlike the circumstance in the centralized mode where data only from a single entrance, PromptFL involves several participants. Thus we redeclare that the number of shots for the FL mode implies the overall shot containing the entire participants, which is more practical and accord with the class unbalance scenarios in the real-world. From the result in  3, we observe that as the number of training examples per class increases, the performance of PromptFL enhanced. Furthermore, for the setting of adequate clients with sufficient class number on each client, the accuracy for each setting reveals rather steady.

Next, to eliminate the possible impact caused by different clients, we further study the performance of PromptFL with different clients from 16 to 32 to 64, with the iid mode that each client owns random set of classes. Also, to avoid that the collapse of performance due to deficiency of class on each client, especially for the case with large clients number, we also range the number of shots from 2 to 4 to 8 to 16. We observe that for different number of clients, performance will reach similar optimum for as the classes on each client is sufficient. For example, in caltech101, all settings achieve around 89% with 16 shots.

We also analyse the efficiency of PromptFL with regard to the computation and communication cost during training. We measure the communication cost by the size of uploaded data per round, and the total round to be transmitted. For the computation cost, we calculate the GPU memory utilization and training GPU time for given steps. Fig.  4 shows the comparison between existing finetuning framework and our proposed PromptFL. We observe that PromptFL can save at most $110$ times communication cost per round compared to existing prevailing method, let alone that PromptFL takes half of rounds to reach convergence, which makes a wider disparity in communication cost between them. As for the computation cost, we report the comparison of GPU time as in the same given steps, where PromptFL remains outperform existing framework around $3$ times. Further more, there is huge advantage that PromptFL consumes far less GPU memory during training, which can alleviate the system burden in practical.