기술적 요약: 인맥락 엔티티 검색의 혼합 메커니즘

문제 제기

언어 모델 (LM) 은 인맥락 추론에 탁월한데, 이 능력은 기본적으로 텍스트 내에서 관련 엔티티를 결합(예: "Ann loves pie"에서 "Ann"을 "pie"에 결합) 하고 이후 쿼리에 따라 이를 검색(예: "누가 파이를 좋아할까?"에 답변) 할 수 있는 능력에 의존합니다. 이전 연구는 LM 이 인맥락 시퀀스 내의 인덱스에 기반하여 엔티티를 검색하는 위치 기반 메커니즘을 활용한다고 규명했으나, 이 관점은 주로 짧은 엔티티 그룹 목록 (예: $n=2$ 또는 $3$) 을 다룬 연구에서 도출된 것이었습니다.

본 연구에서 다루는 핵심 문제는 더 복잡하고 현실적인 설정에서 이 위치 기반 메커니즘의 신뢰성입니다. 결합된 엔티티 그룹의 수가 증가함에 따라 LM 은 긴 맥락에 대한 추론에 어려움을 겪으며 "중간 분실 (lost-in-the-middle)" 현상과 같은 현상을 보입니다. 저자들은 위치 기반 메커니즘이 더 긴 시퀀스로 일반화되는지, 아니면 LM 이 검색 정확도를 유지하기 위해 대체 전략을 사용하는지 조사합니다.

방법론

저자들은 LM 의 내부 인과 구조를 탐구하기 위해 기계적 해석 가능성 기법, 구체적으로 교환 개입 (interchange interventions) 을 활용합니다.

1. 작업 설정: 연구는 "엔티티 그룹"(특정 역할을 채우는 엔티티 튜플, 예: 이름 + 음식) 을 포함하는 템플릿 기반 인맥락 추론 작업을 사용합니다. 모델은 특정 그룹 내의 쿼리 엔티티를 기반으로 대상 엔티티를 검색하도록 요청받습니다.
2. 반사실적 설계: 경쟁 가설들을 구분하기 위해 저자들은 원본과 반사실적 입력의 쌍으로 구성된 데이터셋을 구축합니다. 이러한 쌍은 세 가지 가설적 메커니즘—위치 기반, 어휘적, 반사적—이 각각의 중간 변수에 대해 교환 개입이 수행될 경우 서로 다른 출력을 예측하도록 설계됩니다.
   • 위치 기반 메커니즘 ($P$): 반사실적 쿼리의 그룹 인덱스에 있는 엔티티를 검색합니다.
   • 어휘적 메커니즘 ($L$): 반사실적 쿼리 엔티티에 결합된 엔티티를 검색합니다 (예: 쿼리가 "Tim"에서 "Ann"으로 변경되면 "Ann"이 좋아하는 것을 검색합니다).
   • 반사적 메커니즘 ($R$): 대상 자체에서 시작하여 쿼리를 통해 역참조되는 직접 포인터를 통해 대상 엔티티를 검색합니다.
3. 개입 프로토콜: 저자들은 검색이 발생하기 직전으로 식별된 특정 레이어 ($\ell$) 에서 마지막 토큰 잔류 스트림에 교환 개입을 수행합니다. 반사실적 실행에서 잔류 스트림 벡터를 원본 실행에 패치하고 결과적인 다음 토큰 분포를 관찰합니다.
4. 인과 모델링: LM 의 행동을 시뮬레이션하기 위한 인과 모델 ($M$) 이 개발됩니다. 이 모델은 세 가지 메커니즘을 가중치 혼합으로 결합합니다.
   • 위치 기반 항은 쿼리 인덱스를 중심으로 한 가우시안 분포로 모델링되며, 분산은 인덱스에 의존합니다.
   • 어휘적 및 반사적 항은 원-핫 (one-hot) 분포로 모델링됩니다.
   • 모델은 손실 지표로 제이슨-샤논 발산 (JSD) 을 사용하여 LM 의 로그 확률 분포를 예측하도록 훈련됩니다.
5. 범위: 실험은 2B 에서 72B 파라미터에 이르는 세 가지 계열 (Llama, Gemma, Qwen) 의 아홉 가지 모델을 다루며, 열 가지 다른 결합 작업에서 평가됩니다.

주요 기여

1. 세 가지 구별되는 메커니즘의 식별

본 논문은 검색이 오직 위치 기반이라는 지배적인 관점에 도전합니다. LM 이 세 가지 메커니즘의 혼합을 활용함을 입증합니다.

