A Multi-Task Embedder For Retrieval Augmented LLMs

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 notably improves the LLM's performances in various downstream tasks, and outperforms both general and taskspecific retrievers with a substantial advantage. et al., 2022). Many of the challenges can be traced 045 back to the inherent limitations of LLMs in terms 046 of knowledge, memory, and action. Specifically, 047 LLMs cannot internalize the vast and constantly 048 changed world knowledge due to their finite and 049 static parameters. LLMs are incapable of memo-050 rizing and utilizing long-term information because 051 of the limited context length. Finally, LLMs re-052 quire manually in-context examples and tools to 053 accomplish complex real-world tasks. 054 Retrieval augmentation stands as a vital mech-055 anism to address these inherent limitations of the 056 LLM. It brings in useful information from exter-057 nal sources, such as knowledge, memory pieces, 058 in-context examples, and tools, which substantially 059 enhances the LLM for the generation of desired 060 outputs (Gao et al., 2023). The embedding model 061 (a.k.a. embedder) is a critical part of retrieval aug-062 mentation, which bridges the LLM's information 063 needs with external sources. The existing embed-064 ding models can be briefly partitioned into two 065 categories. One is the general-purpose embedders, 066 which aim to be universally applicable for various 067