Active Example Selection for In-Context Learning

In plain words

A language model can often learn a task from a few solved examples placed in its prompt, with no retraining. How well this works depends heavily on which examples are picked. Two sets of the same size can give very different accuracy, and simply reordering a set does not reliably fix that. The paper asks a different question: which examples should go into the prompt?

The authors treat this like active learning. You have a pile of unlabeled texts. You choose one, reveal its label, add the pair to the prompt, and repeat, up to four times. A small neural network learns the choosing policy with reinforcement learning. Each choice is rewarded by how much it raises accuracy on a held-out set of labeled examples. The network is trained on logged random choices, not on live runs of the policy.

The learned choices help on a task the policy saw during training, and help less on a task it never saw. The gains shrink, and become mixed, when the same examples are moved to larger GPT-3 models. The pitfalls section lists where the paper and its code disagree and what the reported numbers do not show.

TL;DR

• The paper frames choosing in-context demonstrations as a sequential decision with four steps over a pool of unlabeled examples §3.1.
• A three-layer MLP Q-network is trained offline with conservative Q-learning on logs of 2,000 random episodes §3.2 §4.1 app. A app. B.
• The reward is the marginal gain in accuracy on a labeled reward set, scored with a calibrated GPT-2 medium model §3.2 §4.1.
• Author-reported: picking from seen examples gives 71.4 average accuracy against 59.6 for random tab. 3. Picking from 1,000 unseen-task examples gives 65.4 tab. 4.
• Transfer to GPT-3 is mixed: small gains on Ada, negative or small results on Babbage and Curie tab. 7 app. C.2.

Problem

With in-context learning, a model is shown a few labeled examples and predicts the label of a new input. The paper shows that the result depends strongly on which demonstrations are chosen. For GPT-2, the best and worst of 30 random four-example sets differ by more than 30% on every task. For three of the four tasks, the worst sets fall below random guessing §2.2 fig. 1.

Two common fixes do not remove the spread. Reordering the examples helps only when the set already has a good ordering, and even the best of the 4! = 24 orders leaves wide gaps §2.2 fig. 1. Calibration raises the average for GPT-2 but does not reduce the variance §2.2 fig. 2. For GPT-3 models, calibration does reduce the variance §2.2 tab. 2.

The paper's setting assumes three things: a pool of unlabeled examples from which labels are revealed one at a time; a labeled reward set used to score choices during training; and a calibrated left-to-right model §3.1 §4.1. Under these assumptions, labels are revealed only for the examples that get picked §3.1.

Method

Overview

The prompt is a sequence of up to k = 4 labeled examples §2.1 §3.1. The state is the prefix of picked pairs. The action set is every unlabeled example plus a special end-of-prompt action, written ⊥ §3.1. Choosing an example reveals its label and appends the pair to the prefix. The episode ends when the budget is used, or when ⊥ is chosen §3.1.

The reward for choosing action a in state s is the change in an accuracy score f: the score of the model with the new prefix, minus the score without it §3.2. Summed along an episode, these gains telescope to the final score minus the zero-shot score. The paper treats this as a shaped reward that leaves the optimal policy unchanged, and uses a discount factor of 1 §3.2 §3.2 footnote.

The paper uses accuracy on the reward set as f §3.1. The features of a candidate are computed without its label. They come from the model's predicted label probabilities given the current prefix, plus the entropy of that prediction §3.2.

Features

• State: the number of examples already picked, |s| §3.2.
• Action: the predicted label probabilities for the candidate given the prefix, and the entropy of that prediction. The paper names these two features, and the shipped training config adds more (see Paper vs code) §3.2.

Q-network

The Q-function is an MLP applied to the concatenated state and action features, Q(s, a) = MLP([s ∥ a]). The paper uses a three-layer MLP with hidden size 16 §3.2 app. B. The paper does not say how the end-of-prompt action is represented; the code uses a learned vector (see Paper vs code).

Training: offline Q-learning with a conservative penalty

Training data are collected once, before learning. A random policy runs for 2,000 episodes. Each episode picks four of the 100 training examples at random, and the transitions and rewards are stored §3.2 §4.1. The learner then uses only this logged data. The paper uses conservative Q-learning (CQL) to limit overestimation of Q-values for actions that are rare in the data §3.2 app. A.

