Mitigating Copy Bias in In-Context Learning through Neuron Pruning

## Abstract

Abstract Large language models (LLMs) have demonstrated impressive few-shot in-context learning (ICL) abilities. Still, we show that they are sometimes prone to a ‘copying bias’, where they copy answers from provided examples instead of learning the underlying patterns. In this work, we propose a novel and simple method to mitigate such copying bias. First, we create a synthetic task and use the Integrated Gradients method to identify neurons that prioritize copying over generalization. We demonstrate that pruning these neurons consistently improves performance across a diverse set of ICL tasks. We also show that our method is applicable across various LLM architectures, including Transformers and State-Space Models, without requiring modifications.
In our analysis, we adopt a task-recognition perspective on ICL and examine task vectors (Hendel et al., 2023 ) induced by the model. We find that pruning enhances the quality of these vectors, suggesting that the pruned neurons previously hindered effective task recognition.

## 1 Introduction

In-Context Learning (ICL) (Brown et al., ) has recently emerged as a powerful and simple alternative to traditional training and fine-tuning. ICL involves presenting a Large Language Models (LLM) with a “context” consisting of several example pairs, each containing an input and its corresponding correct output, followed by a test example for prediction. For instance, consider the following prompt:

New York $\rightarrow$ 2,   Lyon $\rightarrow$ 1,   Florida $\rightarrow$

In this case, the model must leverage the contextual information from the given examples to identify the underlying pattern of mapping words to vowel counts; based on this, it must then predict that the corect answer for “Florida” is “3”.

While ICL has shown considerable effectiveness, its use in few-shot scenarios faces significant challenges (Zhao et al., ; Razeghi et al., ). In these settings, the inherent scarcity of labeled examples becomes a critical bottleneck, as ICL often requires a substantial number of in-context examples to generalize effectively. Moreover, the performance of ICL is highly sensitive to various aspects of the prompt and the presentation of the examples. Factors such as the specific wording of the prompt (Wang et al., ), the order in which the examples are presented (Lu et al., ), and their relevance to the target example can significantly influence the outcome. Consequently, in domains where labeled data is limited, these challenges collectively hinder the reliable application of ICL, emphasizing the need for strategies that can mitigate sensitivity and make the most of the scarce examples available.

Recent research has primarily addressed these challenges by focusing on prompt formulation strategies, including techniques for selecting optimal templates and examples (Zhou et al., ; Hao et al., ; Lu et al., ), as well as calibration methods (Han et al., ; Zhao et al., ). However, existing work has not yet explored how errors in in-context learning (ICL) relate to the internal processes of LLMs or how to correct them through targeted model modifications.

Our study takes a novel approach by investigating neural activation patterns related to a common challenge in ICL: copying errors. Referring back to the vowel-counting example, a copying error would occur if the model were to output “2” or “1” for Florida, instead of the correct answer “3”. In these cases, the model appears to directly copy an answer from the provided examples rather than generating the correct, novel response based on the induced underlying pattern.

In this work, we hypothesize that there is a small subset of neurons in language models that prioritize copying behavior over task recognition. We posit that these mechanisms can be task-agnostic; that is, the same neurons are responsible for this reasoning shortcut across a range of tasks. We further hypothesize that deactivating these neurons will make the model less likely to follow shortcuts and encourage it to focus on recognizing underlying regularities.

To identify these neurons, we employ the vowel-counting task and apply the attribution method, Integrated Gradients (IG) (Sundararajan et al., ), to trace the copying errors to individual neurons. We then select the top contributing neurons as ”Copying Neurons.” The vowel-counting task is particularly interesting, as it appears challenging for a range of models and elicits the copying behavior in these models. We demonstrate that deactivating the neurons identified using this single task improves results across a diverse range of ICL problems, making our method practical – since the neurons do not need to be selected for each individual task – and confirming the existence of a general mechanism prioritizing the shortcut over reasoning.

In summary, our contributions are fourfold:
(1) We identify the copying bias in ICL and demonstrate that LLMs, particularly smaller ones, suffer from a high percentage of these errors.
(2) We introduce a method to identify specific neurons responsible for triggering the copying behavior (copying neurons).
(3) We show that pruning these identified neurons leads to out-of-the-box improvement across a wide range of ICL tasks.
(4) We utilize the task vectors framework introduced by  (Hendel et al., ), quantifying a model’s ability to recognize and adapt to tasks during ICL. Using this framework, we provide evidence that pruning the copying neurons enhances the quality of task vectors, indicating improved task representation. This finding explains the observed performance gains across various ICL tasks.

## 2 Related Work

ICL, first introduced by (Brown et al., ), has attracted significant interest in recent years due to its ability to enable large language models to perform complex downstream tasks without explicit fine-tuning. By leveraging contextual information provided within the input prompts, these models can dynamically adapt their behavior and generate contextually relevant outputs across a wide range of tasks. While it is commonly associated with Transformer architectures, ICL has also been explored in other model architectures, such as State-Space Models and the RWKV model (Grazzi et al., ; Park et al., ). Leading to a wide line of works that seek to improve the effectiveness of ICL mechanisms (Zhao et al., ; Wei et al., ; Chu et al., ; Zhou et al., ; Wei et al., ; Li et al., ), as well as studies that aim to explain the underlying processes and dynamics of how models internalize and utilize context (Min et al., ; Liu et al., ; Xie et al., ; Olsson et al., ; Von Oswald et al., ; Dai et al., ).

Works proposing methods to improve ICL have mainly focused on prompt selection and prompt formulation strategies, Zhou et al. () propose a different prompting strategy that breaks down a complex problem into a series of simpler subproblems and then solves them in sequence. Zhang et al. () reformulate the example selection for ICL as a sequential decision problem, and propose a reinforcement learning algorithm for identifying generalizable policies to select demonstration examples. Sorensen et al. () introduces a method for unsupervised selection of prompt templates by maximizing the mutual information between the input and the corresponding model output. Lu et al. () show that the order in which the samples are provided can make a significant difference and propose a method for finding the optimal permutation.

While the majority of approaches to improving ICL have focused on various methods of prompt engineering, such as prompt selection and formulation strategies, our work takes a fundamentally different approach. To the best of our knowledge, this is the first attempt to enhance ICL through a neuron-level analysis, specifically by identifying and pruning neurons that contribute to copying (which we define in Section 3.1) within large language models.

Neuron-Level analysis involves examining individual neurons within the model to determine their specific roles. Relevant studies in this area have aimed to understand and categorize neurons into groups based on their functional roles. For example, Voita et al. () show that individual neurons in LLMs correspond to different groups such as dead neurons, positional neurons, and N-gram neurons. Furthermore,  Gurnee et al. () discusses an additional group of categories including entropy neurons, alphabet neurons, syntax neurons, and semantic neurons. Chen et al. () identify safety neurons, which are responsible for safety behaviors. Neuron-level analysis was further expanded to study multilingual LLMs.  Tang et al. () detect language-specific and language-agnostic related neurons in multilingual language models. Neuron-Level analysis is typically performed through activation analysis, where one examines the patterns of neuron activations across various inputs (Voita et al., ; Gurnee et al., ; Stolfo et al., ). Attribution methods, such as Integrated Gradients (Sundararajan et al., ), are also employed to quantify the contribution of individual neurons to the model’s output, hence, allowing for discovering different neuron-families (Dai et al., ; Shi et al., ).

## 3 Method

This section presents our method for detecting and mitigating copying in ICL. First, in section 3.1 we revisit the ICL setting, formally define what are the copying errors, and introduce the Integrated Gradients attribution method. In section 3.2 we elaborate on how we create a synthetic dataset that will be used as a proxy for our proposed detection method, and in section 3.3 we present our method for detecting and mitigating copying neurons.

### 3.1 Preliminaries

#### In-Context Learning

ICL enables a model $f$ to adapt to downstream tasks without any parameter updates. This is achieved by forming a prompt $p$ that includes concatenated training examples. In ICL, a prompt $p$ is constructed by linking the task inputs with their corresponding labels as follows:

$$ $p=x_{1}:y_{1},x_{2}:y_{2},\dots x_{n}:y_{n},x_{n+1}:$ (1) $$

Using this prompt, the model $f$ predicts the most probable label $y$ that completes the prefix $p$ according to the function $f$. In this framework, few-shot learning is characterized by the number of examples in the prompt. Furthermore, we denote $S=[y_{1},y_{2},\dots y_{n}]$ as the set of the in-context example answers for prompt $p$, then we say $p$ is an $|S|=n$-shot in-context prompt.

#### Copying Bias

