Learning to Retrieve Iteratively for In-Context Learning

TL;DR

This paper presents an iterative retrieval framework optimized via reinforcement learning to improve in-context learning for semantic parsing, outperforming previous methods and generalizing across different large language models.

Contribution

It introduces a learned iterative retrieval method that enhances exemplar selection in in-context learning, with minimal additional parameters, and demonstrates superior performance and generalization.

Findings

Outperforms previous exemplar selection methods on semantic parsing datasets.

Requires only 4M additional parameters for state encoding.

Generalizes across different inference LLMs beyond training models.

Abstract

We introduce iterative retrieval, a novel framework that empowers retrievers to make iterative decisions through policy optimization. Finding an optimal portfolio of retrieved items is a combinatorial optimization problem, generally considered NP-hard. This approach provides a learned approximation to such a solution, meeting specific task requirements under a given family of large language models (LLMs). We propose a training procedure based on reinforcement learning, incorporating feedback from LLMs. We instantiate an iterative retriever for composing in-context learning (ICL) exemplars and apply it to various semantic parsing tasks that demand synthesized programs as outputs. By adding only 4M additional parameters for state encoding, we convert an off-the-shelf dense retriever into a stateful iterative retriever, outperforming previous methods in selecting ICL exemplars on semantic parsing datasets such as CalFlow, TreeDST, and MTOP. Additionally, the trained iterative retriever generalizes across different inference LLMs beyond the one used during training.