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

왜 중요한가

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

핵심 기여

OV 출력의 정답 언베딩 기여를 측정하는 Logit-Contribution Scoring(LOCOS) 도입

LOCOS는 각 헤드의 위치별 OV 출력에 대해 정답 토큰의 unembedding 벡터와의 스칼라 내적을 계산해 logit 기여를 얻는다. 이 값을 needle 영역과 off-needle 영역에서 길이 정규화해 공간적 대조(spatial contrast)를 수행한 뒤 모든 정답 스텝을 평균해 헤드 점수를 산출한다. 이 절차는 한 번의 forward pass로 각 헤드의 쓰기 방향(write)과 위치별 읽기(read)를 동시에 반영해 비문자적 검색 헤드를 식별하게 한다.

NoLiMa 비문자적 검색 벤치마크에서 LOCOS가 더 높은 인과적 영향력 보임

세 개 모델 계열(Qwen3, Gemma-3, OLMo-3.1)과 여섯 구성에서 상위 LOCOS 헤드들을 평균 대체(mean-ablation)하면 NoLiMa ROUGE-L이 더 낮은 k 값에서 급락했다. 예를 들어 Qwen3-8B에서 상위 50개 LOCOS 헤드를 제거하면 ROUGE-L이 0.401에서 0.000으로 붕괴했다. 같은 실험에서 attention 기반 기준선은 동일 깊이에서 상당한 성능을 유지해 LOCOS가 비문자적 검색에 대해 더 인과적인 헤드를 선택했음이 확인됐다.

OV 프로젝션과 공간 대조의 분리 실험을 통한 기여 성분 규명

OV 프로젝션을 제거하고 attention만으로 동일한 공간 대조를 수행하는 통제 실험을 도입해 OV 성분의 효과를 고립시켰다. 이 통제는 모델-구성에 따라 경쟁력을 보였으나 전반적으로 LOCOS가 더 일관된 인과적 붕괴를 유도했다. 따라서 OV 회로의 정답-정렬 기여는 비문자적 검색 헤드 탐지에서 신뢰도를 높이는 주요 성분으로 확인되었다.

기능적 분리(functonal dissociation) 검증으로 검색 특이성 확보

LOCOS로 선택된 헤드의 제거는 문맥 기반 검색 성능을 크게 떨어뜨리는 반면 파라메트릭 기억(도시-국가 연관, PopQA, 산수)은 거의 유지되었다. 이를 통해 LOCOS가 모델의 전반적 중요도 대신 검색 특이적 회로를 식별함이 확인되었다. 분리 점수(Dissociation Score)를 도입해 특정 k에서 검색 손실 대비 파라메트릭 손실의 격차를 정량화했다.

문체별·레이어별 분포 및 한계 규명

LOCOS 점수는 모델 계열에 따라 늦은 레이어에 더 집중되는 경향을 보였으나 계열별로 차이가 존재했다. Gemma-3-27B 등 일부 설정에서는 attention-only 대조가 더 강한 영향을 보이는 역전 현상이 관찰되어 추가 진단(tuned-lens 등)을 수행했다. 논문은 오프-니들 편향과 아키텍처 범위(예: MoE, encoder–decoder, SSM 계열 미포함)를 한계로 명시했다.

핵심 아이디어 이해하기

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 분해에 의존하므로 원리상 확장 가능하지만 인과적 영향량과 레이어 분포는 아키텍처별로 다를 수 있다고 명시했다.

실무 활용

LOCOS는 단일 forward pass로 위치별 쓰기 기여를 측정하므로 긴 문맥에서 비문자적 검색 회로를 탐지하고 타겟형 수리적 개입이나 디버깅을 수행할 때 실용적이다. 논문에서 제공한 mean-ablation 프로토콜은 특정 헤드 집단을 평균 대체해 인과적 검증을 수행하므로 모델 수정 전후 효과를 정량적으로 확인할 수 있다. 공개된 저장소는 탐지 코드와 재현 실험을 포함해 실무 적용을 용이하게 한다.

장문 질문응답 시스템에서 문맥 검색 회로가 잘못 작동할 때 문제를 영역별로 고립해 패치 우선순위를 정하는 용도
모델 압축 또는 KV cache 관리 시 검색에 핵심적인 헤드를 보존하고 덜 중요한 헤드는 축소하는 기준으로 활용하는 용도
안정성 평가에서 비문자적 검색이 과도하게 작동해 오답 합성이 발생하는 케이스를 진단하고 완화 전략을 설계하는 용도

코드 공개 여부: 공개

코드 저장소 보기

키워드

attention heads(어텐션 헤드)OV circuit(OV 회로)logit-contribution scoring(LOCOS)non-literal retrieval(비문자적 검색)NoLiMa(비문자적 검색 벤치마크)

용어 해설

OV 회로(OV circuit): — Attention head는 QK 회로로 어디를 읽는지 결정하고 OV 회로로 읽은 표현을 residual stream에 쓴다. 본문에서는 OV 회로가 읽은 위치의 표현을 정답 토큰의 unembedding 방향으로 어떻게 투영하는지가 비문자적(non-literal) 검색의 핵심이라고 설정했다. OV 회로는 어떤 정보를 쓰는지(what)와 어디를 읽는지(where)를 분리해 이해하는 근거가 된다.
로그잇 기여 점수(LOCOS)(Logit-Contribution Scoring (LOCOS)): — 각 헤드의 위치별 OV 출력이 정답 토큰의 unembedding 방향으로 기여하는 스칼라를 계산해, needle과 off-needle의 기여를 공간적으로 대조하여 점수를 산출하는 방법이다. 이 방식은 한 번의 forward pass로 비문자적 검색 헤드를 탐지할 수 있도록 설계되었다. LOCOS는 attention 기반 검출이 놓치는 '무엇을 쓰는가'를 직접 측정한다.
공간 대조(Spatial Contrast): — 정답을 생성하는 시점에서 needle(정답을 포함한 구간)과 동일 길이의 off-needle 영역에 대한 logit 기여를 비교해 그 차이를 점수로 사용하는 집계 방식이다. 이 대조는 동일 디코딩 스텝 내부에서 위치별 기여를 비교하므로 토큰 빈도나 전반적 편향을 상쇄한다. LOCOS의 핵심 집계 규칙으로 단일 정답 스텝으로부터 유의미한 신호를 확보한다.
바늘-건초더미(NIAH)(Needle-in-a-haystack (NIAH)): — 짧은 정답 스팬(needle)을 긴 방해 문맥(haystack) 안에 삽입해 모델이 정답을 직접 복사할지 의미를 파악해 합성할지를 평가하는 실험 설정이다. NIAH에서는 정답 토큰이 source에 문자 그대로 존재하므로 attention 기반 탐지가 잘 작동한다. 본문은 NIAH와 달리 의미 기반 복원이 필요한 NoLiMa를 중심으로 비문자적 검색을 평가했다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드