The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866911232813957120 |
|---|---|
| author | Kahardipraja, Patrick Achtibat, Reduan Wiegand, Thomas Samek, Wojciech Lapuschkin, Sebastian |
| author_facet | Kahardipraja, Patrick Achtibat, Reduan Wiegand, Thomas Samek, Wojciech Lapuschkin, Sebastian |
| contents | Large language models are able to exploit in-context learning to access external knowledge beyond their training data through retrieval-augmentation. While promising, its inner workings remain unclear. In this work, we shed light on the mechanism of in-context retrieval augmentation for question answering by viewing a prompt as a composition of informational components. We propose an attribution-based method to identify specialized attention heads, revealing in-context heads that comprehend instructions and retrieve relevant contextual information, and parametric heads that store entities' relational knowledge. To better understand their roles, we extract function vectors and modify their attention weights to show how they can influence the answer generation process. Finally, we leverage the gained insights to trace the sources of knowledge used during inference, paving the way towards more safe and transparent language models. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_15807 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation Kahardipraja, Patrick Achtibat, Reduan Wiegand, Thomas Samek, Wojciech Lapuschkin, Sebastian Computation and Language Information Retrieval Machine Learning Large language models are able to exploit in-context learning to access external knowledge beyond their training data through retrieval-augmentation. While promising, its inner workings remain unclear. In this work, we shed light on the mechanism of in-context retrieval augmentation for question answering by viewing a prompt as a composition of informational components. We propose an attribution-based method to identify specialized attention heads, revealing in-context heads that comprehend instructions and retrieve relevant contextual information, and parametric heads that store entities' relational knowledge. To better understand their roles, we extract function vectors and modify their attention weights to show how they can influence the answer generation process. Finally, we leverage the gained insights to trace the sources of knowledge used during inference, paving the way towards more safe and transparent language models. |
| title | The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation |
| topic | Computation and Language Information Retrieval Machine Learning |
| url | https://arxiv.org/abs/2505.15807 |