FedSEA-LLaMA: A Secure, Efficient and Adaptive Federated Splitting Framework for Large Language Models

TL;DR

FedSEA-LLaMA tackles privacy, efficiency, and adaptability challenges in federated splitting of large language models by introducing Gaussian-noise based privacy preservation on forward hidden states, attention-mask compression, and collaborative KV cache for inference acceleration. It enables adaptive partitioning of the LLaMA2 model, allowing the first and last blocks to reside on clients while the middle blocks stay on the server, significantly reducing memory and communication overhead. Empirical results across NLU, summarization, and conversational QA show performance on par with centralized LLaMA2, with up to 8x speedups and substantial privacy protection against model inversion attacks. The framework demonstrates robustness to partition choices and noise scales, offering a practical pathway for privacy-preserving, resource-efficient distributed LLM deployment in real-world federated settings.

Abstract

Private data holds promise for improving LLMs due to its high quality, but its scattered distribution across data silos and the high computational demands of LLMs limit their deployment in federated environments. To address this, the transformer-based federated split models are proposed, which offload most model parameters to the server (or distributed clients) while retaining only a small portion on the client to ensure data privacy. Despite this design, they still face three challenges: 1) Peer-to-peer key encryption struggles to secure transmitted vectors effectively; 2) The auto-regressive nature of LLMs means that federated split learning can only train and infer sequentially, causing high communication overhead; 3) Fixed partition points lack adaptability to downstream tasks. In this paper, we introduce FedSEA-LLaMA, a Secure, Efficient, and Adaptive Federated splitting framework based on LLaMA2. First, we inject Gaussian noise into forward-pass hidden states to enable secure end-to-end vector transmission. Second, we employ attention-mask compression and KV cache collaboration to reduce communication costs, accelerating training and inference. Third, we allow users to dynamically adjust the partition points for input/output blocks based on specific task requirements. Experiments on natural language understanding, summarization, and conversational QA tasks show that FedSEA-LLaMA maintains performance comparable to centralized LLaMA2 and achieves up to 8x speedups in training and inference. Further analysis of privacy attacks and different partition points also demonstrates the effectiveness of FedSEA-LLaMA in security and adaptability.

Figures (5)

• Figure 1: Model architecture of FL-GLM and FedSEA-LLaMA. 1FedSEA-LLaMA employs Gaussian noise addition to protect intermediate vectors. 2FedSEA-LLaMA minimizes the transmission of intermediate data and further accelerates long-context inference. 3FedSEA-LLaMA adaptively adjusts model partition points based on task requirements and system load.
• Figure 2: Overall Framework of FedSEA-LLaMA. Model Split with adaptive partition points and Gaussian noise on forward-passing vectors to preserve privacy for LoRA fine-tuning strategy, and achieving attention-mask compression and KV Cache collaborations to accelerate inference process. The local client stores the first ${p}$ blocks ($Blocks_A$) and the last ${q}$ blocks ($Blocks_C$), while the server stores the remaining ${k}$ blocks ($Blocks_B$). During training, 1$h_A$ is transmitted to the server-side model $Blocks_B$; 2$h_B$ is transmitted to the client-side model $Blocks_C$; 3 the gradient with respect to $h_B$ is backpropagated to the server; 4 the gradient with respect to $h_A$ is backpropagated to the client-side model.
• Figure 3: (a)Threat model attack in Multi-client Training. 1Clients upload hidden states; 2Server processes via $Blocks_B$; 3Attack decoder learns from malicious client's data; 4Normal outputs are sent back to all clients. (b) Gaussion noise injection. Adding noise to $h_1$ provides protection for both the forward and backward passes.
• Figure 4: Computation load on client side of centralized LLaMA2-7B and FedSEA-LLaMA on ReCoRD dataset.
• Figure 5: (a)Different noise injection on MultiRC dataset. (b)Ablation study on attention-mask. (c)Ablation study on KV cache.