PaperlayerView exact revision on arXiv

Source-linked AI summary

Active Example Selection for In-Context Learning

Yiming Zhang, Shi Feng, Chenhao Tan

arXiv:2211.04486v1cs.CLcs.AI

TL;DR

In-context learning can vary substantially with the selected demonstration examples, raising questions about how reliably language models acquire new information. The paper formulates active example selection as a sequential decision problem and trains reinforcement-learning policies, improving GPT-2 accuracy and generalizing to new tasks while gains diminish on larger GPT-3 models.

  • Problem

    In-context learning performance is highly sensitive to demonstration-example selection, while existing reordering and calibration methods do not reliably stabilize GPT-2.

  • Method

    The paper treats selecting, annotating, and adding unlabeled examples as sequential decision making and uses reinforcement learning to train selection policies.

  • Results

    GPT-2 policies improve performance by 5.8% on average on unseen tasks, while selected examples yield a small improvement on GPT-3 Ada that diminishes on larger models.

  • Takeaways & Limitations

    Example selection can stabilize and improve GPT-2 in-context learning, but transfer across model sizes is limited and reflects differing model capabilities.

  • Takeaways & Limitations

    Results from GPT-2 should not be assumed to generalize to GPT-3, whose behavior differs across model sizes.

Abstract

from arXiv · show

With a handful of demonstration examples, large-scale language models show strong capability to perform various tasks by in-context learning from these examples, without any fine-tuning. We demonstrate that in-context learning performance can be highly unstable across samples of examples, indicating the idiosyncrasies of how language models acquire information. We formulate example selection for in-context learning as a sequential decision problem, and propose a reinforcement learning algorithm for identifying generalizable policies to select demonstration examples. For GPT-2, our learned policies demonstrate strong abilities of generalizing to unseen tasks in training, with a $5.8\%$ improvement on average. Examples selected from our learned policies can even achieve a small improvement on GPT-3 Ada. However, the improvement diminishes on larger GPT-3 models, suggesting emerging capabilities of large language models.

1 Introduction

The paper revisits whether demonstration-example selection substantially affects in-context learning and investigates reinforcement learning for selecting examples. Learned policies improve GPT-2 performance and generalize across tasks, while effects differ across GPT model sizes.

  • Motivation: In-context learning performance remains highly variable across demonstration sets, even with reordering and calibration, especially for GPT-2.The authors interpret this variance as evidence that models do not yet acquire new information in-context efficiently and reliably.
  • Approach: The paper formulates example selection as sequential decision making and uses reinforcement learning to discover generalizable selection policies.It focuses on active selection from unlabeled datasets rather than fine-tuning from an existing labeled set.
  • Results: 12.1% improvement over a max-entropy active learning baseline is observed for GPT-2 on a seen task with labeled data.This evaluation tests the learned policy within a task represented during training.
  • Results: 5.8% improvement on average is achieved when the GPT-2 learned policy generalizes to new tasks.The result suggests that the policy captures systematic biases in how GPT-2 acquires information.
  • Results: Examples selected using GPT-2 policies produce a small improvement on GPT-3 Ada, but the improvement diminishes on larger GPT-3 models.The paper presents this divergence as related to emerging abilities in large language models.

2 The Effect of Example Selection

The study defines few-shot text classification through prompts containing labeled examples, then shows that example choice creates substantial instability. Reordering and calibration help in some settings but do not reliably resolve the problem, motivating selection beyond ordering.

  • Definition: A prompt consists of k labeled examples, and predictions choose the label token with the highest language-model likelihood.The formulation concatenates the prompt with a test input and restricts predictions to the label set or proxy tokens.
  • Sensitivity: Across all four tasks, GPT-2 performance ranges by more than 30% between the best and worst randomly sampled demonstration sets.For AGNews, SST-2, and TREC, the worst set performs below random guessing; TREC reaches 10.0% versus 16.7% random accuracy.
  • Reordering: 8.1% average improvement over random demonstration examples results from selecting the best of 24 orderings using a 100-example validation set.This estimate is optimistic for a true few-shot setting because such a validation set may not be available.
  • Reordering: Best orderings still leave wide performance ranges, including 79.6% versus 32.7% on AGNews and below-random performance for 9 of 30 TREC sets.These results indicate that considerable proportions of demonstration sets lack a good ordering.
  • Calibration: Calibration improves average performance on three tasks but leaves GPT-2 variance unresolved and decreases performance marginally on Amazon.On calibrated AGNews, accuracy still spans 79.5% to 26.1%, a gap exceeding 53.4%.
  • Model comparison: GPT-3 models also show instability, especially on multi-label tasks, although calibration diminishes their variance.Binary tasks exhibit much smaller variance, and calibration may no longer reflect the model’s innate information-acquisition ability.

