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

Foundation models (FMs) adapt surprisingly well to downstream tasks with fine-tuning.However, their colossal parameter space prohibits their training on resource-constrained edge-devices.For federated fine-tuning, we need to consider the smaller FMs of few billion parameters at most, namely on-device FMs (ODFMs), which can be deployed ondevice.Federated fine-tuning of ODFMs has unique challenges non-present in standard fine-tuning: i) ODFMs poorly generalize to downstream tasks due to their limited sizes making proper fine-tuning imperative to their performance, and ii) devices have limited and heterogeneous system capabilities and data that can deter the performance of fine-tuning.Tackling these challenges, we propose HET-LORA, a feasible and effective federated finetuning method for ODFMs that leverages the system and data heterogeneity at the edge.HETLORA allows heterogeneous LoRA ranks across clients for their individual system resources, and efficiently aggregates and distributes these LoRA modules in a data-aware manner by applying rank self-pruning locally and sparsity-weighted aggregation at the server.It combines the advantages of high and low-rank LoRAs, achieving improved convergence speed and final performance compared to homogeneous LoRA.Furthermore, HET-LORA has enhanced computation and communication efficiency compared to full finetuning making it more feasible for the edge.