기술 요약: 어텐션 전용 트랜스포머에서의 간접 목적어 식별을 위한 최소 회로의 출현

문제 정의

기계론적 해석 가능성(Mechanistic interpretability)은 대규모 언어 모델(LLM)을 인간이 이해할 수 있는 계산 회로로 역공학하는 것을 목표로 한다. 그러나 현대 LLM의 규모, 잔차 연결(residual connections), 그리고 비선형성은 특정 추론 작업을 수행하는 데 필요한 최소한의 메커니즘을 종종 가려버린다. 이전 연구(예: Wang et al., 2023)는 GPT-2 small과 같은 사전 학습된 모델에서 간접 목적어 식별(Indirect Object Identification, IOI)을 위한 복잡한 다중 홉(multi-hop) 어텐션 회로를 확인했다. 하지만 이러한 메커니즘은 자연어 및 다중 작업 압력이 존재하는 복잡한 최적화 환경 내에서 발생한다. 본 연구의 과제는 이러한 추론 능력이 더 단순하고 제약된 아키텍처에서도 나타날 수 있는지, 그리고 그 과정에 관여하는 근본적인 계산 모티프를 분리해낼 수 있는지 결정하는 것이다.

방법론

이를 해결하기 위해 저자들은 순수하게 기호화된 버전의 IOI 작업에 대해 처음부터(from scratch) 학습된 최소한의 어텐션 전용 트랜스포머 모델을 훈련시켰다. 이 접근 방식은 언어적 혼란 요소를 제거하고 핵심적인 관계 추론 과제를 격리한다.

• 작업 구성: IOI 작업은 주어(S)와 간접 목적어(IO)가 도입되고 주어가 반복되는 문장에서 간접 목적어를 예측하는 것이다. 모델은 <MID> 토큰 이전에 반복되지 않은 이름을 예측해야 한다. 저자들은 이를 6개 토큰의 기호 시퀀스 <BOS> IO S1 S2 <MID> ?로 추상화했다. 모델은 <MID> 위치의 토큰을 예측하며, 이는 <MID> 토큰 이전에 반복되지 않은 이름이어야 한다 (예: "BABA"와 "BAAB" 템플릿을 구별하는 것과 같음).
• 모델 아키텍처: 본 연구는 어텐션 메커니즘을 고립시키기 위해 MLP와 레이어 정규화를 생략하고 절대적 위치 임베딩을 사용하는 어텐션 전용 트랜스포머를 활용했다.
  • 구성 1: 2개의 어텐션 헤드를 가진 단일 레이어 모델.
  • 구성 2: 각 레이어에 1개의 어텐션 헤드를 가진 2개 레이어 모델.
  • 어휘집(Vocabulary): 총 8개 토큰 (6개의 이름 + <BOS> + <MID>).
  • 훈련: 모델은 모든 60개의 고유 시퀀스에 대해 크로스 엔트로피 손실, AdamW 옵티마이저, OneCycle 학습률 스케줄러를 사용하여 훈련되었다.
• 분석 기법: 저자들은 잔차 스트림 분해(residual stream decomposition), QK(Query-Key) 및 OV(Output-Value) 전이 행렬의 스펙트럼 분석, 임베딩 개입(LogitLens), 그리고 표적 제거 실험(위치 임베딩 또는 특정 합성 경로 제거)을 사용했다.

주요 결과

1. 단일 헤드 모델의 실패

제로 레이어 모델(bigram)과 단일 레이어/단일 헤드 모델은 과제를 해결하는 데 실패했다. 단일 헤드 모델은 종속 절에서 두 이름에 거의 동일한 확률을 할당했다. 분석 결과, 단일 헤드는 "참조 탐지"(올바른 토큰 식별)와 "복사"(해당 정보를 예측 위치로 전달)라는 역할을 동시에 수행할 수 없었으며, 이 두 역할은 단일 어텐션 헤드 내에서 기능적으로 상충하는 것으로 나타났다.

2. 2헤드 1레이어 모델의 성공

1개의 레이어에 2개의 어텐션 헤드를 가진 모델은 100% 정확도를 달성했다. 기계론적 분석 결과, 매우 해석 가능한 역할 분담이 드러났다:

• 헤드 1 (가산적, Additive): 초기 종속 절의 두 이름에 균등하게 어텐션을 준다. 이 헤드의 출력은 이름 임베딩의 합(sum) 방향(정답 + 오답)과 일치한다. 이는 참조 신호의 집합체 역할을 한다.
• 헤드 2 (대조적, Contrastive): 주절의 주어와 종속 절의 반복되지 않은 이름에 어텐션을 준다. 이 헤드의 출력은 차이(difference) 방향(정답 - 오답)과 일치한다. 이는 오답 후보에 대한 대조적 억제(contrastive suppression)를 수행한다.
• 메커니즘: 최종 로짓(logits)은 이 두 헤드의 가산적 결합으로 계산된다. 헤드 1의 가산 신호와 헤드 2의 대조 신호가 결합하여 오답 토큰의 로짓은 상쇄시키고 정답 토큰의 로짓은 증폭시킨다.

3. 스펙트럼 분석

QK 및 OV 행렬의 스펙트럼 분석은 기능적 특성화를 확인해주었다:

• 헤드 1: 증폭 및 억제 고윳값이 혼재되어 비교적 중립적인 어텐션 역학을 보이며, "패스스루(passthrough)" 또는 복사 메커니즘 역할을 한다.
• 헤드 2: 강한 억제 역학(QK에서 지배적인 음의 고윳값)을 보이며, OV 회로에서 가산적 및 회전적(대조적) 기여가 혼재되어 나타나 오답 토큰의 기여를 효과적으로 차감한다.

4. 위치 임베딩의 역할

위치 임베딩을 제거했을 때 모델은 여전히 약 70%의 정확도를 달성하며 의미론적 문맥 관계에 의존했다. 그러나 위치 임베딩이 있을 때 모델은 100% 정확도에 도달했으며, 가산-대조 역할로 완벽하게 분리되었다. 위치 단서(positional cues)는 최적화 지형을 단순화하여 헤드들이 더 깔끔하게 특화될 수 있도록 도왔다.

5. 2레이어 1헤드 모델

1개의 헤드를 가진 2개 레이어 모델 역시 과제를 해결했다. 병렬적 특화가 나타나는 2헤드 모델과 달리, 이 아키텍처는 레이어 간의 **합성(composition)**에 의존했다. Q, K, V 합성 실험에 대한 제거 실험(ablation) 결과, 모델은 Q-합성(제거 시 정확도 약 100% 하락)과 V-합성(제거 시 약 93% 하락)에 크게 의존함을 보여주었다. 이는 두 번째 레이어가 과제 해결을 위해 첫 번째 레이어로부터 집계된 정보를 읽고 변형한다는 것을 나타낸다.

기여

본 논문은 세 가지 주요 기여를 한다:

1. 최소 충분성(Minimal Sufficiency): 1레이어 2헤드 어텐션 전용 모델이 고정된 템플릿 상에서 IOI 과제를 완벽하게 해결하는 데 충분하다는 것을 입증했다. 이는 이전에 더 큰 모델에서 복잡한 다중 홉 회로가 필요하다고 생각되었던 과제이다.
2. 기계론적 회로 발견: 솔루션이 특화된 두 어텐션 헤드 출력(하나의 가산 헤드와 하나의 대조 헤드)의 직접적인 가산적 결합을 통해 계산되는 최소 회로를 밝혀냈다. 이는 GPT-2 small에서 발견된 복잡한 다중 홉 파이프라인과 대조된다.
3. 훈련 역학에 대한 통찰: 복잡한 회로는 다중 작업 압력에서 기인할 수 있는 반면, 과제 제약적 훈련은 더 간결하고 해석 가능한 메커니즘을 드러낼 수 있음을 주장한다.

의의 및 주장

저자들은 본 연구가 트랜스포머 추론의 계산적 기초를 조사하기 위한 통제된 테스트베드를 제공한다고 주장한다. 복잡성을 제거함으로써 다음과 같은 점을 입증했다:

• 과제 특화 훈련은 매우 해석 가능한 최소 회로를 유도한다.
• LLM의 "블랙박스" 복잡성은 추론 과제 자체의 본질적 특성이라기보다, 대규모 다중 작업 설정에서의 최적화 지형에 따른 부산물일 수 있다.
• 식별된 특정 계산 모티프(가산-대조 회로 및 레이어별 합성)는 더 큰 모델이 유사한 추론을 어떻게 구현하는지에 대한 기준점을 제공하며, IOI 해결의 근본 논리가 기존의 사전 학습된 모델에서 관찰된 것보다 더 단순할 수 있음을 시사한다.

본 논문은 연구 범위에 대해 겸허한 태도를 유지하며, 이러한 발견이 엄격한 시퀀스 구조에 국한되어 있으며, 가변적인 시퀀스 길이, 동적 구문, 또는 훈련 중 회로 출현의 발전 역학으로 어떻게 확장되는지는 아직 탐구하지 않았음을 인정한다.