FedEAT: A Robustness Optimization Framework for Federated LLMs

TL;DR

Federated LLMs confront robustness challenges from data heterogeneity and adversarial clients in privacy-preserving settings. FedEAT addresses this by integrating embedding-space adversarial training with a robust aggregation rule based on the geometric median, offering resilience against both heterogeneous data and malicious updates. Empirical results show FedEAT improves adversarial robustness with only minor utility loss across multiple architectures and tasks, demonstrating practical viability. This work provides a scalable, privacy-preserving approach for deploying robust Federated LLMs in sensitive domains by combining embedding-space defenses with robust aggregation.

Abstract

Significant advancements have been made by Large Language Models (LLMs) in the domains of natural language understanding and automated content creation. However, they still face persistent problems, including substantial computational costs and inadequate availability of training data. The combination of Federated Learning (FL) and LLMs (federated LLMs) offers a solution by leveraging distributed data while protecting privacy, which positions it as an ideal choice for sensitive domains. However, Federated LLMs still suffer from robustness challenges, including data heterogeneity, malicious clients, and adversarial attacks, which greatly hinder their applications. We first introduce the robustness problems in federated LLMs, to address these challenges, we propose FedEAT (Federated Embedding space Adversarial Training), a novel framework that applies adversarial training in the embedding space of client LLM and employs a robust aggregation approach, specifically geometric median aggregation, to enhance the robustness of Federated LLMs. Our experiments demonstrate that FedEAT effectively improves the robustness of Federated LLMs with minimal performance loss.

Figures (4)

• Figure 1: This figure illustrates the training framework for federated large language models(federated LLMs). Due to limited computational resources on the client side, the client's LLM utilizes LoRA to train only a subset of its parameters.
• Figure 2: This figure depicts the process of generating adversarial samples in the embedding space. First, the input prompt is tokenized and then embedded into a vector space using the LLM's embedding layer. Next, a random perturbation is generated and iteratively updated using gradient ascent to maximize its adversarial effect. Finally, the perturbation is added to the original embedding vector to obtain the adversarial sample.
• Figure 3: The figure presents a comparison of accuracy on a benign dataset between models trained using the FedAVG and FedEAT methods within a federated LLMs framework based on the ZEPHYR-7B. Accuracy is used as a metric to assess the utility of these methods.
• Figure 4: Ablation Study on Robustness Improvement: The Impact of Embedding-space Adversarial Training and Geometric Median Aggregation in SST2 and QNLI Tasks. The ASR for both EAT-only and GM-only is lower than that of FedAVG, indicating that both embedding-space adversarial training and geometric median aggregation methods are effective in enhancing the robustness of federated LLMs.