A retrieval conditioned rebinding circuit for dynamic entity tracking in large language models

기술 요약: 대규모 언어 모델의 동적 엔티티 추적을 위한 검색 조건부 재바인딩 회로 (A Retrieval Conditioned Rebinding Circuit for Dynamic Entity Tracking in Large Language Models)

문제 정의
대규모 언어 모델(LLM)은 컨텍스트를 올바르게 해석하고, 엔티티를 그 속성에 바인딩하며, 상태가 변화함에 따라 이러한 바인딩을 업데이트함으로써 관련 정보를 검색해야 한다. 선행 연구들은 LLM이 엔티티 추적을 수행할 수 있음을 보여주었으나, 자연스러운 맥락에서 명시적인 상태 전이를 처리하는 방식에 있어 모델의 행동은 설정에 따라 일관되지 않으며, 그 기저의 계산 메커니즘은 완전히 이해되지 않았다. 특히, 모델이 모든 상태 변화(예: 스왑 연산) 후에 전역적으로 업데이트된 "상황 모델(situation model)"을 유지하는지, 아니면 검색이 필요할 때만 선택적으로 바인딩을 업데이트하는 더 타겟팅된 메커니즘을 사용하는지는 불분명하다.

방법론
저자들은 Kim과 Schuster (2023)를 변형한 박스와 물체에 관한 동적 상태 추적 과제를 사용하여 이 문제를 조사한다. 본 연구는 네 가지 인스트럭션 튜닝 모델(Gemma2-9B-it, Gemma3-12B-it, Llama3.2-3B-it, Llama3.1-8B-it)의 활성화 함수에 대한 인과적 개입(causal interventions)을 사용한다. 분석은 네 단계로 진행된다:

1. 인과적 매개 분석 (Causal Mediation Analysis): 인터체인지 개입(interchange interventions)을 사용하여 토큰 수준의 정보 흐로를 추적한다. 저자들은 참조 박스(reference boxes), 스왑 대상 박스(swap target boxes), 그리고 대상 물체(target objects)와 관련된 정보가 최종 예측에 어떻게 영향을 미치는지 결정하기 위해 반사실적(counterfactual) 입력을 구성한다.
2. 경로 패칭 (Path Patching): 어텐션 헤드 수준에서 메커니즘을 국소화하기 위해, 저자들은 경로 패칭(V-patching 및 Q-patching)을 적용하여 압축된 회로를 식별한다. 이들은 손상되었을 때 성능에 가장 부정적인 영향을 미치는 경로를 반복적으로 선택하여, 과업을 담당하는 어텐션 헤드의 하위 집합을 격리한다.
3. 역할 특화 인터체인지 개입 (Role-Specific Interchange Interventions): 저자들은 어텐션 패턴을 원본과 반사실적 실행 간에 패칭함으로써 식별된 헤드 그룹의 기능적 역할을 테스트한다. 이를 통해 주소 선택(address selection, routing)을 매개하는 헤드와 객체의 내용(object content)을 전달하는 헤드를 구분한다.
4. 바인딩 ID 개입 (Binding ID Interventions): Feng와 Steinhardt (2024)의 프레임워크를 응용하여, 저자들은 특정 헤드 그룹의 쿼리(Q) 및 키(K) 서브스페이스(subspace)에서 인과적 개입을 수행한다. 이들은 바인딩 ID 방향과 시프트 벡터(shift vectors)를 추정하여, 검색 라우팅이 바인딩 식별자의 Q/K 매칭에 의존하는지 테스트한다.

주요 기여 및 결과