Copying bias, as we define it, refers to the phenomenon where a language model $f$ returns an incorrect answer that is one of the labels $S$ of the in-context samples provided in the prompt. In other words, given a prompt $p$ under the $n$-shot ICL settings, a prediction $y_{n+1}$ is called a copying prediction if (1) : $y_{n+1}$ is a wrong prediction and (2) : $y_{n+1}\in S$.

We hypothesize that there exists a small number of neurons, which we call copying neurons, that trigger the model to copy responses from the prompt examples $S$ of the prompt $p$. The identification of these neurons is, therefore, crucial for understanding how LLMs balance copying and generalization, and for enhancing the reliability and interpretability of these models. We further hypothesize that pruning these neurons by setting their weights to zero would encourage the model to reason rather than solely rely on copying, thereby improving its ability to generalize under few-shot ICL.

#### Integrated Gradients (IG)

Integrated Gradients (IG) is a popular technique in explainable AI (Sundararajan et al., ) used to elucidate the relationship between a model’s predictions and its input features. It applies to any differentiable model and is computationally efficient. IG works by generating attributions based on integrating the gradients as the input varies between a baseline and the final input of interest (the path):

$$ $\textbf{IG}(f,\tilde{x},x,i)=(x_{i}-\tilde{x}_{i})\int_{\alpha=0}^{1}\frac{ \partial f(\hat{x})\cdot d\alpha}{\partial f(\hat{x_{i}})}\Bigg{|}_{\hat{x}= \tilde{x}+\alpha(x-\tilde{x})}\,,$ (2) $$

where $f({\cdot})$ denotes the prediction of the model, $x$ is the input vector of interest that we want to attribute, $\tilde{x}$ is the baseline input vector and $i$ is an index denoting the indices of features of interest. The baseline represents a reference point against which the input of interest is compared. More specifically, the baseline is an input vector that is supposed to reflect a neutral, missing, or reference state. The idea behind using a baseline is to measure how the model’s output changes as we transition from this baseline state to the actual input of interest. This change is quantified by integrating the gradients along this path. In our experiments, we use the constant zero baseline as proposed in the original paper of IG. This baseline is straightforward to implement and often provides meaningful attributions. The integral is practically computed using the Riemann sum approximation:

$$ $\textbf{IG}(f,\tilde{x},x,i)\approx\frac{(x_{i}-\tilde{x}_{i})}{m}\sum_{r=1}^{ m}\frac{\partial f(\tilde{x}+\frac{r}{m}(x-\tilde{x}))}{\partial x_{i}},$ (3) $$

where $m$ is the total number of Riemann steps.

### 3.2 Proxy ICL Dataset Generation

The Integrated Gradients (IG) method operates by backpropagating the probability of the prediction. To effectively utilize IG within our framework, we require a set of in-context examples with ground-truth labels. We avoid relying on access to target task data for neuron selection, as this would limit the practicality of our approach. Instead, we utilize a synthetic proxy dataset to identify neurons and prune the same set of neurons on evaluation tasks. This approach also aligns with our research hypothesis that copying neurons are largely task-agnostic. This synthetic dataset is employed within the Integrated Gradients framework to identify copying neurons.