The learner uses a separate target network and a replay buffer §3.2. The conservative penalty adds, for each state, the log-sum-exp of Q over all actions minus the Q-value of the action in the data, weighted by alpha app. A.

# Offline update with a conservative penalty  [§3.2] [app. A]
# Input:  a batch of logged transitions (B of them)
# Output: one gradient step on the Q-network

def optimize_step(batch):
    # Q-values of every candidate plus the end action  [§3.2]
    # (B, A + 1)
    q_all = policy_net(states, actions_with_end)
    # Q-value of the candidate that was actually picked  [§3.2]
    q_taken = q_all[row, picked_index]
    # Bellman target, discount 1: reward plus the best next value  [§3.2]
    # Terminal rows keep only their reward
    # [code: src/rl/agents/dqn_agent.py#L410-L414]
    next_best = max(target_net(next_states, next_actions))
    target = reward + non_terminal * next_best
    # Shipped code uses mean absolute error
    # [code: src/rl/agents/dqn_agent.py#L416]
    # The appendix writes a squared Bellman error  [app. A]
    td_loss = mean(abs(q_taken - target))
    # Penalty: log-sum-exp over actions minus the taken Q-value  [app. A]
    # [code: src/rl/agents/dqn_agent.py#L349-L351]
    penalty = mean(logsumexp(q_all, dim=1) - q_taken)
    # Weighted sum; the shipped config uses alpha = 0.2  [app. B]
    loss = td_loss + alpha * penalty
    return loss

Picking at test time

At test time the policy scores every candidate, plus the end action, and takes the highest value. It stops after four picks, or earlier if the end action wins §3.1 §4.1. Candidate features are computed without labels, so the policy never sees a label before it chooses §3.2.

# Greedy selection of demonstrations  [§3.1] [§3.2]
# Input:  unlabeled pool, trained policy Q, budget k = 4
# Output: ordered list of picked examples

def select_demonstrations(pool, Q, k=4):
    # Start with an empty prefix  [§3.1]
    prefix = []
    for step in range(k):
        # Features for each remaining candidate, without its label  [§3.2]
        candidates = [features(prefix, x) for x in pool]
        # Score every candidate and the end action  [§3.1] [§3.2]
        scores = Q(prefix, candidates + [end_feature])
        # Stop if the end action scores highest  [§3.1]
        if argmax(scores) == len(pool):
            break
        # Reveal the label of the chosen example and append it  [§3.1]
        prefix.append(pool[argmax(scores)])
    return prefix

Evaluation settings

The paper evaluates in three settings §4.1:

• Seen examples, same task: the policy picks from the labeled training set of 100 examples, on a task it trained on §4.1.
• New examples, same task: the policy picks from 100 or 1,000 unseen unlabeled examples of the training task §4.1.
• New task: the policy trains on three datasets and is tested on a fourth, held-out dataset, picking from 100 or 1,000 unlabeled examples of that dataset §4.1.

Baselines are random selection, max-entropy active learning, and reordering by global entropy §4.1. Two oracles, best-of-10 and greedy selection, use labels and a reward set. The paper does not compare them in the new-task setting §4.1.

Hyperparameters

Values from the paper are cited to their table or section. Values marked "code" appear only in the official configuration files and are not in the paper's text agnews-same-task.yaml. "Not stated" means neither the paper nor the code gives the value in the places checked.

Grid values are taken from tab. 5. The paper says it chose values by grid search on the validation scheme app. B. It does not say which grid point was reported.

Results

Author-reported. Setup: GPT-2 medium (the paper says 345M) with calibration, 4-shot prompts, datasets AGNews, Amazon, SST-2 and TREC §2.1 §4.1. Tables give means with 95% confidence intervals as subscripts. The paper does not say whether these are half-widths tab. 3.

Same task, seen examples (average and per task, accuracy %) tab. 3

The paper reports that the seen-examples policy beats random, max-entropy and reordering by 11.8%, 12.1% and 7.9% on average §4.2. These differences match the average column of tab. 3.

New task, held-out dataset tab. 4

The table caption names SST-2 twice where it should name TREC. The average column has no confidence interval.

GPT-2 variance versus GPT-3 models (accuracy %, subscript = standard deviation) tab. 2

(C) means calibrated. The caption of tab. 2 gives standard deviations, whereas tab. 3, tab. 4 and tab. 7 give confidence intervals.

