검색 헤드는 동적이다

초록

본 논문은 대형 언어 모델(LLM)에서 “검색 헤드”가 고정된 집합이 아니라 토큰 생성 시점마다 동적으로 변한다는 사실을 실증한다. 동적 검색 헤드가 정적 헤드로 대체될 수 없으며, 모델의 은닉 상태가 향후 검색 헤드 패턴을 예측한다는 내부 계획 메커니즘을 발견한다. Needle‑in‑a‑Haystack와 다중 홉 QA 실험을 통해 이러한 주장을 검증하고, 동적 검색 헤드를 활용한 Retrieval‑Augmented Generation(RAG) 프레임워크가 정적 헤드보다 성능이 우수함을 보였다.

상세 분석

논문은 먼저 “검색 헤드”를 기존 연구와 동일하게 정의한다. 즉, 현재 토큰이 생성될 때 해당 헤드가 가장 높은 어텐션 가중치를 부여받은 토큰과 동일하고, 그 토큰이 사전 정의된 “needle” 영역에 있을 경우 해당 헤드의 복사‑붙여넣기 점수를 1로 설정한다. 이 정의를 기반으로 저자들은 자동 회귀 생성 과정 전체에 걸쳐 각 헤드의 점수를 시계열적으로 추적한다. Figure 1에서 확인할 수 있듯이, 특정 헤드가 높은 점수를 보이다가도 다음 토큰에서는 급격히 낮아지는 등 급격한 전환이 빈번히 일어난다. 이를 정량화하기 위해 Jaccard 유사도와 인접 Jaccard, 엔트로피 세 가지 지표를 제시한다. “Jaccard w/ Static” 값이 0.180.46 수준으로 낮아 정적 헤드와 동적 헤드의 겹침이 제한적임을 보여준다. “Adj. Jaccard”가 0.280.51 수준으로, 연속 토큰 간에 활성 헤드가 빠르게 교체됨을 의미한다. 엔트로피가 3.0 이상으로, 20개 이상의 헤드가 동적으로 참여한다는 사실을 뒷받침한다.

다음으로 저자들은 동적 헤드의 irreplaceability를 검증하기 위해 두 단계 어블레이션 실험을 설계한다. 첫 번째는 매 토큰 단계마다 동적 헤드를 식별하고, 이를 마스크한 뒤 동일 입력으로 재생성해 성능 저하를 측정한다. 정적 상위 20개 헤드와 무작위 헤드를 동일 수만큼 마스크한 경우와 비교했을 때, 동적 헤드 마스크는 정확도(Exact Match)와 ROUGE‑L 모두에서 현저히 큰 감소를 보였다(Figure 2). 두 번째 실험에서는 동적 헤드를 점진적으로 마스크하면서 모델이 새롭게 활성화하는 보상 헤드가 정적 상위 20개에 포함되는 비율을 측정했다(Figure 3). 보상 헤드가 정적 헤드와 겹치더라도 전체 성능 회복은 미미했으며, 이는 동적 헤드가 특정 컨텍스트와 시점에 특화된 역할을 수행한다는 강력한 증거가 된다.

세 번째 주장은 은닉 상태와 미래 검색 헤드 패턴 간의 상관관계이다. 저자들은 각 토큰의 최종 은닉 벡터를 입력으로 사용해 다음 토큰에서 활성화될 검색 헤드 집합을 예측하는 간단한 선형 분류기를 학습시켰다. 예측 정확도는 무작위 대비 2배 이상 높았으며, 이는 모델이 현재 내부 상태를 통해 향후 어떤 헤드가 검색 역할을 할지 미리 계획하고 있음을 시사한다. 이러한 계획 메커니즘은 기존 “정적 회로” 해석과 달리, 시간에 따라 재구성되는 동적 회로망을 제안한다.

마지막으로, 동적 검색 헤드를 활용한 Retrieval‑Augmented Generation(D‑RAG) 파이프라인을 구축했다. 여기서는 매 생성 단계마다 현재 은닉 상태를 기반으로 가장 가능성이 높은 검색 헤드를 선택하고, 해당 헤드가 가리키는 컨텍스트를 외부 지식 베이스와 결합한다. 실험 결과, D‑RAG은 정적 헤드 기반 RAG 대비 Needle‑in‑a‑Haystack와 다중 홉 QA 모두에서 정확도와 F1 점수가 평균 7%p 상승했다. 이는 동적 헤드 선택이 실제 응용에서도 의미 있는 성능 향상을 제공함을 입증한다.

전반적으로 이 논문은 LLM 내부 메커니즘을 정적 관점에서 해석하던 기존 패러다임을 깨고, 시간에 따라 변하는 동적 어텐션 회로가 핵심적인 정보 검색을 담당한다는 새로운 통찰을 제공한다. 이는 모델 해석, 효율적인 헤드 프루닝, 그리고 동적 지식 통합 방식 설계에 중요한 시사점을 제공한다.