Learning to Retrieve In-Context Examples for Large Language Models
This addresses the bottleneck of example selection in in-context learning for LLM users, offering an incremental improvement over existing retrieval methods.
The paper tackles the problem of selecting high-quality in-context examples for large language models (LLMs) by proposing a framework to train dense retrievers, resulting in significant performance enhancements on 30 tasks and generalization to unseen tasks.
Large language models (LLMs) have demonstrated their ability to learn in-context, allowing them to perform various tasks based on a few input-output examples. However, the effectiveness of in-context learning is heavily reliant on the quality of the selected examples. In this paper, we propose a novel framework to iteratively train dense retrievers that can identify high-quality in-context examples for LLMs. Our framework initially trains a reward model based on LLM feedback to evaluate the quality of candidate examples, followed by knowledge distillation to train a bi-encoder based dense retriever. Our experiments on a suite of $30$ tasks demonstrate that our framework significantly enhances in-context learning performance. Furthermore, we show the generalization ability of our framework to unseen tasks during training. An in-depth analysis reveals that our model improves performance by retrieving examples with similar patterns, and the gains are consistent across LLMs of varying sizes. The code and data are available at https://github.com/microsoft/LMOps/tree/main/llm_retriever .