Enhancing Large Language Model Fine-Tuning with Sharpness-Aware Minimization Under Split Federated Learning
摘要
Large Language Models (LLMs) have achieved impressive performance in Natural Language Processing (NLP) tasks through large-scale pre-training. However, fine-tuning these models for downstream tasks can be challenging due to overfitting, high computational costs, and distributed data scenarios. This paper presents a novel split federated learning framework that uses Sharpness-Aware Minimization (SAM) with Low-Rank Adaptation (LoRA) to solve the challenges of fine-tuning LLMs. By splitting the model between client and server components, we enable efficient federated fine-tuning without excessively large parameter updates, where the limited parameters refer to the low-rank LoRA components introduced to reduce the number of trainable parameters in the model. Specifically, LoRA approximates large attention weight with low-rank matrices, reducing computational overhead on both the clients and the server. Building on the existing SplitLoRA framework, we integrate SAM into the optimization process on the central server during training, ensuring convergence toward flatter minima for improved generalization to downstream tasks. Experiments using E2E and WebNLG datasets with GPT-2 demonstrate that our approach outperforms baseline methods such as SplitLoRA and FedLoRA, achieving a BLEU score of 70.27 and CIDEr of 2.5309 at \(\rho\) = 0.05 on E2E dataset and 55.28 BLEU score on WebNLG dataset. These findings highlight the potential of SAM-based federated fine-tuning for more reliable and scalable deployment of LLMs in domain-specific natural language generation (NLG) tasks, such as customer service chatbots and healthcare applications.