• 위치 기반: 맥락의 시작과 끝에 있는 엔티티에는 효과적이지만, 그룹 수 ($n$) 가 증가함에 따라 중간 위치에서는 노이즈가 많고 확산되는 경향이 있습니다.
• 어휘적: 쿼리 엔티티의 결합된 대응물을 검색합니다 (예: "pie" $\to$ "Ann"). 이 메커니즘은 날카로운 원-핫 신호를 제공합니다.
• 반사적: 쿼리를 통해 검색되는 대상 엔티티에서 자체로 향하는 직접 포인터를 사용합니다. 이는 대상이 절 내에서 쿼리 앞에 나타나는 경우 (예: "Tim loves tea" $\to$ "Who loves tea?") 특히 중요하며, 이러한 시나리오에서 어휘적 복사는 자기회귀 모델에게 구조적으로 어렵습니다.

2. "U 자형" 신뢰성 패턴

광범위한 개입 실험을 통해 저자들은 일관된 패턴을 발견합니다.

• 첫 번째 및 마지막 위치: 검색은 매우 집중적이고 신뢰할 수 있는 위치 기반 메커니즘에 의해 지배됩니다.
• 중간 위치: 위치 기반 신호가 크게 저하됩니다. 이를 보상하기 위해 LM 은 올바른 엔티티를 구별하기 위해 더 날카로운 신호를 제공하는 어휘적 및 반사적 메커니즘에 크게 의존합니다.
• 상호작용: 메커니즘들은 "경쟁적 시너지"를 보입니다. 인덱스가 일치할 때 신호는 증폭되고, 불일치할 때 서로를 억제합니다.

3. 다음 토큰 예측의 인과 모델

저자들은 이 세 가지 메커니즘을 결합한 인과 모델을 구축합니다. 이 모델은 다양한 작업과 모델 크기에 걸쳐 LM 의 실제 다음 토큰 분포와 95% 일치(제이슨-샤논 유사도로 측정) 를 달성합니다. 반면, 위치 기반 메커니즘에만 의존하는 모델 (지배적인 관점) 은 훨씬 낮은 점수 (~45%) 를 기록하며, 일부 구성에서는 균일한 기준선보다 더 낮은 성능을 보입니다.

4. 자연스러운 입력으로의 일반화

본 연구는 엔티티 그룹이 "채움" 텍스트 (열린-ended 문장) 와 교차되는 더 자연스러운 설정에서도 이러한 발견이 유효함을 검증합니다. 채움 토큰의 추가는 노이즈를 증가시키고 어휘적 메커니즘을 약간 약화시키지만, 세 가지 메커니즘 혼합에 대한 근본적인 의존성은 지속되며, 이는 "중간 분실" 현상에 대한 기계적 설명을 제공합니다.

결과

• 메커니즘 보편성: 긴 맥락의 중간 위치에서 위치 기반 메커니즘은 모델 행동의 약 20% 만 설명하는 반면, 어휘적 및 반사적 메커니즘이 올바른 검색의 주요 동인이 됩니다.
• 모델 견고성: 혼합 패턴은 모델 계열 (Llama, Gemma, Qwen) 과 규모 (2B 에서 72B 파라미터) 전반에 걸쳐 일관됩니다.
• 반사적 검증: 반사실적 답변이 원본 맥락에 존재하지 않는 특정 반사실적 설계를 통해 저자들은 답변 토큰 자체와 구별되는 반사적 포인터 메커니즘의 존재를 입증합니다.
• 결손 연구: 인과 모델에서 단일 메커니즘을 제거하면 성능이 크게 저하되어, LM 행동의 정확한 추상화를 위해 세 가지 모두 필요함을 확인합니다.

중요성 및 주장

본 논문은 LM 이 인맥락에서 엔티티를 결합하고 검색하는 방식에 대한 보다 완전한 그림을 확립했다고 주장합니다. 단순한 위치 기반 관점을 넘어 저자들은 긴 맥락 설정에서 LLM 의 강점(다중 신호를 통한 견고한 검색) 과 취약점(중간 위치의 노이즈로 인한 "중간 분실" 오류) 에 대한 기계적 설명을 제공합니다.

이 연구는 "중간 분실" 현상이 단순히 주의 메커니즘의 실패가 아니라, 어휘적 및 반사적 메커니즘에 의해 부분적으로만 보상되는 위치 기반 신호의 특정 저하임을 시사합니다. 저자들은 인맥락 추론의 인과 메커니즘에 대한 추가 연구를 촉진하기 위해 코드와 데이터를 공개합니다.