Transfer of GPT-2 policies and examples to GPT-3 (accuracy %, CI as subscript) tab. 7

The Curie rows use one random set of demonstrations and have no confidence intervals tab. 7.

Linear policies (seen examples and 1,000 new examples, average) tab. 6

Linear policies average 65.6 (seen) and 65.9 (1,000 new), against 71.4 and 69.0 for the MLP in the same settings tab. 6 tab. 3. The paper calls them clearly worse than the MLP app. B.

Analysis of demonstration properties §4.3 fig. 3 fig. 4 fig. 7

The paper reports that, for Amazon, all-positive demonstration sets averaged 87.8% against 83.3% for perfectly balanced sets §4.3 fig. 3. For TREC, performance peaks when exactly two labels are covered §4.3 fig. 4. For length, there is no evidence that longer demonstrations help; one Amazon correlation is significant but weak app. C.4 fig. 7.

Paper vs code

Official code: ChicagoHAI/active-example-selection at commit f51558ef675cdb59c317728ca26431fd0065f376, MIT licence, linked from the paper's introduction §1 footnote. The licence names the Chicago Human+AI Lab, and the paper's acknowledgments thank that lab; the paper lists the authors at the University of Chicago. I read the RL code and the RL configs closely. The baseline code under src/prompting/strategies/ was not compared line by line with the paper.

• Features. The paper lists the state as the count of picks, and the action as predicted label probabilities plus entropy §3.2. The shipped training config uses the entropy, the minimum, maximum and standard deviation of the probabilities, and the perplexity of the prompt as action features, and the count plus the perplexity as state features agnews-same-task.yaml misc_utils.py models.py. The paper does not mention perplexity or the probability statistics.
• Entropy normalization. Code normalizes entropy by log2 of the number of labels misc_utils.py. The paper does not state a normalization.
• Loss. The appendix writes the Bellman term as one half of its square app. A. The code uses mean absolute error for the Bellman part dqn_agent.py. The conservative penalty in the code is computed for the sampled action only, not as an expectation under the behaviour policy dqn_agent.py.
• Target network. The paper cites double Q-learning for the target network §3.2. The code's target is the maximum over the target network's own values, with no action selection by the online network dqn_agent.py.
• Unit of the target update. The paper gives target update values of 100, 200 and 400 steps app. B. In code, the target is synchronised every 400 train steps dqn_agent.py, and each train step runs 16 optimizer updates dqn_agent.py agnews-same-task.yaml.
• Number of updates. Eight thousand train steps (paper) with 16 updates each (code) means 128,000 gradient steps in the shipped config app. B agnews-same-task.yaml agnews-same-task.yaml.
• Offline phase. The shipped config sets offline steps to 12,000, which is more than the 8,000 train steps. So no rollouts run during training, and training uses only the logged data dqn_agent.py agnews-same-task.yaml. This matches the paper's offline setting §3.2.
• End action. The paper describes ⊥ as a special action §3.1. The code uses a learned vector, appended to the candidate features network.py dqn_agent.py. When the logs are loaded, the code adds synthetic end-action transitions with reward 0 at probability 0.05 per transition. The random logger never picks the end action replay.py random_agent.py. The paper mentions neither.
• Reward-set size in logging. The paper says each set has 100 examples §4.1. The logging config uses 80 reward examples random-agent.yaml. The training config uses 100 agnews-same-task.yaml.
• Selected hyperparameters. The paper gives a grid app. B but not the selected point. The shipped config picks lr 3e-4, alpha 0.2, target update 400 and dropout 0.25 agnews-same-task.yaml.
• Calibration. The paper cites contextual calibration and the N/A anchor §2.2. The code calibrates by dividing the predicted probabilities by their mean over calibration prompts, then renormalizing. The calibration prompts are content-free "N/A" inputs models.py data_utils.py.
• Test set size. The paper does not state the size of the test set. The code caps the test split at 10,000 examples, or 1,000 in labeled-gpt3 mode data_utils.py.
• Seeds. The paper reports five random seeds tab. 3. The shipped configs fix seed 42 agnews-same-task.yaml. The seed also controls the data split through a deterministic shuffle data_utils.py. The listed files contain no script that loops over five seeds.
• Multi-task transfer. The code gives the held-out task the task identifier 0. That is the same task-embedding row the first training task (Amazon) uses environment.py environment.py network.py. The paper does not say whether this sharing is intended.
• Stopping on long prompts. If a prompt exceeds the model's context, the code drops the last pick, gives reward 0, and ends the episode environment.py. The paper says k = 4 avoids overflow §6.
• GPT-3 engine names. The code maps Ada, Babbage and Curie to text-ada-001, text-babbage-001 and text-curie-001 models.py. The paper does not name engines.
• Naming of the base model. The paper says GPT-2 345M §2.1. The shipped configs load gpt2-medium agnews-same-task.yaml. This appears to be the same model under two names.
• Code-only details. Input normalization with running statistics network.py, Kaiming-uniform initialization with a = 0.01 network.py, and leaky ReLU activations network.py are not in the paper.

