[논문 리뷰] Retrieval Head Mechanistically Explains Long-Context Factuality

https://arxiv.org/pdf/2404.15574

2025 ICLR Orals

• 2024, Apr 24
• Peking University, Washington, MIT etc.

간만에 RAG로 돌아왔습니다 ~~

Abstract & Introduction

• Retrieval Head: long context에서 relevant information을 찾는 데 크게 관여하는 attention head

복습

Attention Head란?

⇒ Retrieval에 집중하는 Head가 있다는 것~

CoT 추론(질문 + 이전 답변 사용)에 큰 영향을 미침.

• Context 사용 하는 경우 → 영향 큼
• Intrinsic Knowledge 사용하는 경우 → Retrieval Heads마스킹해도 영향 적음

2. Detecting Retrieval Head

어케 찾았니? ⇒ Retrieval head detection algorithm !

Retrieval Score

• q: query
• k: answer(gt = needle)
• x: given context (haystack)
• ⁍: random index range at x
• a: attention score
• w: current token

위처럼 x의 ⁍에 k 를 넣어 각 attention head의 retrieval score를 계산

• Retrieval Score: decoding중 copy-paste 빈도를 측정한 score

  • copy-paste: 입력 토큰을 출력 토큰으로 복사 = 검색 결과를 답변에 사용 copy-paste 기준 (AND 조건으로 만족)

    1. ⁍
    2. ⁍ ⇒ 답변이 생성될 때, Head 들의 Attention이 가장 높은 token(= j) 을 찾는다.

    그 토큰이 w(current token)와 일치한다 & w가 k에 속한다

    Retrieval Score = Needle Sentence의 token-level Recall이다

    Head 단위로 R-S(내 맘대로 줄임) 가 나옴

    Retrieval score for head ⁍

Retrieval Score가 0.5가 넘는 것은 희소함

→ 모델마다 Head의 수가 다르지만 3~6% 사이의 비율을 유지한다

Basic Properties of Retrieval Heads

Retrieval Head의 속성

1. Universal : Model, Scale, FT 여부에 상관 없이 Retrieval Heads를 가짐
2. Model: Llama, Yi, Qwen, Mistral 전부 해당
3. Scale: 6B, 14B, 34B, 8X7B 전부 해당
4. FT: Base, Chat, MoE에 전부 해당
5. Sparse : 전체 Attention Heads중 5% 미만
6. Intrinsic : pretrained model에 이미 존재. 32K ~ 128K로 확장할 때도 정보 검색은 동일한 Retrieval Heads가 수행

4. Dynamically activated
5. (Llama-2 7B)기준 12개의 Head: Context에 상관 없이 activated. 필요한 정보에만 집중
6. 나머지 Head: 다양한 Context에 따라 activated
7. Causal : Retrieval Heads를 마스킹하면 정보 검색 X, 환각 현상 발생. Random Heads를 마스킹하면 차이 없음

Influence on Downstream Tasks

→ Retrieval Head의 Monosemantic함 증명

→ COT를 쓸 때에는 Retrieval Heads가 더 중요함

Bobb’s 한줄평

Retrieval Head를 Masking 하면 멍청해진다 ⇒ MonoSemantic한 Featrue를 찾았다 👍

더 멍청해지는 실험은 했는데, 더 좋아지는 건 불가능한가? 더 Activate를 할 수 있는 방법은 없나? 🤔