Mitigating Copy Bias in In-Context Learning through Neuron Pruning

TL;DR

This paper introduces a neuron pruning technique to reduce copying bias in large language models during in-context learning, improving their ability to learn underlying patterns across diverse tasks.

Contribution

The authors propose a novel neuron pruning method using Integrated Gradients to mitigate copying bias in LLMs without architectural modifications.

Findings

Pruning neurons identified as copying-prioritized improves ICL performance.

Method is effective across different LLM architectures.

Pruning enhances task recognition quality in models.

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.