SA-FedLora: Adaptive Parameter Allocation for Efficient Federated Learning with LoRA Tuning

2024-05-15

arXiv.org

AI Summary

To address the high communication overhead in fine-tuning large models for IoT-based federated learning and the reliance of existing parameter-efficient fine-tuning (PEFT) methods on manual hyperparameter tuning, this paper proposes an adaptive LoRA parameter allocation method inspired by simulated annealing. The method operates in two phases—initialization and annealing—dynamically regulating the number of trainable parameters per client and progressively shrinking the parameter budget during training to jointly optimize convergence speed, global optimality, and communication efficiency. This work is the first to incorporate simulated annealing into federated LoRA optimization, synergistically integrating LoRA’s low-rank adaptation, federated aggregation regularization, and dynamic budget allocation to mitigate client drift and overfitting. Experiments demonstrate significant improvements over FedAvg: up to 93.62% reduction in transmitted parameters, accelerated convergence, and enhanced global generalization performance.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Sampling/Simulation-based Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Fine-tuning large-scale pre-trained models via transfer learning is an emerging important paradigm for a wide range of downstream tasks, with performance heavily reliant on extensive data. Federated learning (FL), as a distributed framework, provides a secure solution to train models on local datasets while safeguarding raw sensitive data. However, FL networks encounter high communication costs due to the massive parameters of large-scale pre-trained models, necessitating parameter-efficient methods. Notably, parameter efficient fine tuning, such as Low-Rank Adaptation (LoRA), has shown remarkable success in fine-tuning pre-trained models. However, prior research indicates that the fixed parameter budget may be prone to the overfitting or slower convergence. To address this challenge, we propose a Simulated Annealing-based Federated Learning with LoRA tuning (SA-FedLoRA) approach by reducing trainable parameters. Specifically, SA-FedLoRA comprises two stages: initiating and annealing. (1) In the initiating stage, we implement a parameter regularization approach during the early rounds of aggregation, aiming to mitigate client drift and accelerate the convergence for the subsequent tuning. (2) In the annealing stage, we allocate higher parameter budget during the early 'heating' phase and then gradually shrink the budget until the 'cooling' phase. This strategy not only facilitates convergence to the global optimum but also reduces communication costs. Experimental results demonstrate that SA-FedLoRA is an efficient FL, achieving superior performance to FedAvg and significantly reducing communication parameters by up to 93.62%.

Problem

• Efficient continual federated learning with parameter dropout
• Adaptive dropout rate for target performance in FL
• Reducing communication overhead in resource-constrained IoT FL

Innovation

• Stepwise parameter dropout for efficiency
• Adaptive dropout via gradient consistency
• Continual learning for performance stability