FedHFT: Efficient Federated Finetuning with Heterogeneous Edge Clients

TL;DR

FedHFT tackles privacy-preserving, personalized finetuning of pretrained LLMs in federated settings with data and resource heterogeneity. It introduces a mixture of masked adapters and a bi-level optimization pipeline that clusters clients by weight-update similarity using Gaussian Mixture Modeling and communicates only important adapter dimensions via Fisher-information masking, while preserving a global initialization through low-rank updates $\Delta W \sim B A$. Empirically, FedHFT yields up to $\sim$3.1x memory savings and $\sim$136.9x communication reductions with competitive or superior accuracy across multiple tasks and architectures (e.g., BERT, T5) compared to FedAvg and other heterogeneous-FL baselines. The approach enables resource-constrained edge devices to participate in collaborative finetuning with improved personalization and efficiency, albeit at some central-server workload for managing cluster adapters and assignments.

Abstract

Fine-tuning pre-trained large language models (LLMs) has become a common practice for personalized natural language understanding (NLU) applications on downstream tasks and domain-specific datasets. However, there are two main challenges: (i) limited and/or heterogeneous data for fine-tuning due to proprietary data confidentiality or privacy requirements, and (ii) varying computation resources available across participating clients such as edge devices. This paper presents FedHFT - an efficient and personalized federated fine-tuning framework to address both challenges. First, we introduce a mixture of masked adapters to handle resource heterogeneity across participating clients, enabling high-performance collaborative fine-tuning of pre-trained language model(s) across multiple clients in a distributed setting, while keeping proprietary data local. Second, we introduce a bi-level optimization approach to handle non-iid data distribution based on masked personalization and client clustering. Extensive experiments demonstrate significant performance and efficiency improvements over various natural language understanding tasks under data and resource heterogeneity compared to representative heterogeneous federated learning methods.

Figures (8)

• Figure 1: System architecture of FedHFT. At each round, the central server computes the adapter weights for each client based on cluster assignment scores. After receiving the adapter weights, participating edge clients initiate local finetuning over their local database and masks based on Fisher information before uploading them back. Lastly, the central server receives the updated adapter weights, updates the adapter weights of each cluster and cluster assignment scores.
• Figure 2: Three main stages in our framework: (i) we perform federated finetuning without clustering for $T_w$ warmup rounds, (ii) we continue the process with a mixture of adapters, where each adapter tracks the updates of the corresponding cluster, (iii) each client can perform a final finetuning on its local data using the initialization point ($\theta+\sum_{c} p_{kc} B_c^{(T)}A_C^{(T)}$) before deployment for inference.
• Figure 3: Performance of various resource-efficient finetuning methods and FedAVG as baseline under data heterogeneity with BERT-base, $K=50$, and $\alpha \in [1, 5, 20, 50]$. Lower values of $\alpha$ indicate higher data heterogeneity among clients.
• Figure 4: Performance of FedHFT on various resource heterogeneity settings with BERT-base, $K=50$, $\alpha=5$ and the ratio of clients with adapter rank $r \in [4, 8, 16, 32]$. We plot the average performance of each client group corresponding to the same adapter rank with FedHFT (grey scattered stars) and the average performance of all clients for FedHFT (blue) and other resource-efficient techniques LoRA (yellow) and HetLoRA (green).
• Figure 5: Performance of various federated finetuning methods with BERT-base, $K=50$, $\alpha=5$ under varying client availability ratios $r_a \in [10\%, 25\%, 50\%, 100\%]$.