Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

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.

한 줄 요약

LOCOS는 비직관적 검색 헤드를 식별하기 위해 헤드의 OV 회로 출력을 정답 토큰의 unembedding 방향으로 투영하는 새로운 방법이다.

핵심 기여도

기존 토큰 매칭 기반 방법과 달리, LOCOS는 헤드의 OV 회로 출력을 정답 토큰의 unembedding 방향으로 투영하여 '무엇을 쓰는지'를 평가한다.
Qwen3-8B에서 상위 50개 LOCOS 헤드를 평균 제거하면 ROUGE-L이 0.401에서 0.000으로 급락하는 반면, 기존 기반선은 0.292를 유지한다.
MuSiQue와 BABI-Long에서 LOCOS 헤드 제거는 각각 0.55 → 0.08, 0.62 → 0.20으로 성능이 급격히 하락한다.
파라메트릭 추론 및 산술 추론은 동일한 제거 조건에서도 기준선 수준을 유지한다.

핵심 아이디어

기존 토큰 매칭 기반의 헤드 탐지 방법은 헤드가 '어디를 읽는지'만 평가하지만, LOCOS는 헤드가 '무엇을 쓰는지'를 평가한다. 이는 헤드의 OV 회로 출력을 정답 토큰의 unembedding 벡터에 투영하여 측정한다. 이 접근은 토큰이 일치하지 않아도 의미를 바탕으로 정답을 생성하는 비직관적 검색 헤드를 효과적으로 식별할 수 있다.

LOCOS는 단일 순방향 패스에서 '바늘' 위치와 '비바늘' 위치의 투영 값을 대비하여 공간적 대비(spatial contrast)를 사용한다. 이는 헤드가 특정 위치에서만 정보를 쓰는지 여부를 확인하는 데 기여한다.

기술적 접근법

**LOCOS 점수 계산**: 헤드의 OV 회로 출력을 정답 토큰의 unembedding 벡터 $ \mathbf{u}_{y} $에 투영하여 스칼라 값을 계산한다.
**공간적 대비**: 바늘 위치와 비바늘 위치의 투영 값을 비교하여 헤드가 특정 위치에서만 정보를 쓰는지 확인한다.
**실험 모델**: Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B), OLMo-3.1 (32B)의 6개 설정에서 평가.
**제거 실험**: 상위 LOCOS 헤드를 평균 제거(mean-ablation)하여 성능 변화를 분석.

주요 결과

**NoLiMa 벤치마크**: Qwen3-8B에서 상위 50개 LOCOS 헤드 제거 시 ROUGE-L이 0.401 → 0.000으로 급락.
**기존 기반선 대비**: 동일 조건에서 기존 최고 기반선은 0.292를 유지.
**MuSiQue**: LOCOS 헤드 제거 시 정확도가 0.55 → 0.08로 감소.
**BABI-Long**: 0.62 → 0.20으로 감소.
**랜덤 헤드 제거**: 기준선에서 0.05 이내로 유지됨.

의의 및 한계

LOCOS는 기존 토큰 매칭 기반 방법이 포착하지 못한 비직관적 검색 헤드를 효과적으로 식별할 수 있다. 이는 긴 문맥에서 의미 기반의 정답 생성을 이해하는 데 중요한 기여를 한다. 또한, 파라메트릭 추론 및 산술 추론은 LOCOS 헤드 제거에도 영향을 받지 않아 검색에 특화된 헤드임을 입증한다.

한계로는 (i) 문맥 내 관련 정보가 포함된 경우, LOCOS가 넓은 의미 매칭을 수행하는 헤드를 누락할 수 있다. (ii) Mixture-of-Experts, 인코더-디코더 스택, state-space 하이브리드 아키텍처는 포함되지 않았으며, 이러한 아키텍처에의 적용 가능성은 추가 검증이 필요하다.

실용적 활용

LOCOS는 LLM의 긴 문맥 처리 메커니즘을 해석하는 데 활용될 수 있으며, 특히 의미 기반의 정답 생성이 필요한 의료, 법률, 고객 지원 등 산업 분야에서 모델 해석도를 높이는 데 유용하다. 또한, 모델 최적화 및 추론 과정에서 불필요한 헤드를 제거하여 성능을 개선하는 데도 활용 가능하다.