Federated Learning of Large Language Models with Parameter-Efficient Prompt Tuning and Adaptive Optimization (2023)

Abstract: Federated learning (FL) is a promising paradigm to enable collaborative model training with decentralized data. However, the training process of Large Language Models (LLMs) generally incurs the update of significant parameters, which limits the applicability of FL techniques to tackle the LLMs in real scenarios. Prompt tuning can significantly reduce the number of parameters to update, but it either incurs performance degradation or low training efficiency. The straightforward utilization of prompt tuning in the FL often raises non-trivial communication costs and dramatically degrades performance. In addition, the decentralized data is generally non-Independent and Identically Distributed (non-IID), which brings client drift problems and thus poor performance. This paper proposes a Parameter-efficient prompt Tuning approach with Adaptive Optimization, i.e., FedPepTAO, to enable efficient and effective FL of LLMs. First, an efficient partial prompt tuning approach is proposed to improve performance and efficiency simultaneously. Second, a novel adaptive optimization method is developed to address the client drift problems on both the device and server sides to enhance performance further. Extensive experiments based on 10 datasets demonstrate the superb performance (up to 60.8% in terms of accuracy) and efficiency (up to 97.59% in terms of training time) of FedPepTAO compared with 9 baseline approaches. Our code is available at https://github.com/llm-eff/FedPepTAO.

