The Illustrated Transformer를 tensor shape으로 다시 읽기

The Illustrated Transformer는 Transformer를 처음 볼 때 가장 도움이 됐던 자료다. 논문 그림은 한 장에 너무 많은 것을 담고 있는데, 이 글은 encoder와 decoder를 열고 self-attention을 한 단계씩 보여준다.

문제는 그림을 보고 난 뒤였다. 색깔 상자가 움직이는 동안에는 알 것 같았지만, 코드를 보면 다시 view, transpose, matmul, masked_fill이 쏟아졌다. 그림에서 코드로 넘어가려면 중간에 tensor shape로 번역하는 단계가 필요했다.

이 글에서는 작은 문장이 Transformer block 하나를 지나는 과정을 따라간다. 원 해설의 직관을 참고하되, 계산과 구조는 원 논문과 대조했다.

먼저 고정할 기호

실제 구현은 batch 차원이 앞에 붙는다. 앞으로 다음 기호를 사용한다.

입력 token id의 shape은 [B, T]다. Embedding을 통과하면 각 token id가 길이 C인 벡터로 바뀐다.

token ids:       [B, T]
token embedding: [B, T, C]
position info:   [1, T, C] 또는 Q/K에 적용되는 회전
block input x:   [B, T, C]

예시는 작지만 실제 모델도 같은 축을 쓴다. 숫자만 커진다.

1단계: 토큰을 벡터로 바꾼다

문장 나는 사과를 먹었다가 tokenizer를 거쳐 네 개의 token id가 됐다고 가정하자.

[나는] [사과를] [먹었] [다]
   0       1       2      3   ← position

Tokenizer가 문장을 어디서 나누는지는 모델마다 다르다. 여기서는 attention 흐름만 보기 위해 네 토큰으로 단순화했다.

Token embedding table에서 각 id에 해당하는 벡터를 꺼내면 X는 [T, C]가 된다. Batch까지 포함하면 [B, T, C]다. 아직 이 벡터만으로는 나는 사과를과 사과를 나는의 위치 차이를 충분히 알 수 없다. 그래서 위치 정보를 넣는다.

원 Transformer는 sinusoidal positional encoding을 embedding에 더했다. GPT-2와 minGPT는 학습형 position embedding을 더한다. 현대 LLM에서 자주 보는 RoPE는 Q와 K에 회전을 적용한다. 셋은 위치를 표현한다는 목적은 같지만 계산 위치가 다르다.

2단계: 하나의 입력에서 Q, K, V를 만든다

Self-attention에서는 같은 입력 X를 세 번 선형 투영한다.

Q = XW_Q
K = XW_K
V = XW_V

구현에서는 세 개의 Linear를 따로 호출하기보다 한 번에 3C차원으로 투영한 뒤 나누기도 한다.

x                 [B, T, C]
linear(x)          [B, T, 3C]
split into q,k,v   each [B, T, C]

Q, K, V를 설명할 때 검색 시스템 비유를 많이 쓴다.

• Query는 현재 토큰이 찾는 조건
• Key는 다른 토큰이 가진 비교 표지
• Value는 선택됐을 때 전달할 내용

이 비유가 유효한 지점은 QKᵀ가 관련도를 만들고 그 결과로 V를 섞는다는 데까지다. 실제 벡터에 사람이 읽을 수 있는 질문이나 키 이름이 저장되는 것은 아니다. 학습이 세 projection의 기준을 만든다.

왜 K와 V를 나눴을까

Key와 Value가 같아도 가중합 자체는 계산할 수 있다. 하지만 분리하면 어떤 기준으로 선택할지와 선택된 뒤 무엇을 전달할지를 다른 표현 공간에서 학습할 수 있다.

도서 검색을 예로 들면 제목·저자·태그를 매칭 기준으로 쓰고, 검색 결과에는 본문 요약을 반환할 수 있다. 이것도 정확한 내부 동작의 복제는 아니지만 K와 V를 분리하는 이유를 생각하는 데는 도움이 됐다.

3단계: head 축을 꺼낸다

