LLM : 칼럼 리뷰 : A Mathematical Framework for Transformer Circuits

A Mathematical Framework for Transformer Circuits

이 글의 목적은 Transformer 내부에서 실제로 어떤 계산이 일어나는지를 회로(circuit) 단위로 분해해서 이해하는 것이다.

거대한 LLM을 처음부터 해석하기보다는, MLP를 제거한 0~2 layer attention-only Transformer부터 분석한다. 구조를 단순화한 뒤 weight와 activation을 뜯어보면서, Transformer의 동작을 사람이 이해할 수 있는 작은 계산들의 조합으로 표현하려고 한다.

1. Residual Stream

Transformer를 볼 때 가장 중심이 되는 것은 residual stream이다.

각 Attention Head나 MLP는 서로 직접 연결되어 있는 것이 아니라,

Residual Stream에서 정보를 읽고 → 계산하고 → 다시 Residual Stream에 결과를 더한다.

즉 residual stream을 Transformer 내부의 공용 통신 채널처럼 볼 수 있다.

각 component가 자신의 결과를 residual stream에 더하는 구조이기 때문에, 서로 다른 head의 contribution을 따로 떼어서 분석할 수도 있다.

---

2. Attention Head = QK Circuit + OV Circuit

Attention Head 하나의 역할은 크게 두 부분으로 나눌 수 있다.

QK Circuit

[
W_Q^T W_K
]

→ 어디를 볼 것인가

현재 token이 어떤 다른 token에 attention을 줄지를 결정한다.

OV Circuit

[
W_O W_V
]

→ 거기서 무엇을 가져올 것인가

선택한 token의 정보를 어떻게 변환해서 residual stream에 써줄지를 결정한다.

따라서 Attention Head를 단순히 Q, K, V vector들의 계산으로 보는 대신,

QK = 어디에서 정보를 가져올지 결정
OV = 가져온 정보를 어떻게 사용할지 결정

하는 하나의 작은 회로로 볼 수 있다.

---

3. Path Expansion

Residual connection은 계속 더하기 구조이기 때문에 Transformer 전체 계산을 전개하면

여러 개의 end-to-end 계산 경로(path)의 합

으로 표현할 수 있다.

예를 들어

[
Embedding \rightarrow Head_1 \rightarrow Head_2 \rightarrow Unembedding
]

처럼 특정 정보가 모델 안에서 어떤 component들을 거쳐 최종 logit까지 전달되는지를 하나의 path로 볼 수 있다.

이렇게 Transformer 전체를 하나의 거대한 계산으로 보는 대신 여러 개의 작은 circuit/path의 합으로 분해해서 분석하자는 것이 이 framework의 중요한 아이디어다. Attention pattern을 고정하면 이 경로들은 대부분 linear operation으로 다룰 수 있다.

---

4. Attention Head Composition

깊이가 2 layer 이상이 되면 뒤의 Attention Head가 앞 Attention Head가 residual stream에 써놓은 정보를 읽을 수 있다.

즉 Head들이 서로 composition되기 시작한다.

크게

• Q-composition: 이전 Head의 output을 Query 계산에 사용
• K-composition: 이전 Head의 output을 Key 계산에 사용
• V-composition: 이전 Head의 output을 Value 계산에 사용

으로 나눌 수 있다.

이 composition 때문에 단일 Attention Head보다 훨씬 복잡한 알고리즘을 만들 수 있다.

---

5. 모델 깊이에 따라 무엇을 할 수 있는가

저자들은 실제 작은 Attention-only Transformer를 뜯어보면서 모델의 깊이가 증가할수록 어떤 계산이 가능해지는지 분석한다.

0 Layer Transformer

Attention 자체가 없기 때문에 현재 token만 보고 다음 token을 예측한다.

결국 사실상

[
A \rightarrow B
]

형태의 bigram model처럼 동작한다.

---

1 Layer Transformer

Attention으로 이전 token을 볼 수 있기 때문에

[
A ; ... ; B \rightarrow C
]

형태의 skip-trigram 같은 패턴까지 사용할 수 있다.

즉 현재 token뿐만 아니라 앞쪽 token의 정보를 가져와 다음 token 예측에 사용할 수 있다.

---

2 Layer Transformer

여기부터 Attention Head끼리 composition이 가능해진다.

그리고 이 과정에서 저자들이 중요하게 발견한 것이 Induction Head다.

Induction Head는 대략

[
[A][B]; ...; [A] \rightarrow [B]
]

패턴을 수행한다.

현재 token과 동일한 token이 과거에 등장한 위치를 찾고, 그 token 다음에 무엇이 나왔는지를 가져와 현재 다음 token으로 예측한다.

예를 들어 context에

... A B ... A

가 있다면 과거의 A를 찾아 그 뒤에 있던 B를 현재 위치로 복사하는 식이다.

이것은 단순한 언어 통계가 아니라 현재 context 내부에서 발견한 패턴을 이용하는 알고리즘이기 때문에, 저자들은 이를 작은 Transformer에서 나타나는 in-context learning의 핵심 mechanism으로 본다.

---

핵심

이 글의 핵심은 Transformer를 단순히

Embedding → Attention → MLP → Attention → ...

처럼 layer 단위로 보는 것이 아니라,

Residual Stream을 통해 서로 정보를 주고받는 작은 circuit들의 집합

으로 보자는 것이다.

특히

Residual Stream → QK / OV Circuit → Head Composition → Path → Induction Head

라는 관점을 만들면서 이후 Transformer Mechanistic Interpretability 연구에서 사용하는 기본적인 분석 틀을 제시했다.

다만 이 글은 작은 Attention-only Transformer를 중심으로 분석하기 때문에 실제 대규모 Transformer 전체를 설명한 것은 아니다. 특히 MLP를 제대로 해석하는 문제는 이 글에서는 거의 해결하지 못했으며, 저자들도 이를 주요 한계로 인정한다.

 

 

 

 

 

 

• 2-layer attention-only model에서 중요한 건 head composition
• 특히 저자들이 본 모델에서는 K-composition이 핵심
• 1층의 previous-token head가 “바로 전 token이 뭐였는지” 정보를 residual stream에 써놓고
• 2층의 induction head의 K가 그 정보를 읽음
• 그래서 현재 token과 같은 token이 예전에 나온 위치를 찾고, 그 다음 token을 복사하는 induction mechanism이 만들어짐
• 반대로 그 모델에서는 Q-composition과 V-composition은 상대적으로 중요하지 않았음
  
  
  3점 / 5점 : 층마다 교환하는 정보의 흐름을 어렵지만 직관적으로 잘 써둠