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

Abstract: We investigate the mechanism of in-context learning (ICL) on sentence classification tasks with semantically-unrelated labels (“foo”/“bar”). We find intervening in only 1% heads (named “in-context heads”) significantly affects ICL accuracy from 87.6% to 24.4%. To understand this phenomenon, we analyze the value-output vectors in these heads and discover that the vectors at each label position contain substantial information about the corresponding labels. Furthermore, we observe that the prediction shift from “foo” to “bar” is due to the respective reduction and increase in these heads’ attention scores at “foo” and “bar” positions. Therefore, we propose a hypothesis for ICL: in 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. The query and key matrices can be considered as two towers that learn the similarity metric between the last position’s features and each demonstration at label positions. Using this hypothesis, we explain the majority label bias and recency bias in ICL and propose two methods to reduce these biases by 22% and 17%, respectively.

## introduction in-context learning (icl) is an emergent ability @xcite of large language models @xcite @xcite . by using some demonstration-label pairs as prompts, icl performs well without updating parameters on many tasks, such as machine translation @xcite , complexity reasoning @xcite , compositional generalization @xcite and information extraction @xcite . because the mechanism of icl remains unclear, many studies focus on understanding how icl works. @xcite find that icl can be disentangled into task recognition (tr) and task learning (tl). tr does not rely on the demonstration-label mappings because the roles of demonstrations and labels are helping the model know "what is the task". in this situation, the model have similar predictions when the mappings are wrong @xcite , because the predictions are based on pre-trained priors. on the other hand, tl relies on the demonstration-label mappings because the semantic priors are removed. for example, in an icl sentiment classification task, if the labels are "positive/negative", the task is tr. if the labels are "foo/bar", the task is tl because the labels are semantically-unrelated @xcite . @xcite analyze the information flow by averaging all attention heads and find the label words are anchors to merge the semantic information of corresponding demonstrations in shallow layers, and information is extracted from label words to the final prediction in deep layers. although these studies are important for understanding icl, the exact mechanism of icl remains a mystery for several reasons. firstly, the information flow is typically observed as an average across each head, but understanding icl requires exploring the precise importance of each head. secondly, each head has a query matrix, key matrix, value matrix, and output matrix; it is essential to study the role of each matrix in detail. lastly, icl is plagued by issues such as majority label bias and recency bias, and how to explain and mitigate these biases has not yet been thoroughly investigated. in this paper, we address these issues by identifying important heads for icl and studying the roles of each matrix within these heads. using two methods, we identify 12 important heads (named in-context heads) that significantly affect icl accuracy across five datasets, reducing it from 87.6% to 24.4% on average. intervening in 6 heads (fooheads) decreases the probabilities of "foo", while intervening in the other 6 heads (barheads) reduces the probabilities of "bar". to explore the reason of this phenomenon, we study these heads' value-output vectors computing by value-output matrices, and find that the vectors on label positions contain much information about the corresponding labels. moreover, we observe the attention scores in the in-context heads when predictions shift from "foo" to "bar", and find that the attention scores at "foo" positions decrease, while the attention scores at "bar" positions increase. based on these observations, we propose a hypothesis for icl, as shown in figure 1 : in in-context heads, value-output matrices extract label information ("foo"/"bar") from corresponding labels, and query-key matrices compute the similarity between the last position and each label position. the query and key matrices can be regarded as two towers for learning the similarity between the features at last position and each demonstration at label positions. the greater the similarity, the higher the probability of the corresponding label. based on this hypothesis, we explore the reason why icl has majority label bias @xcite and recency bias @xcite . the existing of majority label bias matches our hypothesis: query and key matrices compute the attention weights between the last position and each demonstration, so the sum of one label's attention weights is larger when this label is related to more demonstrations. about recency bias, we hypothesize that it is caused by the influence of positional embedding during attention score computation in both shallow and deep layers. based on our analysis, we propose two methods for reducing these biases. for majority label bias, we increase the attention weight of the imbalanced label's position in in-context heads, and the majority label bias reduces 22%. for recency bias, we remove the affect of position embedding in in-context heads, and the recency bias reduces 17%. our code and data will be released on https://github.com/zepingyu0512/ in-context-mechanism . ## hypothesis for icl mechanism our hypothesis is motivated by a case study in section 3.1. we find that icl performance can be affected much by only 1% heads, where some can enhance the probabilities for "foo" and others for "bar" (section 3.2). to understand why this happens, we analyze the value-output vectors and attention scores in section 3.3 and find that valueoutput matrices extract the label information and attention scores computed by query-key matrices control the label information flow. at last, we discuss our hypothesis for icl in section 3.4. ## understanding majority label bias and recency bias in icl there are several phenomena of icl that haven't been explained. @xcite demonstrate that models tend to predict majority labels and the labels near the input text. @xcite also find that changing the demonstration order can affect predictions a lot. based on our hypothesis, we explore why icl has majority label bias (in section 4.1) and recency bias (in section 4.2). ## conclusion we identify the important heads for icl and analyze the value-output vectors and attention scores in these heads. we propose a hypothesis for the mechanism of icl. in shallow layers, the demonstrations and input text is captured by the label positions and the last position. in in-context heads, the value-output matrices project the label features into value-output vectors. the query and key matrices can be regarded as two towers learning the similarity between the last position's features and each label position's features. if the similarity score is high, the corresponding label's probability is enlarged. based on this hypothesis, we interpret why icl has majority label bias and recency bias. furthermore, we propose two methods to reduce these biases by 22% and 17%. overall, our study provides a new method and a reasonable hypothesis for understanding the mechanism of in-context learning.