Adaptive Rank Allocation for Federated Parameter-Efficient Fine-Tuning of Language Models

TL;DR

FedARA introduces an adaptive rank allocation framework for Federated Parameter-Efficient Fine-Tuning of language models to address non-IID data and fixed-parameter configurations. It combines truncated SVD adaptation, dynamic rank masks with server arbitration, and rank-based module pruning to reduce communication and local computation on edge devices. Empirical results across multiple datasets, models, and edge platforms show average accuracy gains of around 7–9% under non-IID data, and communication reductions of about 2.4x, alongside substantial improvements in training time and energy consumption. These findings demonstrate the practical viability of FedPEFT for efficient federated fine-tuning in resource-constrained settings and set a path for scaling the approach to larger models and tasks.

Abstract

Pre-trained Language Models (PLMs) have demonstrated their superiority and versatility in modern Natural Language Processing (NLP), effectively adapting to various downstream tasks through further fine-tuning. Federated Parameter-Efficient Fine-Tuning (FedPEFT) has emerged as a promising solution to address privacy and efficiency challenges in distributed training for PLMs on resource-constrained local devices. However, our measurements reveal two key limitations of FedPEFT: heterogeneous data across devices exacerbates performance degradation of low-rank adaptation, and a fixed parameter configuration results in communication inefficiency. To overcome these limitations, we propose FedARA, a novel adaptive rank allocation framework for federated parameter-efficient fine-tuning of language models. Specifically, FedARA employs truncated Singular Value Decomposition (SVD) adaptation to enhance similar feature representation across clients, significantly mitigating the adverse effects of data heterogeneity. Subsequently, it utilizes dynamic rank allocation to progressively identify critical ranks, effectively improving communication efficiency. Lastly, it leverages rank-based module pruning to automatically remove inactive modules, steadily reducing local computational cost and memory usage in each federated learning round. Extensive experiments show that FedARA consistently outperforms baselines by an average of 6.95% to 8.49% across various datasets and models under heterogeneous data while significantly improving communication efficiency by 2.40$ \times$. Moreover, experiments on various edge devices demonstrate substantial decreases in total training time and energy consumption by up to 48.90% and 46.95%, respectively.

Figures (17)

• Figure 1: The complete training process of Federated Parameter-Efficient Fine-Tuning (FedPEFT), where local updates utilize Low-Rank Adaptation (LoRA), while keeping the pre-trained model frozen.
• Figure 2: The preliminary measurement results of FedPEFT in a homogeneous FL setting: (a) Comparison of overall training time on different edge devices (BERT; 20News; Rank=12; Batch size=4). (b) Accuracy of FedPEFT under IID and non-IID scenarios (DistilBERT; Rank=24; Adapter-h (A-h) size=108; Adapter-p (A-p) size=216). (c) Impact of inserted positions on accuracy (DistilBERT; Rank=24). (d) Communication and computation time percentage across models and devices using FedLoRA (News Category; Rank=12; Batch size=4).
• Figure 3: System architecture of FedARA, illustrating the interaction between the server and client, including ClientsSel (Client Selection), CommPru (Communication Pruning), FedArb (Federated Arbitration), MaskGen (Rank Mask Generation), and RankDet (Local Rank Detection).
• Figure 4: Illustration of truncated SVD adaptation, where a diagonal matrix $E$ is inserted between the LoRA matrices.
• Figure 5: Magnitude (left) and direction (right) discrepancies between global and local models under different FL structures. FedLoRA and FedSVD correspond to fine-tuned models utilizing LoRA and truncated SVD adaptation, respectively, with a rank of 12.