The specific task we choose is that of vowel counting since the mapping from a word to such structural attributes requires reasoning. LLMs can occasionally make errors on this task and similar tasks,(^1^11See, e.g., [https://community.openai.com/t/incorrect-count-of-r-characters-in-the-word-strawberry/829618](https://community.openai.com/t/incorrect-count-of-r-characters-in-the-word-strawberry/829618).) potentially outputting copying responses. The synthetic samples we utilize simply map an arbitrary word to its corresponding vowel counts (e.g., apple: $2$). To construct a diverse set of examples, we extract words from a dictionary and calculate their respective vowel counts. This results in two sets: a training set used to identify the copying neurons and a validation set employed to select the optimal layer and percentage of neurons for intervention through pruning.

### 3.3 Copying Neurons Detection

Figure: Figure 1: A high-level depiction of our proposed method of detecting copying neurons. First, in (a) we feed ICL prompts from the synthetic dataset. In this phase, we are only interested in the prompts where the model outputs a wrong response which also appears in the prompt examples. Second, in (b) we use these prompts and calculate the sum of the probabilities over predicted responses that appear in the prompt. This sum is used within the IG framework to attribute it to neurons in the targeted layer.
Refer to caption: extracted/5898011/latex/fig/method.jpg

Denote $V$ as the vocabulary space, $p$ as the in-context prompt of interest containing $n$ in-context examples, and $S_{p}=[y_{1},\dots y_{n}]$ the set of labels of the different examples in the prompt $p$. In our detection process, we are only interested in prompts $p$ on which the model outputs a wrong prediction $y\in S_{p}$ (hence the copying) and denote $\hat{y}\notin S_{p}$ as the ground-truth answer. Let $w^{l}\in\mathrm{R}^{d_{1}\times d_{2}}$ be the weight matrix of the linear layer at block $l$ on which our detection process operates. Furthermore, we define the model output $\mathrm{P}_{p}(y|\hat{w}^{l}_{j})$ as the probability of predicting a certain answer $y\in\mathrm{V}$.

$$ $\mathrm{P}(y|\hat{w}^{l}_{j},p)=\mathrm{P}(y|w^{l}_{j}=\hat{w}^{l}_{j},p),$ (4) $$

where $w^{(l)}_{j}$ denotes the $j$-th intermediate neuron in the $l$-th layer of interest (Figure 1),
$\hat{w}^{l}_{j}$ is a given constant that $w^{l}_{j}$ is assigned to. We define $l_{u}=\sum_{i}^{n}\mathrm{P}(y=y_{i}\in S_{p}|w^{l}_{j}=\hat{w}^{l}_{j},p)$ as the probability of predicting an answer which is provided in the prompt examples (in $S$), we also define $l_{v}=\mathrm{P}(y=\hat{y}\notin S_{p}|w^{l}_{j}=\hat{w}^{l}_{j},p)$ as the probability of predicting the ground-truth answer $\hat{y}$. Lastly, we define $\Delta L=l_{v}-l_{u}$ as the prediction shift.

Copying, by definition, occurs when the model’s prediction shifts from the true answer to one of the responses provided in the prompt. Thus, copying neurons are those that drive $\Delta L$.

By leveraging IG, we can attribute $\Delta L$ to individual components. This approach enables us to identify specific neurons responsible for copying within the LLM.

To quantify the contribution of a neuron $w_{j}^{l}$ to the prediction shift ($\Delta L$), we gradually change $w^{(l)}_{j}$ from 0 0 (the baseline) to its original value $\hat{w}^{(l)}_{j}$ computed by the model and integrate the gradients:

$$ $\displaystyle\text{Attr}(w^{l}_{j},p,S_{p})=\mathrm{IG}(w^{l}_{j})=(\hat{w}^{l }_{j}-0)\int_{\alpha=0}^{1}\frac{\partial|\Delta L|}{\partial w^{l}_{j}}d\alpha$ $\displaystyle\qquad\qquad\qquad=\sum_{i}^{n}\hat{w}^{(l)}_{j}\int_{\alpha=0}^{ 1}\frac{\partial\left|\mathrm{P}(y_{i}|\alpha\hat{w}^{l}_{j},p)-\mathrm{P}( \hat{y}|\alpha\hat{w}^{l}_{j},p)\right|}{\partial w^{l}_{j}}d\alpha,$ (5) $$

$$ $\text{Attr}(w^{l}_{j},p,S_{p})=\mathrm{IG}(w^{l}_{j})=\sum_{i}^{n}\frac{\hat{w }^{l}_{j}}{m}\sum_{r=1}^{m}\frac{\partial\left|\mathrm{P}(y_{i}|\frac{r}{m} \hat{w}^{l}_{j},p)-\mathrm{P}(\hat{y}|\frac{r}{m}\hat{w}^{l}_{j},p)\right|}{ \partial w^{l}_{j}},$ (6) $$

where $m=20$ is the number of approximation steps we use in our experiments, following (Sundararajan et al., ).

Finally, we compute the final attribution scores for the neurons $[w_{j}^{l}]$ by averaging $\text{Attr}(w_{j}^{l})$ across all samples in the synthetic dataset, resulting in a relevance score that quantifies the extent to which a neuron contributes to copying.

$$ $\text{R}(w_{j}^{l})=\frac{1}{|D|}\sum_{k=1}^{|D|}\text{Normalize}(\text{Attr}( w_{j}^{l},p_{k},S_{p_{k}})),$ (7) $$

where $D$ is the synthetic dataset and $p_{k}$ is the $k-$th sample of the synthetic dataset. and Normalize is the min-max normalizing function.

To mitigate the copying bias, we propose to suppress the weights of the detected copying neurons as follows:

$$ $w^{l}_{i}=\left\{\begin{aligned} 0,&\qquad&\text{R}(w^{l}_{i})\geq\sigma,\\ w^{l}_{i},&\qquad&\text{R}(w^{l}_{i})<\sigma.\end{aligned}\right.$ (8) $$

where $\sigma$ is the filtering threshold and is tuned using a validation dataset.

## 4 Experiments

Figure: Figure 2: Percentage of total errors and copying errors for both the pruned and un-pruned models, results are shown for 3 ICL tasks across 3 different models: GPT2-Small, BLoom-560M, and OPT-1.3B. The dack bar in each diagram represents the unpruned version while the lighter bar represents the pruned version; the entire bar height represents the total error of the model and the shaded part represents the copying error rate.
Refer to caption: x1.png

Our approach offers a generic method for detecting copying neurons, applicable to any language model. To demonstrate the generalizability of our method across different models and tasks, we conduct extensive experiments on a diverse set of LLMs. This includes the recent state-space models, such as Mamba (Gu & Dao, ) as well as a broad spectrum of transformer-based models, including OPT (Zhang et al., ), GPT2 (Radford et al., ), Bloom (Le Scao et al., ) and LLaMA Touvron et al. (); Dubey et al. ().

#### Data

In all of our experiments, we follow Hendel et al. () and Grazzi et al. () and study 18 tasks in 3 different categories including algorithmic (to lowercase, to uppercase, list first, list last, list max, list min, and next letter), linguistic (present to past, present to gerund, singular to plural, antonyms, and past to perfect), and knowledge (landmark, currency, country to capital, person to language, religion, and continent), the algorithmic tasks are generated automatically, for the linguistic we use the GitHub Repositries(^2^22[https://github.com/Drulac/English-Verbs-Conjugates](https://github.com/Drulac/English-Verbs-Conjugates))(^3^33[https://github.com/sindresorhus/irregular-plurals](https://github.com/sindresorhus/irregular-plurals)) and the knowledge data is taken from Meng et al. (). Additionally, we also incorporate real-world datasets like sentiment classification, including SST2, SST5, and subsets from the BIG-Bench Tasks (Suzgun et al., ). For more information about the data, refer to Appendix B.

#### Implementation Details

For each model, we use the synthetic validation dataset introduced in Section 3.2 to identify the optimal block and the number of copying neurons to prune as follows. IG, as defined in Equation 7, is applied across the layers of interest (summarized in Appendix. A) in all of the blocks in the model. As described in Section. 3.3, this procedure quantifies which neurons contribute most significantly to the copying errors. Furthermore, we use the validation set introduced in Section 3.2 in order to find the optimal pruning rate and the optimal block that maximizes the validation accuracy over the proxy ICL validation set, we apply multiple pruning rates ranging from 1% to 10%, in 1% increments (i.e., [1%, 2%, 3%, …, 10%]) over each layer in all of the blocks in the model, and use the best validation accuracy performing configuration to use for the unseen ICL tasks. This procedure allows us to determine the optimal block and the optimal number of neurons to prune for maximizing the accuracy on the validation proxy ICL set. The layers we choose to apply the detection and pruning procedures are summarized in Appendix A.

Figure: Figure 3: Summary of the results over the synthetic ICL tasks, for more information on the tasks and the exact numbers, refer to Appendix C.
Refer to caption: x2.png

Figure: Figure 4: Results of Llama-2 and Llama-3 over SST2, SST5, and Object Counting task from BBH benhmark
Refer to caption: x3.png

### 4.1 Validating the Prevalence of Copying Errors

To validate the significance of copying errors, we first conduct an error analysis for three representative ICL tasks (to lower, singular to plural, and present simple to gerund) across a range of large language models (GPT2-Small, Bloom-560M, and OPT-1.3B). In this analysis, we show the percentage of copying errors as defined in Section 3.1 out of the total number of errors. The results are presented in Figure 2. In this figure, we present a comparative analysis of error rates across different models and tasks. The dark bars show error rates for the baseline (unpruned) model, while the light bars display results for the pruned model. Each bar is composed of two elements: an unfilled portion representing the total error rate, depicted by the full height of the bar with a black outline, and a filled portion below indicating the copying error rate, which is a subset of the total error rate. Evidently, most of the errors in these few-shot ICL tasks stem from copying, these models tend to replicate responses based on examples provided in the prompt, rather than generalizing to new contexts. Additionally, our pruning method significantly reduces the number of copying errors thus also reducing the total error rate.

### 4.2 Tasks Evaluation

To demonstrate the generalizability of our approach across model architectures, we include a range of models: (1) Transformer-based: OPT, GPT-2, BLOOM, LLaMA, and (2) Mamba state space models of various sizes.

For each synthetic ICL task, we utilize the test sets introduced by Hendel et al. (). We report the mean accuracy across these test sets for each model and task configuration over the different shots (1, 2, 3, and 4) evaluated using various seeds. Figure. 3 presents a scatter plot comparing the performance of our pruned model against the baseline (non-pruned) model. A diagonal line representing equal performance is included for reference. Data points falling above this line indicate instances where our pruned model achieves higher accuracy than the baseline, while points below the line represent cases where the baseline model outperforms. As evident from the distribution of points that is dominated by those above the diagonal, our pruned model consistently demonstrates superior accuracy across a wide range of ICL tasks for different shot instances, underscoring the effectiveness of our targeting neuron pruning strategy. We believe that the fact our technique rarely leads to a performance drop – and when it does, the impact is only marginal – makes it particularly appealing for practical applications. For the exact numbers, we refer the reader to Appendix C.

In order to test our approach beyond this benchmark ICL tasks, we also test it on tasks that are based on three datasets of collected data: SST-2, SST-5, and the object counting sub-task from the BBH benchmark (Suzgun et al., ). In these cases, to allow a comparison with previous work, we use Llama-2 (Touvron et al., ) and Llama-3 (Dubey et al., ). The same synthetic dataset of vowels mapping is used for copying neurons detection.

We include two recent baselines. The first is Weighted In-Context Learning (WICL) by  Yang et al. (), which improves the performance of LLMs on downstream tasks by assigning and applying optimal weights to demonstration examples during ICL. The second, Automatic Prompt Engineer (APE) by Zhou et al. (), automatically generates and selects optimal instructions for large language models to improve their performance on various ICL tasks without relying on human-crafted prompts.

The results for these three benchmarks are presented in Figure 4. Evidently, our method significantly improves over the baseline LLM as well as over the two baselines.

### 4.3 Tasks-Vector Analysis

Next, we build upon the recent task-vectors framework of Hendel et al. () to study the relationship between copying neurons pruning and the quality of the emerged task vectors in ICL. By comparing the task vectors generated by pruned and unpruned models across various ICL scenarios, we seek to understand if our proposed targeted pruning can enhance a model’s ability to distill task-relevant information from demonstrations, specifically under the few shots settings. Furthermore, we follow the setup of Hendel et al. () and report ICL accuracy using the standard ICL promoting (denoted by ICL), the accuracy obtained by using the emerged Task-Vectors without pruning the copying neurons (denoted as Task-Vectors) and the accuracy obtained by using Task Vectors obtained from the model with the pruned copying neurons (denoted as Tak-Vectors-Pruned).

In Figure. 5, we show the results of OPT-2.7B and Bloom-560M models over the “Singular Plural” and “Country Capital” ICL tasks. As can be seen from the results, pruning the copying neurons indeed yields better Task-Vectors for ICL. This suggests that our pruning strategy may be effectively identifying and removing neurons that were interfering with the model’s ability to infer the underlying task. For more results on additional models and tasks, refer to Appendix D
.

Figure: Figure 5: Task-Vectors accuracies over OPT-2.7B and Bloom-560M models tested on (1) Singular Plural and (2) Country Capital ICL tasks. We show the Task-Vectors accuracies with and without pruning the detected copying errors, as can be seen, pruning the copying errors improves the quality of the extracted Task-Vectors across the different shots $\in[1,2,3,4]$ for the two models and ICL tasks.
Refer to caption: x4.png

### 4.4 Ablation Studies

We present multiple ablation studies to evaluate and understand the different components of our proposed detection and pruning methods. These ablation studies were conducted with the OPT-350M, GPT2-Small, and Bloom-560M models, over the Linguistic Antonyms, and Letter to Upper ICL tasks.

Our first experiment focuses on using “Prediction Shift” within the IG framework. We aim to determine whether applying IG to the prediction shift, as defined in Section 3.3, is essential for our proposed method’s effectiveness. To this end, we compare our approach with an alternative where IG is applied to the predicted probability (specifically, the maximum probability) instead. Results are presented in Tab.1 (“Max IG” row) clearly shows that the prediction shift is essential for the success of our proposed method.

We further check the effect of min-max normalizing the IG scores across the samples from the proxy ICL task we use in the detection process. The results without this normalization are reported under “w/o Norm”. As can be seen, normalizing the scores across the samples can significantly enhance the detection process of the copying neurons.

Additionally, we explore a baseline case where we randomly prune the same percentage of neurons as in the best-performing version of our method. This experiment, labeled as “Random” in Tab.1, shows a degradation in ICL accuracy for some shot settings, underscoring the importance of our targeted pruning strategy.

**Table 1: Ablation studies for the different components in our method over OPT-350M, GPT2-Small, and Bloom-560M models applied on the Linguistic Antonyms, and Letter to Upper ICL tasks**
|  | OPT-350M | GPT-Small | Bloom-560M |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Shot | ICL | Ours | Max IG | w/o Norm | Random | ICL | Ours | Max IG | w/o Norm | Random | ICL | Ours | Max IG | w/o Norm | Random |
| Linguistic Antonyms |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| 1 | 0.20 | 0.29 | 0.18 | 0.22 | 0.20 | 0.06 | 0.19 | 0.17 | 0.15 | 0.08 | 0.57 | 0.61 | 0.55 | 0.59 | 0.57 |
| 2 | 0.32 | 0.38 | 0.30 | 0.34 | 0.30 | 0.10 | 0.29 | 0.29 | 0.21 | 0.12 | 0.68 | 0.68 | 0.68 | 0.68 | 0.67 |
| 3 | 0.33 | 0.37 | 0.33 | 0.34 | 0.30 | 0.11 | 0.27 | 0.25 | 0.23 | 0.11 | 0.71 | 0.71 | 0.70 | 0.70 | 0.70 |
| 4 | 0.29 | 0.33 | 0.27 | 0.31 | 0.27 | 0.11 | 0.25 | 0.25 | 0.22 | 0.13 | 0.73 | 0.77 | 0.73 | 0.77 | 0.75 |
| Letter to Uppercase |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| 1 | 0.24 | 0.51 | 0.30 | 0.37 | 0.23 | 0.24 | 0.53 | 0.37 | 0.45 | 0.34 | 0.76 | 0.81 | 0.77 | 0.78 | 0.76 |
| 2 | 0.72 | 0.87 | 0.72 | 0.81 | 0.72 | 0.96 | 0.96 | 0.95 | 0.96 | 0.96 | 0.92 | 0.96 | 0.91 | 0.95 | 0.90 |
| 3 | 0.85 | 0.94 | 0.85 | 0.90 | 0.84 | 1.00 | 1.00 | 0.98 | 0.99 | 1.00 | 0.96 | 0.98 | 0.96 | 0.96 | 0.96 |
| 4 | 0.92 | 0.97 | 0.91 | 0.94 | 0.92 | 1.00 | 1.00 | 0.99 | 1.00 | 1.00 | 0.96 | 0.98 | 0.95 | 0.95 | 0.95 |

## 5 Conclusions and Discussion

We presented a novel method to mitigate copying bias in few-shot In-Context Learning by pruning neurons that are linked to this behavior according to the integrated gradients interoperability method. Our approach consistently improved performance across a variety of ICL tasks and model architectures. These findings highlight the potential of targeted neuron pruning as an effective strategy for optimizing the capabilities of large language models.

The “out-of-the-box” improvements provided by our method, without the need for task-specific data or fine-tuning, have significant practical implications for deploying more reliable few-shot learning systems. Our approach allows for the enhancement of LLM performance across a wide range of tasks using only a simple, synthetic dataset for neuron identification. Moreover, the consistent improvements observed across different model architectures suggest that this method could be broadly applicable, potentially becoming a standard post-processing step in LLM deployment pipelines.

The success of our pruning method in improving performance across various tasks indicates that “copying neurons” may be acting as a form of shortcut, inhibiting the model’s ability to engage in more sophisticated reasoning processes. This observation aligns with recent work on shortcut learning in neural networks (Yom Din et al., ; Belrose et al., ) and suggests that in-context learning quality could potentially be improved by carefully modulating the influence of different neuron groups or pathways.

Our results suggest that by pruning copying neurons, we enhance the model’s ability to distill task-relevant information from demonstrations, leading to more effective task vectors. This raises interesting questions about the relationship between neuron-level representations and the higher-level task embeddings captured by task vectors. Specifically, it may be useful to consider the representation in a way that disentangles multiple activation pathways.

### Acknowledgements

Ivan Titov is supported by the Dutch National Science Foundation (NWO Vici VI.C.212.053). This work was supported by a grant from the Tel Aviv University Center for AI and Data Science (TAD). This research was also supported by the Ministry of Innovation, Science & Technology, Israel (1001576154) and the Michael J. Fox Foundation (MJFF-022407). The contribution of the first author is part of a PhD thesis research conducted at Tel Aviv University.

## Appendix A Target Layers

This section outlines the exact layers targeted by our detection and pruning techniques. We concentrate on specific linear layers within each model block, with GPT2 being an exception where we focus on a CNN layer. Our approach encompasses both transformer-based and Mamba-based architectural designs. For a detailed breakdown of the exact layers our method operates on across various model families, refer to Table 2.

**Table 2: A summary of the specific layers on which we apply our detection and pruning method for different model families**
| Family | Layer | type |
| --- | --- | --- |
| OPT | fc1 | linear |
| GPT2 | mlp.c_fc | cnn |
| Bloom | mlp.dense_h_to_4h | linear |
| Llama | mlp.gate_proj | linear |
| Mamba | mixer.in_proj | linear |

## Appendix B Tasks Datasets

In all of our experiments, we follow Hendel et al. () and Grazzi et al. () and study 18 tasks in 4 different categories including algorithmic, translation, linguistic, and knowledge, the algorithmic tasks are generated automatically, for the linguistic we use the GitHub Repositries(^4^44[https://github.com/Drulac/English-Verbs-Conjugates](https://github.com/Drulac/English-Verbs-Conjugates))(^5^55[https://github.com/sindresorhus/irregular-plurals](https://github.com/sindresorhus/irregular-plurals)) and the knowledge data is taken from Meng et al. (). More details on the datasets are shown in Table B.

| Category | Reference | Task | Example |
| --- | --- | --- | --- |
| Algorithmic | T1 | To Lowercase | A $\rightarrow$ a |
| T2 | To Uppercase | a $\rightarrow$ A |  |
| T3 | List First | q,b,e,s $\rightarrow$ q |  |
| T4 | List Last | q,b,e,s $\rightarrow$ s |  |
| T5 | List Max | 2,1,5 $\rightarrow$ 5 |  |
| T6 | List Min | 2,1,5 $\rightarrow$ 1 |  |
| T7 | Next Letter | a,b,c $\rightarrow$ d |  |
| Linguistic | T8 | Present to past | go $\rightarrow$ went |
| T9 | Present to gerund | go $\rightarrow$ going |  |
| T10 | Singular to plural | cat $\rightarrow$ cats |  |
| T11 | Antonyms | happy $\rightarrow$ sad |  |
| T12 | Past to Perfect | catch $\rightarrow$ caught |  |
| Knowledge | T13 | Landmark | Maybach $\rightarrow$ Germany |
| T14 | Currency | Azerbaijan $\rightarrow$ Manat |  |
| T15 | Country to Capital | France $\rightarrow$ Paris |  |
| T16 | Person to Language | Macron $\rightarrow$ French |  |
| T17 | Religion | Muhammad $\rightarrow$ Islam |  |
|  | T18 | Continent | Swanson Mountains $\rightarrow$ Antarctica |

Beyond synthetic ICL tasks, we use sentiment classification datasets like SST2, and SST5, in SST2 the task is to classify text sentences into one of the two sentiments (negative or positive), while in SST5 the task is to classify text sentences into one of five sentiments (very positive, positive, neutral, negative, very negative). Additionally, we also incorporate the object-counting task from the BBH benchmark (Suzgun et al., ), where the task is to find out the total number of objects given in a context sentence, a sample illustration from the dataset is as follows:

“I have a car, and a toaster. How many objects do I have? $\rightarrow$ 2”

## Appendix C More Results

This section presents the full results of the ICL tasks. We compare the performance of various models across 18 ICL tasks. Each entry in the result tables contains two values: the left value represents the performance of the unpruned model, while the right value shows the performance achieved using our proposed method. Results are averaged across 5 different seeds.

**Table 3: Results over the GPT2 models family. The left number is the base model and the right number is with our pruning approach. Results are averaged across 5 different runs.**
| Task | GPT2-Small | GPT2-Medium |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-shot | 2-shot | 3-shot | 4-shot | 1-shot | 2-shot | 3-shot | 4-shot |  |
| T1 | .35$|$.50 | .94$|$.98 | .98$|$1.0 | .99$|$1.0 | .36$|$.72 | .91$|$.98 | .98$|$1.0 | .99$|$1.0 |
| T2 | .24$|$.53 | .96$|$.96 | 1.0$|$1.0 | 1.0$|$1.0 | .61$|$.78 | .99$|$1.0 | .99$|$1.0 | 1.0$|$1.0 |
| T3 | .42$|$.42 | .62$|$.67 | .63$|$.68 | .62$|$.65 | .73$|$.87 | .82$|$.87 | .86$|$.90 | .87$|$.92 |
| T4 | .42$|$.62 | .67$|$.67 | .68$|$.68 | .68$|$.68 | .66$|$.62 | .78$|$.75 | .82$|$.82 | .83$|$.80 |
| T5 | .57$|$.61 | .53$|$.59 | .52$|$.57 | .52$|$.57 | .54$|$.54 | .51$|$.53 | .52$|$.54 | .54$|$.56 |
| T6 | .30$|$.30 | .34$|$.38 | .38$|$.44 | .40$|$.45 | .35$|$.39 | .37$|$.42 | .41$|$.49 | .44$|$.52 |
| T7 | .02$|$.15 | .17$|$.55 | .34$|$.64 | .46$|$.68 | .40$|$.54 | .38$|$.54 | .32$|$.44 | .35$|$.45 |
| T8 | .10$|$.17 | .25$|$.38 | .26$|$.38 | .28$|$.40 | .16$|$.21 | .28$|$.38 | .32$|$.45 | .36$|$54 |
| T9 | .10$|$.30 | .29$|$.49 | .35$|$.52 | .43$|$.57 | .23$|$.30 | .47$|$.55 | .55$|$.64 | .63$|$.72 |
| T10 | .27$|$.30 | .35$|$.39 | .36$|$.40 | .37$|$.42 | .34$|$.38 | .45$|$.48 | .54$|$.54 | .59$|$.59 |
| T11 | .06$|$.19 | .10$|$.29 | .11$|$.27 | .11$|$.25 | .49$|$.52 | .62$|$.65 | .67$|$.67 | .67$|$.67 |
| T12 | .19$|$.26 | .33$|$.44 | .38$|$46 | .40$|$.49 | .26$|$.31 | .33$|$.42 | .36$|$.47 | .38$|$48 |
| T13 | .18$|$.18 | .16$|$.16 | .18$|$.18 | .20$|$.20 | .17$|$.17 | .24$|$.24 | .26$|$.26 | .26$|$.27 |
| T14 | .28$|$.30 | .38$|$.40 | .40$|$.43 | .42$|$.44 | .31$|$.31 | .41$|$.41 | .45$|$.47 | .47$|$.47 |
| T15 | .11$|$.11 | .22$|$.22 | .24$|$.24 | .25$|$.25 | .19$|$.22 | .30$|$.30 | .32$|$.32 | .36$|$.35 |
| T16 | .26$|$.30 | .35$|$.35 | .33$|$.33 | .33$|$.33 | .41$|$.45 | .47$|$.47 | .47$|$.45 | .49$|$.48 |
| T17 | .49$|$.54 | .64$|$.68 | .66$|$.70 | .67$|$.71 | .41$|$.45 | .55$|$.57 | .49$|$.55 | .51$|$.58 |
| T18 | .52$|$.54 | .48$|$.53 | .61$|$.63 | .70$|$.72 | .54$|$.54 | .55$|$.55 | .59$|$.60 | .62$|$.64 |

**Table 4: Results over the BLOOM models family. The left number is the base model and the right number is with our pruning approach. Results are averaged across 5 different runs.**
| Task | Bloom-560M | Bloom-1.1B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-shot | 2-shot | 3-shot | 4-shot | 1-shot | 2-shot | 3-shot | 4-shot |  |
| T1 | .34$|$.57 | .86$|$.92 | .98$|$1.0 | .99$|$1.0 | .21$|$.25 | .66$|$.70 | .85$|$.90 | .93$|$.95 |
| T2 | .76$|$.81 | .92$|$.96 | .96$|$.98 | .96$|$.98 | .65$|$.69 | .99$|$.99 | 1.0$|$1.0 | 1.0$|$1.0 |
| T3 | .37$|$.46 | .71$|$.77 | .83$|$.87 | .87$|$.90 | .43$|$.48 | .83$|$.85 | .92$|$.93 | .93$|$.95 |
| T4 | .43$|$.54 | .62$|$.71 | .68$|$.78 | .72$|$.83 | .52$|$.56 | .75$|$.75 | .77$|$.77 | .83$|$.80 |
| T5 | .51$|$.54 | .51$|$.51 | .56$|$.56 | .57$|$.57 | .58$|$.60 | .62$|$.62 | .60$|$.63 | .60$|$.60 |
| T6 | .22$|$.27 | .33$|$.36 | .38$|$.42 | .40$|$.42 | .24$|$.24 | .28$|$.28 | .33$|$.33 | .34$|$.34 |
| T7 | .15$|$.19 | .24$|$.28 | .28$|$.28 | .33$|$.32 | .21$|$.32 | .41$|$.60 | .44$|$.65 | .51$|$.73 |
| T8 | .36$|$.47 | .44$|$.58 | .49$|$.62 | .58$|$.67 | .48$|$.57 | .65$|$.72 | .74$|$.76 | .81$|$.83 |
| T9 | .35$|$.42 | .45$|$.56 | .56$|$.64 | .57$|$.66 | .53$|$.58 | .73$|$.76 | .81$|$.85 | .84$|$.84 |
| T10 | .30$|$.36 | .42$|$.49 | .41$|$.48 | .45$|$.52 | .40$|$.44 | .52$|$.55 | .56$|$.60 | .62$|$.66 |
| T11 | .42$|$.47 | .56$|$.61 | .60$|$.64 | .65$|$.65 | .57$|$.61 | .68$|$.68 | .71$|$.71 | .73$|$.77 |
| T12 | .39$|$.46 | .38$|$.50 | .45$|$.54 | .42$|$.51 | .40$|$.47 | .48$|$.56 | .62$|$.68 | .68$|$74 |
| T13 | .16$|$.18 | .22$|$.22 | .21$|$.21 | .23$|$.25 | .29$|$.31 | .40$|$.40 | .41$|$.41 | .42$|$.42 |
| T14 | .26$|$.28 | .39$|$.41 | .45$|$.48 | .41$|$.43 | .36$|$.36 | .48$|$.48 | .44$|$.46 | .47$|$.47 |
| T15 | .18$|$.22 | .26$|$.29 | .29$|$.33 | .29$|$.31 | .28$|$.32 | .37$|$.42 | .40$|$.45 | .42$|$.42 |
| T16 | .24$|$.26 | .36$|$.39 | .32$|$.34 | .30$|$.30 | .52$|$.52 | .58$|$.58 | .57$|$.61 | .57$|$.62 |
| T17 | .49$|$.52 | .45$|$.44 | .57$|$.55 | .63$|$.61 | .57$|$.60 | .63$|$.67 | .66$|$.66 | .70$|$.70 |
| T18 | .59$|$.59 | .55$|$.55 | .64$|$.62 | .66$|$.63 | .49$|$.51 | .53$|$.55 | .64$|$.64 | .65$|$.67 |

**Table 5: Results over the OPT models family, the left number is the base model and the right number is with our pruning approach, results are averaged across 5 different runs**
|  | OPT-125M | OPT-350M |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Task | 1 Shot | 2 Shot | 3 Shot | 4 Shot | 1 Shot | 2 Shot | 3 Shot | 4 Shot |
| T1 | .01$|$.18 | .37$|$.59 | .21$|$.35 | .21$|$.31 | .02$|$.15 | .23$|$.55 | .52$|$.81 | .72$|$.90 |
| T2 | .02$|$.17 | .27$|$.49 | .25$|$.28 | .21$|$.24 | .24$|$.51 | .72$|$.87 | .85$|$.94 | .92$|$.97 |
| T3 | .19$|$.34 | .30$|$.38 | .32$|$.38 | .32$|$.39 | .21$|$.29 | .55$|$.64 | .62$|$.64 | .61$|$.60 |
| T4 | .18$|$.37 | .30$|$.40 | .30$|$.40 | .29$|$.39 | .55$|$.55 | .36$|$.34 | .33$|$.30 | .27$|$.27 |
| T5 | .32$|$.39 | .33$|$.33 | .30$|$.29 | .30$|$.30 | .49$|$.51 | .46$|$.46 | .46$|$.46 | .47$|$.49 |
| T6 | .25$|$.35 | .31$|$.35 | .30$|$.34 | .33$|$.35 | .30$|$.29 | .41$|$.39 | .40$|$.38 | .44$|$.41 |
| T7 | .01$|$.07 | .03$|$.12 | .03$|$.09 | .04$|$.08 | .03$|$.05 | .15$|$.30 | .15$|$.28 | .19$|$.31 |
| T8 | .01$|$.01 | .03$|$.05 | .04$|$.07 | .04$|$.06 | .05$|$.07 | .12$|$.15 | .17$|$.21 | .19$|$.30 |
| T9 | .01$|$.06 | .11$|$.16 | .14$|$.18 | .15$|$.21 | .07$|$.07 | .19$|$.22 | .23$|$.25 | .23$|$.28 |
| T10 | .10$|$.19 | .22$|$.25 | .28$|$.30 | .29$|$.32 | .28$|$.33 | .35$|$.37 | .39$|$.41 | .38$|$.40 |
| T11 | .05$|$.10 | .11$|$.15 | .15$|$.17 | .16$|$.18 | .20$|$.29 | .32$|$.38 | .33$|$.37 | .29$|$.33 |
| T12 | .02$|$.09 | .12$|$.16 | .16$|$.16 | .16$|$.16 | .16$|$.19 | .24$|$.26 | .27$|$.29 | .26$|$.28 |
| T13 | .08$|$.12 | .07$|$.12 | .10$|$.15 | .11.16 | .14$|$.16 | .24$|$.27 | .32$|$.34 | .32$|$.32 |
| T14 | .09$|$.25 | .26$|$.34 | .31$|$.33 | .31$|$.34 | .33$|$.37 | .48$|$.50 | .42$|$.44 | .51$|$.54 |
| T15 | .04$|$.09 | .13$|$.16 | .15$|$.18 | .15$|$.17 | .11$|$.13 | .20$|$.20 | .24$|$.25 | .25$|$.26 |
| T16 | .19$|$.24 | .26$|$.30 | .25$|$.28 | .24$|$.26 | .29$|$.28 | .42$|$.41 | .45$|$.44 | .46$|$.45 |
| T17 | .54$|$.54 | .54$|$.57 | .53$|$.55 | .62$|$.62 | .51$|$.50 | .64$|$.64 | .67$|$.67 | .72$|$.72 |
| T18 | .56$|$.56 | .65$|$.63 | .65$|$.65 | .69$|$.69 | .51$|$.53 | .47$|$.50 | .56$|$.58 | .60$|$.67 |
|  | OPT-1.3B | OPT-2.7B |  |  |  |  |  |  |
| Task | 1 Shot | 2 Shot | 3 Shot | 4 Shot | 1 Shot | 2 Shot | 3 Shot | 4 Shot |
| T1 | .01$|$.13 | .24$|$.83 | .76$|$.98 | .94$|$1.0 | .09$|$.13 | .89$|$.98 | .99$|$1.0 | 1.0$|$1.0 |
| T2 | .02$|$.14 | .73$|$.82 | .92$|$.97 | .97$|$.99 | .04$|$.12 | .72$|$.93 | .86$|$.98 | .91$|$1.0 |
| T3 | .78$|$.93 | .91$|$.91 | .92$|$.92 | .89$|$.88 | .84$|$.90 | .93$|$.95 | .97$|$1.0 | .98$|$1.0 |
| T4 | .37$|$.56 | .62$|$.74 | .59$|$.71 | .68$|$.74 | .47$|$.57 | .64$|$.76 | .77$|$.80 | .74$|$.76 |
| T5 | .67$|$.70 | .61$|$.64 | .65$|$.67 | .62$|$.68 | .56$|$.56 | .49$|$.49 | .54$|$.54 | .55$|$.55 |
| T6 | .29$|$.31 | .43$|$.45 | .51$|$.53 | .54$|$.59 | .35$|$.37 | .47$|$.47 | .51$|$.52 | .54$|$.54 |
| T7 | .28$|$.34 | .52$|$.65 | .53$|$.60 | .56$|$.56 | .14$|$.23 | .44$|$.50 | .50$|$.58 | .44$|$.49 |
| T8 | .11$|$.16 | .29$|$.45 | .42$|$.60 | .57$|$.70 | .20$|$.32 | .49$|$.65 | .61$|$.73 | .68$|$.75 |
| T9 | .12$|$.14 | .27$|$.45 | .49$|$.63 | .59$|$.71 | .28$|$.39 | .68$|$.75 | .80$|$.83 | .82$|$.85 |
| T10 | .34$|$.36 | .36$|$.41 | .42$|$.45 | .45$|$.49 | .35$|$.40 | .46$|$.50 | .52$|$.56 | .59$|$.62 |
| T11 | .60$|$.62 | .70$|$.71 | .74$|$.74 | .75$|$.75 | .64$|$.64 | .75$|$.77 | .76$|$.78 | .78$|$.78 |
| T12 | .27$|$.30 | .30$|$.37 | .38$|$.43 | .41$|$.49 | .28$|$.33 | .40$|$.47 | .42$|$.44 | .47$|$.52 |
| T13 | .32$|$.38 | .38$|$.42 | .46$|$.48 | .45$|$.49 | .42$|$.46 | .50$|$.52 | .55$|$.57 | .61$|$.61 |
| T14 | .41$|$.41 | .48$|$.48 | .51$|$.51 | .51$|$.51 | .43$|$.43 | .53$|$.55 | .49$|$.50 | .52$|$.52 |
| T15 | .24$|$.26 | .38$|$.40 | .42$|$.44 | .44$|$.47 | .25$|$.30 | .40$|$.45 | .50$|$.54 | .53$|$.57 |
| T16 | .56$|$.56 | .65$|$.65 | .65$|$.65 | .68$|$.68 | .70$|$.70 | .74$|$.74 | .72$|$.72 | .74$|$.74 |
| T17 | .55$|$.60 | .67$|$.68 | .73$|$.73 | .77$|$.77 | .54$|$.59 | .62$|$.64 | .72$|$.72 | .72$|$.72 |
| T18 | .55$|$.55 | .61$|$.61 | .65$|$.65 | .71$|$.71 | .48$|$.50 | .59$|$.59 | .69$|$.69 | .73$|$.73 |

**Table 6: Results over the Mamba models family. The left number is the base model and the right number is with our pruning approach. Results are averaged across 5 different runs.**
| Task | Mamba-130M | Mamba-370M |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-shot | 2-shot | 3-shot | 4-shot | 1-shot | 2-shot | 3-shot | 4-shot |  |
| T1 | .77$|$.97 | 1.0$|$1.0 | 1.0$|$1.0 | 1.0$|$1.0 | .95$|$.99 | 1.0$|$1.0 | 1.0$|$1.0 | 1.0$|$1.0 |
| T2 | .09$|$.73 | .95$|$1.0 | 1.0$|$1.0 | 1.0$|$1.0 | .55$|$.83 | .88$|$.92 | 1.0$|$1.0 | 1.0$|$1.0 |
| T3 | .39$|$.72 | .82$|$.95 | .89$|$1.0 | .92$|$1.0 | .75$|$.79 | .98$|$.99 | .99$|$1.0 | 1.0$|$1.0 |
| T4 | .73$|$.73 | .86$|$.88 | .92$|$.92 | .93$|$.93 | .49$|$.53 | .80$|$.80 | .86$|$.86 | 0.91$|$.93 |
| T5 | .47$|$.50 | .49$|$.53 | .53$|$.55 | .53$|$.56 | .53$|$.56 | .51$|$.53 | .62$|$.64 | .60$|$.62 |
| T6 | .32$|$.35 | .46$|$.48 | .47$|$.50 | .47$|$.52 | .34$|$.38 | .41$|$.43 | .51$|$.53 | .54$|$.55 |
| T7 | .14$|$.14 | .70$|$.70 | .73$|$.73 | .81$|$.83 | .60$|$.69 | .87$|$.89 | .88$|$.90 | .85$|$.85 |
| T8 | .08$|$.20 | .36$|$.48 | .42$|$.54 | .51$|$.57 | .38$|$45 | .56$|$.59 | .63$|$.63 | .70$|$.70 |
| T9 | .11$|$.25 | .36$|$.50 | .50$|$.55 | .60$|$.65 | .55$|$.59 | .74$|$.77 | .82$|$.85 | .86$|$.88 |
| T10 | .25$|$.25 | .30$|$.30 | .39$|$.42 | .48$|$.50 | .34$|$.41 | .46$|$.52 | .56$|$.58 | .68$|$.70 |
| T11 | .20$|$.20 | .29$|$.29 | .35$|$.35 | .30$|$.30 | .45$|$.51 | .61$|$.64 | .70$|$.73 | .66$|$.69 |
| T12 | .19$|$.22 | .36$|$.43 | .47$|$.52 | .59$|$.62 | .32$|$.40 | .35$|$.42 | .45$|$.49 | .49$|$.52 |
| T13 | .16$|$.20 | .26$|$.26 | .28$|$.30 | .36$|$.36 | .37$|$.41 | .52$|$.55 | .54$|$.57 | .53$|$.55 |
| T14 | .28$|$.33 | .33$|$.36 | .36$|$.40 | .40$|$.42 | .35$|$.35 | .41$|$.41 | .46$|$.48 | .46$|$.48 |
| T15 | .08$|$.10 | .14$|$.14 | .19$|$.19 | .20$|$.20 | .31$|$.36 | .45$|$.49 | .54$|$.54 | .56$|$.56 |
| T16 | .39$|$.44 | .50$|$.50 | .60$|$.60 | .63$|$.63 | .62$|$.66 | .68$|$.72 | .70$|$.73 | .72$|$.75 |
| T17 | .52$|$.55 | .65$|$.67 | .61$|$.65 | .64$|$.70 | .51$|$.57 | .58$|$.61 | .58$|$.61 | .58$|$.60 |
| T18 | .48$|$.50 | .48$|$.50 | .56$|$.56 | .51$|$.53 | .53$|$.56 | .51$|$.55 | .59$|$.59 | .60$|$.60 |

## Appendix D Task Vectors

We provide additional results for the task-vectors experiment, we include two additional models (GPT2-Small and Bloom1.1B) over the Algorithmic Next Letter task and Linguistics Antonyms.
.

Figure: Figure 6: Additional quantitative results for the task-vectors experiment
Refer to caption: x5.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - How does the pruned model perform on downstream tasks, such as reasoning, math, long-context understanding, instruction following, and safety? While it may not be necessary to conduct an exhaustive evaluation due to computational resource limitations, at the very least, relevant experiments should be included.

- Does it matter if the model is a base model or an instruct model? For example, if Bloom-560M is turned into an instruct model or fine-tuned with DPO, will the experimental results change or remain the same?

- Can data engineering help solve this problem? If it can, why do we still need an additional step to prune the model?
soundness: 3
strengths: This paper successfully implemented neural pruning into tackling the copy bias in ICL. The proposed IG method with pruning strategy are relatively reasonable to solve the copy bias problem. This is a significant contribution as previous work primarily focused on prompt engineering and calibration methods without directly addressing the internal model dynamics causing copying errors.

Also, the proposed method is versatile, as it can be applied across LLM architectures, including both Transformersand SSM, without needing task-specific modifications. This general applicability makes the method highly practical for a wide range of ICL tasks and model types. It provides “out-of-the-box” improvements without requiring model retraining or access to large amounts of task-specific data, making it an efficient post-processing step in LLM deployment pipelines.

The authors use a vowel-counting task as a proxy for detecting copying errors, where language models are prone to making copying mistakes. The proxy ICL dataset generation method is clear are easy to follow. The diverse range of tasks helps to demonstrate the robustness and general applicability
summary: The paper presents a novel approach to addressing a key issue in large language models (LLMs) during In-Context Learning (ICL)—the tendency to copy answers from provided examples, known as “copying bias.” Instead of learning and generalizing from the patterns in the context, models often replicate the examples, especially in few-shot learning settings. To tackle this, the authors propose using Integrated Gradients (IG) to identify neurons responsible for this copying behavior. By pruning these “copying neurons,” the model’s ability to reason and generalize is improved.

The method is task-agnostic and applicable across various LLM architectures, including Transformer-based models like GPT, Bloom, OPT, and LLaMA, as well as State-Space Models such as Mamba. Extensive experiments across different tasks show that neuron pruning significantly reduces copying errors, enhancing overall task performance. The paper also demonstrates that this pruning improves the quality of task vectors, leading to better task recognition and generalization.
weaknesses: The primary weakness of this paper is the lack of a comprehensive evaluation of the model’s capabilities after pruning. While the method effectively mitigates copy bias, this paper presents more of a case study rather than demonstrating its real-world applicability. Pruning techniques generally have an impact on the overall performance of the model, and as such, it is critical for the authors to provide results showing how the pruned model performs on downstream tasks. Unfortunately, this paper does not offer that kind of evaluation, which is a necessary component to fully understand the implications of the pruning approach on the model’s general capabilities.

Another issue with the paper is the limited comparison with baselines. The authors provide several ICL-related baselines, but they seem to operate under the assumption that copy bias can only be addressed with ICL-focused solutions. The paper would benefit from a broader exploration of the problem, including a more detailed definition of copy bias and an examination of its resolution under different types of models, such as base and instruct models. Notably, the authors did not experiment with instruct models, which may behave very differently, especially with methods like APE that are known to improve prompt-following capabilities. Additionally, it would be valuable to explore whether fine-tuning or data engineering, such as modifying the SFT dataset or providing higher-quality alignment data, could address the issue. If data engineering alone could solve the problem, one must question the necessity of introducing an additional, complex pruning process, particularly for large models where pruning is a non-trivial engineering task.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the weakness section above where I asked some questions.
soundness: 2
strengths: 1. The author uses illustration to help the reader better understand how the proposed method works.
2. The author considers different ICL tasks for validation.
summary: This paper introduces "copying bias" in large language models (LLMs) in-context learning and presents a method for detecting copying neurons and improves the performance on ICL tasks. The author conducts experiments on ICL tasks and task vectors to verify the claim.
weaknesses: 1. Lacks some implementation detail: in Section 3.3, what is the size of n? The author assumes that $\hat{y}\notin S_p$ and therefore if n is large, there will be many repetitions in $S_p$ since labels in $S_p$ are the count of vowels, which are integers. Also, does the author explicitly held out an integer so that when constructing $S_p$, the words cannot have the number of vowels equal to this held-out integer?
2. Universality of copying neurons detection: This paper only studies how to detect copy neurons on the vowel counting task. What if this proxy ICL task is switched to other tasks? Will the copying neuron detected still work, and if not (e.g., if all elements in $S_p$ are unique), what makes vowel counting special that it works as a proxy ICL task?
3. Continuation of the previous point: method in 3.3 assumes that the output has single token. This is probably fine for the vowel counting task as single-digits are single-tokens. What if we have other ICL tasks as proxy ICL task and the outputs are multi-tokens?
4. How are errors being calculated? There are two ways: 1) probability-based: all probability summed for labels not equal to the true label and averaged across all inputs; 2) accuracy-based: 1 - (accuracy for the correct output) with greedy decoding.
5. In Figure 2 and 3, it seems that only models up to 1.3B are evaluated. In Figure 2 it seems that a large proportion of error comes from the copying error. However, for larger models (e.g., 8B, 14B) and for the singular to plural task, I suspect that larger model will not have this high copying error if each in-context example is unique. For 1-shot and 2-shot setting, the performance might simply be improved by adding a task description at the very front to eliminate task ambiguity.
6. Pruning the neuron is essentially "training" the model on the ICL task. It would be convincing if the author presents other training-based method as baselines (with a sweep of hyperparameter on the validation set) in Figure 2 and 3.
7. For Figure 5, while there is improvement, the author does not show how much improvement comes from less copying bias. The author has larger models like Llama-3-8B in Figure 4 experiment but there lacks evidence that the pruning decreases the copying error.

Minor issues:
1. Example in the introduction can be ambiguous. it could also be the case to count the number of words and this case the correct answer for "Florida" is 1.
2. Line 063-064: the first quotation mark of "Copying Neurons." is not correct.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. In Figure 2, it’s clear that the copying error rate has dropped consistently across different models and tasks, contributing to improved accuracy. However, I’m surprised that copying errors still dominate the error rate in the unpruned model. Do the authors have any insight into why this phenomenon occurs? Additionally, how does the copying error rate compare for the vowel-counting task specifically? I’m curious about the extent to which copying errors can be reduced on the in-distribution task (vowel counting) versus out-of-distribution tasks (the other ICL tasks).”

2. Why are all the tasks restricted to 4 shots? My intuition is that with more few shot examples, the model is more likely to pick up the actual pattern and less likely to copy, therefore, the method will become less effective. This analysis is not  to weaken the method — as authors have pointed out that some task might have limited ICL examples as a restriction, and the proposed method is a great way to mitigate the bias. I am trying to get a better understanding on the limitation of the proposed method. 

3. For the BBH object counting task, my understanding is that these are complex reasoning tasks that benefit from Chain-of-Thought (CoT) reasoning, which ICL examples could potentially demonstrate. Can the authors clarify if CoT prompting is employed in either the baseline models or the copy-error-pruned method?
soundness: 3
strengths: 1. This paper demonstrates that mechanistic interpretability methods like Integrated Gradients (IG), typically used for understanding model behaviors, can have practical implications such as improving in-context learning (ICL). By identifying and pruning neurons responsible for copying behavior, the paper presents an innovative approach to enhancing ICL performance.

2. The authors conduct experiments across a wide range of tasks and model architectures, demonstrating consistent performance improvements. This thorough evaluation reinforces the robustness and generalizability of the proposed method.

3.  I appreciate that authors can first justify the prevalence of copying errors to justify why such a bias is worth targeting across different models and tasks. This rigorous approach highlights the significance of addressing this bias, underscoring the method’s relevance and value in improving model generalization.
summary: This paper addresses the issue of copying bias in LLMs during in-context learning (ICL). To mitigate copying-biase, the authors identify and prune specific "copying neurons". They use Integrated Gradients (IG) on a synthetic vowel-counting task to locate these neurons and then prune them. The approach is shown to generalize across multiple model architectures (e.g., GPT-2, OsPT, BLOOM) and improves the quality of task vectors—representations that encapsulate task-specific information—indicating better task recognition. This method provides an effective way to enhance ICL performance without task-specific fine-tuning, offering a broadly applicable strategy for deploying more reliable LLMs.
weaknesses: 1. **Similarity of Tasks and Limited Generalizability**: The 18 tasks studied in this paper share notable similarities with vowel counting, meaning that the copying neurons identified through Integrated Gradients primarily apply to a set of very related and rather synthetic tasks. While the authors also attempt to show improvements on non-synthetic tasks, I am not sure how to interpret "copying error" on datasets like SST2 and SST5, which are multiple-choice tasks with a limited set of choices (e.g., 2 or 5 options)? In such cases, does the model exhibit a genuine copying bias or if the improvements are driven by other factors? Would the copying bias be easily mitigated by using unbiased ICL samples (for example, one sample for each answer choice A/B/C/D).

2. **Limited Exploration of Trade-offs**: The paper does not extensively address the potential trade-offs introduced by pruning neurons. Authors only claim that the pruned model never has a drop in performance but what about other side effects? For instance, pruning could theoretically affect the model's performance on tasks where copying might be beneficial (e.g., fact-based retrieval tasks). A more detailed analysis of cases where pruning could have a negative impact would be beneficial.

3. **Minor Reproducibility Concerns**: The paper lacks details on certain experimental settings, such as the dataset specifications (e.g., training/validation sizes, number of shots) for the vowel-counting data, as well as the hyperparameter settings used for pruning. Including these details would enhance reproducibility and provide greater clarity on the experimental setup.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: * **Q1**: Is the copying bias merely an artifact of having a limited number of examples (e.g., <= 4)? For instance, what happens in a 64-shot setting? Could this bias indicate that the copying issue arises from an ill-posed in-context learning (ICL) problem rather than a general model flaw?

* **Q2**: Could the authors provide additional details about the proxy dataset? For example, how many shots are used, and could a few more examples be included for completeness?

* **Q3**: How does pruning of these neurons affect the overall model performance beyond ICL tasks? For example, how do the pruned LLMs perform on more general benchmarks such as [Live-Bench](https://livebench.ai/) or other reasoning benchmarks (e.g., MMLU or GSM8K)? Or, what’s the performance on the indirect object identification (IOI) task? The authors suggest that their method “could be broadly applicable, potentially becoming a standard post-processing step in LLM deployment pipelines” (l. 497-498). However, it remains unclear if this would be a sensible approach for LLMs deployed “in the wild.”

* **Q4**: Are the minimum and maximum values of the normalizing functions calculated across the dataset alone, or are they based on both the dataset and all neurons?

* **Q5**: What’s the significance of the meta-optimization (l. 319-369)? Why not greedily remove the neurons based on the “copying score” from Eq. 7 (this may be related to the question in **Q4**)?

* **Q6**: What’s the relation of the real-word experiments in Sec. 4.2 (sentiment analysis and object counting) to the copying behavior that’s the main topic of the present work?
soundness: 2
strengths: * **S1**: Previous work has focused on improving ICL primarily through black-box approaches (e.g., prompt tuning), whereas this work adopts a mechanistic perspective. This shift is valuable, as it has the potential to deepen our understanding of the internal workings of ICL in these models.

* **S2**: The pruning method effectively improves ICL performance across various architectures and tasks. Additionally, it enhances the quality of task vectors, also leading to improvements in ICL performance.
summary: This paper shows that large language models (LLMs) can exhibit a copying bias, often simply copying their answers from the in-context examples. The authors identified task-agnostic neurons responsible for this copying behavior using integrated gradients [1] on a vowel-counting proxy task (Secs. 3.2 & 3.3). They mitigate their influence by pruning them (i.e., simply setting them to 0; Eq. 8). Their neuron pruning method improves in-context learning (ICL) performance across various tasks, including the ICL benchmark proposed by Hendel et al [2], two sentiment analysis tasks and an object counting task, and various models, including GPT-2, Llama variants or Mamba. They also find that pruning these neurons improves the effectiveness of task vectors [2].

---

[1] Sundararajan, M., Taly, A., & Yan, Q. (2017). Axiomatic attribution for deep networks. In ICML.

[2] Hendel, R., Geva, M., & Globerson, A. (2023). In-context learning creates task vectors. In EMNLP.
weaknesses: * **W1**: The overall technical contributions (Sec. 3) are modest. The work applies a well-established attribution method (integrated gradients) to identify neurons responsible for a specific behavior (copying) and then prunes these neurons to achieve the intended outcome (no copying, though that’s unclear to me, see **W2**). Unfortunately, this is highly similar to prior work such as the one by Ghorbani et al [1]. For example, they used Shapley values (another attribution method) to identify neurons that led to gender-specific biases and pruned them to achieve gender fairness. Thus, while the application to copying in ICL is novel, the methodological contributions remain limited (i.e., the use of prediction shift ΔL instead of max probability for integrated gradients, l. 227-230).

* **W2**: It is unclear to me whether the performance improvement results from the reduction of copying neurons or from the pruning of neurons exhibiting different behavior. For example, why would performance improve on *copying* tasks such as “list first” or “list max” in the ICL benchmark of Hendel et al [2], which require copying from the in-context examples? Tabs. 3-6 consistently also show improvements on these *copying* tasks. If the identified neurons indeed suppress the copying behavior, we would expect to see drops in performance on such tasks. This seems to contradict one of the authors’ main claims that they “identify specific neurons responsible for triggering the copying behavior (copying neurons)” (l. 71-73).

* **W3**: It is unclear whether the copying bias is merely an artifact of the ill-posed nature of few-shot ICL problems, and whether it remains relevant for many-shot ICL problems. Does this copying bias persist beyond the few-shot setting? (This point is also raised in the questions below.)

* **W4**: The writing could be improved. There are several typos throughout, and clarity and reproducibility would benefit from including all necessary details, such as those regarding proxy dataset generation.

* **W5**: Code is not provided, making reproducibility challenging.

--- 

[1] Ghorbani, A., & Zou, J. Y. (2020). Neuron shapley: Discovering the responsible neurons. NeurIPS.

[2] Hendel, R., Geva, M., & Globerson, A. (2023). In-context learning creates task vectors. In EMNLP.

Meta-review: The paper proposes pruning neurons responsible for copying bias in few-shot in-context learning, identified via Integrated Gradients. While the authors report some improvements on synthetic tasks, the overall scope and novelty are limited. Key details are missing, and the paper does not situate its contribution well within the broader literature on in-context learning, including theoretical and empirical advances such as Pan et al. Comparisons to alternative interventions (e.g., better prompts, data engineering) are lacking, and real-world applicability is unclear.

Key limitations
- Evaluation focuses on synthetic or simplistic tasks. Real-world relevance is low.
- The approach draws on existing attribution-based pruning methods, so its novelty is modest.
- Missing details and the absence of code make reproducibility difficult.
- The paper neglects important prior work on ICL, which provides different ways to explain low-accuracy scenarios.
- The effect of pruning on tasks that require copying is not fully addressed.

Suggestions
- Improve literature coverage, referencing theoretical and empirical studies on ICL failures and task recognition and task learning.
- Evaluate on more diverse and many-shot tasks.
- Provide thorough methodology, release code, and compare against other potential solutions.

Overlall, the current version needs broader validation, stronger comparisons, and clearer ties to prior work before acceptance.

Decision: Reject