Multi-head attention을 계산하려면 C를 H × D로 나눈다.

q: [B, T, C]
   → view [B, T, H, D]
   → transpose [B, H, T, D]

K와 V도 같은 shape이 된다.

Q, K, V: [B, H, T, D]

여기서 transpose는 장식이 아니다. 각 batch와 head별로 T × T attention score를 행렬 곱하기 좋게 축을 배치한다.

Q                 [B, H, T, D]
K.transpose       [B, H, D, T]
Q @ K.transpose   [B, H, T, T]

마지막 두 축 T × T가 “각 Query 위치가 각 Key 위치를 얼마나 보는가”를 담는다.

4단계: attention score를 만든다

Query 위치 i와 Key 위치 j의 내적이 raw score다.

score[i, j] = q_i · k_j

내적은 두 벡터가 같은 방향을 향할수록 커진다. 학습된 Q/K 공간에서 방향의 유사도를 관련도 계산에 이용하는 셈이다. 그다음 head dimension의 제곱근으로 나눈다.

scaled_score = score / sqrt(D)

D가 커지면 내적의 분산도 커진다. 큰 값이 softmax를 포화시키지 않도록 scale을 맞추는 단계다.

아주 작은 숫자 예시

학습된 실제 벡터가 아니라 계산을 보기 위한 2차원 예시다.

q = [1, 0]

k1 = [1, 0]
k2 = [0, 1]
k3 = [1, 1]

내적은 [1, 0, 1]이다. D=2이므로 sqrt(2)로 나눈 뒤 softmax를 적용하면 대략 다음 weight가 나온다.

softmax([0.707, 0, 0.707]) ≈ [0.401, 0.198, 0.401]

Value를 다음처럼 두자.

v1 = [1, 0]
v2 = [0, 1]
v3 = [1, 1]

출력은 Value의 가중합이다.

0.401v1 + 0.198v2 + 0.401v3 ≈ [0.802, 0.599]

Q와 K는 weight를 만들 때 쓰이고, 최종 출력에 섞이는 것은 V라는 점이 숫자로 드러난다.

5단계: causal mask로 미래를 가린다

Encoder self-attention은 입력 전체를 볼 수 있다. Decoder-only GPT는 다음 토큰을 예측해야 하므로 현재 위치가 미래 정답을 보면 안 된다.

길이 4인 시퀀스의 causal mask는 아래 삼각형 모양이다.

query 0: [보기, 가림, 가림, 가림]
query 1: [보기, 보기, 가림, 가림]
query 2: [보기, 보기, 보기, 가림]
query 3: [보기, 보기, 보기, 보기]

구현에서는 softmax 전에 허용되지 않은 score를 -inf로 바꾼다.

score = score.masked_fill(mask == 0, -inf)
weight = softmax(score, dim=-1)

exp(-inf)=0이므로 가려진 위치의 softmax weight가 0이 된다. 단순히 attention 결과를 나중에 0으로 만드는 것과 다르다. softmax의 정규화 대상에서 미래 위치를 사실상 제거해야 허용된 위치의 weight 합이 1이 된다.

6단계: score가 아니라 Value를 섞는다

Softmax 이후 attention weight shape은 [B, H, T, T]다. 여기에 V를 곱한다.

attention weight  [B, H, T, T]
V                 [B, H, T, D]
output per head   [B, H, T, D]

각 Query 위치가 모든 Value 위치의 정보를 weight만큼 가져온다. 그 뒤 head들을 다시 합친다.

[B, H, T, D]
→ transpose [B, T, H, D]
→ contiguous + view [B, T, C]
→ output projection [B, T, C]

contiguous()가 코드에 나타나는 이유도 축 이동과 관련 있다. transpose는 메모리를 새로 배열하지 않고 stride를 바꾼 view를 만들 수 있다. 이어지는 view가 기대하는 연속 메모리 배치를 맞추기 위해 contiguous copy가 필요할 수 있다.

7단계: 여러 head를 다시 한 표현으로 합친다

