How do Large Language Models Learn In-Context? Query and Key Matrices of In-Context Heads are Two Towers for Metric Learning

Overview

• The paper investigates the mechanism behind in-context learning (ICL) on sentence classification tasks with unrelated labels.
• They find that intervening in only 1% of the neural network heads (called "in-context heads") significantly impacts ICL accuracy, reducing it from 87.6% to 24.4%.
• To understand this phenomenon, the authors analyze the value-output vectors in these in-context heads and discover that the vectors at each label position contain substantial information about the corresponding labels.
• They observe that the shift in prediction from one label to another is due to changes in the attention scores in the in-context heads.
• Based on these findings, the authors propose a hypothesis for ICL, explaining the majority label bias and recency bias, and suggesting two methods to reduce these biases.

Plain English Explanation

In this research, the authors looked at how large language models learn to perform sentence classification tasks when the labels are not directly related to the content of the sentence (e.g., classifying a sentence as "foo" or "bar"). They found that even small changes to a specific part of the neural network, called the "in-context heads," can have a big impact on the model's performance.

By analyzing the internal workings of these in-context heads, the authors discovered that the model is essentially learning to identify features of the labels themselves, rather than just focusing on the content of the sentence. This means the model can be biased towards predicting the most common label or the most recent label it saw, rather than the label that best fits the sentence.

The authors propose a way to understand this process, which involves the model learning a "similarity metric" between the features of the sentence and the features of each possible label. They also suggest two methods that can help reduce these biases, which could be useful for improving in-context learning in language models.

Technical Explanation

The researchers conducted experiments on sentence classification tasks with semantically-unrelated labels (e.g., "foo" or "bar"). They found that intervening in only 1% of the neural network heads, specifically the "in-context heads," significantly reduced the in-context learning (ICL) accuracy from 87.6% to 24.4%.

To understand this phenomenon, the authors analyzed the value-output vectors in these in-context heads. They discovered that the vectors at each label position contained substantial information about the corresponding labels. Furthermore, they observed that the shift in prediction from one label to another was due to changes in the attention scores of the in-context heads at the label positions.

Based on these findings, the authors proposed a hypothesis for ICL: in the in-context heads, the value-output matrices extract label features, while the query-key matrices compute the similarity between the features at the last position and those at each label position. This process can be seen as the model learning a "similarity metric" between the input features and the label features.

Using this hypothesis, the authors were able to explain the majority label bias and recency bias observed in ICL. They also proposed two methods to reduce these biases by 22% and 17%, respectively. These methods involve modifying the training process to encourage the model to learn a more balanced and robust in-context learning approach.

Critical Analysis

The paper provides a detailed and insightful analysis of the underlying mechanisms of in-context learning in language models. The authors' hypothesis about the role of the query-key and value-output matrices in learning a similarity metric between input features and label features is a compelling explanation for the observed biases.

However, the authors acknowledge that their analysis is limited to a specific set of sentence classification tasks with semantically-unrelated labels. It would be valuable to see if their findings and proposed methods generalize to a wider range of in-context learning tasks, including those with more realistic and semantically-related labels.

Additionally, the authors' experiments focused on intervening in only 1% of the neural network heads. It would be interesting to explore the impact of intervening in different proportions of heads or different subsets of the network to gain a more comprehensive understanding of the underlying mechanisms.

Overall, this paper provides a significant contribution to the understanding of in-context learning in language models and offers promising directions for further research and potential improvements in this area.

Conclusion

This research paper offers valuable insights into the underlying mechanisms of in-context learning (ICL) in language models. The authors' discovery that the value-output and query-key matrices in a small subset of neural network heads, called "in-context heads," play a crucial role in ICL provides a new perspective on how these models learn and make predictions.

By proposing a hypothesis that explains the majority label bias and recency bias in ICL, and suggesting methods to reduce these biases, the authors have made important strides towards improving the robustness and reliability of in-context learning. These findings could have significant implications for the development of more advanced and versatile language models, with the potential to enhance their ability to perform a wide range of tasks in an effective and unbiased manner.