TDR: Task-Decoupled Retrieval with Fine-Grained LLM Feedback for In-Context Learning

TL;DR

TDR introduces a task-decoupled retrieval framework that leverages fine-grained LLM feedback to enhance example selection for in-context learning, leading to improved performance across diverse NLP tasks.

Contribution

The paper presents a novel TDR framework that decouples task-specific retrieval from cross-task data, utilizing LLM feedback to guide high-quality example retrieval for ICL.

Findings

TDR achieves state-of-the-art results on 30 NLP tasks.

TDR consistently improves ICL performance across various datasets.

The method is plug-and-play and compatible with multiple LLMs.

Abstract

In-context learning (ICL) has become a classic approach for enabling LLMs to handle various tasks based on a few input-output examples. The effectiveness of ICL heavily relies on the quality of these examples, and previous works which focused on enhancing example retrieval capabilities have achieved impressive performances. However, two challenges remain in retrieving high-quality examples: (1) Difficulty in distinguishing cross-task data distributions, (2) Difficulty in making the fine-grained connection between retriever output and feedback from LLMs. In this paper, we propose a novel framework called TDR. TDR decouples the ICL examples from different tasks, which enables the retrieval module to retrieve examples specific to the target task within a multi-task dataset. Furthermore, TDR models fine-grained feedback from LLMs to supervise and guide the training of the retrieval module, which helps to retrieve high-quality examples. We conducted extensive experiments on a suite of 30 NLP tasks, the results demonstrate that TDR consistently improved results across all datasets and achieves state-of-the-art performance. Meanwhile, our approach is a plug-and-play method, which can be easily combined with various LLMs to improve example retrieval abilities for ICL. The code is available at https://github.com/Nnn-s/TDR.