不用库实现注意力

한국어 원문으로 표시합니다.

목표

어텐션을 라이브러리 없이 직접 구현합니다. numpy 도 torch 도 쓰지 않습니다 — 표준 라이브러리(math)만 씁니다.

느립니다. 그래서 차원을 작게(d=8–64, n=4–8) 씁니다. 목적은 속도가 아니라 안이 보이는 것입니다.

만들 것

/root/work/tf/model.py 에 다음을 정의합니다.

함수	계약
`softmax(xs)`	합이 1. 큰 입력에서도 터지지 않는다
`attention(Q, K, V, mask=None)`	`(out, weights)` 를 돌려준다. √d_k 로 나눈다
`causal_mask(n)`	`mask[i][j]` 가 참이면 i 가 j 를 볼 수 있다
`multi_head(Q, K, V, h, mask=None)`	마지막 차원을 h 등분해 각각 어텐션 후 이어 붙인다
`pos_encoding(n, d)`	자리마다 다른 `[-1,1]` 범위 벡터
`layer_norm(row, eps=1e-5)`	평균 0, 분산 1
`block(X, h, mask=None)`	`X + multi_head(LN(X), ...)`

행렬은 전부 파이썬 리스트의 리스트([[float]])입니다. 자리 × 차원.

확인

cd /root/work/tf
python3 -c "import model; print(model.softmax([1,2,3]))"

단계

softmax — 오버플로 안 남
attention + √d 측정 → 02-scale.txt
causal_mask
multi_head (h=1 이면 attention 과 동일)
pos_encoding + 순열 실험 → 05-perm.txt
layer_norm
block (pre-LN + 잔차)
정리 → 08-notes.md

참고

채점기는 참조 구현과 1e-6 이내로 비교합니다. 어텐션은 구현이 달라도 같은 숫자가 나오는 연산이라, 값이 어긋나면 대개 스케일링이나 마스크 순서가 틀린 것입니다.

터지지 않는 소프트맥스

/root/work/tf/model.py 에 softmax(xs) 를 만드세요. 합이 1이어야 하고, [1000, 1001, 1002] 처럼 큰 입력에서도 inf·nan 없이 동작해야 합니다.

mkdir -p /root/work/tf. 표준 라이브러리만 씁니다(import math). 최댓값을 빼고 exp 하세요 — exp(x - max) 는 수학적으로 같은 값이지만 오버플로가 나지 않습니다. 이 한 줄이 없으면 큰 값이 들어오는 순간 터집니다.

어텐션 세 줄과 √d

attention(Q, K, V, mask=None) 을 만들어 (출력, 가중치) 를 돌려주세요. 점수는 반드시 √d_k 로 나눕니다. 그리고 나눴을 때와 안 나눴을 때 최대 가중치가 얼마나 다른지 재서 02-scale.txt 에 남기세요.

Q·K·V 는 [[float]] (자리 × 차원). score[i][j] = dot(Q[i], K[j]) / sqrt(len(K[j])), w[i] = softmax(score[i]), out[i] = Σ w[i][j] * V[j]. 측정은 d=64 짜리 난수 벡터로 하세요 — 나누지 않으면 최대 가중치가 1.0 에 붙습니다. 그게 '가중 평균이 아니라 하나 고르기' 가 된 상태입니다.

미래를 가린다

causal_mask(n) 을 만드세요. mask[i][j] 가 참이면 i 번째 질의가 j 를 볼 수 있다는 뜻입니다. attention 이 이 마스크를 받으면 가려진 자리의 가중치가 정확히 0 이어야 합니다.

i >= j 면 볼 수 있습니다. 구현할 때 소프트맥스 전에 점수를 -inf(또는 아주 작은 값)로 만드세요. 소프트맥스 뒤에 0을 곱하면 남은 가중치의 합이 1이 아니게 됩니다 — 그게 흔한 버그입니다.

쪼개고 다시 붙인다

multi_head(Q, K, V, h, mask=None) 을 만드세요. 마지막 차원을 h 등분해 각각 어텐션을 돌리고 다시 이어 붙입니다. 출력 모양은 입력과 같아야 합니다.

각 헤드의 차원은 d // h 이고, 스케일링도 그 작은 차원 기준이어야 합니다. h=1 이면 attention 과 정확히 같은 결과가 나와야 합니다 — 그게 가장 좋은 자체 검증입니다.

어텐션이 순서를 모른다는 것을 증명한다

pos_encoding(n, d) 를 만드세요. 그리고 입력 순서를 섞었을 때 위치 인코딩이 없으면 출력이 그대로 따라 섞이고, 더하면 그렇지 않다는 것을 확인해 05-perm.txt 에 남기세요.

사인/코사인이든 다른 방식이든 좋습니다. 값은 [-1,1] 범위이고 자리마다 달라야 합니다. 증명 방법: X 를 섞은 X' 로 어텐션을 돌린 결과가 원래 결과를 같은 순서로 섞은 것과 같은지 비교하세요. 이게 permutation equivariance 입니다.

LayerNorm 은 특징 축으로

layer_norm(row, eps=1e-5) 를 만드세요. 한 토큰 벡터를 평균 0, 분산 1로 맞춥니다.

배치가 아니라 그 벡터 안에서 평균과 분산을 냅니다. 그래서 배치 크기나 시퀀스 길이에 영향받지 않습니다 — BatchNorm 과의 결정적 차이입니다. 상수 벡터가 들어와도 0으로 나누지 않도록 eps 를 쓰세요.

블록을 조립한다

block(X, h, mask=None) 을 만드세요. pre-LN + 잔차 구조입니다: X + multi_head(LN(X), ...). 출력 모양은 입력과 같아야 합니다.

rows = [layer_norm(r) for r in X] 한 뒤 멀티헤드를 돌리고, 원래 X 를 더합니다. 잔차가 없으면 깊이 쌓았을 때 기울기가 입력까지 못 갑니다. 요즘 모델이 post-LN 이 아니라 pre-LN 인 이유는 워밍업 없이도 학습이 되기 때문입니다.

세 가지를 숫자로 정리한다

08-notes.md 에 세 줄 이상. 2단계에서 잰 두 최대 가중치, 5단계가 보여 준 성질, 그리고 시퀀스 길이를 2배로 늘리면 점수 행렬이 몇 배가 되는지.

본문에 스케일링, 순서, 제곱 이 들어가야 합니다. 세 번째가 실무에서 가장 자주 부딪히는 것입니다 — 컨텍스트를 4k 에서 8k 로 늘리면 4배입니다.