Head마다 다른 projection을 사용하므로 같은 문장에서도 다른 score matrix를 학습할 수 있다. 하나는 가까운 토큰 관계에 민감할 수 있고, 다른 하나는 더 먼 의존성에 반응할 수 있다. 중요한 것은 역할 이름을 사람이 지정하지 않는다는 점이다.

각 head의 결과를 이어 붙인 뒤 W_O projection을 거쳐 다시 C차원으로 돌린다.

Concat(head_1, ..., head_H)W_O

왜 굳이 W_O가 필요한지 처음에는 놓쳤다. Concatenate만 하면 head별 subspace가 옆으로 붙어 있을 뿐이다. Output projection은 그 결과를 다음 residual stream에서 사용할 하나의 표현으로 다시 섞는다.

8단계: Residual stream에 변화량을 더한다

Attention 출력은 원래 입력과 같은 [B, T, C]다. Shape을 유지하는 이유 중 하나가 residual addition이다.

GPT-2 스타일의 pre-LN block을 단순화하면 다음과 같다.

x = x + attention(layer_norm_1(x))
x = x + mlp(layer_norm_2(x))

x는 block 사이를 계속 흐르는 residual stream이다. Attention과 MLP는 이 stream을 완전히 교체하기보다 수정값을 계산해 더한다. 이 관점으로 코드를 보면 block이 덜 복잡해 보인다.

원 Transformer 논문의 post-LN은 순서가 다르다.

x = layer_norm(x + attention(x))
x = layer_norm(x + ffn(x))

그림을 코드에 대입하기 전에 어떤 normalization 배치인지 확인해야 한다.

9단계: FFN이 각 위치를 따로 가공한다

Attention에서 토큰 사이 정보 이동이 끝나면 MLP 또는 FFN이 각 토큰 위치에 같은 네트워크를 적용한다.

[B, T, C]
→ Linear(C, 4C)
→ activation
→ Linear(4C, C)
→ [B, T, C]

원 논문은 ReLU를 사용했고, GPT 계열은 GELU를 주로 사용한다. 4C는 널리 쓰이는 비율이지 반드시 지켜야 하는 수학 법칙은 아니다.

FFN에는 T × T 연산이 없다. 위치끼리 섞는 일은 attention이 맡고, FFN은 각 위치의 channel을 넓혔다 줄이며 비선형 변환한다.

Attention: token mixing
FFN:       channel transformation

이 구분은 Transformer 최적화를 볼 때도 이어진다. 긴 context에서는 attention의 T² 성분이 문제가 되고, 큰 hidden dimension에서는 MLP 행렬 곱의 연산량과 파라미터가 큰 비중을 차지할 수 있다.

위치를 넣는 세 방식 다시 비교하기

Sinusoidal positional encoding

원 논문 방식이다. 위치마다 서로 다른 주기의 sine/cosine 벡터를 만들고 token embedding에 더한다. 파라미터가 없고 훈련 길이보다 긴 위치에도 값을 계산할 수 있다.

Learned position embedding

GPT-2와 minGPT에서 확인할 수 있다. 위치 id를 별도 embedding table에서 조회해 token embedding에 더한다. 구현은 단순하지만 학습한 최대 위치를 넘는 외삽은 자연스럽지 않다.

RoPE

RoPE는 token embedding에 위치 벡터를 더하지 않고 주로 Q와 K에 위치별 회전을 적용한다. 두 위치 m, n에서 회전한 Query와 Key를 내적하면 관계가 절대 위치 자체보다 각도 차이, 즉 m-n에 의존하는 형태를 갖는다.

2차원 성분 쌍 하나만 보면 다음 회전 행렬을 생각할 수 있다.

R(θ) = [[cos θ, -sin θ],
        [sin θ,  cos θ]]

위치 m의 Query에는 R(mθ), 위치 n의 Key에는 R(nθ)를 적용한다. 회전 행렬의 성질 때문에 두 벡터의 내적에는 상대 각도 (n-m)θ가 남는다. 이것이 RoPE를 상대 위치 관점으로 설명하는 핵심이다.

