LLM 수학 04 — Transformer를 식으로 한 바퀴

선형대수, 확률, 미분과 최적화를 따로 봤다. 이제 decoder-only Transformer 한 layer의 forward, training loss, autoregressive generation을 한 경로로 연결한다.

2017년 원 논문의 Transformer는 encoder-decoder, sinusoidal position encoding, post-LayerNorm 구조다. 오늘날 decoder-only LLM은 pre-norm, RMSNorm, RoPE, SwiGLU, GQA 같은 변형을 자주 사용한다. “Transformer 공식” 하나가 모든 model의 정확한 구현은 아니다.

전체 계산 지도

text
→ tokenizer
→ token ids                         (B, T)
→ token embedding                   (B, T, D)
→ position information              (B, T, D) 또는 Q/K rotation
→ N × Transformer block             (B, T, D)
   ├─ normalization
   ├─ causal self-attention
   ├─ residual
   ├─ normalization
   ├─ feed-forward / gated MLP
   └─ residual
→ final normalization               (B, T, D)
→ vocabulary projection             (B, T, V)
→ shifted cross-entropy              scalar

기호는 다음과 같다.

$B$ — batch size
$T$ — sequence length
$V$ — vocabulary size
$D$ — model dimension
$H_q$ — query head 수
$H_{kv}$ — key-value head 수
$d_h$ — head dimension
$L$ — layer 수

일반 multi-head attention에서는 $D=H_qd_h$ 이고 $H_q=H_{kv}$ 다.

Tokenization은 수학 이전의 계약이다

text를 token id sequence로 바꾸는 tokenizer는 model 밖의 전처리처럼 보이지만 다음을 바꾼다.

sequence length $T$
vocabulary size $V$
숫자·코드·한국어의 분해 방식
input/output embedding parameter 수
context window에 들어가는 실제 text 양
perplexity의 token 단위

같은 문장이 tokenizer A에서는 10 token, B에서는 16 token일 수 있다. “8K context”는 8K 글자나 단어가 아니라 해당 tokenizer의 token 수다.

Token embedding

embedding table은

E\in\mathbb{R}^{V\times D}

이다. token id $x_t$ 의 row를 고른다.

e_t=E[x_t]\in\mathbb{R}^{D}

batch 전체는

X^{(0)}\in\mathbb{R}^{B\times T\times D}

이다. lookup은 one-hot $o_t\in\mathbb{R}^{V}$ 와의 곱으로도 쓸 수 있다.

e_t=o_t^\top E

구현은 one-hot을 만들지 않고 index로 row를 가져온다.

Position이 필요한 이유

self-attention의 content 연산만 보면 token row를 같은 permutation으로 섞었을 때 output도 같은 방식으로 섞이는 permutation equivariance가 있다. 순서를 model에 알려줄 별도 signal이 필요하다.

Sinusoidal position encoding

원 논문의 식은 짝수·홀수 dimension에 sine과 cosine을 사용한다.

PE_{(pos,2i)} = \sin\left(pos/10000^{2i/D}\right)

PE_{(pos,2i+1)} = \cos\left(pos/10000^{2i/D}\right)

입력 embedding에 더한다.

h_{pos}=e_{token}+PE_{pos}

서로 다른 주파수로 position을 표현한다. 원 논문은 fixed sinusoidal과 learned positional embedding이 비슷한 결과를 냈고, 더 긴 length extrapolation 가능성을 보고 sinusoidal을 선택했다. 이것이 모든 후속 model에서 sinusoidal이 최선이라는 뜻은 아니다.

Learned absolute position embedding

h_t=e_{token,t}+P[t]

position table $P\in\mathbb{R}^{T_{max}\times D}$ 를 학습한다. training에서 보지 않은 position으로 확장하는 방법이 자동으로 정해지지 않는다.

Rotary Position Embedding

RoPE는 query와 key의 dimension pair를 position-dependent angle로 회전한다. 2차원 pair를 단순화하면

R_{\theta,t} = \begin{bmatrix} \cos(t\theta)&-\sin(t\theta)\\ \sin(t\theta)&\cos(t\theta) \end{bmatrix}

