Large Language Models Know What Makes Exemplary Contexts

TL;DR

This paper introduces a unified, reinforcement learning-based framework enabling large language models to self-select, rank, and optimize in-context examples, significantly improving their few-shot learning performance.

Contribution

It proposes a parameter-efficient retrieval method for LLMs to self-optimize demonstration selection and ordering, enhancing in-context learning without extensive retraining.

Findings

Improved ICL performance with self-selected demonstrations

Effective identification of representative and diverse examples

Validation of the method's effectiveness through experiments

Abstract

In-context learning (ICL) has proven to be a significant capability with the advancement of Large Language models (LLMs). By instructing LLMs using few-shot demonstrative examples, ICL enables them to perform a wide range of tasks without needing to update millions of parameters. This paper presents a unified framework for LLMs that allows them to self-select influential in-context examples to compose their contexts; self-rank candidates with different demonstration compositions; self-optimize the demonstration selection and ordering through reinforcement learning. Specifically, our method designs a parameter-efficient retrieval head that generates the optimized demonstration after training with rewards from LLM's own preference. Experimental results validate the proposed method's effectiveness in enhancing ICL performance. Additionally, our approach effectively identifies and selects the most representative examples for the current task, and includes more diversity in retrieval.