기술 요약: 트랜스포머는 어떻게 토큰 간 연관성을 학습하는가: 경사 선도 항이 기계적 해석 가능성을 제공한다

문제 제기

자기 주의 (self-attention) 기반의 대규모 언어 모델 (LLM) 은 사실적 지식 포착과 일관된 텍스트 생성에서 놀라운 능력을 입증해 왔으며, 이는 주로 의미적 연관성(예: "새"와 "날았다" 간의 연결, 또는 "자동차"와 "트럭"의 상호 교환 가능성) 을 학습하는 능력에 기인합니다. 최근 연구들이 인두 (induction) 헤드나 선형 의미 관계와 같은 특정 메커니즘을 규명해 왔지만, 실제 자연어 데이터로 주의 기반 트랜스포머를 훈련하는 과정에서 이러한 연관성이 어떻게 발현되는지에 대한 원칙적이고 이론적 근거를 갖춘 설명은 여전히 부재합니다.

이전의 이론적 연구들은 종종 비현실적인 가정에 의존해 왔습니다. 예를 들어, 합성 구조화된 언어, 위치 인코딩이나 잔차 연결이 없는 단순화된 아키텍처, 또는 비표준 훈련 체계(예: 순차적 구성 요소별 훈련) 등이 그것입니다. 이러한 실제와의 괴리는 현대 LLM 에 대한 통찰의 일반화 가능성을 제한합니다. 본 논문은 다음과 같은 질문을 통해 이 격차를 해소합니다: 자연어 데이터로 훈련되는 주의 기반 언어 모델의 훈련 과정에서 의미적 연관성은 어떻게 발현되는가?

방법론

저자들은 훈련 역학에 기반한 이론적 프레임워크를 제안하며, 특히 훈련 초기 단계에서의 경사 선도 항 (leading-term) 근사를 분석합니다.

1. 모델 설정: 이 연구는 위치 인코딩 (상대적), 인과적 마스킹, 잔차 스트림을 갖춘 표준 $L$-레이어 주의 기반 트랜스포머에 초점을 맞추며, 자연어 말뭉치에 대해 표준 전체 배치 경사 하강법으로 훈련합니다.
2. 선도 항 근사: 트랜스포머가 핵심 동작(의미적 관계 포함) 을 훈련 초기에 획득한다는 점을 인식하여, 저자들은 경사 업데이트를 근사함으로써 가중치 행렬 (출력 $W_O$, 값 $V^{(l)}$, 쿼리 - 키 $W^{(l)}$, 위치 인코딩 $P^{(l)}$) 에 대한 폐형식 (closed-form) 표현식을 유도합니다. 그들은 충분한 수의 단계 $s$(여기서 $s \leq \eta^{-1} \min(\dots)$) 에 대해 가중치가 경사 선도 항에 가깝게 유지됨을 보여줍니다.
3. 기저 함수 분해: 이론적 분석은 학습된 가중치가 말뭉치 통계에서 유도된 세 가지 근본적인 기저 함수의 단순한 합성으로 표현될 수 있음을 밝힙니다:
   • 빅그램 매핑 ($\bar{B}$): 다음 토큰 의존성 ($P(e_j|e_i)$) 을 포착합니다.
   • 교환 가능성 매핑 ($\Sigma \bar{B} = \bar{B}^\top \bar{B}$): 이전 토큰 분포의 유사성에 기반한 토큰 간의 기능적 유사성(예: 동의어 또는 공유된 문법적 역할) 을 포착합니다.
   • 문맥 매핑 ($\bar{\Phi}$): 토큰이 나타나는 평균적인 문맥을 나타내는 더 긴 범위의 접두어 - 접미어 공발생을 인코딩합니다.

주요 기여

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