q'_t=R_{\theta,t}q_t,qquad k'_s=R_{\theta,s}k_s

회전된 dot product에는 상대 position $t-s$ 가 들어간다.

(q'_t)^\top k'_s =q_t^\top R_{\theta,s-t}k_s

실제 RoPE는 여러 frequency를 dimension pair마다 적용한다. position vector를 hidden state에 더하는 대신 attention의 Q와 K에 회전을 적용한다. context extension에서는 base frequency scaling과 training length 조건을 확인해야 한다.

한 layer의 입력

layer $\ell$ 입력을

X^{(\ell)}\in\mathbb{R}^{B\times T\times D}

라 하자. modern pre-norm block을 단순화하면

U=X^{(\ell)}+\operatorname{Attention}(\operatorname{Norm}(X^{(\ell)}))

X^{(\ell+1)}=U+\operatorname{FFN}(\operatorname{Norm}(U))

이다. dropout, bias, parallel residual 같은 세부는 model마다 다르다.

Q, K, V projection

일반 MHA부터 본다.

Q=XW_Q,qquad K=XW_K,qquad V=XW_V

W_Q,W_K,W_V\in\mathbb{R}^{D\times D}

projection 뒤

Q,K,V\in\mathbb{R}^{B\times T\times D}

이고 head를 나누면

Q,K,V\in\mathbb{R}^{B\times H\times T\times d_h}

이다.

이름을 정보 검색과 연결해 읽는다.

query — 현재 position이 찾는 조건
key — 각 position이 match될 주소
value — match weight로 섞을 내용

$K$ 와 $V$ 가 같지 않은 이유는 “어디를 볼지”와 “무엇을 가져올지”에 다른 projection을 학습하게 하기 위해서다.

Scaled dot-product attention

head 하나에 대해

S=\frac{QK^\top}{\sqrt{d_h}}

이다.

shape는

(B,H,T,d_h)(B,H,d_h,T) \rightarrow(B,H,T,T)

다. $S_{bhij}$ 는 batch $b$ , head $h$ 에서 query position $i$ 와 key position $j$ 의 compatibility score다.

왜 $\sqrt{d_h}$ 로 나누는가

$q_i,k_i$ 가 mean 0, variance 1이고 대략 independent라고 단순 가정하면 dot product는

q\cdot k=\sum_{i=1}^{d_h}q_i k_i

이고 variance는 대략 $d_h$ 다.

\operatorname{Var}(q\cdot k)\approx d_h

standard deviation은 $\sqrt{d_h}$ 이므로 나누면 score scale을 order 1로 유지한다.

\operatorname{Var}\left(\frac{q\cdot k}{\sqrt{d_h}}\right)\approx1

큰 logit은 softmax를 포화시켜 gradient가 작아질 수 있다. 이 유도는 independence와 unit variance를 둔 직관이다. 실제 trained Q, K가 정확히 그 distribution이라는 주장은 아니다.

Causal mask

decoder-only model이 position $i$ 에서 미래 $j>i$ 를 보지 못하게 한다.

M_{ij}= \begin{cases} 0,&j\le i\\ -\infty,&j>i \end{cases}

P=\operatorname{softmax}(S+M)

$e^{-\infty}=0$ 이므로 미래 key probability가 0이 된다. 구현은 dtype에서 충분히 작은 finite value를 사용할 수 있다.

padding mask는 batch마다 실제 token이 아닌 padding position을 막는다. causal mask와 목적이 다르며 broadcasting shape를 검증해야 한다.

모든 key가 masked된 row에 softmax를 적용하면 undefined한 $0/0$ 또는 NaN 문제가 생길 수 있다. library kernel의 all-masked behavior를 확인한다.

Softmax axis

P_{ij} = \frac{\exp(S_{ij})}{\sum_{j'}\exp(S_{ij'})}

query $i$ 를 고정하고 key axis $j$ 로 합이 1이다.

\sum_jP_{ij}=1

$(B,H,T_q,T_k)$ 의 마지막 axis다. query axis로 softmax하면 다른 의미의 연산이 된다.

Value의 가중합

O=PV

shape는

(B,H,T,T)(B,H,T,d_h) \rightarrow(B,H,T,d_h)

