어텐션 가중치를 직접 꺼내 본다
앞의 일곱 편은 층을 쌓고 걸음을 재는 이야기였다. 이번 편은 구조가 다른 층 하나를 본다. 어텐션은 결국 가중평균이고, 특이한 점은 그 가중치가 학습된 상수가 아니라 매번 입력에서 계산된다는 것이다.
그 가중치는 화면에 나오지 않는다. 꺼내서 그려 보자.
토큰 여섯 개
주제 세 개를 두 토큰씩 나눠 가진 임베딩을 만든다. t0과 t3이 한 주제,
t1과 t4가, t2와 t5가 각각 한 주제다.
import numpy as np
X = 2.6 * np.array([
[1.0, 0.2, 0.0, 0.0], # t0
[0.0, 1.0, 0.3, 0.0], # t1
[0.0, 0.0, 1.0, 0.2], # t2
[0.9, 0.3, 0.0, 0.0], # t3 - t0 과 같은 주제
[0.1, 0.9, 0.2, 0.0], # t4 - t1 과 같은 주제
[0.0, 0.1, 0.9, 0.3], # t5 - t2 와 같은 주제
])
d = X.shape[1]
def softmax(a):
a = a - a.max(-1, keepdims=True)
e = np.exp(a)
return e / e.sum(-1, keepdims=True)
A = softmax(X @ X.T / np.sqrt(d)) # 어텐션 가중치
out = A @ X # 그 가중치로 만든 가중평균Q, K, V 를 만드는 가중치 행렬은 일부러 뺐다. 세 개를 다 넣으면 무엇이
구조에서 나오고 무엇이 학습에서 나오는지가 섞인다. 여기서는 Q = K = V = X
로 두고 구조만 본다.
행 하나가 한 토큰의 시선이다
행 t0 을 보면 자기 자신에 0.510, 같은 주제인 t3 에 0.389, 나머지에는
0.015 에서 0.039 를 준다. 짝과 무관한 토큰의 비가 0.389 / 0.015, 약 26배다.
여섯 행이 모두 같은 모양이다. 모든 행에서 상위 두 칸이 자기 자신과 자기 주제
짝이고, 그 둘이 가져가는 몫이 0.81 에서 0.90 이다. 나머지 네 칸을 다 합쳐도
0.2 가 안 되고, 그중 가장 큰 값이 0.063 이다.
다만 상위 두 칸의 순서는 행마다 다르다. t3, t4, t5 는 자기 자신보다
짝에게 더 준다 - 예를 들어 t3 은 자신에게 0.382, t0 에게 0.468 이다.
내적은 방향뿐 아니라 크기도 보기 때문이다. t0 의 노름이 2.651, t3 이
2.467 이라 t3 입장에서는 t0 과의 내적이 자기 자신과의 내적보다 크다.
어텐션이 “나는 나를 본다” 는 규칙을 갖고 있지 않다는 뜻이다. 규칙은 유사도
하나뿐이고, 크기도 유사도에 들어간다.
행의 합은 1이다. 어텐션이 하는 일이 선택이 아니라 배분이기 때문이다.
어디에 얼마나 볼지를 정하고, 그 비율로 값을 섞는다. 이 행렬을 곱한 결과
A @ X 가 층의 출력이다.
여기까지에 학습된 것은 하나도 없다. 임베딩이 비슷하면 내적이 크고, 내적이 크면 소프트맥스가 큰 몫을 준다.
다만 Q = K = X 로 둔 대가가 하나 있고, 이건 짚어 둬야 한다. 그렇게 두면
로짓 행렬이 정확히 대칭이다. 재 보면 max|L - L.T| 가 0.00e+00 이다.
즉 “t0 이 t3 을 보는 정도” 와 “t3 이 t0 을 보는 정도” 가 로짓 수준에서
같도록 강제된다. 위 행렬에서 두 값이 0.389 와 0.468 로 다른 것은 행마다
따로 정규화하는 소프트맥스 때문이고, 그 아래 로짓은 같은 값이다.
실제 트랜스포머는 X 대신 XW_q, XW_k 를 쓴다. W_q 와 W_k 가 다르면
로짓이 비대칭이 된다 - 무작위 W 200쌍으로 재면 |L - L.T| / |L| 이 중앙값
0.953 이다. 즉 그 두 행렬은 어떤 유사도를 볼지 고르기만 하는 게 아니라,
보는 방향에 앞뒤를 만든다. 이 토이는 어텐션이 무엇을 계산하는지는 보여
주지만 그 방향성은 못 보여 준다.
sqrt(d) 로 나누는 한 줄
X @ X.T 뒤에 / np.sqrt(d) 가 붙어 있다. 이게 없으면 어떻게 되는지 차원을
올려 가며 재 본다. 무작위 벡터 여덟 개로 8×8 어텐션을 만든다.
for d in (4, 16, 64, 256, 1024):
rng = np.random.default_rng(0)
for _ in range(200): # 200회 평균
q = rng.standard_normal((8, d))
k = rng.standard_normal((8, d))
L = q @ k.T
raw, scaled = softmax(L), softmax(L / np.sqrt(d)) d 로짓 표준편차 최대 확률(그대로) 최대 확률(나눔) 엔트로피(그대로) 엔트로피(나눔)
4 1.93 0.526 0.342 1.290 1.750
16 3.88 0.755 0.359 0.662 1.732
64 7.94 0.872 0.357 0.323 1.728
256 15.81 0.937 0.361 0.157 1.719
1024 31.89 0.968 0.363 0.078 1.716여덟 개에 고르게 배분했을 때의 엔트로피가 2.079 다.
로짓의 표준편차가 차원을 따라 커진다. 성분이 분산 1인 두 벡터의 내적은 항이
d 개 더해진 합이라 표준편차가 sqrt(d) 다. 측정값이 1.93, 3.88, 7.94, 15.81, 31.89 로, 2, 4, 8, 16, 32 위에 그대로 앉는다.
소프트맥스는 입력 차이가 커지면 뾰족해진다. 그래서 d=1024 에서 최대 확률이
0.968 로 올라가고 엔트로피가 0.078 로 떨어진다. 고르게 볼 때가 2.079
니까 거의 한 곳만 보고 나머지는 버린다. 가중평균이라기보다 그냥 선택이 된다.
sqrt(d) 로 나누면 로짓의 표준편차가 차원과 무관하게 1 부근으로 돌아온다.
표에서 차원이 256배 늘어나는 동안 최대 확률은 0.342에서 0.363, 엔트로피는
1.750에서 1.716 으로 사실상 움직이지 않는다.
진짜 문제는 기울기다
뾰족해지는 게 왜 나쁜가. 3편의 관점으로 보면 답이 나온다. 소프트맥스의 야코비
행렬은 대각이 p_i(1-p_i), 비대각이 -p_i p_j 다. 확률이 한 곳에 몰려 p 가
0이나 1에 붙으면 두 항이 모두 0으로 죽는다. 아래는 대각합 sum p(1-p) 를 잰
것이고, 이게 작다는 건 야코비 전체가 작다는 뜻이다.
d sum p(1-p) (그대로) sum p(1-p) (나눔)
4 0.6052 0.7747
16 0.3397 0.7691
64 0.1809 0.7705
256 0.0908 0.7668
1024 0.0465 0.7651나누지 않으면 이 값이 차원을 따라 계속 줄어서 d=1024 에서 0.0465 가 된다.
나눈 쪽은 0.77 근처에서 꿈쩍도 안 한다. 16.5배 차이다. 4편에서 본 것과
같은 종류의 사고다 - 순전파 값이 포화하면 그 자리의 미분이 사라진다.
sqrt(d) 는 성능 튜닝이 아니라 학습이 시작되게 하는 조건이다.
어텐션은 순서를 모른다
마지막으로 하나. 위 행렬은 토큰의 내용만 보고 만들어졌다. 위치는 어디에도 안 들어갔다.
perm = [3, 1, 5, 0, 4, 2]
A2 = softmax(X[perm] @ X[perm].T / np.sqrt(d))
np.allclose(A2 @ X[perm], (A @ X)[perm]) # True입력 순서를 섞으면 출력도 정확히 같은 순서로 섞인다. 값이 달라지는 게 아니라 자리만 바뀐다. 즉 어텐션에게 “앞” 과 “뒤” 는 없다.
문장에서 순서는 의미다. 그래서 트랜스포머는 위치 정보를 따로 넣는다. 위치 인코딩이 왜 필요한지는 이 한 줄의 실험이 전부 설명한다 - 넣지 않으면 모델이 단어 순서를 볼 방법이 아예 없다.
그래서
- 어텐션은 가중평균이고, 가중치는 입력에서 계산된다. 행의 합이 1이다
- 그 행렬은 꺼내 볼 수 있다. 이 예에서 모든 행의 상위 두 칸이 자기 자신과 주제
짝이고, 같은 주제 짝
0.389대 무관한 토큰0.015로 26배 차이가 난다 sqrt(d)는 로짓이 차원을 따라sqrt(d)로 커지는 것을 되돌린다. 없으면d=1024에서 최대 확률0.968, 엔트로피0.078로 포화한다- 포화의 대가는 기울기다.
d=1024에서0.0465대0.7651, 16.5배 차이가 난다 - 어텐션은 순서에 대해 아무것도 모른다. 위치 인코딩은 장식이 아니라 누락된 정보를 채우는 것이다
다음 편은 방금 남긴 구멍을 메운다. 어텐션이 순서를 모른다면 순서를 어떻게 집어넣는지, 위치 인코딩을 재 본다.
Comments