[Paper Review] FLoRA: Federated Fine-Tuning Large Language Models with Heterogeneous Low-Rank Adaptations

Ziyao Wang, Zanming Shen+5 more|arXiv (Cornell University)|Sep 9, 2024

Topic Modeling4 citations

TL;DR

FLoRA proposes a noise-free, stacking-based aggregation method for federated fine-tuning of large language models using heterogeneous low-rank adapters (LoRAs). By stacking local LoRA matrices instead of averaging them, FLoRA eliminates mathematical aggregation errors and supports variable LoRA ranks across clients, achieving superior performance over SOTA methods in both homogeneous and heterogeneous settings.

ABSTRACT

The rapid development of Large Language Models (LLMs) has been pivotal in advancing AI, with pre-trained LLMs being adaptable to diverse downstream tasks through fine-tuning. Federated learning (FL) further enhances fine-tuning in a privacy-aware manner by utilizing clients' local data through in-situ computation, eliminating the need for data movement. However, fine-tuning LLMs, given their massive scale of parameters, poses challenges for clients with constrained and heterogeneous resources in FL. Previous methods employed low-rank adaptation (LoRA) for efficient federated fine-tuning but utilized traditional FL aggregation strategies on LoRA adapters. These approaches led to mathematically inaccurate aggregation noise, reducing fine-tuning effectiveness and failing to address heterogeneous LoRAs. In this work, we first highlight the mathematical incorrectness of LoRA aggregation in existing federated fine-tuning methods. We introduce a new approach called FLORA that enables federated fine-tuning on heterogeneous LoRA adapters across clients through a novel stacking-based aggregation method. Our approach is noise-free and seamlessly supports heterogeneous LoRA adapters. Extensive experiments demonstrate FLORA' s superior performance in both homogeneous and heterogeneous settings, surpassing state-of-the-art methods. We envision this work as a milestone for efficient, privacy-preserving, and accurate federated fine-tuning of LLMs. Our code is available at https://github.com/ATP-1010/FederatedLLM.

Motivation & Objective

• Address the mathematical inaccuracy in existing federated LoRA aggregation methods, which introduces noise due to independent averaging of LoRA matrices A and B.
• Enable effective federated fine-tuning in realistic, heterogeneous client environments where clients have varying data distributions and computational resources.
• Support heterogeneous LoRA ranks across clients, allowing clients to use different adapter ranks based on local data and hardware constraints.
• Develop a theoretically sound aggregation mechanism that preserves model update integrity and accelerates convergence in federated learning.
• Achieve state-of-the-art performance in both homogeneous and heterogeneous LoRA settings for LLM fine-tuning while maintaining privacy and efficiency.

Proposed method

• Propose a stacking-based aggregation mechanism that concatenates local LoRA matrices A and B from all clients along the rank dimension, forming a global LoRA matrix without averaging.
• Theoretical analysis proves that stacking preserves the mathematical correctness of the global model update, eliminating the intermediate noise term present in FedIT’s averaging approach.
• Support heterogeneous LoRA ranks by allowing clients to use different ranks for their A and B matrices, with stacking naturally accommodating these differences.
• Introduce a privacy-preserving mechanism by splitting LoRA matrices into rank-1 sub-modules and shuffling their order before stacking to prevent reconstruction by malicious clients.
• Integrate with standard privacy mechanisms such as encryption and differential privacy, ensuring compatibility with the privacy guarantees of federated learning.
• Use a standard FedAvg-like FL training loop, where clients fine-tune their local LoRA modules on-device and upload them to the server, which then applies stacking to form the global model.

Experimental results

Research questions

• RQ1 What is the mathematical source of aggregation noise in existing federated LoRA methods like FedIT, and how can it be eliminated?
• RQ2 Can a federated fine-tuning method support heterogeneous LoRA ranks across clients while maintaining model accuracy and convergence?
• RQ3 Does a stacking-based aggregation mechanism for LoRA matrices lead to faster convergence and improved performance compared to averaging-based methods?
• RQ4 How does FLoRA perform in realistic, heterogeneous client settings with varying data and hardware resources?
• RQ5 Can the proposed stacking method be secured against privacy leakage from global model updates without sacrificing communication efficiency?

Key findings

• FLoRA eliminates aggregation noise by replacing averaging with stacking of LoRA matrices, which is mathematically correct and leads to faster convergence.
• The stacking mechanism naturally supports heterogeneous LoRA ranks across clients, enabling broader client participation in federated fine-tuning.
• FLoRA outperforms the SOTA method FedIT across multiple benchmarks, including GSM8K, HumanEval, and WizzAR, on LLaMA, Llama2, and TinyLlama models.
• LoRA rank 8 consistently yields strong performance across models and datasets, while extreme ranks (e.g., 64) show model-dependent performance, suggesting a correlation between optimal rank and model capacity.
• Communication overhead in FLoRA is only slightly higher than FedIT and remains negligible compared to full fine-tuning, making it practical for large-scale deployment.
• FLoRA’s privacy-preserving design, including rank-1 sub-module shuffling and compatibility with encryption and differential privacy, prevents reconstruction of local LoRA matrices by malicious clients.