Retrieval Head Mechanistically Explains Long-Context Factuality

Despite the recent progress in long-context language models, it remains elusive how transformer-based models exhibit the capability to retrieve relevant information from arbitrary locations within the long context. This paper aims to address this question. Our systematic investigation across a wide spectrum of models reveals that a special type of attention heads are largely responsible for retrieving information, which we dub retrieval heads. We identify intriguing properties of retrieval heads:(1) universal: all the explored models with long-context capability have a set of retrieval heads; (2) sparse: only a small portion (less than 5\%) of the attention heads are retrieval. (3) intrinsic: retrieval heads already exist in models pretrained with short context. When extending the context length by continual pretraining, it is still the same set of heads that perform information retrieval. (4) dynamically activated: take Llama-2 7B for example, 12 retrieval heads always attend to the required information no matter how the context is changed. The rest of the retrieval heads are activated in different contexts. (5) causal: completely pruning retrieval heads leads to failure in retrieving relevant information and results in hallucination, while pruning random non-retrieval heads does not affect the model's retrieval ability. We further show that retrieval heads strongly influence chain-of-thought (CoT) reasoning, where the model needs to frequently refer back the question and previously-generated context. Conversely, tasks where the model directly generates the answer using its intrinsic knowledge are less impacted by masking out retrieval heads. These observations collectively explain which internal part of the model seeks information from the input tokens. We believe our insights will foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.

한 줄 요약

LLaMA-2 7B 등 다양한 모델에서 정보 검색을 담당하는 'retrieval head'가 존재하며, 이는 hallucination과 추론 성능에 직접적인 영향을 미친다.

핵심 기여도

**LLaMA-2 7B** 등 4개 모델 가족에서 **5% 미만의 attention head**가 정보 검색을 담당함.
**retrieval head**는 **short-context로 사전 학습된 모델에도 존재**하며, **context 길이 확장 후에도 동일 head가 유지됨**.
**12개 retrieval head**는 **context 변화에 관계없이 항상 필요한 정보에 집중**하며, 나머지 head는 **동적 활성화**됨.
retrieval head를 **완전 제거하면 hallucination 발생**, **랜덤 non-retrieval head 제거는 성능에 영향 없음**.

핵심 아이디어

기존 연구에서 CopyNet과 Induction Head의 아이디어를 결합하여, **transformer 내부의 특정 attention head가 조건부 복사 알고리즘을 구현**함을 가정하였다. 이 연구는 **LLaMA-2 7B, Yi, QWen, Mistral** 등 다양한 모델에서 **retrieval head**가 존재함을 밝히고, 이 head가 **Needle-in-a-Haystack 테스트 성능**에 핵심적인 역할을 함을 보여준다. retrieval head는 **context 길이와 학습 방식에 관계없이 유지**되며, **12개의 head는 항상 활성화**, 나머지는 **context에 따라 동적으로 활성화**된다. 이는 모델이 **입력 정보를 참조하는 메커니즘**을 명확히 설명한다.

기술적 접근법

**4개 모델 가족**, **6개 모델 규모**, **3가지 fine-tuning 방식**에 걸쳐 실험 수행.
**Needle-in-a-Haystack 테스트**를 통해 retrieval head를 탐지.
**attention head 활성도 분석**을 통해 retrieval head 식별.
**retrieval head 제거 실험**을 통해 **causal 영향** 확인.
**LLaMA-2 7B**에서 **12개 retrieval head**가 항상 활성화됨을 확인.
**context 길이 32-128K로 확장**한 모델에서도 동일 head가 유지됨.

주요 결과

**LLaMA-2 7B**에서 **12개 retrieval head**가 **context 변화에 관계없이 항상 활성화**됨.
retrieval head를 **완전 제거하면 hallucination 발생**, **랜덤 non-retrieval head 제거는 성능에 영향 없음**.
**Chain-of-Thought (CoT) 추론**에서 retrieval head가 **강한 영향**을 미침.
**LLaMA-2 7B 80K** 모델에서 retrieval head는 **context 확장 후에도 동일하게 유지**됨.
**Qwen Chat** 등 fine-tuned 모델에서도 retrieval head가 **기본 모델과 동일하게 유지**됨.

의의 및 한계

이 연구는 **transformer 내부의 특정 attention head가 정보 검색을 담당**함을 체계적으로 밝혀내며, **mechanistic interpretability 분야의 중요한 진전**을 이룬다. retrieval head의 존재는 **KV cache 압축**, **hallucination 감소**, **추론 성능 향상** 등 실용적 문제 해결에 기여할 수 있다. 그러나 **모든 retrieval head가 동일한 역할을 하는지**, **다양한 task에서의 영향 차이**는 추가 연구가 필요하다. 또한, **retrieval head가 활성화되지 않아도 hallucination이 발생하는 경우**가 있어, **활성화 조건에 대한 이해**가 부족하다는 한계가 있다.

실용적 활용

이 연구는 **long-context 모델의 KV cache 압축**, **hallucination 감소**, **추론 성능 향상**을 위한 **모델 최적화**에 활용될 수 있다. 특히, **LLaMA-2 7B 80K**와 같은 **대규모 모델의 효율적 배포**와 **CoT 추론 성능 개선**에 기여할 수 있다.