기술적 요약: 검색 헤드는 동적입니다

문제 제기

최근 기계적 해석 가능성 연구는 Large Language Models(LLM) 내에서 입력 컨텍스트에서 정보를 추출하는 특정 어텐션 헤드를 식별했는데, 이를 "검색 헤드 (retrieval heads)"라고 부르며 종종 "복사 - 붙여넣기" 행동을 보입니다. 그러나 이전 연구들 (예: Wu et al., 2024; Zhang et al., 2025; Fu et al., 2024) 은 주로 데이터셋 전체에 걸쳐 집계된 정적 통계 (static statistics) 에 의존하여 평균적으로 검색을 수행하는 고정된 헤드 서브셋을 식별해 왔습니다.

이러한 정적 관점은 autoregressive 생성의 세밀한 시간적 역동성 (fine-grained temporal dynamics) 을 간과합니다. 이는 헤드의 평균적 행동이 실시간 운영을 위한 충분한 휴리스틱이라고 가정합니다. 저자들은 이 접근 방식이 모델의 행동을 지나치게 단순화할 위험이 있다고 주장합니다. 통계적으로 지배적인 헤드가 중요한 시간 단계에서 활성화되지 않을 수 있는 반면, "덜 중요한" 헤드가 특정 컨텍스트에서 대체 불가능한 역할을 할 수 있기 때문입니다. 다루어진 핵심 문제는 검색 헤드가 고정된 집합인지, 아니면 생성 컨텍스트에 조건부로 결정되는 동적 집합인지, 그리고 동적이라면 그것이 예측 가능하고 정적 대응물과 기능적으로 구별되는지 여부입니다.

방법론

본 논문은 두 가지 주요 작업에 걸친 검색 헤드의 체계적 분석을 수행합니다: Needle-in-a-Haystack (NIAH) 작업 (복사 - 붙여넣기 검색용) 과 HotpotQA 데이터셋 (다중 홉 추론용).

1. 정의 및 지표

• 정적 검색 헤드 (Static Retrieval Heads): Wu et al., 2024 에 따라 코퍼스 전체에 집계되었을 때 빈번하게 검색 작업을 수행하는 헤드로 정의됩니다.
• 동적 검색 헤드 (Dynamic Retrieval Heads): 토큰 수준에서 정의됩니다. 특정 시간 단계 $t$에서 헤드가 검색 헤드로 간주되려면 다음 조건을 만족해야 합니다:
  • 생성된 토큰이 해당 헤드로부터 가장 높은 어텐션 가중치를 받는 토큰과 일치해야 합니다.
  • 해당 최고 어텐션 토큰은 목표 컨텍스트 ("바늘") 내에 있어야 합니다.
  • 추론 작업 (HotpotQA) 의 경우, 정의는 어텐션 싱크와 로컬 윈도우를 제외하고 유효한 컨텍스트에 대한 지원 사실에 할당된 어텐션의 비율을 측정하도록 일반화됩니다.
• 지표: 본 연구는 동적 집합과 정적 집합을 비교하기 위해 자카드 유사도 (Jaccard similarity) 를, 분포의 확산을 측정하기 위해 엔트로피를, 그리고 은닉 상태와 미래 검색 패턴 간의 관계를 측정하기 위해 정규 상관 분석 (Canonical Correlation Analysis, CCA) 을 활용합니다.

2. 실험 절차

• 헤드 제거 (Head Ablation): 대체 불가능성을 테스트하기 위해 저자들은 각 생성 단계에서 동적 검색 헤드를 마스킹하고, 동일한 수의 상위 순위 정적 헤드나 무작위 헤드를 마스킹했을 때의 성능 저하와 비교합니다.
• 보상 분석 (Compensation Analysis): 본 연구는 모델이 제거된 동적 헤드를 정적 헤드를 활성화하여 보상할 수 있는지 조사하며, 새로 활성화된 헤드와 상위 20 개 정적 집합 간의 중첩을 측정합니다.
• 예측 모델링:
  • CCA: 시간 단계 $n$의 최종 은닉 상태와 미래 시간 단계 $n+k$의 검색 점수 간의 선형 상관관계를 측정합니다.
  • MLP 프로브: 비선형 예측 가능성을 테스트하기 위해 최종 은닉 상태에서 검색 헤드 패턴 (NIAH 의 경우 이진 점수, HotpotQA 의 경우 연속 점수) 을 예측하도록 훈련됩니다.
• 동적 RAG 사례 연구: 저자들은 동적 검색 헤드를 DRAGIN 프레임워크에 통합합니다. 각 단계에서 MLP 프로브가 활성 동적 헤드를 예측하고, 어텐션 마스크가 모델의 시야를 이러한 헤드의 상위-$k$ 어텐션 위치로 제한하여 동적 검색 과정을 시뮬레이션합니다.

주요 기여 및 주장

본 논문은 검색 헤드의 성격에 관한 세 가지 핵심 주장을 확립합니다:

1. 동적성 (Dynamism): 검색 헤드는 시간 단계에 따라 동적으로 변화합니다. 데이터셋 전체에 걸쳐 데이터를 집계하는 통계적 방법은 개별 시간 단계에서 나타나는 헤드의 고유한 패턴을 포착하지 못합니다. 활성 검색 헤드의 집합은 크게 변동하며, 많은 고유한 헤드가 시간에 따라 활성화되는 "긴 꼬리 (long-tail)" 분포를 보입니다. 이는 전통적인 정적 서브셋 (예: 상위 20 개) 의 범위를 훨씬 초과합니다.
2. 대체 불가능성 (Irreplaceability): 특정 시간 단계의 동적 검색 헤드는 기능적으로 구체적이며 정적 검색 헤드로 효과적으로 대체될 수 없습니다. 동적 헤드를 제거하면 정적 헤드나 무작위 헤드를 동일한 수만큼 제거했을 때보다 훨씬 심각한 성능 저하가 발생합니다. 모델이 정적 헤드를 활성화하여 보상을 시도하지만, 이 보상은 검색 정확도를 유지하기에 충분하지 않습니다.
3. 상관관계 (예측 메커니즘, Correlation): 모델의 은닉 상태는 미래 검색 헤드 패턴에 대한 예측 신호를 인코딩합니다. 은닉 상태와 미래 검색 활성화 간에는 강한 상관관계가 있습니다 (예: $k=0$에서 CCA 계수 0.966). 이는 모델이 실행되기 전에 어떤 헤드가 필요한지 "알고" 있는 내부 계획 메커니즘을 나타냅니다.

결과

• 정량적 통계: 모델들 (Llama 3.1/3.2, Qwen, Phi) 전반에 걸쳐 동적 검색 헤드와 상위 20 개 정적 헤드 간의 자카드 유사도는 낮습니다 (0.18~0.46 범위). 이는 정적 정의가 활성 헤드의 상당 부분을 놓치고 있음을 확인시켜 줍니다. 엔트로피 값은 3.0 을 초과하여 헤드가 소수의 고정된 서브셋에 국한되지 않고 광범위하게 관여함을 나타냅니다.
• 제거 성능: NIAH 및 HotpotQA 작업에서 동적 헤드를 마스킹하면 거의 완전한 정확도 하락 (예: NIAH 에서 1.0 에서 0.0 으로) 이 발생했지만, 정적 헤드를 마스킹할 때는 미미한 저하만 발생했습니다.
• 예측 정확도: MLP 프로브는 은닉 상태로부터 검색 패턴을 성공적으로 예측하여 NIAH 에서는 높은 F1 점수 (0.80–0.86) 를, HotpotQA 에서는 높은 $R^2$ 점수 (최대 0.81) 를 달성했습니다. 이는 은닉 상태의 예측 가능성을 검증합니다.
• 동적 RAG 유용성: DRAGIN 프레임워크에서 프로브에 의해 예측된 동적으로 선택된 헤드를 사용하는 것은 HotpotQA 에서 정적 헤드 선택 및 무작위 기준선보다 일관되게 우수한 성능을 보였습니다. Llama 3.1-8b 의 경우, 동적 전략은 정적 전략의 0.5098 대비 F1 점수 0.5586 을 달성했습니다.

의의

본 논문은 검색 헤드를 정적 정의보다 더 정확한 LLM 내부 메커니즘 이해를 제공하는 동적이고 컨텍스트에 조건부인 집합으로 보는 관점이 중요하다고 주장합니다. 이러한 발견의 의의는 다음과 같습니다:

• 해석 가능성: 검색이 고정된 회로가 아니라 모델이 현재 생성 상태에 따라 어텐션 자원을 동적으로 할당하는 유동적인 과정임을 드러냅니다.
• 효율성 및 압축: 이러한 발견은 동적 KV 캐시 압축을 향한 길을 제시합니다. 평균 중요도에 기반하여 헤드를 제거하는 것 (희귀하지만 대체 불가능한 "긴 꼬리" 헤드를 버릴 위험이 있음) 대신, 향후 시스템은 경량 프로브를 사용하여 현재 단계에서 활성인 특정 헤드에 대해서만 KV 쌍을 동적으로 유지하거나 가져올 수 있습니다.
• 할루시네이션 감지 및 정밀 RAG: 검색 헤드 활성화의 유무는 사실 생성 중 검색 신호가 부족할 때 할루시네이션을 감지하거나, 모델이 진정한 필요를 나타낼 때만 외부 검색을 트리거하는 내재적이고 해석 가능한 지표로 작용할 수 있으며, 이를 통해 컨텍스트 오염을 최소화할 수 있습니다.

저자들은 이러한 동적 관점이 LLM 에 대한 세밀한 이해를 제공하여 더 정밀하고 상태 인식적인 개입 기술 개발을 가능하게 한다고 결론지었습니다.