각 query output은 허용된 value vector의 convex combination이다.

o_i=\sum_jP_{ij}v_j

attention weight이 explanation과 같다는 보장은 없다. output projection, residual, 다음 layer의 nonlinear transformation을 거친 최종 prediction 영향과 한 head의 weight는 다르다.

손으로 계산하는 causal attention

두 token, head dimension 2라고 하자.

Q= \begin{bmatrix} 1&0\\ 1&1 \end{bmatrix},quad K= \begin{bmatrix} 1&0\\ 1&2 \end{bmatrix},quad V= \begin{bmatrix} 2&0\\ 0&4 \end{bmatrix}

raw score는

QK^\top = \begin{bmatrix} 1&1\\ 1&3 \end{bmatrix}

$\sqrt2$ 로 나누고 causal mask를 더한다.

S+M = \begin{bmatrix} 1/\sqrt2&-\infty\\ 1/\sqrt2&3/\sqrt2 \end{bmatrix}

첫 row의 probability는 $[1,0]$ 이다. 첫 token은 자기 자신만 본다.

둘째 row는 두 logit 차이가

\frac{3-1}{\sqrt2}=\sqrt2\approx1.414

이므로

P_{2,\ast}\approx[0.196,0.804]

이다. output은

o_1=[2,0]

o_2=0.196[2,0]+0.804[0,4] \approx[0.392,3.216]

이다.

Multi-head attention

head마다 다른 projection으로 다른 subspace의 관계를 계산한다.

\operatorname{head}_h = \operatorname{Attention}(XW_Q^{(h)},XW_K^{(h)},XW_V^{(h)})

\operatorname{MHA}(X) = \operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_H)W_O

concat 뒤 shape는 $(B,T,Hd_h)=(B,T,D)$ 이고 $W_O\in\mathbb{R}^{D\times D}$ 가 head 정보를 다시 섞는다.

head 수를 늘려도 $D$ 를 고정하면 각 head dimension $d_h=D/H$ 는 작아진다. head 수 증가가 parameter와 표현력을 단순히 비례 증가시키지는 않는다.

MHA, MQA, GQA

query head 수를 $H_q$ , KV head 수를 $H_{kv}$ 라 하자.

MHA — $H_{kv}=H_q$
MQA — $H_{kv}=1$
GQA — $1<H_{kv}<H_q$

GQA는 여러 query head가 한 KV head를 공유하는 group을 만든다. query 표현 수는 유지하면서 KV cache와 decode memory bandwidth를 줄인다.

Q shape는

(B,H_q,T,d_h)

K, V shape는

(B,H_{kv},T,d_h)

다. attention 계산 때 KV head를 query group에 대응시킨다. 단순히 H_q == H_kv라고 가정한 코드는 GQA model에서 shape가 틀린다.

Output projection과 residual

head output을 합친 뒤

A=\operatorname{Concat}(O_1,\ldots,O_H)W_O

A\in\mathbb{R}^{B\times T\times D}

입력과 같은 shape이므로 residual 덧셈이 가능하다.

U=X+A

residual은 새 계산 $A$ 에 원래 representation $X$ 를 보존하는 경로를 제공한다. shape가 같아야 한다.

Feed-forward network

attention이 token axis의 정보를 섞는다면 FFN은 각 token position에 같은 nonlinear function을 독립 적용해 feature를 변환한다.

원 논문 형태는

\operatorname{FFN}(x) =W_2\phi(W_1x+b_1)+b_2

이다.

W_1\in\mathbb{R}^{D\times D_{ff}},qquad W_2\in\mathbb{R}^{D_{ff}\times D}

shape는

(B,T,D)\rightarrow(B,T,D_{ff})\rightarrow(B,T,D)

이다. token끼리 섞이지 않고 각 token row에 동일 weight를 쓴다.

SwiGLU

gated FFN의 한 형태는

\operatorname{SwiGLU}(x) = \operatorname{SiLU}(xW_g)\odot(xW_v)

\operatorname{FFN}(x) = \operatorname{SwiGLU}(x)W_o

