기술 요약: 언어 모델은 상태 변화를 거치며 개체를 추적하는가?

문제 제기

개체 추적 (Entity Tracking, ET) 은 복잡한 추론에 필수적인 기술로, 작업 (예: 이동, 추가, 제거) 이 발생할 때 문맥 내 개체의 상태를 유지하고 업데이트하는 능력을 포함합니다. 이전 연구는 트랜스포머 언어 모델 (LM) 이 정적 개체 바인딩 (상태 변화 없이 개체를 위치와 연관 짓는 것) 을 어떻게 처리하는지 조사하거나, 단순화된 합성 언어와 '토이 (toy)' 모델을 사용하여 ET 문제를 해결한 바 있으나, 비토이 (non-toy) 사전 훈련된 LLM 이 자연어로 표현된 여러 상태 변경 작업을 포함하는 현실적인 ET 작업을 어떻게 처리하는지에 대한 이해는 제한적입니다. 구체적으로, 이러한 모델이 문맥이 전개됨에 따라 점진적으로 세계 상태를 구축하는지, 아니면 다른 전략을 사용하여 쿼리를 해결하는지는 여전히 불분명합니다.

방법론

저자들은 '박스 (box)' 벤치마크 (Kim & Schuster, 2023) 에서 파생된 데이터셋을 사용하여 비토이 모델 (Gemma-2-2B, CodeLlama-13B, Llama-3.1-70B) 의 ET 기저 메커니즘을 조사했습니다. 이 데이터셋은 상자 내 개체에 대한 초기 설명, 상태 변경 작업 (PUT, REMOVE, MOVE) 의 시퀀스, 그리고 특정 상자의 내용을 묻는 쿼리로 구성됩니다.

본 연구는 행동 분석과 기제적 해석 가능성 (mechanistic interpretability) 을 결합한 이중 접근법을 사용합니다:

1. 행동 프로빙 (Behavioral Probing):

   • 전역 vs. 로컬 프로브: 최종 토큰의 잔류 스트림 (residual stream) 에 선형 프로브를 훈련시켜 모델이 모든 개체의 위치인 전역 상태를 인코딩하는지, 아니면 쿼리된 상자의 내용인 로컬 상태만 인코딩하는지 확인합니다.
   • 이전 상태 프로빙: 프로브를 훈련하여 이전 상태 (최종 작업 이전의 상태) 가 레이어를 가로질러 순차적으로 디코딩되는지 탐지함으로써, 모델이 상태 변화를 점진적으로 처리하는지 병렬로 처리하는지 테스트합니다.

2. 기제적 분석 (Mechanistic Analysis):

   • 회로 추적 (Path Patching): PUT 작업에 관여하는 특정 어텐션 헤드를 식별하고, 이를 정적 개체 바인딩 (DESCRIPTION) 을 위한 회로와 비교합니다.
   • 부분공간 분석: 기호 기반 반사실적 마스킹 (DCM) 과 부분공간 패칭을 사용하여 PUT 및 DESCRIPTION 회로가 위치 정보를 전달하기 위해 동일한 통신 부분공간을 공유하는지 확인합니다.
   • REMOVE 메커니즘 조사: '부재 (absence)' 신호 분석의 어려움으로 인해, 저자들은 로그/순위 분석, 삼진 프로브 (상태를 {존재, 제거됨, 존재하지 않음} 으로 분류), 그리고 인과적 개입 (프로브 가중치 무효화) 을 사용하여 REMOVE 작업이 어떻게 구현되는지 특성화합니다.

주요 발견

1. 비점진적 추적 전략

토큰이나 레이어를 처리함에 따라 LLM 이 점진적으로 세계 상태를 구축한다는 가설과 달리, 결과는 비점진적 전략을 지지합니다:

• 토큰 간: 전역 프로브는 최종 토큰의 잔류 스트림에서 전체 세계 상태를 디코딩하지 못하지만, 로컬 프로브 (쿼리된 상자를 대상으로 함) 는 높은 정확도를 달성합니다. 이는 모델이 전역 상태를 유지하지 않고 쿼리가 제시될 때만 관련 정보를 동적으로 통합함을 시사합니다.
• 레이어 간: 이전 상태를 프로빙한 결과, 모델이 레이어를 가로질러 상태 변화를 순차적으로 처리하지 않는다는 것이 드러났습니다. 대신, 이전 상태는 동일한 레이어에서 최종 상태와 유사한 정확도로 디코딩되어, 모든 관련 작업이 최종 토큰에서 병렬로 처리됨을 나타냅니다.

