Heterogeneous L o RA for Federated Fine-tuning of On-Device Foundation Models (2024)

Abstract: 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 on-device. 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 HetLoRA, a feasible and effective federated fine-tuning 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, HetLoRA has enhanced computation and communication efficiency compared to full fine-tuning making it more feasible for the edge.

## introduction the emerging foundation models (fms), easily ranging from few to hundreds of billions of parameters (gemini @xcite @xcite @xcite @xcite @xcite , have remarkable zero/few shot learning capabilities. they often perform surprisingly well on a variety of tasks such as text/image generation, language translation, and conversing in natural language with fine-tuning. such standard fms, however, demand costly resources for directly fine-tuning their entire parameter space. to tackle this issue, many works have proposed different parameter-efficient fine-tuning (peft) methods such as prompt tuning @xcite , utilizing adapters @xcite , or lowrank adaptation (lora) of the original model @xcite which freezes the original pre-trained parameters of the fm and train additional, smaller number of parameters instead. unfortunately, such peft methods rely on the assumptions that i) fms are deployed to and trained with the data of a single client for adaptation to the downstream task and ii) the client has enough resources to even fit a standard fm of hundred billion size for, at least, inference. in practice, more often than not we are interested in fine-tuning fms for on-device private data that is distributed across multiple devices, and these clients at the edge rarely have such resources to even fit standard fms for inference. one of the examples is to use fms for downstream tasks that requires training with private, non-sharable data such as medical or law-related documents distributed across clients @xcite @xcite . in such cases, fine-tuning of the fms needs to be brought to the edge that has limited system resources. federated fine-tuning & odfms. we consider such federated fine-tuning scenarios, where we train a set of parameters collaboratively across clients to obtain a global set of parameters that can be plugged in to the fm for the targeted downstream task. note that federated fine-tuning is orthogonal to personalization of fms in federated learning (fl) that has been explored in other previous work @xcite , which aims to train parameters that perform well for individual clients rather than general downstream tasks. we also define on-device fms (odfms) as models with few billion parameters at max that are able to fit into memory on limited capacity clients considering current hardwares. challenges. federated fine-tuning of odfms entails two major challenges non-present in neither standard peft of fms nor federated training of non-fms. first, odfms poorly generalize to downstream tasks. unlike odfms, standard fms have their zero/few-shot learning capability by their large parameter space trained on massive data. however, as we show in second, devices have limited and heterogeneous system capabilities and data distributions @xcite @xcite . without suitable peft methods that flexibly adapts to such heterogeneity across devices, we would only have limited performance output from federated fine-tuning of odfms. for instance, with homogeneous rank deployment in lora that is agnostic of system heterogeneity, the client with the least system resource becomes the bottleneck, forcing a smaller rank to be deployed across all clients despite the other clients being able to handle higher ranks (see t 1 9 p j e @xmath0 with lora modules , contributions. in our work, we push the limits of federated fine-tuning of odfms by proposing hetlora that is aware of heterogeneity residing at the edge. our main flagship results (in parameter-efficient fine tuning. most peft methods either train a subset of parameters within the existing fm whilst other parameters are freezed or introduce an additional set of trainable parameters whilst keeping the original fm freezed. for the former, methods such as head or bias finetuning @xcite @xcite has been explored, and for the latter, methods such as adapters @xcite , prompt @xcite or prefixtuning (li and liang, 2021), and low-rank approximation @xcite has been proposed. while these number of methods has been proven to perform as well as full model fine-tuning with just a few number of parameters for the centralized setting, it has not been thoroughly explored how these methods would perform for a much smaller fm such as odfms, in the decentralized setting where clients' system-capacities and data can be heterogeneous and much limited. federated fine-tuning. several recent work has proposed to combine the peft methods devised for the centralized setting to fl @xcite such as training prompts or adapters collaboratively with fl @xcite @xcite @xcite . another line of work has proposed to perform a few-shot or nearly zero-shot training of fms with fl for improved communication-efficiency @xcite . however, these work either overlooks that most devices do not have the resource to fit standard fms @xcite even for inference or does not consider the heterogeneous system capacities of the clients. it is detrimental to consider these factors since fms that actually fits to the devices in fl are much smaller, making them weaker in the general intelligence capabilities, and also heterogeneous system capacities can prohibit deploying same sized peft parameters across clients. lora and fl. there has been a number of variations of lora @xcite @xcite ) from its first proposal @xcite . however, only a few number of recent work has looked in to using lora for fl. in @xcite , the importance of the initialization for the lora modules is evaluated where they propose to train the lora modules with fl and then perform singular value decomposition (svd) to gain a good initialization of the lora modules. however, the training process of lora itself is not altered to adapt to heterogeneous system capabilieis of devices. another recent work @xcite has evaluated lora in the context of personalized fl, but other than applying lora to personalization, the lora method itself is, again, not changed. to the best of our knowledge, our work is the first method to consider both data and system heterogeneity for improving training speed, communication/computation efficiency, and final performance of odfm's federated fine-tuning. contemporary approaches @xcite have proposed to use svd to decompose the full matrix to deploy heterogeneous ranks across clients. we compare our proposed method with using svd for heterogeneous lora in our work and show that the performance is underwhelming compared to our proposed heterogeneous lora. ## 5 discussions and concluding remarks in our work, we investigated federated fine-tuning for odfms by proposing hetlora that enables utilizing the power of fms at the edge by taking device system and data heterogeneity into consideration. hetlora is not only practical but also achieves better training speed, communication/computation efficiency, and final performance compared to other relevant baselines such as svdbased lora, homogeneous lora, or full finetuning. our findings in this work opens up several questions worth investigating. for instance, if the settings allow us to assign specific ranks to clients what will be the effective way to assign the ranks across clients for better convergence and performance? another important next step of our work includes pursuing the theoretical convergence and generalization of heterogeneous lora.