Personalized Wireless Federated Learning for Large Language Models

In the pre-training stage, an LLM is deployed on a server and undergoes self-supervised learning using a substantial volume of unlabelled data. Numerous client devices transmit their unlabelled data that does not contain sensitive information to the server for learning purposes over the wireless network. The goal of pre-training is to teach patterns, grammar, semantics, and world knowledge to LLMs by predicting missing or masked words in sentences. For instance, in the BERT model, it is done using a technique called Masked Language Modelling (MLM), where a certain percentage of words in a sentence are randomly replaced with a special mask token, and the model has to predict the original word [3]. During the pre-training stage, the LLM learns to capture contextual relationships and build a general understanding of language. A large number of parameters in the model enables it to encode a vast amount of information from the training corpus, resulting in a rich language representation.

However, FL may not be necessary for the pre-training stage of LLMs due to the following reasons:

• •

  High Resource Requirements: Pre-training of LLMs involves adjusting all their parameters, which requires significant computational resources. In FL, the model is trained across many devices or servers, with each device obtaining updates for the model based on its computation and storage resource [4]. Additionally, this distributed computation increases the communication overhead, as the update from each device needs to be aggregated to obtain the global model. Centralized training, on the other hand, can leverage powerful servers and optimized infrastructure to train LLMs more efficiently.

• •

  Privacy Concerns: Pre-training of LLMs typically involves using a large, publicly available corpus of text, such as books, websites, and other online resources. Since this data is already publicly available and does not contain sensitive personal information, there are no data privacy concerns that necessitate the use of FL. FL is more applicable in scenarios where data is sensitive and decentralized, such as in healthcare or personal mobile devices, where privacy regulations or concerns prevent the data from being shared directly.

Once the LLM completes pre-training on the server side, it undergoes further fine-tuning on a smaller, more specific dataset. This dataset typically consists of local private data. The clients download the pre-trained parameters of the LLM from the server and perform fine-tuning on an extremely small subset of parameters to enhance the performance of the LLM. To ensure security and privacy of local data, the fine-tuning process is conducted locally on the client side, and the updated subset of parameters are transmitted back to the server through wireless networks. This local dataset is often labelled, meaning that it comes with correct answers that the LLM should learn to predict. Fine-tuning requires less computational resources and smaller amount of data compared to pre-training. Fine-tuning of LLMs can be categorized into the following types:

RAG is an approach that combines LLMs with information retrieval techniques to enhance performance of LLMs[5].

Due to the enormous training cost of LLMs, the learned knowledge has a temporal lag. For instance, the training data for GPT-3.5 is current up until January 2022, implying that it lacks knowledge of any events transpiring after January 2022. During the generation process, instead of relying solely on the pre-trained or fine-tuned model’s knowledge, retrieval mechanisms are employed to retrieve the latest relevant information from external sources such as Internet or local knowledge bases. These knowledge bases are often deployed at the edge, and clients retrieve the latest local information by querying knowledge bases. The retrieved local information, along with the client’s request, is then sent to the server-side LLM by wireless networks. Moreover, LLMs can leverage the Internet to retrieve the latest relevant public information. Then, this retrieved information is incorporated into the generated output of the LLM, ensuring that the generated content is contextually relevant and factually accurate.

However, RAG may also not be suitable for FL due to the following reasons:

• •

  Additional Data Exposure: Sharing sensitive retrieval queries or accessing external resources across multiple clients may compromise the privacy and security of the distributed data.

• •

  No Weight Update: RAG enhances the performance of LLMs by incorporating the retrieved data into the prompt for in-context learning. It does not require updating parameters of LLMs, thus eliminating the need for local gradient descent in FL.

LLMs not only need to align with human values during the learning process, but also interact with external knowledge, which represents a more complex learning paradigm. Due to the massive amount of insensitive data and extensive model parameters involved in the pre-training stage, it is more suitable for centralized cloud-based learning. In contrast, during the RAG stage, only local data embedding and querying are required, making it more suitable for local execution. Therefore, we primarily focus on designing PFLs for the fine-tuning stage.

Federated instruction tuning involves FL for instruction tuning of LLMs [4], which necessitates each client to possess sufficient computational resources to fine-tune a larger number of parameters compared with task fine-tuning methods. Additionally, it requires transmitting more model parameters over the network, which can be bandwidth-intensive and time-consuming. Furthermore, it is challenging to define a clear, algorithmic loss to measure the quality of the LLM’s output for local instructions. The following techniques hold the potential to address these challenges.

Federated task tuning employs FL for fine-tuning downstream tasks of LLMs [8], which is relatively efficient as it requires less data and computational resources of cilents than that of the instruction tuning. It allows for task-specific adaptation but also risks overfitting or forgetting the original knowledge of local LLMs. PEFT is a set of techniques designed to adapt LLMs to specific downstream tasks with minimal changes to the original parameters. The following PEFT methods can help mitigate overfitting and catastrophic forgetting in LLMs.

There have been several joint research efforts on wireless FL and LLMs. The authors in [3] combined split learning and FL to pretrain the BERT model. Reference [4] proposed a federated instruction learning called Shepherd, which utilizes LoRA for fine-tuning LLMs through instruction data. Similarly, the authors in [8] presented a low-parameter federated learning based on LoRA for task fine-tuning of LLMs. Reference [10] addressed co-tuning and off Site-tuning of LLMs through a comprehensive FL open-source framework called Fate-LLM. However, all of these studies aim to train a unified LLM in a distributed manner, overlooking device variations, user preferences, and characteristics. As a result, they fail to achieve device-adaptive and user-centric LLMs.

