Heterogeneous LoRA for Federated Fine-tuning of On-Device Foundation Models

TL;DR

This paper introduces HetLoRA, a novel federated fine-tuning method for on-device foundation models that uses heterogeneous low-rank approximations to address data and system heterogeneity, improving convergence and efficiency.

Contribution

HetLoRA allows heterogeneous LoRA ranks across devices, enabling efficient federated fine-tuning of foundation models with improved convergence and performance.

Findings

HetLoRA outperforms homogeneous LoRA in convergence speed.

HetLoRA achieves better final performance on federated tasks.

HetLoRA enhances computational efficiency compared to full fine-tuning.

Abstract

Foundation models (FMs) adapt well to specific domains or tasks with fine-tuning, and federated learning (FL) enables the potential for privacy-preserving fine-tuning of the FMs with on-device local data. For federated fine-tuning of FMs, we consider the FMs with small to medium parameter sizes of single digit billion at maximum, referred to as on-device FMs (ODFMs) that can be deployed on devices for inference but can only be fine-tuned with parameter efficient methods. In our work, we tackle the data and system heterogeneity problem of federated fine-tuning of ODFMs by proposing a novel method using heterogeneous low-rank approximations (LoRAs), namely HetLoRA. First, we show that the naive approach of using homogeneous LoRA ranks across devices face a trade-off between overfitting and slow convergence, and thus propose HetLoRA, which allows heterogeneous ranks across client devices and efficiently aggregates and distributes these heterogeneous LoRA modules. By applying rank self-pruning locally and sparsity-weighted aggregation at the server, HetLoRA combines the advantages of high and low-rank LoRAs, which achieves improved convergence speed and final performance compared to homogeneous LoRA. Furthermore, HetLoRA offers enhanced computation efficiency compared to full fine-tuning, making it suitable for federated fine-tuning across heterogeneous devices.