다만 RoPE는 원 Transformer나 초기 minGPT의 구성요소가 아니다. minGPT를 읽을 때 wpe를 찾게 되는 이유도 학습형 absolute position embedding을 사용하기 때문이다.

Encoder와 Decoder에서는 QKV의 출처가 다를 수 있다

Self-attention은 Q, K, V가 같은 시퀀스 표현에서 나온다. Cross-attention은 출처가 갈린다.

이 표를 이해하면 encoder-only, decoder-only, encoder-decoder의 차이도 선명해진다.

• Encoder-only는 양방향 self-attention으로 입력 표현을 만든다.
• Decoder-only는 causal self-attention으로 다음 토큰을 예측한다.
• Encoder-decoder는 source를 읽는 encoder와 target을 생성하는 decoder를 cross-attention으로 연결한다.

GPT는 세 번째 cross-attention을 사용하지 않는다. Prompt와 생성 중인 텍스트가 하나의 token stream 안에 있고, causal language modeling 하나로 학습하기 때문이다.

그림에서 코드로 넘어갈 때 내가 놓쳤던 것

Attention weight가 곧 지식은 아니다

Weight는 해당 층, 해당 head, 해당 입력에서 Value를 섞는 계수다. 모델 전체의 판단 근거나 사실 저장 위치라고 단정할 수 없다. 여러 층의 MLP, residual stream, 다른 head가 함께 출력을 만든다.

QKV는 입력마다 바뀌고 projection weight는 학습된다

W_Q, W_K, W_V는 학습 파라미터다. Q, K, V는 현재 입력 X에 그 파라미터를 적용한 activation이다. 둘을 섞어 부르면 state dict와 forward tensor를 구분하기 어렵다.

학습 병렬화와 생성 병렬화는 다르다

Causal mask가 있어도 학습에서는 정답 target 전체가 있으므로 각 위치의 loss를 한 번에 계산할 수 있다. 생성에서는 새 token을 매번 샘플링한 뒤 다시 입력해야 한다. 이 차이가 KV cache와 decoding 최적화로 이어진다.

Self-attention은 외부 검색이 아니다

RAG의 retrieval은 모델 밖에서 문서를 가져와 context를 구성한다. Self-attention은 주어진 context 안의 token 표현을 섞는다. 둘 다 “관련 정보를 찾는다”는 비유를 쓸 수 있지만 시스템 경계가 다르다.

내가 다시 설명할 수 있어야 하는 문장

지금 단계의 복습 기준이다.

1. 입력 [B,T,C]에서 QKV를 만들고 [B,H,T,D]로 바꾸는 이유를 설명할 수 있는가?
2. QKᵀ가 왜 [B,H,T,T]가 되는지 손으로 적을 수 있는가?
3. scaling과 masking이 왜 softmax 전에 들어가는지 말할 수 있는가?
4. attention weight와 V의 곱이 어떤 정보를 만드는지 설명할 수 있는가?
5. head 결과를 합친 뒤 output projection이 필요한 이유를 말할 수 있는가?
6. Attention과 FFN의 역할을 token mixing과 feature transformation으로 나눌 수 있는가?
7. 원 논문의 post-LN과 GPT 계열 pre-LN 코드를 구분할 수 있는가?
8. sinusoidal PE, learned position embedding, RoPE를 같은 것으로 부르지 않는가?
9. GPT가 원 Transformer decoder에서 cross-attention을 뺀 구조라는 점을 설명할 수 있는가?

아홉 문장을 실제 tensor와 코드로 답할 수 있어야 그림이 내 것이 됐다고 판단할 수 있을 것 같다.

다음 확인

다음 단계는 minGPT의 CausalSelfAttention.forward()에 hook 또는 print를 붙여 다음 shape를 직접 기록하는 것이다.

x
q, k, v
attention score
attention weight
head output
concatenated output
logits
loss

아직 실행 결과는 없다. 다음 글에는 실행 전 알아야 할 개념과 코드를 읽을 순서를 먼저 정리했다.

참고한 자료

• Jay Alammar, The Illustrated Transformer
• Vaswani et al., Attention Is All You Need
• Su et al., RoFormer
• Karpathy, minGPT model.py