로짓 기여도 점수를 활용한 비문자적 정보 검색 헤드 식별

Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

대규모 언어 모델은 긴 문맥에서 답변을 생성할 때, 관련 문맥의 의미를 종합하는 방식으로 작동하며, 단순히 내용을 그대로 복사-붙여넣기 하는 경우는 드뭅니다. 어떤 어텐션 헤드가 이러한 정보 통합 과정을 수행하는지 파악하는 것은 긴 문맥 모델의 동작 방식을 이해하는데 중요합니다. 그러나 기존의 탐지 방법들은 이러한 헤드를 제대로 식별하지 못합니다. 이는 기존 방법들이 생성된 토큰과 어텐션을 받은 토큰이 일치하는 경우에 높은 점수를 부여하도록 설계되었기 때문인데, 이는 모델이 어디에서 정보를 읽어들이는지만 파악할 뿐, 어떤 방식으로 정보를 출력하는지(특히, 비문자적 정보 검색을 담당하는 출력-값 회로)를 파악하지 못합니다. 본 연구에서는 '로짓 기여도 점수 (LOCOS)'라는 새로운 탐지 방법을 제안합니다. LOCOS는 각 헤드가 생성한 답변 토큰의 방향으로 투영된 출력-값 회로의 영향을 측정하여, 비문자적 정보 검색을 수행하는 헤드를 식별합니다. 세 가지 모델 패밀리(Qwen3, Gemma-3, OLMo-3.1)에서 LOCOS를 사용하여 NoLiMa 비문자적 정보 검색 벤치마크 상위 헤드들을 제거했을 때, 기존의 어텐션 기반 탐지 방법보다 적은 수의 헤드를 제거했을 때 ROUGE-L 점수가 현저히 감소했습니다. 예를 들어, Qwen3-8B 모델에서 50개의 헤드를 제거하면 ROUGE-L 점수가 0.401에서 0.000으로 떨어지는 반면, 기존 방법론의 가장 좋은 성능을 보이는 기준 모델은 여전히 0.292의 점수를 유지합니다. 선택된 헤드들은 정보 검색과 관련된 특정 기능(예: 파라메트릭 회상, 산술 추론)에는 영향을 미치지 않습니다. 동일한 방식으로 Qwen3-8B 모델에서 LOCOS를 사용하여 헤드를 제거했을 때, MuSiQue 점수는 0.55에서 0.08로, BABI-Long 점수는 0.62에서 0.20으로 감소하는 반면, 무작위로 선택된 헤드들을 사용하는 경우 기준 성능과 0.05 이내의 차이를 보입니다.

In long-context use, large language models frequently synthesize answers from the meaning of a relevant context span rather than literally copy-pasting them. Identifying which attention heads perform this synthesis matters for interpreting long-context model behavior. Yet existing detectors miss these heads by construction: they reward heads whose attended token matches the generated token, a literal-copy criterion that captures where a head reads but not what it writes through its output-value (OV) circuit, the very mechanism that carries non-literal retrieval. We introduce Logit-Contribution Scoring (LOCOS), a write-aware detector that scores each head by the projection of its OV-circuit output onto the answer-token unembedding direction, contrasting needle and off-needle source positions in a single forward pass. Across three model families (Qwen3, Gemma-3, OLMo-3.1), mean-ablating the top LOCOS heads on the NoLiMa non-literal retrieval benchmark collapses ROUGE-L at lower head counts than prior attention-based detections; on Qwen3-8B, ablating 50 heads drives ROUGE-L from 0.401 to 0.000 while the strongest baseline still retains 0.292. The selected heads are retrieval-specific: parametric recall and arithmetic reasoning stay at baseline under the same ablation. On Qwen3-8B, the same ablation also drops MuSiQue from 0.55 to 0.08 and BABI-Long from 0.62 to 0.20, while a random-heads control stays within 0.05 of baseline.