1. 명시적 가중치 특성화: 다음 토큰 예측 손실 하에서 실제 텍스트 말뭉치로 훈련된 주의 기반 트랜스포머의 가중치에 대한 최초의 명시적이고 폐형식적인 특성화를 제시합니다. 가중치는 위에서 언급한 세 가지 기저 함수의 합성으로 나타남이 입증됩니다.
2. 의미적 연관성의 기계적 해석: 저자들은 가중치에서 학습된 특징들을 빅그램, 교환 가능성, 그리고 문맥 매핑의 특정 합성으로 해석합니다. 이러한 합성들이 모델이 의미적 연관성을 포착하는 방식을 입증합니다(예: "물고기"가 "연못"과 어떻게 문맥 매핑을 통해 관련되는지, 또는 "빨간색"이 "트럭"과 어떻게 빅그램 매핑을 통해 관련되는지).
3. 실증적 검증: 이론적 예측은 TinyStories 데이터셋으로 훈련된 장난감 3-레이어 주의 모델과 OpenWebText 로 훈련된 실용적인 대규모 LLM(Pythia-1.4B) 모두에서 검증되었습니다.

결과

• 이론적 정합성: TinyStories 로 훈련된 3-레이어 트랜스포머에 대한 실험은 100 에포크 훈련 후에도 이론적으로 예측된 선도 항과 실제로 학습된 가중치 간의 코사인 유사도가 매우 높게 유지됨을 보여줍니다 (예: 주의 및 값 행렬의 경우 >0.99).
• 의미적 구조: 정성적 분석은 기저 함수가 의미 있는 언어 구조를 포착함을 확인시켜 줍니다. 예를 들어, 빅그램 매핑 $\bar{B}$는 형용사를 그것이 수식하는 명사와 상관관계가 있으며, 문맥 매핑 $\bar{\Phi}$는 개체들을 그들의 전형적인 환경(예: "물고기"와 "연못") 과 상관관계가 있습니다.
• 실제 LLM 에 대한 일반화: Pythia-1.4B 에 적용되었을 때, 이론적 선도 항 특징은 학습된 임베딩 및 주의 상관관계와 특히 훈련 초기 단계에서 강한 합의를 보입니다. 분석은 훈련이 진행됨에 따라 가중치가 더 풍부한 지식 표현으로 이동하지만, 상당 기간 동안 이론이 예측한 근본적인 연관성 특징을 유지함을 보여줍니다.
• 구성 요소 협력: 전체 모델 계산 (식 12) 에 대한 분석은 잔차 스트림이 평균 다음 토큰 예측을 제공하는 반면, 자기 주의 블록($\bar{Q}$와 $\bar{\Phi}^\top \Sigma \bar{B}$에 의해 주도됨) 이 말뭉치 통계에 기반하여 다음 토큰을 가장 잘 나타내는 토큰을 선택적으로 집중함으로써 이 예측을 정제하는 방식을 보여줍니다.

의의

본 논문은 딥러닝 훈련 역학과 언어 이론을 연결함으로써 대규모 언어 모델의 표현 학습에 대한 기계적 기초를 제공한다고 주장합니다. 자연주의적 텍스트 분포와 표준 훈련 절차를 분석의 기반으로 삼음으로써, 이 작업은 공식적 이론 분석과 실용적 LLM 사용 간의 격차를 크게 줄였습니다.

저자들은 그들의 발견이 각 트랜스포머 구성 요소의 기능을 문맥화하는 엄격한 방법을 제공한다고 주장합니다. 모델을 블랙박스로 보는 대신, 이 이론은 모델의 동작이 먼저 이러한 근본적인 연관성 특징(빅그램, 교환 가능성, 문맥) 을 확립함으로써 발달하며, 이것이 더 복잡한 의미 관계를 학습하기 위한 앵커 역할을 한다고 시사합니다. 이 작업은 가중치 행렬이 다양한 구성 요소에서 공유된 말뭉치 통계의 단순한 함수로 분해될 수 있음을 시사함으로써 해석 가능성 연구에 새로운 길을 열어주며, 특정 메커니즘을 넘어 모델에서 개념이 어떻게 발생하는지에 대한 더 광범위한 가설을 가능하게 합니다.