A Multi-Task Embedder For Retrieval Augmented LLM s (2024)

Abstract: LLMs confront inherent limitations in terms of its knowledge, memory, and action. The retrieval augmentation stands as a vital mechanism to address these limitations, which brings in useful information from external sources to augment the LLM. However, existing retrieval methods encounter two pressing issues. On one hand, the general retrievers are not properly optimized for retrieval augmentation hence exhibit limited effectiveness; on the other hand, the task-specific retrievers excel in the targeted retrieval augmentation scenario, while lack the versatility to handle diverse scenarios. In this work, we propose LLM-Embedder for the unified support of diverse retrieval augmentation scenarios. Our method presents three technical contributions. Firstly, we introduce a new reward formulation, namely rank-aware reward. It exploits the ranking position of the desired output among N sampled outputs from the LLM, which leads to fine-grained and robust computation of reward from the LLM’s feedback. Secondly, we design a novel distillation objective, called graded distillation. It incorporates both the absolute value and the relative order of the reward for more sufficient utilization of the LLM’s feedback. Thirdly, we systematically optimize the multi-task learning, which effectively unifies the multiple retrieval functionalities into one model. In our experiment, LLM-Embedder substantially improves the LLM’s performances in various downstream tasks, while introducing superior retrieval augmentation’s effect over both general and task-specifc retrievers. Our data, code, and model have been released at https://github.com/FlagOpen/FlagEmbedding.

## introduction large language models (llms) present a unified foundation to support general artificial intelligence applications @xcite @xcite . despite the substantial improvement over the last-gen methods, llms still face many severe problems, such as hallucination @xcite , limited memory @xcite , misfollowing of instructions @xcite . many of the challenges can be traced back to the inherent limitations of llms in terms of knowledge, memory, and action. specifically, llms cannot internalize the vast and constantly changed world knowledge due to their finite and static parameters. llms are incapable of memorizing and utilizing long-term information because of the limited context length. finally, llms require manually in-context examples and tools to accomplish complex real-world tasks. retrieval augmentation stands as a vital mechanism to address these inherent limitations of the llm. it brings in useful information from external sources, such as knowledge, memory pieces, in-context examples, and tools, which substantially enhances the llm for the generation of desired outputs @xcite . the embedding model (a.k.a. embedder) is a critical part of retrieval augmentation, which bridges the llm's information needs with external sources. the existing embed-ding models can be briefly partitioned into two categories. one is the general-purpose embedders, which aim to be universally applicable for various retrieval tasks @xcite @xcite . despite their popularity, they are not properly optimized for retrieval augmentation, and are thus prone to an inferior effectiveness in the corresponding task. the other one is the task-specific embedders, which are tailored for one specific retrieval augmentation scenario, e.g., knowledge retrieval @xcite and example retrieval @xcite . however, these methods lack versatility across different scenarios. as the llms require assistance from diverse external sources in solving real-world problems, it becomes imperative to develop an effective and versatile embedding model to support the diverse retrieval augmentation needs. in this paper, we present llm-embedder, a unified embedding model to support a broad range of retrieval augmentation scenarios, including knowledge retrieval, memory retrieval, example retrieval, and tool retrieval. training such a versatile embedding model presents multiple challenges in terms of 1) how to learn from the llm, and 2) how to harmonize different retrieval tasks. in llm-embedder, the following technical contributions are presented. to summarize, llm-embedder stands as a pioneering work for the uniform support of the diverse retrieval augmentation scenarios of llms. it makes threefold technical contributions, and brings valuable inspirations on how to learn from llm's feedback and how to harmonize different retrieval tasks. in our experiment, llm-embedder achieves a superior performance, where it notably improves the llm's performance in a variety of downstream tasks. meanwhile, its retrieval augmentation's effect is superior to both general and task-specific retrieval methods. our model and code will be publicly available to facilitate future research. ## llm-embedder in this section, we will present the retrieval augmentation scenarios with llm-embedder ( §3.1), and introduce its training methodology ( §3.2). four seasons. seven seasons. rick&morty has five ... how many seasons is rick&morty ⋯ ⋯ rick and morty is an american ... it has seven seasons until now ... rank 𝑟 𝑎 = 𝟖 reward 𝑅 𝐶 𝑖 = 𝟕 user input 𝑈 retrieval candidate 𝐶 𝑖 seven seasons. six seasons. rick&morty is a ... how many seasons is rick&morty ⋯ ⋯ user input 𝑈 rank 𝑟 𝑏 = 𝟏 ranked outputs ranked outputs figure 2: the rank-aware reward for each retrieval candidate. it measures the improvement of the rank of the desired output among multiple sampled outputs. ## experiment the experimental studies aim to investigate three research questions. rq 1. can llm-embedder support the llm's diverse retrieval augmentation need? ( §4.2) rq 2. what is llm-embedder's impact on each retrieval augmentation scenario? ( §4.3) rq 3. what is the individual contribution of each technique in llm-embedder? ( §4.4) 4.1 settings ## conclusion in this work, we present llm-embedder, a unified embedding model to support the llm's diverse retrieval augmentation needs, including knowledge retrieval, memory retrieval, example retrieval, and tool retrieval. we propose three key techniques to facilitate the training of llm-embedder, spanning from reward formulation, distillation objective, and multi-task learning recipe. our experiments show llm-embedder's empirical advantages over both general and task-specific embedding models across all evaluation scenarios. this highlights its effectiveness as a foundational building block to support the retrieval augmentation of the llm.