LoraRetriever: Input-Aware LoRA Retrieval and Composition for Mixed Tasks in the Wild

TL;DR

LoraRetriever is proposed, a retrieve-then-compose framework that adaptively retrieves and composes multiple LoRAs according to the input prompts, highlighting its practical effectiveness and versatility.

Abstract

Low-Rank Adaptation (LoRA) provides an effective yet efficient solution for fine-tuning large language models (LLM). The modular and plug-and-play nature of LoRA enables the integration of diverse domain-specific LoRAs to enhance the capabilities of LLMs. Previous research on exploiting multiple LoRAs either focuses on specific isolated downstream tasks or fixes the selection of LoRAs during training. However, in real-world scenarios, LLMs receive diverse prompts covering different tasks, and the pool of candidate LoRAs is often dynamically updated. To bridge this gap, we propose LoraRetriever, a retrieve-then-compose framework that adaptively retrieves and composes multiple LoRAs according to the input prompts. LoraRetriever contains three main components: firstly, identifying and retrieving LoRAs relevant to the given input; secondly, formulating strategies for effectively integrating the retrieved LoRAs; and thirdly, developing efficient batch inference to accommodate heterogeneous requests. Experimental results indicate that LoraRetriever consistently outperforms the baselines, highlighting its practical effectiveness and versatility.

Figures (7)

• Figure 1: Illustration of serving multiple LoRAs within a dynamically updated LoRA pool for mixed-task scenarios. a) LoRAs from various domains and tasks aimed at enhancing specific capabilities of the LLM can be uploaded to or updated to the LoRA pool. b) The multi-LoRA serving framework aims to leverage the plug-and-play nature of LoRAs to offer comprehensive services. c) The downstream tasks, presented in a mixed-task form, require personalized expert routing.
• Figure 2: The LoraRetriever Framework. This framework, equipped with a pool of candidate LoRAs from various domains/tasks, is designed to offer personalized services tailored to the input provided. It begins by executing an input-aware LoRA retrieval process aimed at identifying LoRAs corresponding to tasks analogous to the input (§\ref{['sec:lora_retrieval']}). Subsequently, it employs a specialized LoRA composition mechanism to efficiently utilize the retrieved LoRAs (§\ref{['sec:lora_composition']}). By constructing a LoRA mapping matrix for batch inputs, the framework facilitates effective batch inference (§\ref{['sec:batch']}).
• Figure 3: LoRA embedding similarity heatmap. Tasks from the same domain are grouped in square brackets.
• Figure 4: The left figure shows the performance of LoraRetriever varying the number of LoRAs. The right figure shows the performance of Throughput varying the batch size.
• Figure 5: Showcasing How the LoRARetrieval Framework Employs Multiple LoRAs for Cooperative Problem Solving.