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

arXiv: https://arxiv.org/abs/2402.02872 · HTML: https://mlsanity.com/paper/2402.02872/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 50 references that are in the map.
- First posted on arXiv 5 Feb 2024.

## Builds on (most cited first, 28 of 28)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Least-to-Most Prompting Enables Complex Reasoning in Large Language Models](https://mlsanity.com/paper/2205.10625/) (2022)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space](https://mlsanity.com/paper/2203.14680/) (2022)
- [Dissecting Recall of Factual Associations in Auto-Regressive Language Models](https://mlsanity.com/paper/2304.14767/) (2023)
- [Larger language models do in-context learning differently](https://mlsanity.com/paper/2303.03846/) (2023)
- [The (Un)reliability of saliency methods](https://mlsanity.com/paper/1711.00867/) (2017)
- [Good Debt or Bad Debt: Detecting Semantic Orientations in Economic Texts](https://mlsanity.com/paper/1307.5336/) (2013)
- [Direct and Indirect Effects](https://mlsanity.com/paper/1301.2300/) (2013)
- [Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning](https://mlsanity.com/paper/2305.14160/) (2023)
- [Demystifying Prompts in Language Models via Perplexity Estimation](https://mlsanity.com/paper/2212.04037/) (2022)
- [What In-Context Learning "Learns" In-Context: Disentangling Task Recognition and Task Learning](https://mlsanity.com/paper/2305.09731/) (2023)
- [The Closeness of In-Context Learning and Weight Shifting for Softmax Regression](https://mlsanity.com/paper/2304.13276/) (2023)
- [ICL-D3IE: In-Context Learning with Diverse Demonstrations Updating for Document Information Extraction](https://mlsanity.com/paper/2303.05063/) (2023)

## Built on by (most cited first, 18 of 18)

- [Attention Heads of Large Language Models: A Survey](https://mlsanity.com/paper/2409.03752/) (2024)
- [Interpreting Arithmetic Mechanism in Large Language Models through Comparative Neuron Analysis](https://mlsanity.com/paper/2409.14144/) (2024)
- [Locate, Steer, and Improve: A Practical Survey of Actionable Mechanistic Interpretability in Large Language Models](https://mlsanity.com/paper/2601.14004/) (2026)
- [Can In-context Learning Really Generalize to Out-of-distribution Tasks?](https://mlsanity.com/paper/2410.09695/) (2024)
- [Understanding Multimodal LLMs: the Mechanistic Interpretability of Llava in Visual Question Answering](https://mlsanity.com/paper/2411.10950/) (2024)
- [How Do Large Vision-Language Models See Text in Image? Unveiling the Distinctive Role of OCR Heads](https://mlsanity.com/paper/2505.15865/) (2025)
- [Unveiling In-Context Learning: A Coordinate System to Understand Its Working Mechanism](https://mlsanity.com/paper/2407.17011/) (2024)
- [From Compression to Expression: A Layerwise Analysis of In-Context Learning](https://mlsanity.com/paper/2505.17322/) (2025)
- [CAFE: Retrieval Head-based Coarse-to-Fine Information Seeking to Enhance Multi-Document QA Capability](https://mlsanity.com/paper/2505.10063/) (2025)
- [From Unstructured Data to In-Context Learning: Exploring What Tasks Can Be Learned and When](https://mlsanity.com/paper/2406.00131/) (2024)
- [ConText: Driving In-context Learning for Text Removal and Segmentation](https://mlsanity.com/paper/2506.03799/) (2025)
- [BASIC: Boosting Visual Alignment with Intrinsic Refined Embeddings in Multimodal Large Language Models](https://mlsanity.com/paper/2508.06895/) (2025)
- [Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis](https://mlsanity.com/paper/2509.24164/) (2025)
- [Understanding the Dynamics of Demonstration Conflict in In-Context Learning](https://mlsanity.com/paper/2603.04464/) (2026)
- [Quantifying LLM Biases Across Instruction Boundary in Mixed Question Forms](https://mlsanity.com/paper/2509.20278/) (2025)
- [Logit Separability-Driven Samples and Multiple Class-Related Words Selection for Advancing In-Context Learning](https://mlsanity.com/paper/2406.10908/) (2024)
- [In-Context Fixation: When Demonstrated Labels Override Semantics in Few-Shot Classification](https://mlsanity.com/paper/2605.08295/) (2026)
- [AttnLink: Turning Attention into Schema Links for Text-to-SQL](https://mlsanity.com/paper/2608.00693/) (2026)

## Abstract (arXiv)

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.