기술 요약: Assign and Add: 구성적 산술에 관한 기계론적 연구

문제 정의

대규모 언어 모델(LLM)은 훈련 과정에서 보지 못한 복잡한 과업을 수행하기 위해 기술을 구성하는 능력을 보여주지만, 이러한 "구성적 일반화(compositional generalization)"를 가능하게 하는 구체적인 메커니즘은 여전히 불분명하다. 일반화가 원자적 메커니즘의 구성에 기인한다는 설명이 많음에도 불구하고, 서로 다른 메커니즘—구체적으로 변수 바인딩(유도 헤드를 통한 연상 회상)과 모듈러 산술(푸리에 특징을 통한 구조적 표현)—사이의 상호작용은 통제된 환경에서 체계적으로 연구된 바 없다.

본 논문은 다음 질문을 다룬다: 트랜스포머는 변수 할당과 모듈러 덧셈 메커니즘을 어떻게 구성하여 훈련 분포를 넘어 일반화하며, 이러한 능력은 훈련 중에 어떻게 발현되는가?

방법론

저자들은 변수 할당과 모듈러 덧셈을 결합한 합성 태스크를 수행하는, 기계론적으로 해석 가능한 작은 트랜스포머를 사용하는 통제된 실험 설정을 채택하였다.

• 모델 아키텍처: $d_{model} = d_{head} = 128$, $d_{mlp} = 512$인 2-레이어, 단일 헤드 트랜스포머를 사용한다. 이 아키텍처는 회로 발견을 단순화하기 위해 편향(bias), 정규화 레이어, 그리고 Layer-1 MLP를 의도적으로 제외하였다. GPT-2 스타일의 절대 위치 임베딩과 인과적 마스킹(causal masking)을 사용한다.
• 태스크 정의: 모델은 변수가 상수이거나 컨텍스트 내에서 동적으로 할당된 변수인 경우(예: c=1, b=17, a=42, b+19=?) 모듈러 덧셈($N=59$)을 수행하도록 훈련된다.
• 데이터 분포: 시퀀스는 절대적 위치에 의존하는 것을 방지하기 위해 무작위 패딩과 함께 토큰화된다. 훈련 세트는 일반화를 평가하기 위해 다음과 같이 분할된다:
  • 0-변수 (0-var): 순수 상수 덧셈.
  • 1-변수 (1-var) 및 2-변수 (2-var): 할당된 변수를 포함하는 덧셈.
  • 홀드아웃 세트 (Held-out Sets): 덧셈 쌍의 약 30%가 훈련에서 제거된다 (예: 만약 1+3이 학습되지 않았다면, 해당 시퀀스를 포함하는 모든 시퀀스가 제외됨). 또한, 특정 변수가 훈련 세트의 특정 위치에 오는 것을 제한하여 위치 일반화를 테스트한다.
• 분석 기법: 본 연구는 어텐션 패턴, 잔차 스트림 벡터, 가중치 행렬을 분석하기 위해 TransformerLens 라이브러리를 통한 기계론적 해석 도구를 활용한다. 또한 선형 프로브와 이론적 그래디언트 분석을 사용하여 훈련 역학을 추적한다.

주요 기여

1. 구성적 산술을 위한 다루기 쉬운 모델: 저자들은 변수 할당이 모듈러 산술과 어떻게 구성되는지를 기계론적으로 해석할 수 있게 해주는 단순화된 모델을 제시한다.
2. 분리된 훈련 역학: 분석 결과, 모듈러 산술은 변수 할당보다 먼저, 그리고 독립적으로 학습된다는 것이 밝혀졌다. 모델은 변수 바인딩이 완전히 확립되기 전에도 0-var 태스크에서 완벽에 가까운 정확도를 달성한다.
3. 일반화에 대한 기계론적 설명: 본 논문은 일반화가 두 가지 특정 회로의 구성으로부터 발생함을 식별한다:
   • 할당된 변수를 검색하고 그 값을 라우팅하는 Layer 1의 수정된 유도 헤드(modified induction head).
   • 검색된 값들에 대해 산술을 수행하는 Layer 2의 모듈러 덧셈 회로(구체적으로 Output-Value 행렬과 Layer-2 MLP).
