기술적 요약: 트랜스포머는 선형적으로 고도로 구조화된 세계 모델을 표현한다

문제 제기

이 논문은 순차적 추론 흔적으로 학습된 트랜스포머가 작업 도메인의 표면적 표현이 아닌, 그 도메인의 대수적 구조를 반영하는 내부 "세계 모델"을 개발하는지 여부를 조사한다. 이전 연구 (예: Li et al. [2023], Nanda et al. [2023]) 는 오델로 (Othello) 로 학습된 모델이 보드 점유 상태의 선형 표현 (셀 하나당 하나의 특징) 을 자발적으로 학습함을 보였으나, 이러한 현상이 다른 구조적 제약을 가진 작업에서도 유효한지는 명확하지 않다. 구체적으로, 저자들은 유효성이 개별 셀이 아닌 겹치는 하위 구조 (행, 열, 3×3 박스) 에 의해 결정되는 조합적 제약 충족 작업인 스도쿠에서, 모델이 상태가 셀 단위로 표현되는지 아니면 하위 구조 단위로 표현되는지 질문한다.

방법론

저자들은 백트래킹이 포함된 Norvig 스타일 솔버로 생성된 270 만 개의 스도쿠 해결 흔적으로 8 레이어 인과적 디코더 전용 트랜스포머 (8 개 어텐션 헤드, $d_{model}=576$) 를 학습시켰다. 모델은 홀드아웃 테스트 세트에서 셀 단위 정확도 98.4%, 그리드 단위 정확도 97.5% 를 달성했다.

내부 계산을 분석하기 위해 저자들은 "프로브 및 패치 (probe-and-patch)" 방법론을 사용했다:

1. 선형 프로빙: [clues_end] 토큰 (초기 단서 끝을 표시) 에서의 잔류 스트림 활성화에 로지스틱 회귀 분류기를 훈련하여 특정 개념을 디코딩했다. 세 가지 프로브 계열을 테스트했다:
   • 셀 상태: 특정 셀의 숫자를 예측 (81 개 프로브).
   • 셀 후보: 특정 셀의 유효한 후보 숫자를 예측 (729 개 프로브).
   • 하위 구조 상태: 특정 하위 구조 (행, 열, 또는 박스) 에 숫자가 존재하는지 예측 (243 개 프로브).
2. 인과적 개입 (활성화 패치): 디코딩된 표현이 인과적으로 사용되는지 확인하기 위해, 저자들은 잔류 스트림에서 특정 프로브 방향을 빼고 타겟 토큰 로짓의 감소량을 측정했다.
3. 어텐션 분석: 저자들은 어텐션 헤드를 쿼리 - 키 (QK) 및 출력 - 값 (OV) 회로로 분해하여 헤드가 어떤 정보를 주시하는지 그리고 잔류 스트림을 어떻게 수정하는지 결정했다.
4. 로짓 렌즈 및 직접 로짓 귀속 (DLA): 이러한 도구들은 예측이 레이어를 거쳐 어떻게 진화하는지 추적하고, 로짓 변화를 특정 구성 요소 (헤드 또는 뉴런) 에 귀속시키는 데 사용되었다.
5. 뉴런 제거: 최종 MLP 레이어의 특정 뉴런을 식별하고 제거하여 의사 결정에서의 인과적 역할을 테스트했다.

주요 기여 및 결과

1. 하위 구조 세계 모델의 출현

오델로 모델에서 발견된 "셀 단위" 표현과 대조적으로, 스도쿠 트랜스포머는 퍼즐의 제약 대수학을 중심으로 조직된 세계 모델을 구축한다.

• 선형 분해 가능성: "하위 구조 $S$에 숫자 $d$가 존재하는가?"에 대한 선형 프로브는 모든 243 개의 (하위 구조, 숫자) 쌍에서 중간 레이어 (레이어 4~7) 에서 완벽한 정확 일치 정확도 (1.0) 를 달성한다.
• 셀 수준의 한계: 반면, 개별 셀의 상태를 예측하는 프로브는 약 80% 정확도에서 정체된다. 셀 후보 프로브는 100% 정확도에 도달하지만 구조적 정렬을 보인다: 하위 구조를 공유하는 셀들의 후보 벡터는 매우 상관관계가 높아, 이들이 하위 구조 수준의 특징의 선형 결합임을 시사한다.
• 안정성: 이 하위 구조 표현은 레이어 4 에서 나타나며 교차 레이어 및 교차 위치 전수 실험을 통해 확인된 바와 같이 해결 흔적 전체에 걸쳐 기하학적으로 안정적으로 유지된다.
• 인과적 사용: 활성화 패치는 이러한 방향이 인과적으로 필수적임을 확인한다; 하위 구조에 숫자가 존재함을 나타내도록 이를 덮어쓰면 해당 배치에 대한 모델의 예측이 붕괴된다.

