Robustness of Named-Entity Replacements for In-Context Learning (2023)

Abstract: A key feature of modern large language models (LLMs) is their ability to perform in-context learning, a prompting technique where query- answer demonstrations are shown before the final query. This allows for generalization to novel distributions at inference time where the LLM can learn new rules without parameter updates. However, the choice of demonstrations and their relationship to a particular query can have a profound impact on model accuracy, raising concerns about the true in-context generalization capabilities (Zhao et al., 2021). In this work, we explore the robustness of the in-context learning paradigm by focusing on entities. In particular, we seek to understand the robustness of LLM in-context learning with respect to named entity replacements. We discover a significant variance in downstream performance based on the choice of the named entities, across three popular reasoning tasks and two popular LLMs. Specifically, model accuracy on the test sets can fluctuate between -2.7 to +8.0 points depending on the choice of named entity replacements. Our analysis exposes the sensitivity of LLM in-context learning with respect to named entities, and offers a simple recipe to improve test performance by hyper-parameter tuning the named entities for a given dataset. Code and datasets for reproducing the results are publicly available.

## introduction modern large language models (llms) routinely demonstrate in-context learning, where the model can solve a wide range of tasks from a few demonstration examples (or exemplars) @xcite . in-context learning, or few-shot prompting, is a powerful learning paradigm as it allows for fast adaptation to new tasks without the need to train example from flan-t5 on gsm8k @xcite . here we replace all male entities (e o ) in demonstrations with douglas, and observe a change in the query answer prediction (from the incorrect answer 40 to the correct answer 10). only 1 out of @xmath0, with the remaining demonstrations denoted by "[...]". model parameters. however, model performance in the in-context setting can vary significantly with the choice of demonstrations. previous work has uncovered performance instability when altering the demonstrations, particularly when the demonstrations change semantically @xcite @xcite . naturally, a question arises whether llms are robust to semantically preserving, minimal alterations to the demonstrations. in this paper, we seek to address this question using simple, semantically preserving alterations: named entity replacements. because changing named entities should not, in principle, change the correct response to the query, the model is expected to remain invariant to such modifications. for example, given the demonstration "alice has the ball. alice went to the kitchen. where is the ball? kitchen.", the model should be invariant to the choice of the named entity "alice"-i.e., a model should behave the same if we swap all instances of "alice" with another named entity, e.g. "eve". however, we observe that the model behavior is not always robust to such alterations (figure 1 )prompting us to critically evaluate named entity invariance within in-context learning. we conduct experiments to investigate the robustness of llm in-context learning with respect to named entity replacements in three popular named entity heavy reasoning tasks, namely, babi, @xcite @xcite and gpt-j @xcite . our key findings are: ## experimental setup the goal of this study is to investigate the robustness of llms' in-context learning performance by changing entities in the demonstrations and queries. we consider a fixed set of prompt demonstrations aside from the named entities. by replacing the named entities and observing any accuracy changes, we can better understand the effect of certain entities on a model's accuracy. while many prior works have studied the robustness of llms with respect to perturbations @xcite @xcite , we focus our analysis on in-context learning. a more in-depth analysis of related work is provided in appendix subsection a.1. ## results performance across bootstrap sets. we find a noticeable change in accuracy with the best and worst performing entity compared to the baseline or median performing entity. figure 2 shows flan-t5-xl results on all validation datasets (analogous gpt-j results are in figure 5 in the appendix, and follow a similar pattern). this indicates that specific entities can noticeably alter in-context learning performance. furthermore, the standard deviations for the median and best performing entity across bootstraps are low, suggesting that the selection of different name sets will likely yield similar median and best performance. however, the variation in the worst performing entities is large. ## analysis why are the models not invariant to entity replacements? we conducted several analyses in attempts to find the root cause of this result, but they were largely inconclusive. does the frequency of names affect how much they shift model performance? the lack of robustness could well be explained by the frequency effects of named entities in the training corpus of the pre-trained model @xcite . we use the entity frequency data provided by @xcite , and assessed performance changes across names with varying frequencies. we do not find any correlation between name frequency and the model performance post entity replacement (see appendix section b.3). does the token length of names affect performance? another hypothesis of the invariance could be due to the inconsistent tokenization of named entities. rare names can be tokenized into longer tokens based on the learned subwords of the model. we analyzed token length post tokenization and compared it with model performance, but we are unable to find any meaningful relationship between the two (see appendix section b.4). does improved performance result from a decrease in prompt perplexity due to entity replacement? finally, the surprisal of the llm to certain entities could be a promising hypothesis to explain the change in model performance @xcite . we compare prompt perplexity after entity replacement to llm performance, and yet again we are unable to find a statistical correlation (see appendix section b.5). while gonen et al. (2022b) suggests improvement of model performance with decrease in perplexity, we observe certain entities displaying the opposite behavior. ## conclusion in this study, we find that llms exhibit a consistent lack of robustness with respect to named entity replacement for the in-context learning setting, further questioning the process of reasoning employed by these models @xcite . we demonstrate this across three axes: the complexity of the task, the model architecture, and which named entities are to be replaced (in the demonstrations, query, or both). the cause for this effect remains elusive, which opens up exciting research directions. one intriguing result is that the entities that can boost performance are dataset specific. this points to potential avenues for enhancing llm downstream performance by performing hyperparameter search over a range of entities.