PFL allows for designing personalized LLMs that can adapt to the data of individual clients over wireless networks, which may improve user satisfaction and engagement levels. Advantages of applying PFL to LLMs include:

We propose a PFIT method based on RL, in which each client has personalized requirements for helpfulness and safety of the LLM. Helpfulness emphasizes the quality and accuracy of generated content, such as grammatical correctness, logical coherence, and relevance of the responses. Safety, on the other hand, emphasizes the legality and ethicality of the generated content, such as the absence of sensitive or harmful information. For example, in Fig. 2, the LLM with high helpfulness would provide answers to user questions regarding employee information, including sensitive details. However, the LLM with high safety levels would refrain from answering sensitive employee information, prioritizing user privacy and data protection. To achieve PFIT, we introduce three key innovations:

• •

  Double Reward Model: We define two reward models to evaluate the helpfulness and safety of local instruction responses. The quality reward for different clients’ outputs is obtained by linearly weighting these two reward models, enabling personalized instruction fine-tuning of the LLM for different clients’ requirements.

• •

  Personalized Reward Function: We design a personalized reward function that includes both the quality reward from two reward models and a negative regularization reward for global knowledge. The regularization reward is based on the Euclidean distance between local model parameters and global model parameters, promoting the knowledge sharing among clients in the PFL system.

• •

  Sparse Attention Update: To encourage lightweight devices to participate in PFL, we only train the last two layers of the LLM. Additionally, we adopt a sparse attention mechanism to further reduce the computation complexity and communication overhead in self-attention layers during the fine-tuning phase of the LLM.

The specific workflow of the PFIT is as follows:

Step 1: Initialize the pre-trained LLM as the global model on the server and freeze the earlier parts of the LLM.

Step 2: Each client uses the global LLM as the initial local LLM, sets a personalized reward function (In the red dashed box of Fig. 2), and selects its own instruction data for local learning.

Step 3: Based on the current local LLM, the client calculates the regularization reward related to the global LLM, evaluates the helpfulness and safety of responses, computes the quality reward for instructions, and then updates the unfrozen part of the local LLM using the PPO algorithm [11] according to the personalize reward function, including both quality reward and regularization reward.

Step 4: The server aggregates the sparse tunable layers (the unfrozen parts of LLMs) from all clients, obtains the updated global LLM, and sends the global LLM (the unfrozen part) to all clients.

Step 5: Steps 3-4 are repeated until the convergence criteria of the PFL system are met.

We propose a PFTT method based on PEFT, in which a set of clients have different task goals, and each client possesses a set of non-IID task data. For instance, in Fig. 3, different clients are tasked with film classification, but the distribution of labelled data varies across different clients. Some clients may have a higher proportion of science fiction and realistic film labels, while others may have a larger number of comedy and tragedy film labels. As a result, the LLM will excel at providing personalized movie classifications based on the local distribution of labelled data available to it. To achieve PFTT, we introduce three key innovations:

• •

  Universal Adapter: We incorporate adapters into the global pre-trained LLM to share task knowledge among different clients.

• •

  Local LoRA: We introduce LoRA in the local LLMs to enable personalized local knowledge learning. The size of LoRA can be adjusted based on the data volume or computational resources available on each client.

• •

  Partial Aggregation: During global aggregation, only the adapter parameters are aggregated for global knowledge sharing, while the LoRA parameters are not aggregated for maintaining personalization. This approach enables personalized task tuning in heterogeneous devices.

The specific workflow of the PFTT is as follows:

Step 1: Initialize the pre-trained LLM as the global LLM on the server and insert adapters to the LLM.

Step 2: Each client uses the global LLM as the initial local LLM and designs local LoRA parameters based on the data volume or computational resource of the local LLM.

Step 3: Based on the current global LLM and local LoRA parameters, the client fine-tunes the LLM using local task data and updates the adapter and LoRA parameters over the wireless networks.

Step 4: The server aggregates the adapter parameters from all clients, obtains the updated global adapter parameters, and sends them to the clients.

Step 5: Steps 3-4 are repeated until the convergence criteria of the system are met.

Through PFL, both fine-tuning methods can achieve better performance for personalized LLMs.

We consider a PFL system for LLMs. Suppose the system consists of four clients and one server, where each client possesses different local data and model preferences. The server has sufficient resources. The PFL system aims to achieve personalized fine-tuning of the local LLMs for all clients in wireless networks with Rayleigh channel, the communication round is set to 40 and the SNR is set to 5dB.

Fig. 4 presents the evaluation results for PFIT and its contenders, where SFL means to a fine-tuning method that uses only a single reward model (helpfulness metric) and incorporates 20% sparse attention. PFL, on the other hand, represents personalized fine-tuning without the sparse attention. Shepherd is an FL method that employs LoRA for instruction fine-tuning [4]. We can see that PFIT enables the local model to obtain the highest rewards than SFL and Shepherd with the single reward model. Moreover, compared to PFL that does not use sparse attention, PFIT reduces communication overhead by 20%. Shepherd, utilizing LoRA for instruction fine-tuning, results in the lowest communication overhead. However, this approach also affects the LLM’s performance, resulting in lower rewards compared to PFIT.

Fig. 5 presents the evaluation results for PFTT and its benchmarks. In vanilla FL [1], the parameters of both adapters and LoRAs all need to be uploaded. FedBert [3] is a federal split learning method, and FedLora [8] is a federated fine-tuning method that exclusively incorporates LoRA. The results indicate that PFTT achieves the highest accuracy, which highlights the effectiveness of the LoRA-based personalized structure. Similarly, due to the fact that PFTT only requires the transmission of a part of fine-tuned parameters (universal adapters), it incurs the minimum communication overhead compared to other methods.