Learning to Retrieve Iteratively for In-Context Learning
This work addresses the challenge of efficient exemplar selection for in-context learning in semantic parsing, offering a novel method that improves performance and generalization, though it is incremental as it builds on existing dense retrieval techniques.
The paper tackles the combinatorial optimization problem of selecting optimal in-context learning exemplars for semantic parsing by introducing an iterative retrieval framework that uses reinforcement learning with LLM feedback. The result is a stateful iterative retriever that outperforms previous methods on datasets like CalFlow, TreeDST, and MTOP, using only 4M additional parameters and generalizing across different LLMs.
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.