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

About this paper

Zeping Yu, Sophia Ananiadou

A mechanistic hypothesis is proposed for how Large Language Models perform In-Context Learning, characterizing specific 'in-context heads' where query and key matrices operate as 'two towers' for metric learning. This framework provides an explanation for common ICL biases and enables effective strategies to mitigate them, reducing label imbalance accuracy changes by up to 29.1%.

Key insights

• A small subset of 'in-context heads' (approximately 1% of total heads) are disproportionately crucial for ICL, with distinct 'fooheads' and 'barheads' specifically enhancing or reducing label probabilities.
• The value-output matrices within these 'in-context heads' extract distinct label features from demonstration labels, while the query-key matrices function as 'two towers' for metric learning, computing semantic similarity between input and demonstration features.
• This 'two towers' mechanism mechanistically explains ICL biases: majority label bias arises from increased attention to more frequent labels, and recency bias from positional embeddings influencing attention scores for later demonstrations.

Summary by alphaXiv.