--- title: "The Illustrated Transformer를 tensor shape으로 다시 읽기" slug: "illustrated-transformer-tensor-flow" category: "AI / LLM" topic: "ai-llm" subtopic: "transformer" tags: ["Transformer","Tensor Shape","Self-Attention","RoPE","GPT"] status: "published" created: "2026-07-17" updated: "2026-07-21" summary: "The Illustrated Transformer의 직관을 실제 QKV tensor와 attention 계산 순서로 옮겨, 한 문장이 block을 통과하는 흐름을 다시 그렸다." kind: "Concept Walkthrough" evidence: "해설·논문 교차 확인" series: "Transformer from first principles" --- [The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/)는 Transformer를 처음 볼 때 가장 도움이 됐던 자료다. 논문 그림은 한 장에 너무 많은 것을 담고 있는데, 이 글은 encoder와 decoder를 열고 self-attention을 한 단계씩 보여준다. 문제는 그림을 보고 난 뒤였다. 색깔 상자가 움직이는 동안에는 알 것 같았지만, 코드를 보면 다시 `view`, `transpose`, `matmul`, `masked_fill`이 쏟아졌다. 그림에서 코드로 넘어가려면 중간에 **tensor shape로 번역하는 단계**가 필요했다. 이 글에서는 작은 문장이 Transformer block 하나를 지나는 과정을 따라간다. 원 해설의 직관을 참고하되, 계산과 구조는 [원 논문](https://arxiv.org/abs/1706.03762)과 대조했다. ## 먼저 고정할 기호 실제 구현은 batch 차원이 앞에 붙는다. 앞으로 다음 기호를 사용한다. | 기호 | 의미 | 예시 | |---|---|---:| | `B` | batch size | 2 | | `T` | sequence length | 4 | | `C` | model dimension, `d_model` | 8 | | `H` | attention head 수 | 2 | | `D` | head dimension, `C / H` | 4 | 입력 token id의 shape은 `[B, T]`다. Embedding을 통과하면 각 token id가 길이 `C`인 벡터로 바뀐다. ```text 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가 됐다고 가정하자. ```text [나는] [사과를] [먹었] [다] 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`를 세 번 선형 투영한다. ```text Q = XW_Q K = XW_K V = XW_V ``` 구현에서는 세 개의 `Linear`를 따로 호출하기보다 한 번에 `3C`차원으로 투영한 뒤 나누기도 한다. ```text 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`로 나눈다. ```text q: [B, T, C] → view [B, T, H, D] → transpose [B, H, T, D] ``` K와 V도 같은 shape이 된다. ```text Q, K, V: [B, H, T, D] ``` 여기서 `transpose`는 장식이 아니다. 각 batch와 head별로 `T × T` attention score를 행렬 곱하기 좋게 축을 배치한다. ```text 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다. ```text score[i, j] = q_i · k_j ``` 내적은 두 벡터가 같은 방향을 향할수록 커진다. 학습된 Q/K 공간에서 방향의 유사도를 관련도 계산에 이용하는 셈이다. 그다음 head dimension의 제곱근으로 나눈다. ```text scaled_score = score / sqrt(D) ``` `D`가 커지면 내적의 분산도 커진다. 큰 값이 softmax를 포화시키지 않도록 scale을 맞추는 단계다. ### 아주 작은 숫자 예시 학습된 실제 벡터가 아니라 계산을 보기 위한 2차원 예시다. ```text q = [1, 0] k1 = [1, 0] k2 = [0, 1] k3 = [1, 1] ``` 내적은 `[1, 0, 1]`이다. `D=2`이므로 `sqrt(2)`로 나눈 뒤 softmax를 적용하면 대략 다음 weight가 나온다. ```text softmax([0.707, 0, 0.707]) ≈ [0.401, 0.198, 0.401] ``` Value를 다음처럼 두자. ```text v1 = [1, 0] v2 = [0, 1] v3 = [1, 1] ``` 출력은 Value의 가중합이다. ```text 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는 아래 삼각형 모양이다. ```text query 0: [보기, 가림, 가림, 가림] query 1: [보기, 보기, 가림, 가림] query 2: [보기, 보기, 보기, 가림] query 3: [보기, 보기, 보기, 보기] ``` 구현에서는 softmax 전에 허용되지 않은 score를 `-inf`로 바꾼다. ```text 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를 곱한다. ```text attention weight [B, H, T, T] V [B, H, T, D] output per head [B, H, T, D] ``` 각 Query 위치가 모든 Value 위치의 정보를 weight만큼 가져온다. 그 뒤 head들을 다시 합친다. ```text [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`차원으로 돌린다. ```text 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을 단순화하면 다음과 같다. ```text 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은 순서가 다르다. ```text x = layer_norm(x + attention(x)) x = layer_norm(x + ffn(x)) ``` 그림을 코드에 대입하기 전에 어떤 normalization 배치인지 확인해야 한다. ## 9단계에서 FFN이 각 위치를 따로 가공한다 Attention에서 토큰 사이 정보 이동이 끝나면 MLP 또는 FFN이 각 토큰 위치에 같은 네트워크를 적용한다. ```text [B, T, C] → Linear(C, 4C) → activation → Linear(4C, C) → [B, T, C] ``` 원 논문은 ReLU를 사용했고, GPT 계열은 GELU를 주로 사용한다. `4C`는 널리 쓰이는 비율이지 반드시 지켜야 하는 수학 법칙은 아니다. FFN에는 `T × T` 연산이 없다. 위치끼리 섞는 일은 attention이 맡고, FFN은 각 위치의 channel을 넓혔다 줄이며 비선형 변환한다. ```text 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차원 성분 쌍 하나만 보면 다음 회전 행렬을 생각할 수 있다. ```text 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은 출처가 갈린다. | 위치 | Q 출처 | K/V 출처 | 볼 수 있는 범위 | |---|---|---|---| | Encoder self-attention | encoder 이전 층 | encoder 이전 층 | source 전체 | | Decoder masked self-attention | decoder 이전 층 | decoder 이전 층 | 현재까지의 target | | Encoder-decoder cross-attention | decoder | encoder output | source 전체 | 이 표를 이해하면 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를 직접 기록하는 것이다. ```text x q, k, v attention score attention weight head output concatenated output logits loss ``` 아직 실행 결과는 없다. 다음 글에는 실행 전 알아야 할 개념과 코드를 읽을 순서를 먼저 정리했다. ## 참고한 자료 - Jay Alammar, [The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/) - Vaswani et al., [Attention Is All You Need](https://arxiv.org/abs/1706.03762) - Su et al., [RoFormer](https://arxiv.org/abs/2104.09864) - Karpathy, [minGPT `model.py`](https://github.com/karpathy/minGPT/blob/37baab71b9abea1b76ab957409a1cc2fbfba8a26/mingpt/model.py)