• 검색 조건부 재바인딩의 증거 (Evidence for Retrieval Conditioned Rebinding): 인과적 매개 분석 결과, 참조 박스와 스왑 대상 박스에 대한 정보는 초기 언급에서 직접 흐르는 것이 아니라, 스왑 연산 이후에 나타나 예측을 향해 전파됨을 보여준다. 반면, 물체 정보는 더 직접적으로 라우팅된다. 이러한 해리(dissociation)는 "검색 조건부 재바들딩" 가설을 뒷받침한다. 즉, 모델은 스왑 후의 전체 상태를 전역적으로 재인코딩하지 않는다. 대신, 원래의 바인딩을 보존하고 검색 시에만 바인딩 ID에 대한 재매핑(remapping)을 적용한다.
• 압축된 회로의 식별 (Identification of a Compact Circuit): 경로 패칭을 통해 전체 모델 구성 요소(예: Gemma-9B의 38개 헤드)의 3~10%에 해당하는 특정 어텐션 헤드 회로를 식별하였으며, 이는 전체 모델 성능(전체 모델 1.00 대비 0.89 정확도)에 근접하게 회복한다. 이 회로는 다섯 가지 기능적 그룹으로 구성된다:
  • 그룹 E (Binding Anchor): 박스 정체성 정보를 쓴다(writes).
  • 그룹 D (Swap Position Transmitter): 스왑 정보를 전파한다.
  • 그룹 C (Position Updater): 업데이트된 바인딩 위치를 향해 어텐션을 형성한다.
  • 그룹 B (Dereferencer): 포인터 역할을 하며, 주소 선택을 매개한다.
  • 그룹 A (Answer Retriever): 대상 물체 토큰을 검색한다.
• 주소 라우팅의 메커니즘 (Mechanism of Address Routing): 인터체인지 개입은 그룹 B가 주소 선택에 민감한 포인터로서 기능함을 확인시켜 준다. 이의 어텐션 패턴을 패칭하면 모델의 초점이 반사실적 물체 내용을 전달하지 않고도 올바른 소스 위치로 이동한다.
• Q/K 서브스페이스 시그니처 (Q/K Subspace Signatures): 바인딩 ID 개입은 포인터 메커니즘이 Q/K 매칭에 의존함을 밝혀냈으나, 구현 방식은 모델 제품군에 따라 다르다:
  • Gemma 모델: 바인딩 정보가 쿼리와 키 서브스페이스 모두에 표현되는 명확한 Q/K 매칭 시그니처를 보인다. 쿼리 벡터를 이동시키면 어텐션을 효과적으로 재지정할 수 있다.
  • Llama 모델: 주로 키 벡터에 의존하여 바인딩 정보를 나타낸다. 키 측면의 개입이 가장 큰 어텐션 변화를 일으키며, 쿼리 측면의 효과는 더 약하다.
• 일반화 가능성 (Generalizability): 검색 조건부 재바인딩 메커니즘은 Gemma와 Llama 제품군 모두에서 관찰되며, 이는 구체적인 회로 시그니처의 차이에도 불구하고 이것이 동적 상태 추적을 위한 공유된 계산 전략임을 시사한다.

의의 및 주장
본 논문은 LLM이 맥락 의존적 상태 추적을 수행하는 방식에 대해 해석 가능한 회로 수준의 설명을 제공한다고 주장한다. 연구 결과는 LLM이 반드시 모든 상태 변화 후에 세계에 대한 완전하고 전역적으로 업데이트된 정신적 시뮬레이션을 유지하는 것은 아님을 시사한다. 대신, 모델은 스왑 관련 정보를 보존하고 검색 시점에 로컬하게 적용하여 업데이트된 바인딩 ID를 통해 올바른 객체로 어텐션을 재지정하는 "재바인딩" 전략을 활용한다.

저자들은 계산 전략(검색 조건부 재바인딩)은 모델 제품군 전반에 걸쳐 일관되게 나타나지만, 회로 수준의 구현(특히 Q 대 K 서브스페이스에 대한 의존도)은 다르다는 점에 주목한다. 이는 해당 메커니즘이 동적 추적 문제를 해결하기 위한 공유된 솔루션일 수 있으나, 그 실현 방식은 특정 모델 아키텍처나 사전 학습 분포에 따라 달라질 수 있음을 의미한다.

한계점
저자들은 본 연구의 과제 설정이 매우 통제되어 있으며(단일 스와프, 단일 토큰 물체, 알파벳 순서의 박스 이름), 자연스러운 담화의 전체 범위를 포착하지 못할 수 있음을 인정한다. 또한, 개입에 사용된 바인딩 ID 방향은 근사적인 선형 표현 가정을 바탕으로 추정되었으므로, 만약 바인딩 정보가 비선형적으로 표현되거나 여러 서브스페이스에 분산되어 있다면 실제 메커니즘을 과소평가할 수 있다.