로그잇 기여 점수(LOCOS)를 이용한 비문자적 검색 헤드 식별

대형 언어 모델이 긴 문맥에서 정답을 단어 그대로 복사하지 않고 문맥 의미를 합성해 응답하는 경우가 빈번하다. 이러한 비문자적(non-literal) 검색을 수행하는 attention head를 정확히 찾는 일은 모델 내부 동작을 해석하고 장기 문맥 성능 문제를 고치는 데 핵심적이다. 본 논문은 OV 회로의 쓰기(output-value) 방향이 정답 unembedding에 어떻게 기여하는지를 직접 측정해 기존의 attention 기반 검출이 놓치는 헤드를 발견할 수 있음을 보였다.

Transformer의 attention head는 QK 회로로 어디를 읽을지(positions)를 정하고 OV 회로로 읽은 값을 residual stream에 쓰는 두 단계로 구성된다. 기존의 attention 기반 검색 헤드 탐지는 주로 QK가 읽은 위치와 생성 토큰의 일치 여부에 의존했는데, 이 접근은 읽은 위치가 실제로 무엇을 쓰는지(what)와 분리되어 있다는 점을 반영하지 못한다. 비문자적 검색에서는 읽은 구간의 의미를 OV 회로가 변환해 정답 방향으로 쓰므로 '어디를 읽었는가'만으로는 인과적 기여를 포착할 수 없다.

LOCOS의 핵심 산식은 각 헤드의 위치별 logit 기여 φ_{t,j}^{(l,h)} = u_{y_t}^⊤ o_{t,j}^{(l,h)} = α_{t,j}^{(l,h)} ⋅ u_{y_t}^⊤ W_O^{(l,h)} v_{t,j}^{(l,h)} 로 정의된다. 이 식은 [어떤 값을 입력으로] → [어떤 연산을 수행해] → [어떤 결과를 얻고] → [그 값의 의미] 패턴으로 해석될 때, 입력은 위치 j의 value 벡터 v_{t,j}이고 수행 연산은 OV 투영 W_O와 attention 스칼라 α의 곱이며 결과는 정답 토큰의 logit에 더해지는 스칼라 φ이고 그 의미는 해당 위치가 정답에 쓰기 기여를 했는지 여부다. 공간 대조는 동일 디코딩 스텝에서 needle(정답 스팬)과 off-needle 영역의 φ 합을 길이 정규화해 Φ^{+}_t와 Φ^{-}t를 구하고 그 차이를 각 정답 스텝에서 집계해 헤드별 최종 점수 S{l,h}를 산출한다.

NoLiMa 벤치마크에서 여섯 모델 구성에 대해 상위 LOCOS 헤드 평균 대체는 모든 경우에서 ROUGE-L을 더 빠르게 붕괴시켰다. Qwen3-8B에서 상위 50개 헤드를 제거하면 ROUGE-L이 0.401에서 0.000으로 떨어진 반면 attention 기반 최강 기준선은 동일 깊이에서 0.292를 유지했다. OV 투영을 제거한 attention-only 공간 대조 통제는 모델별로 성능 차이를 보였으나 전체적으로 LOCOS가 더 일관된 인과적 영향을 유도했다.

모델 구성은 L 레이어, 각 레이어당 H 헤드, 헤드 차원 d_h, 모델 차원 d=H·d_h 표기법을 사용한다. 각 헤드의 위치별 출력은 o_{t,j}^{(l,h)} = α_{t,j}^{(l,h)} ⋅ W_O^{(l,h)} v_{t,j}^{(l,h)} 이고 logits는 최종 residual stream을 W_U로 투영해 얻는다. LOCOS는 이 per-position OV 출력을 정답 unembedding u_{y_t}와 내적해 logit 기여 φ를 계산하며, 이 값을 동일 길이의 영역별 합으로 정규화해 공간 대조를 수행한다.

첫 번째 한계는 문맥 내에 정답과 의미적으로 관련된 방해 문구가 다수 존재하면 off-needle 기여 Φ^{-}가 올라가 넓은 의미 매칭을 수행하는 헤드를 낮게 평가할 수 있다는 점이다. 두 번째 한계는 평가 대상 모델군에 제한이 있어 Mixture-of-Experts, encoder–decoder 아키텍처, 또는 state-space hybrid 같은 구조로의 일반화는 추가 검증이 필요하다는 점이다. 논문은 LOCOS의 원리 자체는 OV/QK 분해에 의존하므로 원리상 확장 가능하지만 인과적 영향량과 레이어 분포는 아키텍처별로 다를 수 있다고 명시했다.