Retrieval Head Mechanistically Explains Long-Context Factuality

Meta info.

• Authors: Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, Yao Fu
• Paper: https://arxiv.org/pdf/2404.15574
• Published: April 24, 2024
• Code: https://github.com/nightdessert/Retrieval_Head

TL; DR

특정 attention head가 retrieval을 담당한다

Problem States

long input에 대해 어떻게 Retrieval이 동작하고 있는지 밝혀진 바 없음

Suggestion

retrieval이 어떻게 동작하는지 attention head로 확인

Effects

• long context에 대해서는 항상 retrieval-head가 존재함
• 보통은 5% 미만의 head가 이에 해당
• short context로 학습된 PLM은 이미 이 retrieval head가 있고, 긴 context로 학습한대도 retrieval-head 개수가 증가하는건 아님
• LLaMA-2 7B 모델 기준으로, context가 어떻든지간에 정답에 attend하는 retrieval-head는 12개로 고정, 나머지는 다른 context를 봄
• retrieval-head를 제거하면 검색하지 못해 hallucination 발생. random하게 non-retrieval head를 제거하면 딱히 검색에는 지장 없음
• CoT에서도 영향 큰것으로 보임
• 모델이 intrinsic knowledge로 답변하는 경우 retireval-head를 마스킹해도 영향 덜받음