2. 작업의 메커니즘

• PUT 작업: PUT 메커니즘은 이전 연구 (Prakash 등, 2024) 에서 발견된 개체 바인딩 회로와 유사합니다. 이는 특정 어텐션 그룹 (A, B, C, D) 이 위치 정보 (순서 ID) 를 사용하여 새 개체를 상자 ID 에 바인딩하는 '되돌아보기 (look-back)' 메커니즘을 활용합니다. 중요하게도, PUT 과 DESCRIPTION 회로는 위치 정보를 전송하기 위해 기능적 부분공간을 공유합니다.
• REMOVE 작업: REMOVE 작업은 '전역 억제 태그 (global suppression tag)'를 통해 구현됩니다.
  • 로그 분석에 따르면, 개체가 제거될 때 해당 개체의 로그가 억제되지만, 이 억제는 엄격히 로컬적이지 않으며 상자 문맥과 관계없이 개체에 전역적으로 영향을 미칩니다.
  • 삼진 프로브는 주로 개체 토큰 (상자 ID 토큰이 아님) 에 위치하는 '제거 태그' 신호를 식별합니다.
  • 인과적 개입은 개체 토큰에서 제거 태그를 무효화하면 모델이 제거된 개체를 예측하지 못하게 하지만, 상자 ID 토큰에서의 개입은 효과가 없음을 확인합니다. 이는 모델이 정밀한 상자별 바인딩 해제 메커니즘 대신 취약한 전역 휴리스틱을 사용하여 개체를 억제함을 의미합니다.

3. 예측 및 관찰된 실패 모드

'전역 제거 태그'에 대한 기제적 통찰은 표준 행동 테스트에서는 명확하지 않은 특정 실패 모드를 예측합니다:

• 무작위 제거 (No-op Remove): 문맥에서 언급된 적이 없는 상자에서 개체를 제거하거나 (존재하지 않는 개체를 제거하는 경우), 해당 개체 유형의 전역 억제를 초래하여 다른 상자에서 오류를 발생시킵니다.
• 공유 레이블 개체: 동일한 개체 유형이 여러 상자에 존재하는 경우, 한 상자에서 이를 제거하면 모델이 이를 전역적으로 억제하여 다른 상자에 대한 잘못된 예측을 초래합니다.
• 재도입: 개체를 제거한 후 다른 상자에 다시 넣는 작업은 전역 억제 태그가 지속되기 때문에 실패합니다.
  이러한 특정 시나리오에 대한 행동 테스트는 높은 퇴화율을 확인하여 기제적 가설을 검증했습니다.

기여

1. 비점진적 추적의 증거: 본 논문은 LLM 이 토큰이나 레이어를 통한 누적적 세계 구축이 아닌, 최종 쿼리 단계에서 병렬적으로 대상 토큰을 검색함으로써 ET 작업을 해결한다는 실증적 증거를 제공합니다.
2. 작업의 기제적 특성화: 본 연구는 PUT 을 정적 개체 추적과 유사한 바인딩 회로로 특성화하고, REMOVE 를 개체 수준 태그에 의해 주도되는 전역 억제 메커니즘으로 규명합니다.
3. 행동 및 기제적 분석의 연결: 기제적 통찰 (전역 제거 태그) 이 새로운 실패 모드를 예측하고, 이를 새로운 행동 테스트를 통해 검증하는 피드백 루프를 보여줍니다. 반대로, 행동 결과는 기제적 가설의 형성을 안내합니다.
4. 개입적 해결책: 저자들은 개입을 통해 전역 제거 태그를 무효화하면 이러한 실패 모드를 부분적으로 수정할 수 있음을 보여주어, 기제적 이해를 통한 모델 견고성 개선의 개념 증명을 제공합니다.

의의

본 논문은 현재 트랜스포머 LLM 의 근본적인 아키텍처적 한계를 강조합니다: 즉, 시간에 따른 상태 변경 추적이라는 본질적으로 순차적인 작업을 비순차적, 병렬적 전략으로 해결한다는 점입니다. 지속적인 세계 상태를 유지하는 대신 마지막 토큰에서 관련 문맥을 집계하는 이러한 의존성은 작업이 모호하거나 개체가 여러 문맥에 공유될 때 취약한 행동을 초래합니다. 본 연구는 행동 평가와 기제적 해석 가능성을 결합함으로써 숨겨진 실패 모드를 발견하고 언어 모델의 더 견고한 추론 능력 개발을 안내할 수 있음을 보여줍니다.