"다음 토큰 예측"은 8×4 행렬로도 재현된다

"The cat is sitting on the mat"에서 "mat"을 예측하는 학습은 이 문장 하나만으로 6개의 예측 문제를 동시에 푼다. 입력을 한 칸 왼쪽, 정답을 한 칸 오른쪽에 놓고(teacher forcing), 인과 마스크 덕분에 위치 t의 출력은 위치 0~t만 보고 t+1을 맞힌다. 아래 숫자는 어휘 V=8, 폭 d=4, head 1개, 층 1개, 시퀀스 T=3인 toy 모델을 NumPy로 실제로 돌린 출력이다. 어휘는 ["the","cat","is","sitting","on","mat","dog","<eos>"], 입력 [0,1,2]("the cat is"), 정답 [1,2,3]("cat is sitting")이다.

초기화 직후의 attention은 단순 평균이다

std 0.02로 초기화하면 Q·K가 10⁻³ 규모라 점수 QKᵀ/√4는 모두 ≈0이고, 마스킹 후 softmax를 하면 attention 가중치는 완전 균등이 된다.

A = [[1.      0.      0.    ]    ← "the"는 자기 자신만 본다
     [0.5     0.5     0.    ]    ← "cat"은 the/cat을 반반
     [0.3333  0.3333  0.3333]]   ← "is"는 셋을 균등 평균 (행 합 = 1)

logits에 softmax를 취하면 모든 확률이 1/8=0.125 근처다. 정답 토큰의 확률은 [0.1251, 0.1250, 0.1251], −log의 평균은 2.079017. ln 8 = 2.079442이니 6강의 "초기 손실 ≈ ln V"가 그대로 재현된다. 마지막 위치가 "mat"에 주는 확률은 0.1249이고 argmax는 엉뚱하게 "is"다. 학습되지 않은 모델은 이렇게 생겼다. 참고로 std를 0.5로 키우면 평균 손실이 1.8710으로 우연히 ln 8보다 낮아지는데, 이게 "초기 손실이 ln V보다 훨씬 낮으면 초기화 스케일을 의심하라"는 검사의 근거다.

기울기는 (softmax − onehot)/T에서 시작한다

dL/dlogits = [[ 0.0417 -0.2916  0.0417  0.0416  0.0417  0.0417  0.0417  0.0417]
              [ 0.0417  0.0418 -0.2917  0.0416  0.0417  0.0417  0.0416  0.0416]
              [ 0.0417  0.0417  0.0417 -0.2916  0.0417  0.0417  0.0416  0.0417]]

각 행에서 오답 토큰은 +p/3 ≈ +0.0417(확률을 낮추라는 신호), 정답 토큰은 (p−1)/3 ≈ −0.2917(높이라는 신호)이다. 정답 쪽 기울기가 오답 7개의 합과 정확히 상쇄돼 행 합이 0인 것이 softmax의 성질이다. 이 텐서를 hᵀ와 곱하면 출력 헤드의 기울기가 나오고, 실제 모델에서는 같은 기울기가 W_o·W_v·W_k·W_q·MLP·임베딩까지 연쇄 법칙으로 전파된다.

AdamW 한 step의 숫자

lr=3e-4, β₁=0.9, β₂=0.95, t=1로 적용하면 편향 보정 때문에 m̂/(√v̂+ε) = sign(g)가 되고, 첫 step의 갱신은 Δw = −lr·(sign(g) + 0.1·w)다. "cat" 행의 기울기 부호가 [−,−,−,+]이니 Δw = [+0.0003, +0.0003, +0.0003, −0.0003]. 출력 헤드만 갱신하고 손실을 다시 재면 2.079017 → 2.078987로 아주 조금 내려간다. 이게 한 step이고, 프론티어 모델은 이걸 10⁵~10⁶번, 수십조 토큰에 대해 반복한다. 이 예에서 실무 규칙 셋이 확인된다. 첫 Adam step은 부호 step이라 워밍업이 필요하고, 정답 확률을 올리는 힘과 오답을 내리는 힘이 같은 텐서에서 나오며, step당 이동 거리는 배치가 아니라 lr이 정한다.

import numpy as np; rng=np.random.default_rng(0); V,d,T=8,4,3; std=0.02
softmax=lambda x:(lambda e:e/e.sum(-1,keepdims=True))(np.exp(x-x.max(-1,keepdims=True)))
E=rng.normal(0,std,(V,d)); P=rng.normal(0,std,(T,d))
Wq,Wk,Wv=[rng.normal(0,std,(d,d)) for _ in range(3)]; Wo=rng.normal(0,std,(d,d))/np.sqrt(2*1); Wlm=E.copy()
ids,tg=[0,1,2],[1,2,3]; X=E[ids]+P; Q,K,Vv=X@Wq,X@Wk,X@Wv
S=Q@K.T/np.sqrt(d); mask=np.triu(np.ones((T,T)),1).astype(bool); A=softmax(np.where(mask,-np.inf,S))
H=X+(A@Vv)@Wo; logits=H@Wlm.T; probs=softmax(logits); loss=-np.log(probs[np.arange(T),tg]).mean()
onehot=np.zeros_like(probs); onehot[np.arange(T),tg]=1; dlogits=(probs-onehot)/T; dW=dlogits.T@H
lr,b1,b2,eps,wd=3e-4,0.9,0.95,1e-8,0.1; m=(1-b1)*dW; v=(1-b2)*dW**2
Wnew=Wlm-lr*((m/(1-b1))/(np.sqrt(v/(1-b2))+eps)+wd*Wlm)

같은 문장이 Llama-3-8B를 지나가면

실제 모델과의 차이는 규모뿐이다계산. 문장이 약 7개 토큰으로 잘리면 입력은 [1, 7], 임베딩 후 [1, 7, 4096]. 각 층에서 Q는 [1, 32, 7, 128], K와 V는 GQA 때문에 [1, 8, 7, 128]이고 32개 쿼리 head가 4개씩 하나의 KV head를 공유한다. SwiGLU는 4096 → 14336 → 4096. 32층을 지나 logits는 [1, 7, 128256], 즉 897,792개의 실수다. 학습 시 이 문장은 EOD 토큰으로 다른 문서와 이어져 8,192 토큰 시퀀스에 패킹되고, 16M 토큰 배치에서 이 7개 토큰이 기울기에 기여하는 몫은 7/16,777,216이다. 추론 시 KV 캐시는 토큰당 2×32층×8head×128차원 = 65,536개 값, BF16으로 128KB다. 8,192 토큰 문맥이면 약 1.07GB. MHA(KV head 32)였다면 4배인 512KB/토큰이고, 이게 3강에서 Meta가 GQA를 택한 이유다.

출처: nanoGPT · Llama-3-8B config · FlashAttention-3 · toy 계산은 원 조사자가 NumPy(seed 0)로 실행한 값