이다. projection 두 개를 element-wise 곱하므로 hidden width와 parameter 수를 vanilla FFN과 비교할 때 projection 개수를 같이 세야 한다.

Normalization과 block 순서

LayerNorm은 mean을 빼고 variance로 나눈 뒤 trainable scale과 bias를 적용한다.

\operatorname{LN}(x) = \gamma\odot\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta

RMSNorm은 mean subtraction 없이 RMS로 나눈다.

\operatorname{RMSNorm}(x) = \gamma\odot\frac{x}{\sqrt{\frac1D\sum_i x_i^2+\epsilon}}

원 논문의 post-norm은

y=\operatorname{LN}(x+F(x))

modern pre-norm은

y=x+F(\operatorname{Norm}(x))

형태다. checkpoint architecture를 구현할 때 normalization 종류와 위치가 모두 맞아야 weight를 올바르게 사용할 수 있다.

Layer를 쌓는다

X^{(L)}=\operatorname{Block}_{L-1}(\cdots\operatorname{Block}_0(X^{(0)}))

self-attention 한 layer에서 모든 token이 이전 token의 정보를 섞을 수 있어도 여러 layer가 불필요한 것은 아니다. attention, nonlinear FFN, residual이 반복되며 representation을 단계적으로 바꾼다.

causal model에서 information은 미래에서 과거로 흐르지 않는다. 어느 layer에서도 position $i$ 는 $j>i$ 의 input token을 볼 수 없다.

Vocabulary projection

마지막 hidden을 vocabulary logit으로 바꾼다.

Z=HW_U+b

H\in\mathbb{R}^{B\times T\times D},qquad W_U\in\mathbb{R}^{D\times V}

Z\in\mathbb{R}^{B\times T\times V}

input embedding $E$ 와 output weight를 공유하는 weight tying을 쓰면 $W_U=E^\top$ 형태가 된다. 모든 model이 tying을 쓰는 것은 아니다.

각 position $t$ 의 next-token distribution은

p_\theta(x_{t+1}=i\mid x_{\le t}) = \operatorname{softmax}(Z_t)_i

이다.

Shifted language-model loss

입력 token이

[BOS, 나는, 학교에, 간다]

라면 target은 한 칸 왼쪽으로 민

[나는, 학교에, 간다, EOS]

다. 식은

\mathcal{L}_{LM} =-\sum_{t=1}^{T}\log p_\theta(x_{t+1}\mid x_{\le t})

이다.

teacher forcing에서는 training 시 position $t$ 의 context에 model이 전에 생성한 token이 아니라 dataset의 실제 prefix가 들어간다. inference에서는 model output이 다음 input이 되어 error가 누적될 수 있다.

instruction tuning에서 prompt token loss를 mask하고 answer token만 학습할 수 있다.

\mathcal{L} =-\frac{\sum_t m_t\log p_\theta(x_{t+1}\mid x_{\le t})} {\sum_t m_t}

mask 정책이 system, user, assistant token 중 어디에 loss를 주는지 확인한다.

Backpropagation 경로

한 target token의 gradient는

cross-entropy
→ vocabulary logits
→ output projection
→ final hidden
→ 모든 Transformer block
→ attention/FFN weights
→ token embedding

로 흐른다. shared parameter는 sequence의 여러 position과 batch sample에서 받은 gradient가 합쳐진다.

causal mask는 forward information flow뿐 아니라 gradient path도 제한한다. 미래 token의 loss는 과거 token representation과 weight에 영향을 줄 수 있지만 과거 position output이 미래 input을 직접 본 것은 아니다.

Parameter 수를 근사한다

bias와 norm을 잠시 제외한다.

Attention projection

MHA에서 Q, K, V, O가 모두 $D\times D$ 면

N_{attn}\approx4D^2

Vanilla FFN

N_{ffn}\approx2DD_{ff}

$D_{ff}=4D$ 면

N_{ffn}\approx8D^2

따라서 한 block은 대략 $12D^2$ 다. embedding은 $VD$ , layer $L$ 개면 주요 parameter는

N\approx VD+12LD^2

이다. SwiGLU, GQA, bias, tying 여부에 따라 coefficient가 달라진다. model card의 정확한 architecture config로 다시 계산해야 한다.

