Review

TL; DR

Summary

• 연구진: 위스콘신 대학, 시드니 공과대학

• 인용수 : 2

연구 동기

• Self-attention 기반 LLM은 factual 지식과 word knowledge를 모두 학습

  ⇒ 모델 내부에서 어떤 구조가 만들어지고, 어떻게 학습되는걸까?
  ⇒ 아래와 같은 패턴 발견

  • induction heads (패턴 복사)
  • linear semantic relations (선형 의미 관계)
  • topic clustering (주제별 묶임)

• 결론: Semantic association (단어들 간 의미적 연결)이 LLM의 핵심 능력이다!
  • 정의: 토큰들 사이의 통계적(얼마나 자주?) + 기능적(문장 안에서 비슷한 역할을 하는지?) 관계
    • bird ↔ flew ⇒ 같이 자주 등장
    • car ↔ truck ⇒ 서로 대체하기 쉬움
    • country ↔capital ⇒ 비슷한 의미로 묶임

    → 이런 의미적 연관성이 있어야 문장 생성과 일반화가 가능

“그렇다면 트랜스포머는 어떻게 단어 간 의미적 연관(semantic association)을 학습하는가?”

⇒ gradient descent 최적화로부터 자연스럽게 드러남
⇒ 이러한 구조가 학습 시 어떻게 만들어지는지 파악하기 위한 연구들이 진행됨

기존 연구의 한계

트랜스포머의 학습 방식이 매우 복잡해서, 기존 연구는 다음과 같은 비현실적인 가정을 채택

• Synthetic structured language
  • 반복적인 패턴 or 문법이 완전히 규칙적인 toy language
  • 실제 자연어와 괴리 큼

• Positional encoding이나 residual connection이 제거된 단순화된 모델 구조
  • 자연어는 단어의 순서가 중요한데, 순서 정보 전달 기능을 상실함 → bag-of-words처럼 행동
  • Residual connection 없애면 초기 feature 유지가 안되고 gradient 흐름이 깨짐

• 비현실적인 학습 방식
  • end-to-end가 아닌 component별 학습 / 일부 가중치를 freeze
  • 실제 LLM 학습 방식과 다름 (gradient 자체가 전체 network를 통해 전달되기 때문)

제안 아이디어

• 학습 초기 gradient leading term (주요 항) 분석을 통해 트랜스포머 가중치를 수학적으로 해석하자!
  • 트랜스포머는 의미 관계 학습을 초반부에 대부분 진행되고 이후에도 유지되기 때문
  • 초기에는 gradient가 단순해서 가중치를 closed-form으로 근사 가능

⇒ 발견: 학습된 가중치는 단순히 co-occurence가 아닌, 3가지 basis function의 조합으로 표현된다!

• basis function이란?

  • 의미를 구성하는 가장 기본적인 변환 단위

    ⇒ 토큰 → 다른 토큰들과의 관계로 바꿔주는 규칙 / 함수

    • e.g., "fish" → [pond, water, lake, swim, ...]

• Bigram mapping : 인접 토큰 간 의존관계

• Interchangeability mapping: 토큰 간 유사도 반영 (e.g., 유의어, 문법적 역할)

• Context mapping: 고차 의미 관계 (long-range)

Theoretical Analysis

• 어텐션 기반 트랜스포머의 가중치가 gradient의 가장 중요한 부분(leading term)으로 대부분 설명됨

  

  • 작은 초기값(Gaussian initialization)에서 시작해서, learning rate와 step 수가 너무 크지 않은 초기 학습 구간 가정

  • gradient descent를 몇 번 진행하면, 트랜스포머의 각 가중치들이 특정한 형태에 가깝게 됨

    

  ⇒ 모든 layer가 동일한 구조를 배운다!

  • 3가지 basis function

    • Bigram mapping $\bar{B}$

      • $\bar{B}_{ij}=\mathcal{P}_t(e_i)\mathcal{P}_t(e_j|e_i) - \mathcal{P}_t(e_i)/|\mathcal{V}|$
        • $\mathcal{P}_t(e_i)$: 토큰 $e_i$가 전체 데이터에서 얼마나 자주 나오는지
        • 첫번째 항 : i→j로의 등장 확률
        • 두번째 항: random하게 다음 토큰이 등장할 확률

      ⇒ 토큰 $e_i$ 뒤에 토큰 $e_j$가 얼마나 자주 오는지를 기반으로 두 토큰 사이의 “연관성”을 나타냄

    • Interchangeability mapping $\sum_{\bar{B}}$

      • $\sum_{\bar{B}}=\bar{B}^T\bar{B}$ (토큰 간 상관관계 행렬)
        • 두 토큰이 비슷한 방식으로 등장하는지 파악

      • 이를 풀어쓰면 아래 수식과 같음

        

      • 자주 나오는 토큰일수록 중요하며, 두 토큰이 동일한 이전 토큰을 공유하는 지 여부를 고려
        토큰	앞에 오는 토큰
        dog	the, big
        cat	the, small

      ⇒ 이전에 등장하는 토큰 분포가 비슷하면 그 토큰들은 비슷한 기능/역할을 함

    • Context mapping $\bar{\phi}$

      

      • 토큰 $e_i$ 주변 context에 토큰 $e_j$가 얼마나 자주 등장하는 지?
        • 토큰 $e_i$의 등장 위치를 기준으로, 이전 prefix를 모두 살펴봐서 그 중 $e_j$가 등장하는 지 확인

      ⇒ 이 context 정보가 Attention & Value 행렬을 생성하는 핵심 요소가 됨

  • 트랜스포머 각 가중치 행렬이 무슨 의미를 가지나? (데이터 통계 관점)

    • Output matrix $W_O$

      

      ⇒ Bigram mapping이 leading term! → output layer는 bigram 모델이다

      ⇒ 트랜스포머는 next-token 분포부터 학습

    • Value matrix $V^{(l)}$

      

      ⇒ 토큰이 주어지면 context를 요약한 후, next-token에 반영

      ⇒ 단순히 값을 전달하는 것이 아닌, 의미 표현을 생성

    • Attention matrix $W^{(l)}$

      

      • 토큰 간 상관관계를 담고있음

      • query-key가 크게 나온다는 건 단순히 “비슷하다”가 아니라 저 토큰을 보면 지금 필요한 다음 단어 예측이 더 잘 된다는 뜻

      • 그러면 Q가 어떻게 만들어지나?
        • 먼저 각 이전 토큰이 현재 목표 출력(next token)과 얼마나 의미적으로 연결되는지 score를 매김
        • Score에서 미래 토큰은 제거하고, 위치에 따라 정리
        • Score를 “출력 토큰과의 관계”가 아니라 어텐션이 실제로 활용하는 “쿼리 토큰과의 관계”로 변환