2. 하위 구조 정렬 어텐션 메커니즘

중간 레이어 어텐션 헤드 (레이어 4~6) 는 특정 하위 구조에서 작동하도록 전문화되어 있다.

• 라우팅: 어텐션 헤드는 특정 행, 열, 박스, 밴드, 또는 스택에 집중한다.
• 제약 강제: 이러한 헤드의 OV 회로는 대칭적인 제약 확인 메커니즘을 구현한다. 만약 숫자가 주시된 하위 구조에 존재하면, 헤드는 관련 셀에서 해당 숫자의 로짓을 억제한다; 존재하지 않으면 숫자를 촉진한다.
• 국소화: 특정 헤드 (예: 열 4~6 을 주시하는 헤드) 의 인과적 제거는 해당 특정 하위 구조 내의 불법 배치 로짓만 증가시키고, 다른 하위 구조에는 영향을 미치지 않는다.

3. 네이키드 - 싱글 회로

저자들은 최종 MLP 레이어에서 "네이키드 싱글" (정확히 하나의 유효한 후보를 가진 셀) 배치를 확정하는 데 책임 있는 희소하고 단의적 (monosemantic) 인 회로를 식별했다.

• MLP 이전 상태: 최종 MLP 이전에도 잔류 스트림은 고유한 네이키드 싱글 사례의 98.77% 에서 올바른 숫자를 1 위 순위로 이미 선정한다.
• 전용 뉴런: 최종 MLP 는 해당 셀이 정확히 하나의 후보를 가질 때만 발화하는 91 개의 전용 뉴런을 포함한다. 이러한 뉴런은 단의적이며 다른 후보 수에는 활성화되지 않는다.
• 증폭 메커니즘: 네이키드 싱글 뉴런이 발화하면, 해당 특정 셀의 모든 배치 토큰에 균일하고 거대한 보강 (+3.00 로짓) 을 적용한다. 올바른 숫자는 이미 중간 레이어에 의해 최상위 순위로 선정되었으므로, 이 균일한 증폭은 우위를 유지하고 소프트맥스 확률을 올바른 숫자에 집중시킨다.
• 인과적 확인: 이러한 뉴런을 제거하면 타겟 네이키드 싱글의 확률이 약 0.59 감소하지만, 같은 보드 위의 다른 동시 네이키드 싱글에는 영향을 미치지 않아, 이를 최종 확정 게이트로서의 역할을 확인한다.

중요성 및 주장

이 논문은 출현하는 세계 모델의 기하학이 도메인의 표면적 분해가 아닌 제약 대수학에 의해 형성됨을 입증한다고 주장한다. 스도쿠에서는 유효성이 행, 열, 박스의 교차로 정의되므로, 모델은 이러한 하위 구조의 관점에서 상태를 학습하여 표현한다.

또한, 이 연구는 트랜스포머가 어떻게 조합적 추론 작업을 해결하는지에 대한 완전한 기계적, 엔드 - 투 - 엔드 알고리즘적 설명을 제공한다:

1. 표현: 중간 레이어는 선형적이고 하위 구조 정렬된 세계 모델을 구축한다.
2. 제약 전파: 전문화된 어텐션 헤드는 해당 하위 구조 내의 무효 후보를 억제함으로써 제약을 강제한다.
3. 의사 결정: 최종 레이어의 희소한 단의적 뉴런 집합이 특정 논리 상태 (네이키드 싱글) 를 감지하고 사전 계산된 올바른 답을 증폭한다.

저자들은 기계적 해석 가능성 도구가 이러한 복잡한 추론 과정을 성공적으로 역공학하여, 결과적인 의사 결정 회로가 희소하고, 단의적이며, 완전히 해석 가능함을 드러냈다고 결론지었다.

한계

저자들은 이러한 발견이 단일 모델 아키텍처 (8 레이어, 8 헤드) 와 단일 퍼즐 유형 (스도쿠) 으로 학습된 모델에 국한됨을 지적한다. 그들은 다른 크기의 모델이나 백트래킹이 아닌 흔적으로 학습된 모델에서 하위 구조 세계 모델이나 네이키드 - 싱글 회로가 반드시 나타날 것이라고 주장하지 않는다. 또한, 어텐션 패턴이 어떻게 선택되는지 (QK 회로) 에 대한 분석은 완전히 상세하지 않으며, 인과적 증거는 특정 패치 프로토콜에 의존한다.