Pitfalls

• Variance claim. The paper says the policy reduces variance on all four tasks §4.2. In tab. 3, the AGNews confidence interval for the seen-examples policy is 7.8, which is larger than the 6.8 for reordering. The claim holds against random and max-entropy on AGNews, but not against reordering.
• Missing average intervals. The new-task text calls the gains over random and max-entropy significant §4.2. tab. 4 gives no confidence interval for the average, and the paper reports no test for it.
• Interval meaning. Subscripts are confidence intervals in tab. 3, tab. 4 and tab. 7, but standard deviations in tab. 2. The paper does not say whether the confidence intervals are half-widths.
• 100-example setting. The shipped configs use the 1,000-example pool only. The code supports 100 through dataset_mode: unlabeled-custom-100 data_utils.py. The paper's 100-example numbers therefore need that setting set by hand.
• Logging reward set. The paper says each reward set has 100 examples §4.1. The logging config uses 80 random-agent.yaml. Whether the logs behind the reported numbers used 80 is not stated.
• Target update units. A target update of 400 counts train steps, which is 6,400 gradient steps in the shipped config dqn_agent.py. Reading "400 steps" as gradient steps gives a different schedule.
• Learning rate. The shipped config uses a cyclic scheduler with equal base and maximum rates, so the learning rate is effectively constant dqn_agent.py.
• Skipped updates. An update is skipped when fewer than two sampled rows have a next state. The skip is only logged at info level dqn_agent.py.
• Normalization statistics. The target network is also in training mode, so its running normalization statistics update on its own forward passes, and they are copied at each target sync dqn_agent.py network.py.
• Compute cost. Reading the code, each step scores every remaining candidate and the reward set again, which is roughly 200 model prompts per step environment.py environment.py. The paper does not report run time or hardware. This is a reading of the code, not a measurement.
• Oracle cost. The paper says greedy selection has exponential runtime §4.2. Best-of-10 needs labels and a reward set, so it is not usable when labels are unavailable §4.1.
• Transfer to GPT-3 is mixed. Ada gains about 1 point tab. 7 app. C.2. Babbage loses 1.6 points with the uncalibrated policy tab. 7. The Curie rows use one random set tab. 7. The paper says it draws no firm conclusion app. C.2.
• Demonstration balance. The paper finds that balanced demonstration sets do not consistently perform better or vary less §4.3 fig. 3.
• Features beyond the paper's list. A reimplementation that follows only the paper's features §3.2 will differ from the released configuration, which adds perplexity and probability statistics agnews-same-task.yaml.
• Multi-task embedding. The held-out evaluation uses the same learned row as Amazon, so the transfer numbers depend on that row. Whether the paper's runs did this is not stated, and the paper does not describe the task identifier.

Lineage

• Builds on: Q-learning §3.2 (Mnih et al., 2013); conservative Q-learning for offline data app. A (Kumar et al., 2020); target networks and replay §3.2 (van Hasselt et al., 2010; Lin, 1992); reward shaping that keeps optimal policies §3.2 (Ng et al., 1999).
• Connects to: reinforcement learning for active learning §5 (Fang et al., 2017; Liu et al., 2018); contextual calibration as the anchor for GPT-2 §2.2 (Zhao et al., 2021).
• Compared against: random, max-entropy, reordering (Lu et al., 2022), best-of-10 and greedy oracles §4.1; linear policies tab. 6.
• Related analysis: demonstrations matter less than expected for in-context learning (Min et al., 2022) §1.