4. 이론적 프레임워크: 저자들은 그래디언트 역학이 어떻게 변수 바인딩과 모듈러 덧셈 회로의 구성을 촉진하는지 보여주는 이론적 분석을 제공하며, Layer-2 Query-Key 행렬에 대한 음의 그래디언트가 자연스럽게 연상 메모리 블록의 형성을 유도함을 보여준다.

결과

훈련 역학 및 단계

훈련 과정은 세 가지 뚜렷한 단계를 보인다:

1. 모듈러 덧셈 학습: 모델은 먼저 변수가 없는 시퀀스(0-var)에 대한 덧셈을 학습하며, 변수 바인딩이 완전히 확립되기 전에 높은 정확도를 달성한다.
2. 변수 할당 학습: 모델은 이후 변수 할당에 필요한 구조를 학습한다. 정확도의 급격한 "스파이크"(약 5,300~5,400 스텝)는 변수 할당 라우팅 회로의 출현과 일치하며, 이때 Layer-1 어텐션이 상수를 할당된 변수로 올바르게 라우팅한다.
3. 정교화 단계: 두 번째 단계(14,500~16,000 스텝)는 "어려운" 시퀀스(예: 홀드아웃된 위치의 특정 변수)에 대한 라우팅을 정교화하며, 특정 변수 위치에 주의를 기울이는 데 발생하는 초기 실패를 수정한다.

일반화 능력

• 위치 일반화: 모델은 새로운 변수 순서(예: 14+b를 본 후 b+14를 예측)에 대해 일반화한다.
• 분포 외(OOD) 일반화: 모델은 훈련 중에 어떤 형태로도 본 적 없는 홀드아웃 덧셈 쌍(예: 1+3)을 성공적으로 예측하며, 큐레이션된 테스트 세트에서 99.4%의 정확도를 달성한다.
• 데이터 요구량: 보지 못한 상수 쌍에 대한 일반화에는 전체 분포의 약 25%($f \approx 0.25$)에 대한 훈련이 필요하다. 2-변수 시퀀스에 대한 일반화에는 훈련 데이터 내 2-var 시퀀스의 상대적 빈도가 약 20%($r \approx 0.2$)가 필요하다.

기계론적 발견

• 공유 표현: 변수가 할당된 덧셈과 순수 상수 덧셈의 잔차 스트림 벡터는 Layer-2 MLP 이전에서 거의 동일하며(코사인 유사도 $\approx 1$), 이는 공유된 계산 메커니즘을 나타낸다.
• 푸리에 특징: Layer-2 MLP 뉴런의 프리액티베이션(preactivations)은 이론적인 푸리에 특징의 형태($\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$)와 일치하며, 이는 모듈러 산술을 위한 확립된 주기적 메커니즘의 사용을 확인해 준다.
• 유도 헤드 수정: Layer-1 어텐션은 상수에 대한 "이전 토큰 헤드(previous-token head)"로서 역할을 하며, 변수 바로 앞에 오는 상수에 주의를 기울인다. Layer-2 Query-K 행렬은 변수 임베딩 상에서 강한 항등 패턴(identity pattern)을 형성하여, = 토큰이 변수와 관련된 올바른 상수에 직접 주의를 기울일 수 있게 한다.

의의

본 논문은 트랜스포머의 구성적 일반화가 내부 메커니즘의 구성성(compositionality)의 자연스러운 결과일 수 있다고 주장한다. 서로 다른 메커니즘(유도 헤드와 푸리에 기반 산술)이 독립적으로 학습될 수 있고, 어텐션 라우팅을 통해 구성될 수 있음을 보여줌으로써, 본 연구는 트랜스포머가 복잡한 추론 과업으로 어떻게 일반화하는지에 대한 기계론적 기초를 제공한다.

저자들은 분석을 용이하게 하기 위해 작고 통제된 환경으로 연구가 제한되어 있음을 언급하며 연구의 범위를 겸손하게 유지한다. 이들은 이 결과가 모든 LLM의 행동을 직접 설명한다고 주장하는 것이 아니라, 트랜스포머에서 일반화 능력이 어떻게 발생하는지에 대한 이해를 돕는 프레임워크를 제공하며, 향후 더 복잡한 과업과 스케일링 법칙에 대한 조사를 촉발할 수 있음을 시사한다.