Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive LoRA Experts Allocation and Selection for Federated Fine-Tuning

Lei Wang, Jieming Bian, Letian Zhang, Jie Xu

TL;DR

The paper tackles the problem of privately fine-tuning large language models under data heterogeneity in federated settings by introducing FedLEASE, a framework that jointly optimizes LoRA expert allocation and adaptive expert usage. It clusters clients based on representation similarity of LoRA adapters, selects the optimal number of experts with silhouette analysis, and employs an adaptive top-$M$ router that guarantees inclusion of each client’s assigned expert while flexibly recruiting others as needed. Empirical results on NLU (GLUE with RoBERTa-Large) and NLG (FLAN with LLaMA-2-7B) show consistent improvements over strong federated LoRA baselines, highlighting gains in both accuracy and generation metrics across task-heterogeneous and label-heterogeneous settings. The approach achieves these benefits with modest communication overhead, demonstrating practical applicability for privacy-preserving, domain-specific fine-tuning of LLMs.

Abstract

Large Language Models (LLMs) have demonstrated impressive capabilities across various tasks, but fine-tuning them for domain-specific applications often requires substantial domain-specific data that may be distributed across multiple organizations. Federated Learning (FL) offers a privacy-preserving solution, but faces challenges with computational constraints when applied to LLMs. Low-Rank Adaptation (LoRA) has emerged as a parameter-efficient fine-tuning approach, though a single LoRA module often struggles with heterogeneous data across diverse domains. This paper addresses two critical challenges in federated LoRA fine-tuning: 1. determining the optimal number and allocation of LoRA experts across heterogeneous clients, and 2. enabling clients to selectively utilize these experts based on their specific data characteristics. We propose FedLEASE (Federated adaptive LoRA Expert Allocation and SElection), a novel framework that adaptively clusters clients based on representation similarity to allocate and train domain-specific LoRA experts. It also introduces an adaptive top-$M$ Mixture-of-Experts mechanism that allows each client to select the optimal number of utilized experts. Our extensive experiments on diverse benchmark datasets demonstrate that FedLEASE significantly outperforms existing federated fine-tuning approaches in heterogeneous client settings while maintaining communication efficiency.

Adaptive LoRA Experts Allocation and Selection for Federated Fine-Tuning

TL;DR

The paper tackles the problem of privately fine-tuning large language models under data heterogeneity in federated settings by introducing FedLEASE, a framework that jointly optimizes LoRA expert allocation and adaptive expert usage. It clusters clients based on representation similarity of LoRA adapters, selects the optimal number of experts with silhouette analysis, and employs an adaptive top- router that guarantees inclusion of each client’s assigned expert while flexibly recruiting others as needed. Empirical results on NLU (GLUE with RoBERTa-Large) and NLG (FLAN with LLaMA-2-7B) show consistent improvements over strong federated LoRA baselines, highlighting gains in both accuracy and generation metrics across task-heterogeneous and label-heterogeneous settings. The approach achieves these benefits with modest communication overhead, demonstrating practical applicability for privacy-preserving, domain-specific fine-tuning of LLMs.

Abstract

Large Language Models (LLMs) have demonstrated impressive capabilities across various tasks, but fine-tuning them for domain-specific applications often requires substantial domain-specific data that may be distributed across multiple organizations. Federated Learning (FL) offers a privacy-preserving solution, but faces challenges with computational constraints when applied to LLMs. Low-Rank Adaptation (LoRA) has emerged as a parameter-efficient fine-tuning approach, though a single LoRA module often struggles with heterogeneous data across diverse domains. This paper addresses two critical challenges in federated LoRA fine-tuning: 1. determining the optimal number and allocation of LoRA experts across heterogeneous clients, and 2. enabling clients to selectively utilize these experts based on their specific data characteristics. We propose FedLEASE (Federated adaptive LoRA Expert Allocation and SElection), a novel framework that adaptively clusters clients based on representation similarity to allocate and train domain-specific LoRA experts. It also introduces an adaptive top- Mixture-of-Experts mechanism that allows each client to select the optimal number of utilized experts. Our extensive experiments on diverse benchmark datasets demonstrate that FedLEASE significantly outperforms existing federated fine-tuning approaches in heterogeneous client settings while maintaining communication efficiency.

Paper Structure

This paper contains 23 sections, 1 theorem, 27 equations, 11 figures, 15 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminary and Motivation
  4. LoRA and MoE Integration
  5. Heterogeneous Federated Fine-tuning Scenario
  6. Heterogeneous Federated Fine-tuning Scenario
  7. Observations
  8. Proposed Method
  9. Adaptive LoRA Experts Allocation
  10. Adaptive top-M LoRA Experts
  11. Algorithm Workflow
  12. Experiment
  13. Training Details
  14. Natural Language Understanding
  15. Natural Language Generation
  16. ...and 8 more sections

Key Result

Theorem F.5

With the assumptions of client-level smoothness, client-level strong convexity, limited inter-cluster influence, and cluster assignment stability, we can derive: If $\beta M < 1$ and the local training at each round converges to a neighborhood of the optimal solution, then the sequence of cluster models $\{\Theta_j^t\}$ generated by FedLEASE converges to a stable point for each cluster $j$.

Figures (11)

  • Figure 1: Illustration of LoRA-MoE mechanism.
  • Figure 2: Performance comparison between FedIT and FedLORA-MoE under two scenarios with different clients’ task heterogeneity.
  • Figure 2: Ablation on adaptive expert allocation.
  • Figure 3: Visualization result of cosine similarity among clients using different LoRA matrices. Note that client 0 and 1 hold data from SST-2 and client 2 and 3 hold data from QNLI.
  • Figure 3: Performance on FLAN dataset (LLaMA-2-7B).
  • ...and 6 more figures

Theorems & Definitions (1)

  • Theorem F.5: Convergence of FedLEASE