Stabilized Fine-Tuning with LoRA in Federated Learning: Mitigating the Side Effect of Client Size and Rank via the Scaling Factor
This work is significant for researchers and practitioners deploying large language models in privacy-preserving federated learning environments, as it resolves a critical instability issue with high-rank LoRA adaptations, which are often necessary for performance.
The paper addresses the instability of integrating Low-Rank Adaptation (LoRA) with Federated Learning (FL), where aggregating updates from multiple clients introduces statistical variance that scales with client count, leading to gradient collapse with high-rank adapters. The authors propose Stabilized Federated LoRA (SFed-LoRA), which derives an optimal scaling factor to mitigate this aggregation error, restoring the efficacy of high-rank adaptation and achieving significantly improved stability and faster convergence compared to state-of-the-art baselines.
Large Language Models (LLMs) are pivotal in natural language processing. The impracticality of full fine-tuning has prompted Parameter-Efficient Fine-Tuning (PEFT) methods like Low-Rank Adaptation (LoRA), optimizing low-rank matrices A and B. In distributed scenarios where privacy constraints necessitate Federated Learning (FL), however, the integration of LoRA is often unstable. Specifically, we identify that aggregating updates from multiple clients introduces statistical variance that scales with the client count, causing gradient collapse when using high-rank adapters. Existing scaling factor candidates, such as the one used by Rank-Stabilized LoRA, ignore the interaction caused by the aggregation process. To bridge this gap, this paper introduces Stabilized Federated LoRA (SFed-LoRA), a framework that theoretically characterizes the interaction between adapter rank and federated aggregation. We derive an optimal scaling factor designed to effectively mitigate the aggregation error accumulating across N clients. By correcting the scaling mismatch inherent in previous approaches, SFed-LoRA restores the efficacy of high-rank adaptation without altering the original model architecture or increasing inference latency. Extensive experiments in diverse tasks, model architectures, and heterogeneous data distributions are conducted to validate our results. We demonstrate that SFed-LoRA prevents high-rank collapse, and achieves significantly improved stability and faster convergence compared with state-of-the-art baselines for high-rank adaptation.