• 가중치들이 실제로 같이 어떻게 작용하는가?

  • 기존 트랜스포머 single-layer 수식

    

  • 이 논문의 근사치 대입

    

  ⇒ 트랜스포머는 통계적 관계(B, Φ, ΣB)를 조합하는 구조다

Experiments

• 실험 개요

  • 실험 목표: 이론이 실제로 맞는지 검증하고, 가중치 안에 의미적 구조가 실제로 존재하는지 확인
    • 가중치가 leading term으로 근사가 되는지?
    • 3가지 basis function 구조가 실제로 나타나는지?

  • 데이터셋: TinyStories
    • 행렬이 크면 분석이 어렵기에 3000 단어로 제한

• 실험 결과 1: 3-Layer Transformer 이론 검증

  “SGD로 학습된 가중치와 이론으로부터 유도된 leading term 가중치 간 유사도 비교”

  

  • 모든 행렬에 대해 거의 완벽하게 일치한다! → neural net이 아니라 단순한 통계로도 설명이 가능함

  “학습 epoch 변화에 따라서는 유사도가 어떻게 변할까?”

  

  • 초기 학습은 이론 그대로 진행되며, 모델이 많이 학습된 이후에도 어느정도 유지

• 실험 결과 2: 3-Layer Transformer 의미 해석 검증

  “각 토큰마다 가장 관련된 토큰 top 30을 뽑아서 실제로 의미가 맞는지 확인”

  

  • $\bar{B}$ → 다음에 나올 가능성이 높은 단어들

  • $\sum_{\bar{B}}$ → 비슷한 역할을 하는 단어들

  • $\bar{\phi}$ → 같은 문맥에서 등장하는 단어들

• 실험 결과 3: 실제 LLM에서도 적용 가능한 지 검증

  • Pythia-1.4B 활용 (학습 중간마다 checkpoint 제공해서 layer 별 분석 가능)

  • 하지만 실제 LLM은 MLP, multi-head attention 등 추가적인 component를 포함하기에 가중치 해석이 불가

  ⇒ 가중치를 직접 보지 말고 임베딩 기반으로 토큰 간 상관관계를 간접적으로 추출하자!

  1. 트랜스포머에 각 토큰을 input으로 부여

  2. Layer 통과 전 임베딩, Layer 통과 후 임베딩, 어텐션 통과 후 임베딩 각각 계산

  3. 이를 기반으로 임베딩 행렬 구성

  ⇒ 근데 이론에서는 토큰 간 상관관계인데, 실제로는 토큰 ↔ 임베딩 간 상관관계라서 직접 비교가 어려움
  1. Leading term 계산 (OpenWebTest 데이터에서 실제 통계 계산)

  2. 정규화 후 covariance 행렬 계산 (임베딩을 공분산으로 변환하면 토큰 행렬로 변환됨)

  ---

  

  실제 LLM도 초기에는 이론과 거의 동일하게 학습되며,
  이 구조를 기반으로 점점 더 복잡한 표현을 쌓아감

  • 첫번째 레이어는 아직 충분한 context가 없기에 어텐션에 대한 코사인 유사도는 낮음
    • MLP 레이어는 임베딩 공간을 형성하는 역할 수행
    • 이후에는 어텐션 영향이 커짐

  • 이후 점점 복잡한 context를 배우면서 유사도 상승

  ---

  “각 어텐션 head가 이론적으로 계산된 구조와 얼마나 닮았을까?”

  

  • 초기 레이어는 semantic association 늦게 학습 → 아직 코사인 유사도는 낮음

  • 중간 레이어에서 head들이 서로 다른 역할로 나뉘기 시작

  • 말단 레이어는 variance 감소 → 각 head의 역할이 명확히 구분됨