## introduction as a promising paradigm to handle decentralized data, federated learning (fl) @xcite enables collaborative model training without transferring the raw data across a massive number of devices. as a bunch of legal restrictions (official journal of the european union, 2016; californians for consumer privacy, 2020) have been implemented, aggregating the decentralized data into a central server or data center becomes complicated or even impossible @xcite . fl generally exploits a parameter server module @xcite to manage the distributed model updates in devices, which only exchanges the parameters of the updated models instead of the raw data, between the parameter server and devices. large language models (llms) @xcite @xcite have achieved major advances in natural language processing (nlp) tasks. the scale of llms can range from 110 million parameters to 175 billion parameters, which correspond to huge communication and computation costs to update parameters during the pre-training process @xcite or fine-tuning process @xcite . both pre-training and fine-tuning update the whole set of parameters of the language model. thus, the application of fl in the pre-training or the fine-tuning process is almost impossible due to significant communication burden brought by large amount of parameters. prompt design @xcite can lead to excellent performance while freezing the original llms. when the llms are frozen, only the prompts or prefix are updated during the tuning process, which can significantly reduce the number of parameters to update. for instance, for a sample from a sentiment analysis task (e.g., "beautiful place!"), a discrete prompt "it was @xcite ." for prompt tuning @xcite and continuous task-specific vectors for prefix tuning @xcite can be concatenated to be sent to a llm, which generates the label of the sample to be "terrible" or "great". numerous parameter-efficient prompt or prefix tuning approaches have been proposed to tune the large language models through updating a few train-able parameters while achieving comparable performance compared with fine-tuning. in order to avoid human involvement in the prompt design, prompt tuning methods @xcite are proposed to search proper prompts within a discrete space of words, which corresponds to inferior performance compared with fine-tuning. continuous prompts, i.e., prefix tuning, can be updated to achieve better performance @xcite . however, this approach leads to sub-optimal performance for the models with less than 10 billion parameters. although p-tuning v2 @xcite achieves comparable performance compared with fine-tuning, it introduces more parameters, which may correspond to heavier communication costs in the setting of fl compared with other parameterefficient tuning approaches. some other parameterefficient prompt tuning methods either suffer from low performance with the focus on low-rank hypercomplex adapter layers @xcite or prompt with a single layer @xcite , or introduce extra computation costs with attentions @xcite . in addition, existing fl techniques for fine-tuning large language models typically incur performance degradation or low efficiency due to huge communication costs @xcite @xcite . adaptive optimization methods, e.g., adaptive moment estimation (adam) and stochastic gradient descent (sgd) with momentum (sgdm) @xcite , have been utilized either on server side @xcite or on device side @xcite @xcite to achieve superior performance in fl. however, the direct application of the adaptive optimization methods may incur problems of the convergence within the training process @xcite . furthermore, the application of adaptive optimization on a single side, i.e., either device or server, may correspond to poor performance. however, when the adaptive optimization is exploited on both sides @xcite ) may incur heavy communication costs. in addition, client drift (karimireddy et al., 2020b) may exist in terms of the adaptive optimization due to non-independent and identically distributed (non-iid) data among devices. in this paper, we propose a parameter-efficient prompt tuning approach with adaptive optimization, i.e., fedpeptao, to tune large language models with fl. as transferring the whole set of pa-rameters in all the prompt layers corresponds to heavy communication costs, we propose an efficient and effective method to choose proper layers of prompts based on the importance of each layer. we design a scoring method to identify the importance of each layer according to the tuning impact of the layer on the final convergence accuracy. in addition, we propose an adaptive optimization method on both server side and device side with control measures on each device to achieve superb accuracy. we summarize out major contributions as follows: ## related work as updating all the parameters of a pre-trained llm consumes a large amount of memory and computation resources, prompt tuning @xcite or prefix tuning @xcite is proposed to update a few parameters with a frozen language model while achieving comparable performance compared with fine-tuning. while prompt tuning may correspond to inferior performance with discrete space of words @xcite , prefix tuning @xcite can deal with continuous prompts to achieve better performance. adapter modules @xcite are exploited to tune large language models with prompts, which may incur heavy computation costs due to the calculation of feed-forward project and non-linearity or attention mechanism @xcite . although efficient low-rank hypercomplex mechanism @xcite can be utilized to reduce parameters to update, the performance may degrade. p-tuning v2 achieves comparable performance compared with fine-tuning @xcite , which the prompts added into each layer of the large language model. prompts can be added at a single layer to further reduce the computation costs @xcite , which may incur performance degradation and depends on multiple trials with each layer. in addition, the selection of the layer may incur long execution time to verify the impact on the final accuracy. although the naswot algorithm @xcite can be exploited to analyze the performance of a neural network architecture, it is only compatible with the neural networks based on the relu activation function @xcite . parallel, distributed, and federated learning have been extensively studied in recent years @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite . some existing fl techniques have been proposed to finetuning large language models, which may suffer from performance degradation or low efficiency due to huge communication costs @xcite @xcite . fedbert @xcite exploits split learning to split a model into two parts, i.e., one with transformer and the other one with head and embedding. as the transformer is shared on the server, fedbert may correspond to inferior performance and huge communication costs compared with prompt-tuning or prefix-tuning. some other fl methods only fine-tune a part of the model weights @xcite , which still suffer from heavy communication costs for big language models. fedprompt @xcite enable fl based on prompt tuning, which communicates the whole set of parameters in prompts corresponding to huge communication costs. adaptive moment estimation (adam) and stochastic gradient descent (sgd) with momentum (sgdm) @xcite are exploited within fl on server side @xcite or on device side @xcite @xcite to address the client drift problem brought by the non-iid data in fl @xcite . however, the direct application of adaptive optimization on devices may lead to convergence problem @xcite . in addition, the application of adaptive optimization on both server and device sides may incur heavy communication costs @xcite . different from the previous work, we propose a general scoring method to analyze the correlation of each layer and the output of the large language model, which can represent the importance of each layer. then, we select the prompt parameters of proper layers to be updated with fl while leaving other prompt parameters of other layers to be adjusted locally with a lossless method so as to achieve superb performance with limited communication costs. in addition, we introduce control measures on each device to alleviate the client drift problem and propose a novel adaptive optimization method on both server and device sides to further improve the performance. ## problem formulation the problem to address in this paper is how to efficiently tune a large language model based on prompt tuning in fl. given a large language model m with l layers, we add prompts for each layer in m and denote the set of parameters to generate prompts by p l with l representing the number of layer in m. the whole set of prompts is denoted by p during the tuning process, the parameters in m are frozen and cannot be updated while the parameters in p are updated to improve the performance of m. we consider a fl setting with a parameter server and m devices. we assume that the data for the tuning process of m is distributed among multiple devices. on each device i, a dataset d @xmath0, m i } n i is located with s i , m i , and n i representing a sample, the corresponding label of s i , and the number of samples in d i . we denote the total number of the samples on all the devices by n , the set of all the samples by s and that of labels by m . due to the limited computation capacity, each device i can only perform the inference of m while updating p i with p i representing the prompt parameters in device i. in order to reduce communication costs, we selected prompts for communication layer 1 prompts layer 2 prompts layer n prompts       e(pretty) e(dull) e(.) prediction (with linear head)             prompts uploading and aggregation ① selected prompts downloading ② upload 𝜁 1 and 𝑟 𝑖 to the server ( before efficient tuning phase) ③ download and distribute 𝑆𝐿 to all clients (before efficient tuning phase) ④ prompts tuning update enable the exchange the parameters of the prompts within a subset of selected layers sl i between device i and the parameter server while the other prompt parameters are only updated within each device. we denote the set of prompt parameters in all the devices by p. the problem to address in this paper can be formulated as how to efficiently generate p such that the global loss is minimized: for nlp tasks, each sample s @xmath1. each sample s k is composed of multiple tokens, i.e., s @xmath2, s 2 k , ..., s t k }, where t represents the length of the input. the prompt p consists of multiple tokens @xmath3, p 2 , ..., p h }, and the corresponding prompt parameters can be trained. the prompts differ according to layers. we denote the template by t (•), which defines how to concatenate the input tokens with the prompt. for instance, s p @xmath4, p) represents the sample combined with the prompt, which contains one [mask] token. the output of the large language model with the prompts predicts the label m k , which corresponds to the [mask] token after applying a verbalizer v(•), i.e., @xmath5. in this section, we first present the system model of fedpeptao. then, we propose parameterefficient prompt tuning method and the adaptive optimization method, respectively. ## experiments in this section, we present the experimental results over 9 baselines and 10 commonly-used tasks to demonstrate the advantages of fedpeptao. ## conclusion in this paper, we propose an original parameterefficient prompt tuning with adaptive optimization approach for large language models with fl, i.e., fedpeptao. we dynamically calculate the score of each layer to choose proper layers of prompts for fl without accuracy degradation. in addition, we provide a novel adaptive optimization method with control variate on the server to achieve superb performance without extra communication costs. we carry out extensive experiments based on 10 tasks and 9 state-of-the-art baseline approaches, which demonstrate significant advantages of fed-peptao in terms of accuracy (up to 60.8% higher) and efficiency (up to 97.59% faster).