3 Active Example Selection by RL

The paper formulates active example selection as sequential decision making and uses reinforcement learning to select demonstration examples from unlabeled pools. It trains policies to maximize in-context performance while targeting generalization across tasks and models.

  • Problem formulation: Unlabeled examples are selected and annotated sequentially to construct prompts, with selection continuing until k examples are chosen or early termination occurs.The objective is to maximize expected accuracy on unseen test examples despite the exponentially large prompt space.
  • Problem formulation: The MDP state is the current prompt prefix, actions are unlabeled examples plus an end-of-prompt action, and rewards are based on language-model performance.In practice, the reward uses accuracy on a labeled validation set, called the reward set.
  • Generalization: The learned policies aim to identify generalizable qualities of demonstration examples and select useful unlabeled examples on tasks without observed labeled examples.The setup differs from standard reinforcement learning because test-time actions and classification tasks may be unseen during training.
  • RL solution: Reinforcement learning policies are trained with Q-learning, using conservative Q-learning to reduce overestimation for actions absent from the training data.The implementation uses off-policy data generated by a random policy over 2,000 episodes.
  • RL solution: Marginal-utility reward shaping assigns each action its additional gain in the objective, and with γ = 1 the trajectory rewards telescope to the final prompt performance.This shaped reward preserves the optimal policy while avoiding the difficulty of sparse rewards.
  • Implementation: Minimal features are used so policies can be evaluated across GPT-2 and GPT-3, with a learned Q-function approximated by an MLP over state and action representations.The paper notes that other features, including language-model embeddings, could also be used.

4 Results

The experiments evaluate learned example-selection policies across same-task and new-task settings, finding strong same-task gains, reduced variance, and smaller but significant cross-task improvements.

  • Evaluation settings: The evaluation covers seen examples on the same task, new examples on the same task, and new examples on an unseen task.New-task evaluation trains on three tasks and tests on a held-out fourth task, using unlabeled candidate examples.
  • Same-task results: Across all 4 tasks, the method outperforms random, max-entropy, and reordering baselines by average gains of 11.8%, 12.1%, and 7.9%, respectively.The comparison uses four-shot selection evaluated across five random runs.
  • Same-task results: The method consistently decreases variance across all 4 tasks compared with the baselines.The paper uses 95% confidence intervals as a proxy for variance.
  • New examples, same task: Picking from 100 or 1000 new examples largely retains performance gains and variance reductions, with higher overall performance observed for 100 examples.The paper attributes the higher performance with 100 examples to large variance when picking from 1000.
  • Oracle comparisons: The learned method performs relatively close to best-of-10, while greedy-oracle significantly outperforms other methods but incurs exponential runtime.The authors use simple features to promote generalizable selection strategies, which may explain why the method does not outperform oracles.
  • New examples, new task: On new tasks, the method yields smaller but significant improvements over random and max-entropy baselines, while gains over reordering remain limited.The result suggests learnable insights about good demonstrations generalize across tasks, but finding good examples in an unknown task remains challenging.
  • Transfer results: Transferring GPT-2-trained policies and examples produces small gains of approximately 1% on GPT-3 Ada but mixed results on Babbage and Curie.The transfer evaluation applies GPT-2-trained policies because generating GPT-3 training episodes is costly.
  • What Makes Good Examples?: Example-selection behavior is idiosyncratic: balanced demonstrations do not consistently improve performance or reduce variance, and optimal label coverage differs across tasks.TREC performs best with exactly two covered labels, whereas AGNews performs best when all four labels are covered.

5 Related Work

Prior work explains and characterizes in-context learning through latent concepts, label independence, and sensitivity to prompt design and demonstration ordering.

  • In-context learning can be explained by inferring a shared latent concept among demonstration examples.
  • Min et al. find that in-context learning success is largely independent of access to gold labels.
  • Known issues include surface form competition and sensitivity to prompts, instructions, and demonstration-example ordering.

6 Discussion