Compute 비용을 분리한다

한 layer, batch 생략, MHA 기준의 큰 항만 본다.

Projection과 FFN

linear projection은 대략 $O(TD^2)$ 다. FFN도 $O(TDD_{ff})$ 다.

Attention score와 value mixing

\operatorname{cost}(QK^\top)=O(T^2d_h)\ \text{per head}

head 전체는 대략

O(T^2D)

이고 $PV$ 도 같은 order다.

attention matrix memory는 $(H,T,T)$ 라 $O(HT^2)$ 다. FlashAttention 같은 exact attention algorithm은 tiling과 recomputation으로 materialized memory traffic을 줄이지만 attention의 수학적 output을 근사 ANN으로 바꾸는 것은 아니다.

짧은 $T$ 에서는 $D^2$ projection과 FFN이 지배할 수 있고, 긴 $T$ 에서는 $T^2D$ attention 항이 커진다. “Transformer 비용은 항상 $T^2$ ” 한 줄만으로 실제 latency를 설명할 수 없다.

Training, prefill, decode

Training

sequence 전체 position을 병렬로 계산하지만 activation을 backward용으로 저장해야 한다. memory에는 parameter, gradient, optimizer state, activation이 들어간다.

Prefill

prompt token 전체의 Q, K, V와 hidden을 병렬 계산한다. 긴 prompt에서는 compute-heavy한 matrix multiplication 비중이 크다.

Decode

매 step 새 token 하나를 계산한다. 이전 token의 K, V를 cache하지 않으면 prefix 전체를 매번 다시 계산해야 한다.

KV cache

layer마다 이전 token의 K와 V를 저장한다. batch당 element 수는

N_{KV} =2\cdot L\cdot T\cdot H_{kv}\cdot d_h

이다. byte 수는 dtype byte와 batch를 곱한다.

\operatorname{bytes} =B\cdot2LTH_{kv}d_h\cdot s_{dtype}

예를 들어

$B=1$
$L=32$
$T=4096$
$H_{kv}=8$
$d_h=128$
BF16, 2 byte

이면

1\cdot2\cdot32\cdot4096\cdot8\cdot128\cdot2 =536{,}870{,}912\ \text{bytes}

약 512 MiB다. allocator overhead, cache metadata, padding은 제외한 값이다.

MHA의 $H_{kv}=H_q=32$ 였다면 같은 조건에서 약 2 GiB다. GQA가 decode memory에 큰 영향을 주는 이유다.

KV cache는 model parameter memory와 별개이며 concurrent sequence 수와 context length에 비례한다.

Cached decode의 attention

새 token query $q_t$ 는 지금까지 cache된 key와 dot product한다.

s_t = \frac{q_tK_{1\ldots t}^\top}{\sqrt{d_h}} \in\mathbb{R}^{t}

o_t=\operatorname{softmax}(s_t)V_{1\ldots t}

한 step의 attention은 history length $t$ 에 선형이지만 $T$ 개 token을 순차 생성한 총합은

\sum_{t=1}^{T}O(t)=O(T^2)

이다. cache는 이전 K, V projection의 재계산을 없애지만 모든 과거 key와의 attention 자체를 없애지는 않는다.

Decoding

마지막 position logit $z$ 로 distribution을 만든다.

p_i = \frac{\exp(z_i/\tau)}{\sum_j\exp(z_j/\tau)}

greedy면 argmax, sampling이면 categorical draw다. top- $k$ , top- $p$ , repetition penalty, constrained decoding은 model forward 뒤 distribution 또는 logit을 변형한다.

생성 loop는

1. prompt prefill, KV cache 생성
2. 마지막 logit에서 token 선택
3. 선택한 token 하나를 입력
4. 새 K/V를 cache에 append
5. stop condition까지 반복

이다. autoregressive dependency 때문에 token step 자체는 순차적이다. batch, speculative decoding과 parallel hardware가 throughput을 높여도 한 sequence의 accepted token dependency는 남는다.

Attention과 retrieval의 차이

둘 다 query-key similarity와 weighted value를 사용하지만 범위가 다르다.

