Retrieve Anything To Augment Large Language Models (2310.07554v2)

Abstract: LLMs face significant challenges stemming from their inherent limitations in knowledge, memory, alignment, and action. These challenges cannot be addressed by LLMs alone, but should rely on assistance from the external world, such as knowledge base, memory store, demonstration examples, and tools. Retrieval augmentation stands as a vital mechanism for bridging the gap between LLMs and the external assistance. However, conventional methods encounter two pressing issues. On the one hand, the general-purpose retrievers are not properly optimized for the retrieval augmentation of LLMs. On the other hand, the task-specific retrievers lack the required versatility, hindering their performance across the diverse retrieval augmentation scenarios. In this work, we present a novel approach, the LLM-Embedder, which comprehensively supports the diverse retrieval augmentation needs of LLMs with one unified embedding model. Training such a unified model is non-trivial, as various retrieval tasks aim to capture distinct semantic relationships, often subject to mutual interference. To address this challenge, we systematically optimize our training methodology. This includes reward formulation based on LLMs' feedback, the stabilization of knowledge distillation, multi-task fine-tuning with explicit instructions, and homogeneous in-batch negative sampling. These optimization strategies contribute to the outstanding empirical performance of the LLM-Embedder. Notably, it yields remarkable enhancements in retrieval augmentation for LLMs, surpassing both general-purpose and task-specific retrievers in various evaluation scenarios. Our checkpoint and source code are publicly available at https://github.com/FlagOpen/FlagEmbedding.

Retrieve Anything To Augment LLMs

LLMs are pivotal in advancing general artificial intelligence, yet they are encumbered by inherent limitations including deficiencies in knowledge storage, memory retention, alignment fidelity, and action coordination. These constraints necessitate reliance on external aids such as knowledge bases and toolkits. Retrieval augmentation is identified as a crucial method to bridge the gap between LLMs and necessary external resources. However, existing retrieval techniques face significant hurdles: general-purpose retrievers are not sufficiently optimized for the specific needs of LLMs, while task-specific retrievers often lack the necessary versatility to operate across diverse scenarios.

This paper presents an innovative solution through the development of the LLM-Embedder, a unified embedding model designed to address the comprehensive retrieval augmentation needs of LLMs. Training such a model involves overcoming the challenge of mutual interference among varied retrieval tasks. Key optimization strategies are employed, including reward formation derived from LLM feedback, stabilization of knowledge distillation, multi-task fine-tuning with explicit instructions, and the utilization of homogeneous in-batch negative sampling. These collectively enhance the empirical performance of LLM-Embedder, enabling it to surpass both general-purpose and task-specific retrievers in various evaluation scenarios.

Key Contributions

1. LLM-Embedder Model: Introduction of the LLM-Embedder, designed to seamlessly integrate LLMs with external resources, marking the first comprehensive support system for all key facets of retrieval augmentation.
2. Systematic Optimization: Robust optimization strategies across multiple dimensions: reward formulation, knowledge distillation, instruction-based fine-tuning, and negative sampling, ensuring the model's effectiveness.
3. Empirical Validation: Extensive experiments demonstrate the LLM-Embedder's superiority over existing embedding models, significantly enhancing the retrieval augmentation impact on critical aspects of LLMs like knowledge enhancement, in-context learning, and memory modeling.

Experimental Evaluations

The performance of LLM-Embedder is rigorously validated through comprehensive experiments:

• Knowledge Enhancement: Incorporating external knowledge significantly improves question answering accuracy in both MMLU and PopQA datasets, emphasizing the retrieval system's pivotal role.
• In-Context Learning: The model improves the LLMs' instruction-following ability across diverse datasets by effectively leveraging retrieved examples.
• Long-Context Modeling: Demonstrates the model's ability to improve memory retention for language generation tasks, outperforming alternative retrieval methods.
• Tool Learning and Conversational Search: LLM-Embedder excels in retrieving appropriate tools and relevant conversational information, showcasing its versatile applicability.

Conclusion and Future Directions

The LLM-Embedder stands out as a unified solution that meets diverse retrieval needs of LLMs, enhancing their performance across multiple application domains. This work marks an advanced step in reconciling specialized retrieval tasks under a single framework. Looking forward, exploration into further optimizing the retrieval-augmentation mechanism with novel training paradigms, such as reinforcement learning and self-supervised methods, could open new avenues for even more efficient and effective LLM integration with external resources. This advancement fosters both theoretical insights and practical efficiencies in leveraging LLMs' full potential.