The discussion addresses GPT-2’s relevance, model-size differences, experimental choices, and the economic limits of extending experiments to larger GPT-3 models.

  • GPT-2 remains relevant because it is public, economically feasible, incompletely understood, and potentially useful when the largest models are inaccessible or inappropriate.
  • Results from GPT-2 should not support generalizing claims because they may not generalize to GPT-3.
  • The study’s goals are to assess in-context information acquisition and test whether reinforcement learning identifies reliable active example-selection strategies.
  • GPT-3-Babbage shows relatively small post-calibration variance, making experiments on larger models economically unsensible.
  • Experiments use k = 4 because context windows constrain longer prompts and prior work reports diminishing improvements beyond four demonstrations.

7 Conclusion

The paper studies active example selection as a reinforcement-learning problem and finds that learned policies improve GPT-2 generalization while gains diminish on larger GPT-3 models.

  • The method iteratively selects, annotates, and adds unlabeled examples to prompts while training reinforcement-learning policies for active example selection.
  • For GPT-2, learned policies stabilize in-context learning and improve accuracy on new example pools and unseen tasks.
  • Useful demonstration examples can have properties that deviate from human intuitions.
  • Examples selected from GPT-2 produce a small improvement on GPT-3 Ada, but the gain diminishes on Babbage and Curie.

Ethics Statement

The paper frames understanding example selection as a way to develop more effective in-context learning strategies and better large-scale language models, while recognizing potential social harms.

  • Understanding how models acquire new information through example selection may support more effective in-context learning strategies and better large-scale language models.
  • These strategies may also be used in applications that could incur harm to society.

A Conservative Q-Learning

This section presents conservative Q-learning as a regularized offline Q-learning objective and reports that linear policies outperform baselines but underperform MLP policies.

  • Standard Q-learning minimizes the Bellman Error (BE).
  • Offline Q-learning can overestimate out-of-distribution action values, causing unfamiliar actions during evaluation and hurting performance.
  • Conservative Q-learning adds a penalty term that regularizes Q-values and encourages policies to avoid unfamiliar evaluation actions.
  • Linear policies outperform baselines but remain clearly worse than MLP policies.

C Additional Results

The additional results examine unlabeled-set size and report SAME TASK accuracy across four datasets using MLP and linear selection policies. Larger unlabeled sets generally improve transfer performance and reduce variance, although gains are not monotonic.

  • SAME TASK evaluation: SAME TASK accuracy is reported for AGNews, Amazon, SST-2, and TREC using MLP and linear policies.Results aggregate five random seeds with 95% confidence intervals.
  • Unlabeled-set size: Larger unlabeled sets improve transfer performance and reduce variance in the NEW TASK setting.The setting trains policies on three datasets and evaluates them on a held-out dataset.
  • Unlabeled-set size: The improvement from larger unlabeled sets is not necessarily monotonic because of large variance.
  • Unlabeled-set size: Policies remain performant when selecting from only 50 unlabeled examples.

C.2 Transfer to GPT-3

Transfer from GPT-2 to GPT-3 produces mixed results: small gains appear for ADA, while BABBAGE often matches or underperforms its base model and CURIE shows marginal gains in limited experiments. Additional analyses find that policy features are interpretable but variable, and example length generally does not predict accuracy.

  • Transfer results: 1.1% and 0.9% improvements occur when transferring policies and examples, respectively, to uncalibrated GPT-3 ADA.For calibrated ADA, the corresponding improvements are 0.4% and 1.9%.
  • Transfer results: Different models may share preferences for demonstration examples, as suggested by the generalization of learned policies.
  • Transfer results: GPT-3 transfer results are mixed, with BABBAGE showing a 1.6% performance drop for an uncalibrated transferred policy.Other BABBAGE scenarios match or underperform the base models.
  • Transfer results: CURIE experiments show marginal gains of 1.8% for an uncalibrated transferred policy and 1.0% for examples transferred to the calibrated model.CURIE results use one random set and therefore omit variance.
  • Policy analysis: Linear policies are less effective than MLP policies but provide greater interpretability, while coefficient variance remains large.Positive entropy coefficients suggest some value for class-balancing strategies, but feature importance varies by dataset.
  • Example-length analysis: There is no evidence that longer examples improve performance: correlations are absent or nonsignificant in most datasets, with only a weak significant Amazon correlation.The Amazon correlation has p-value = 0.019.
Loading 2211.04486v1…