Safely Learning with Private Data: A Federated Learning Framework for Large Language Model (2024)

Abstract: Private data, being larger and quality-higher than public data, can greatly improve large language models (LLM). However, due to privacy concerns, this data is often dispersed in multiple silos, making its secure utilization for LLM training a challenge. Federated learning (FL) is an ideal solution for training models with distributed private data, but traditional frameworks like FedAvg are unsuitable for LLM due to their high computational demands on clients. An alternative, split learning, offloads most training parameters to the server while training embedding and output layers locally, making it more suitable for LLM. Nonetheless, it faces significant challenges in security and efficiency. Firstly, the gradients of embeddings are prone to attacks, leading to potential reverse engineering of private data. Furthermore, the server’s limitation of handling only one client’s training request at a time hinders parallel training, severely impacting training efficiency. In this paper, we propose a Federated Learning framework for LLM, named FL-GLM, which prevents data leakage caused by both server-side and peer-client attacks while improving training efficiency. Specifically, we first place the input block and output block on local client to prevent embedding gradient attacks from server. Secondly, we employ key-encryption during client-server communication to prevent reverse engineering attacks from peer-clients. Lastly, we employ optimization methods like client-batching or server-hierarchical, adopting different acceleration methods based on the actual computational capabilities of the server. Experimental results on NLU and generation tasks demonstrate that FL-GLM achieves comparable metrics to centralized chatGLM model, validating the effectiveness of our federated learning framework.

## introduction existing large language models (llm) have achieved astonishing results by utilizing vast amounts of public data and massive parameters. in comparison to public data, private data holds advantages in both quantity and quality, because private datasets typically encompass more comprehensive and detailed information about individuals or organizations, and the data production process is more rigorous. therefore, private data can undoubtedly further enhance the performance of llm. however, private data is often stored in isolated data silos. for example, mobile users' data is kept locally, involving a significant amount of personal privacy. considering privacy and security, llm cannot store private data in a centralized manner for training. hence, securely leveraging private data for language model training remains a challenging problem. an ideal solution is to utilize the federated learning (fl) @xcite ) framework, which allows data to be retained on the user device for local training and only passes the model parameters to server for model aggregation. this approach achieves the goal of keeping the data stationary while making the model updates. by using fl for llm training, data privacy can be preserved, and the performance of llm can be further improved. unfortunately, traditional fl frameworks, such as fedavg @xcite and fed-prox @xcite , are not suitable for llms because they require each client to have sufficient computational resources to train the entire llm. as an alternative method, transformer with split learning, represented by fedbert @xcite in figure 1 (a), focuses most of the parameters on the server while continuously training the embedding layer and output layer on the local client, making it more suitable for llm. the process involves the client using the embedding layer for data input and forwarding it to the server, which then computes and returns the output states. the client calculates loss, sends gradients back for server updates, and receives updated gradients for the embedding layer. however, this method presents security risks. embedding gradients are vulnerable to attacks @xcite , potentially allowing attackers to reconstruct private data through beam search @xcite or reverseengineered @xcite . additionally, since the server processes one client at a time, it hinders parallel training and reduces efficiency. in this paper, we propose a fl framework called fl-glm for llm, as shown in figure 1(b) . we partition the transformers of the chatglm into three parts: input and output blocks are stored on the client (shared with peers), while the remaining large parameters are kept on the server. during the training process, the client first performs forward propagation on the input data to obtain hidden states. then, these hidden states are encrypted with a secure key and sent to the server. subsequently, the server, either in a client-batch or server-layered approach, receives more client hidden states at a training time and executes forward propagation to send output hidden states back to each client. it is clear that our fl-glm framework can effectively prevent data leakage attacks from both servers and peer-clients while enhancing training efficiency. clients and servers jointly own and utilize the entire model, with certain input and output blocks placed on local clients to thwart embedding gradient attacks from the server. although sharing input and output blocks between all clients can improve results, interception by peers poses risks, which can be resolved through key encryption during client-server communication. to overcome server capacity limitations, we propose various training acceleration methods. for clusters with multiple machines and gpus, a hierarchical server architecture initializes sub-servers for parallel client training, with central server aggregating and distributing models. with single machines and multiple gpus, the client-batch method concatenates client information for training, enabling parallel execution and enhanced efficiency compared to traditional serial execution in split learning. experimental results on nlu and generation tasks demonstrate that fl-glm achieves performance comparable to centralized chatglm-6b models, validating the effectiveness of our framework. further analysis of training costs indicates that our client-batch and server-hierarchical mechanisms can save more than 48% of training time. our code is available at: https://github.com/tapllm/splitfedllm . the innovations in this paper are as follows: ## model in this section, we provide the details of the fl-glm framework, as shown in figure 1(b ). fl-glm consists of three parts: model split, encrypted transmission, and parallel acceleration. firstly, we split llm into three parts, saving the first block 0 and the last block n-1 on the local client and placing the remaining parameters on the server. then, the smashed data is encrypted using keys during client-server transmission. finally, the server employs either client-batch or hierarchical-server methods to achieve parallel acceleration. ## experiments in order to demonstrate the performance of chatglm model within the federated learning framework(fl-glm), we conduct experiments using the same benchmarks as those used in glm model @xcite . ## conclusions to address the challenge of distributed training of llms with limited client computational resources, we propose to utilize the split learning method to segment the generative model. we place the input and output blocks locally on client devices, while the remaining primary model parameters are centralized on a server with ample computational resources. we secure client-server information transfers with encryption methods. to enhance training efficiency, we suggest selecting the client-batch and server-hierarchical acceleration optimization methods based on the server's actual computational capacity, thereby enabling parallel training. this distributed architecture not only ensures that user private data remains on local devices but also effectively reduces the training time, making it more suitable for the scale and complexity of llms. in the future, we contemplate employing more advanced privacy-preserving techniques, such as differential privacy, to safeguard the data transmitted from clients, enabling the application of large language models in privacy-sensitive scenarios.