구분	Self-attention	External retrieval
후보	현재 context의 token	corpus document/chunk
index	매 forward의 K	미리 구축한 search index
normalization	보통 context key에 softmax	top- $k$ score가 확률일 필요 없음
학습	model 내부 end-to-end	retriever와 generator가 분리될 수 있음
목적	representation mixing	외부 evidence 후보 recall

“attention이 있으므로 RAG가 필요 없다”거나 “vector search가 곧 attention”이라고 동일시하면 system 경계를 놓친다.

구현을 읽을 때 확인할 순서

config에서 $V,D,L,H_q,H_{kv},d_h,D_{ff}$ 를 찾는다.
tensor shape를 batch와 head까지 적는다.
normalization 종류와 pre/post 위치를 본다.
position encoding이 absolute, relative, RoPE 중 무엇인지 본다.
causal mask와 padding mask를 구분한다.
softmax axis와 dtype cast를 본다.
FFN activation과 projection 개수를 본다.
input/output embedding tying 여부를 본다.
training target shift와 loss mask를 본다.
inference에서 KV cache update 위치를 본다.

자주 생기는 오해

Attention weight가 model의 설명이다

한 layer 한 head의 local mixing coefficient다. 최종 output의 causal attribution과 같지 않다.

Context window가 길면 모든 token을 잘 사용한다

입력 가능 길이와 long-context retrieval 정확도는 다르다. position method, training length, attention behavior와 evaluation을 확인한다.

KV cache가 attention을 constant time으로 만든다

과거 K/V projection 재계산을 피하지만 새 query는 여전히 과거 key를 본다.

Head 수가 많으면 항상 더 좋다

$D$ 가 고정이면 head당 dimension이 줄고, 원 Transformer 논문 ablation에서도 head 수를 지나치게 늘렸을 때 품질 저하가 있었다.

모든 Transformer가 원 논문 식 그대로다

encoder-decoder와 decoder-only, MHA와 GQA, LayerNorm과 RMSNorm, sinusoidal과 RoPE가 다르다.

Parameter 수가 inference memory 전부다

weight 외에 KV cache, activation, workspace, runtime allocator가 필요하다. concurrency가 커지면 KV cache가 병목이 될 수 있다.

완료 체크

다음 글은 model 밖으로 나가 corpus를 찾고, ranking하고, RAG 답변을 평가하는 수학을 다룬다. LLM 수학 05 — 검색과 RAG.

Reference

Ashish Vaswani 외. Attention Is All You Need, 2017.
Jianlin Su 외. RoFormer — Enhanced Transformer with Rotary Position Embedding, 2021.
Noam Shazeer. GLU Variants Improve Transformer, 2020.
Joshua Ainslie 외. GQA — Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023.
Tri Dao 외. FlashAttention, 2022.

전체 계산 지도

Tokenization은 수학 이전의 계약이다

Token embedding

Position이 필요한 이유

Sinusoidal position encoding

Learned absolute position embedding

Rotary Position Embedding

한 layer의 입력

Q, K, V projection

Scaled dot-product attention

왜 dh\sqrt{d_h}dh​​로 나누는가

Causal mask

Softmax axis

Value의 가중합

손으로 계산하는 causal attention

Multi-head attention

MHA, MQA, GQA

Output projection과 residual

Feed-forward network

SwiGLU

Normalization과 block 순서

Layer를 쌓는다

Vocabulary projection

Shifted language-model loss

Backpropagation 경로

Parameter 수를 근사한다

Attention projection

Vanilla FFN

Compute 비용을 분리한다

Projection과 FFN

Attention score와 value mixing

Training, prefill, decode

Training

Prefill

Decode

KV cache

Cached decode의 attention

Decoding

Attention과 retrieval의 차이

구현을 읽을 때 확인할 순서

자주 생기는 오해

Attention weight가 model의 설명이다

Context window가 길면 모든 token을 잘 사용한다

KV cache가 attention을 constant time으로 만든다

Head 수가 많으면 항상 더 좋다

모든 Transformer가 원 논문 식 그대로다

Parameter 수가 inference memory 전부다

완료 체크

Reference

왜 $\sqrt{d_h}$ 로 나누는가