Source-linked AI summary
Ambiguity-Aware In-Context Learning with Large Language Models
Lingyu Gao, Aditi Chaudhary, Krishna Srinivasan, Kazuma Hashimoto, Karthik Raman, Michael Bendersky
TL;DR
Similarity-based retrieval for in-context learning does not account for the LLM’s task-specific label ambiguity. This paper combines retrieval with ambiguity-aware demonstration selection, finding improvements across three text-classification tasks, particularly when selected examples were misclassified and fell on the test example’s decision boundary.
Problem
Similarity-only demonstration selection does not explicitly account for the LLM’s knowledge of the task-specific label space, even though prompt choice affects ICL.
Source: §1 · p. 1Method
The method retrieves similar examples, identifies the test input’s two most ambiguous labels, and selects demonstrations using their gold labels and model predictions.
Source: §1 · p. 1, §3.4 · p. 5Results
AMBIG-* methods outperform RETR-* by +3.0% (avg.) for Flan-PaLM 2 (M) and +1.2% (avg.) for Flan-PaLM 2 (L), with misclassified demonstrations on the test example’s decision boundary yielding the overall best performance.
Source: §4 · p. 6Takeaways & Limitations
Across three text-classification tasks, the authors find that selecting demonstrations using output-label information is as important as considering input similarity.
Source: §4 · p. 6, §6 · p. 8The method may perform unsatisfactorily when the test example’s gold label often differs from the model’s top-two choices, indicating weak zero-shot performance or flawed ambiguous-label selection.
Source: §7 · p. 8
Abstract
from arXiv · showhide
In-context learning (ICL) i.e. showing LLMs only a few task-specific demonstrations has led to downstream gains with no task-specific fine-tuning required. However, LLMs are sensitive to the choice of prompts, and therefore a crucial research question is how to select good demonstrations for ICL. One effective strategy is leveraging semantic similarity between the ICL demonstrations and test inputs by using a text retriever, which however is sub-optimal as that does not consider the LLM's existing knowledge about that task. From prior work (Lyu et al., 2023), we already know that labels paired with the demonstrations bias the model predictions. This leads us to our hypothesis whether considering LLM's existing knowledge about the task, especially with respect to the output label space can help in a better demonstration selection strategy. Through extensive experimentation on three text classification tasks, we find that it is beneficial to not only choose semantically similar ICL demonstrations but also to choose those demonstrations that help resolve the inherent label ambiguity surrounding the test example. Interestingly, we find that including demonstrations that the LLM previously mis-classified and also fall on the test example's decision boundary, brings the most performance gain.
1 Introduction
The introduction asks whether LLM knowledge of label ambiguity can improve on similarity-only ICL demonstration selection. It proposes combining semantic retrieval with label- and prediction-based constraints, reporting gains across three classification tasks.
1 Introduction
- Similarity-based retrieval selects examples using input text but does not explicitly capture the LLM’s knowledge of the task-specific label space.On SST, Flan-PaLM 2 (L) was reported to confuse Very Negative with Negative more often than Neutral with Very Negative.Source: §1 · p. 1
- The proposed approach combines retrieved examples with the test example’s ambiguous labels and the model’s predictions on candidate demonstrations.It selects examples whose gold labels are ambiguous for the test input, then considers misclassified examples whose predicted labels also fall within that set.Source: §1 · pp. 1–2, Figure 1 · p. 2
- +1.5%, +2.2%, and +2.6% F1 macro improvement over retriever-based ICL followed incrementally from test-label ambiguity, misclassified demonstrations, and training-label ambiguity, averaged across datasets.The improvements correspond to adding each constraint in that order.Source: §1 · p. 2
- +3.9% gain on Flan-PaLM 2 (M), compared with +1.4% on Flan-PaLM 2 (L), was reported for label-based constraints over retriever-based ICL.Source: §1 · pp. 2–3
- The authors attribute the method’s success partly to the ambiguous label set acting as a proxy for the gold test label, while demonstration labels bias predictions.They argue that showing a likely gold label guides the model toward the correct prediction.Source: §1 · p. 3
2 Proposed Method
The method starts with semantically similar training examples, then uses the LLM’s label uncertainty and training-example predictions to select demonstrations intended to address the test example’s ambiguity.
2 Proposed Method
- The method first retains retrieved demonstrations whose gold labels belong to the test example’s ambiguous label set.Source: §2 · p. 3
- For each test input, the method retrieves semantically similar training examples and uses the model’s log-likelihood scores to identify its two most likely labels as the ambiguous label set.The retrieved examples are candidates for demonstration selection.Source: §2 · p. 3
- It further selects demonstrations the model misclassified, restricting them to cases where the model’s prediction falls within the ambiguous set.The authors describe these as lying on the test example’s decision boundary.Source: §2 · p. 3
3 Experimental Setup
The experiments use two Flan-PaLM 2 sizes and three fine-grained English classification datasets. Baselines and proposed methods are compared using retrieved or filtered demonstrations, with variants evaluated at four and eight shots.
3.1 Model
- The experiments use instruction-tuned Flan-PaLM 2 models in medium and large sizes, with demonstrations selected by an off-the-shelf retriever fine-tuned on mT5-base.Demonstration order is shuffled across three random seeds and results are averaged.Source: §3.1 · pp. 3–4
3.2 Data
- The English datasets are selected for fine-grained label spaces and exclusion from the Flan mixture: SST sentiment, GoEmotions with neutral and 27 emotion classes, and four-category EDOS Task-B.Source: §3.2 · p. 4
3.3 Baselines
- The baselines are frequent-label prediction, zero-shot ICL, static N-shot ICL with one demonstration per label, and retriever-based ICL using four or eight nearest neighbors.Source: §3.3 · p. 4
3.4 Proposed Method: AMBIG-ICL
- AMBIG-ICL filters retrieved candidates by whether their gold labels are among the test input’s top-two ambiguous labels, then by misclassification and predicted-label ambiguity.Training-data predictions come from the zero-shot model.Source: §3.4 · p. 5
- The proposed variants use four or eight demonstrations, and Table 2 reports F1 macro comparisons while Table 3 reports average differences from RETR.Source: §3.4 · p. 5, Table 2 · p. 5, Table 3 · p. 5
4 Results and Discussion
Across three tasks, the proposed methods outperform baselines, and label-space-aware selection improves on retriever-based ICL. The analyses link this advantage to more gold-label-matching demonstrations and lower model confusion, even when selection sacrifices semantic similarity.
4 Results and Discussion
- For a GoEmotions example, RETR predicted “Joy” incorrectly, whereas AMBIG-ICL predicted correctly under all three settings.The example compares demonstrations selected by RETR and AMBIG-ICL for n = 4 on Flan-PaLM 2 (L).Source: Table 4 · p. 6
- ZERO outperforms few-shot baselines by 1.4% on Flan-PaLM 2 (M), whereas Flan-PaLM 2 (L) gains +1.4% from demonstrations; on SST and GoEmotions, L performs better with n = 4 than n = 8.Thus, the few-shot benefit and the effect of demonstration count vary across model size and task.Source: §4 · pp. 5–6
- AMBIG-* outperforms RETR-* by +3.0% on Flan-PaLM 2 (M) and +1.2% on Flan-PaLM 2 (L), while proposed methods beat baselines across all three tasks.The comparison uses F1 macro because the tasks are unbalanced.Source: §4 · pp. 5–6
- Adding label-space constraints sacrifices semantic similarity, yet the proposed methods still outperform retriever-based baselines that prioritize similarity.For EDOS Task-B at four shots, selecting examples required up to top-16 retrieved candidates for +GOLD, top-55 for +GOLD+MIS, and more than top-250 for +GOLD+MIS+PRED.Source: §4 · pp. 6–7
- AMBIG demonstrations match the test gold label 44.2% of the time on average, versus 30.9% for RETR, and the ambiguous label set is reported as a close gold-label proxy.Input-similarity-only RETR is 13.3 percentage points behind the proposed methods on average for this label-match measure.Source: §4 · p. 7
- AMBIG-* achieves the lowest entropy across all three datasets and both model sizes, while entropy–F1 correlations are consistently negative only for Flan-PaLM 2 (L).For L, the reported correlations are r=−0.78 for EDOS, −0.48 for SST, and −0.92 for GoEmotions; for M, they are mixed.Source: §4 · pp. 7–8, Table 6 · p. 7
5 Related Work
Prior work selects ICL demonstrations through task-specific or unified retrievers, semantic similarity, uncertainty, diversity, or n-gram overlap. This work adopts an off-the-shelf retriever and uses the LLM’s existing knowledge around test examples to select demonstrations.
5 Related Work
- The authors position their approach as the first, to their knowledge, to use an LLM’s existing knowledge about the test example for demonstration selection.Prior work had typically used model predictions on training data.Source: §5 · p. 8
- The authors adopt an off-the-shelf retriever to examine generalizability across classification tasks and expect their method could also benefit from a task-specific retriever.Source: §5 · p. 8
- Demonstration ordering methods are orthogonal to this work and can be combined with the proposed methods.The paper reports results across three shuffle orders.Source: §5 · p. 8
6 Conclusion and Next Steps
The paper finds that using the LLM’s knowledge of output labels for test examples and demonstrations matters alongside input similarity, and that its methods outperform baselines across three tasks. It identifies expanding the ambiguous label set and applying the techniques to token- and span-level tasks as future directions.
6 Conclusion and Next Steps
- The proposed methods consistently outperform baselines across all three tasks.Source: §6 · p. 8
- The authors find that LLM knowledge of the output label space for test examples and demonstrations is as important as input semantic similarity.Source: §6 · p. 8
- The method uses only the two most ambiguous labels; the authors suggest exploring larger ambiguous sets, especially for fine-grained GoEmotions.Source: §6 · p. 8
- The authors focus on sentence classification and identify NER and POS tagging as token- or span-level tasks for future exploration.Source: §6 · p. 8
7 Limitations
The method’s performance may be unsatisfying when the test gold label often falls outside the model’s top two label choices. The authors identify this pattern as potentially indicating subpar zero-shot performance or flawed ambiguous-label-set selection.
8 Ethics Statement
The paper uses pretrained LLMs for text classification and notes that improved performance on offensive-content classification could be misused. EDOS contains offensive content and was selected for its fine-grained label nuances and to avoid bias from models’ familiarity with the data.
8 Ethics Statement
- The work uses pretrained LLMs for text classification and notes that LLMs exhibit biases.Source: §8 · p. 9
- Improved performance on offensive-content classification could be misused.Source: §8 · p. 9
- EDOS contains offensive content and was selected for its fine-grained label nuances and to avoid bias from models inherently familiar with the data.Source: §8 · p. 9
A.1 Prompt Construction
The appendix presents prompt templates for zero-shot and few-shot ICL and defines the label categories for EDOS, SST, and GoEmotions.
A.1 Prompt Construction
- The prompt templates use a test example and, for few-shot ICL, selected demonstrations with their gold labels.The appendix uses a 4-shot prompt as an example.Source: §A.1 · p. 12
- EDOS classifies sexism into Threat, Prejudiced, Animosity, or Derogation categories.The definitions distinguish threats and justifications of sexism from subtle sexism and explicit derogation.Source: §A.1 · p. 12
- SST assigns movie-review sentences to Great, Good, Okay, Bad, or Terrible sentiment categories.The categories range from extremely positive to extremely negative, with Okay defined as neutral.Source: §A.1 · p. 12
- GoEmotions classifies Reddit comments into 27 emotion categories, including Admiration, Anger, Joy, and Grief.The passage enumerates all 27 labels.Source: §A.1 · p. 12
A.2 Accuracy, Precision, Recall
The appendix reports accuracy, precision, and recall comparisons, analyzes gold-label inclusion in the ambiguity set, and examines entropy-based ordering and an SST example. It also identifies evaluation packages.
A.2 Accuracy, Precision, Recall
- The appendix measures how often the gold label appears in Lambig,t and reports label-wise percentages for the ZERO model.Tables 11 and 12 cover the overall and label-wise percentages, respectively.Source: §A.2 · p. 12, Table 11 · p. 13, Table 12 · p. 14
A.4 Sorting Order with Predicted Probability Distribution Entropy
- Sorting demonstrations by increasing predicted-probability entropy does not consistently improve model performance.Tables 13 and 14 report F1 macro scores and differences from scores averaged over three random seeds.Source: §A.4 · p. 12, Table 13 · p. 14, Table 14 · p. 14
A.5 Example on SST for Comparison between RETR and AMBIG-ICL
- For an SST example, AMBIG-ICL selects demonstrations with which Flan-PaLM 2 (L) predicts correctly, unlike RETR.Table 15 compares the selected demonstrations for n = 4.Source: §A.5 · p. 12, Table 15 · p. 15
A.6 Responsible AI Checklist
- The evaluation uses SciPy and scikit-learn.The passage cites Virtanen et al. (2020) for SciPy and Buitinck et al. (2013) for scikit-learn.Source: §A.6 · p. 12
A.2 Accuracy, Precision, Recall
- The appendix compares proposed methods and baselines on accuracy, precision, and recall with Flan-PaLM 2 (M, L).The supplied passages identify Tables 8–10 as the comparison tables but do not provide their values.Source: §A.2 · p. 12, Table 10 · p. 13, Table 8 · p. 13