√d_k는 어디서 나왔나: 내적의 분산 계산

어텐션 식에는 설명 없이 지나가는 기호가 하나 있습니다.

$\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V$

분모의 $\sqrt{d_k}$ 입니다. 1단원의 지도 글에서 「점수가 너무 커지지 않게 나눈다」고만 적고 넘어갔던 자리인데, 그 문장은 두 가지를 설명하지 않습니다 — 왜 하필 제곱근인가, 그리고 나누지 않으면 정확히 무슨 일이 일어나는가.

이 글은 그 둘을 계산으로 닫습니다. 앞쪽은 분산 계산 세 줄이면 끝나고, 뒤쪽은 softmax의 야코비안이 스케일에 따라 어떻게 변하는지를 재 보면 나옵니다. 그리고 스케일링을 뺀 어텐션이 실제로 학습을 멈추는 것을 400걸음짜리 실험으로 재현합니다.

q·k의 분산은 정확히 d_k다

점수 하나는 질의 벡터와 키 벡터의 내적입니다.

$s = q \cdot k = \sum_{i=1}^{d_k} q_i k_i$

$q$ 와 $k$ 의 성분이 평균 0, 분산 1로 서로 독립이라고 둡니다. 초기화 직후의 표현이 대체로 그렇고, LayerNorm을 지난 뒤라면 더 그렇습니다. 이 가정 아래 $s$ 의 평균과 분산을 계산합니다.

평균부터입니다. 기댓값의 선형성으로 합의 기댓값은 기댓값의 합이고, $q_i$ 와 $k_i$ 가 독립이므로 곱의 기댓값은 기댓값의 곱입니다.

$E[s] = \sum_i E[q_i k_i] = \sum_i E[q_i]E[k_i] = \sum_i 0 \cdot 0 = 0$

분산은 항 하나부터 봅니다. 평균이 0이므로 $\operatorname{Var}(q_ik_i) = E[(q_ik_i)^2] = E[q_i^2]E[k_i^2]$ 이고, $E[q_i^2] = \operatorname{Var}(q_i) + E[q_i]^2 = 1 + 0 = 1$ 입니다. 그러니

$\operatorname{Var}(q_i k_i) = 1 \times 1 = 1$

입니다. 이제 독립인 것들의 합에서 분산이 더해진다는 성질을 씁니다. 서로 다른 $i$ 의 항들은 독립이므로

$\operatorname{Var}(s) = \sum_{i=1}^{d_k}\operatorname{Var}(q_ik_i) = \sum_{i=1}^{d_k} 1 = d_k$

분산이 정확히 $d_k$ 입니다. 표준편차는 그 제곱근인 $\sqrt{d_k}$ 이고, 여기서 그 기호가 나옵니다. 튜닝으로 찾은 값이 아니라 세어서 나온 수입니다.

실제로 난수를 뽑아 이만 번씩 재 보면 이렇습니다.

$d_k$	실측 표준편차	$\sqrt{d_k}$
8	2.829	2.828
64	7.964	8.000
128	11.397	11.314
512	22.669	22.627

그러니 $\sqrt{d_k}$ 로 나누면 분산이 1로 돌아옵니다. 상수 $c$ 를 곱하면 분산은 $c^2$ 배가 되므로

$\operatorname{Var}\!\left(\frac{s}{\sqrt{d_k}}\right) = \frac{1}{d_k}\cdot d_k = 1$

입니다. 차원이 몇이든 softmax에 들어가는 점수의 퍼짐이 같아진다 — 이것이 나누는 목적 전부입니다.

한 가지 짚어 둘 것이 있습니다. 위 계산은 성분끼리 독립이라는 가정 위에 있고, 학습이 진행되면 그 가정은 정확히는 깨집니다. 그래도 $\sqrt{d_k}$ 는 그대로 씁니다 — 정확한 정규화가 목적이 아니라 차원에 따라 스케일이 통째로 어긋나는 것을 막는 것이 목적이기 때문입니다.

스케일이 커지면 분포가 한 점에 몰린다

점수가 커지면 softmax는 어떻게 될까요. $d_k = 64$ 에서 실제로 뽑은 여섯 점수로 봅니다.

$s = (1.585,\; 0.554,\; 4.287,\; -0.693,\; -8.705,\; -9.108)$

8로 나누면 $(0.198, 0.069, 0.536, -0.087, -1.088, -1.138)$ 입니다. 두 벌을 각각 softmax에 넣습니다.

나눈 쪽은 가장 큰 확률이 0.307이고 여섯 자리가 고르게 나뉩니다. 나누지 않은 쪽은 세 번째 자리가 0.911을 가져가고 다섯 번째와 여섯 번째는 $10^{-6}$ 아래로 내려앉습니다.

엔트로피로 재면 숫자 하나로 요약됩니다. 나눈 쪽이 1.643, 나누지 않은 쪽이 0.371입니다. 여섯 자리가 완전히 균등할 때가 $\log 6 = 1.792$ 이므로, 나눈 쪽은 거의 균등에 가깝고 나누지 않은 쪽은 사실상 자리 하나만 남은 것입니다.

이유는 softmax가 차이의 지수를 보기 때문입니다. 점수를 $c$ 배 하면 자리 사이의 차이도 $c$ 배가 되고, 확률의 비는 $e^{c\Delta}$ 로 지수적으로 벌어집니다. 위 예에서 1등과 5등의 차이가 나눈 쪽에서는 1.62, 안 나눈 쪽에서는 12.99이고, $e^{12.99}$ 는 44만입니다.

그래서 기울기가 죽는다

분포가 뾰족한 것 자체는 문제가 아닙니다. 잘 학습된 어텐션은 실제로 뾰족합니다. 문제는 학습이 시작되기도 전에 그렇게 된다는 것이고, 그때 어떤 일이 벌어지는지는 37번에서 유도한 야코비안이 그대로 말해 줍니다.

$J = \operatorname{diag}(p) - p\,p^{\mathsf T}, \qquad J_{ij} = p_i(\delta_{ij} - p_j)$

$p$ 가 원-핫에 가까워지면 어떻게 될까요. 어떤 $p_c \approx 1$ 이고 나머지가 $\approx 0$ 이라면

• $i = j = c$ 성분은 $p_c(1-p_c) \approx 1 \times 0 = 0$
• $i = j \neq c$ 성분은 $p_i(1-p_i) \approx 0 \times 1 = 0$
• $i \neq j$ 성분은 $-p_ip_j$ 인데 둘 중 하나가 0에 가까우므로 $\approx 0$

모든 성분이 0으로 갑니다. 어텐션 가중치에서 점수로 돌아가는 길이 통째로 막히는 것이고, 그러면 $Q$ 와 $K$ 는 기울기를 받지 못합니다.

여기가 마지막 층과 다른 자리입니다. 마지막 층에서는 교차엔트로피의 $1/p_i$ 가 야코비안의 $p_i$ 와 약분되어 $p - y$ 만 남았습니다. 어텐션에는 그 약분을 해 줄 상대가 없습니다 — 어텐션 가중치 위에 붙는 것은 손실이 아니라 $V$ 와의 곱이라서, 야코비안이 소거되지 않고 그대로 곱해집니다.

같은 여섯 점수에 배수를 곱해 가며 세 값을 잽니다. 가로축의 1이 「$\sqrt{d_k}$ 로 나눈 것」이고 8이 「나누지 않은 것」입니다.

배수	가장 큰 확률	엔트로피	야코비안 최대 성분
1	0.3066	1.643	0.2126
2	0.4417	1.397	0.2466
2.52	0.5063	1.280	0.2500
4	0.6670	0.981	0.2221
8	0.9109	0.371	0.0812
16	0.9949	0.034	0.0051
32	0.99998	0.0002	0.0000206

야코비안이 가장 큰 자리는 가장 큰 확률이 0.5 근처일 때이고 그 값이 0.25입니다 — $p(1-p)$ 의 최댓값이므로 앞 글의 sigmoid와 같은 계산입니다. 그 오른쪽부터는 뾰족해질수록 급격히 떨어집니다. 배수 16이면 야코비안이 0.005, 32면 $10^{-5}$ 입니다.

$d_k = 512$ 인 모델에서 나누지 않으면 배수가 22.6이므로 표의 오른쪽 끝보다 더 갑니다. 초기화 직후에 이미 죽어 있는 것입니다.

실험: 스케일링을 빼면 멈춘다

말로 끝내지 않고 재현해 봅니다. 아주 작은 문제 하나를 만듭니다.

• $d_k = 64$, 키와 값이 여덟 쌍. 값의 차원은 8.
• 질의 벡터 $q$ 하나만 학습한다. 목표는 네 번째 값을 그대로 내놓는 것이다.
• 손실은 $L = \tfrac12\lVert o - v_4 \rVert^2$, 여기서 $o = \sum_j p_j v_j$ 이다.

이 손실은 교차엔트로피가 아니므로 야코비안이 소거되지 않습니다. 실제 어텐션이 놓인 자리와 같습니다.

걸음	나눈 것 — 손실 / $p_4$ / $\lVert\partial L/\partial q\rVert$	안 나눈 것 — 손실 / $p_4$ / $\lVert\partial L/\partial q\rVert$
0	1.670 / 0.087 / 7.8e−1	1.597 / 0.0006 / 2.5e+0
10	0.734 / 0.161 / 2.7e−1	5.906 / 0.000 / 6.4e+0
50	0.025 / 0.852 / 4.4e−2	1.893 / 0.000 / 4.2e−9
100	0.008 / 0.918 / 1.5e−2	1.893 / 0.000 / 4.2e−9
399	0.001 / 0.966 / 2.8e−3	1.893 / 0.000 / 4.2e−9

나눈 쪽은 400걸음에 손실 0.0014까지 내려가고 목표 자리의 확률이 0.966이 됩니다. 나누지 않은 쪽은 50걸음쯤에서 손실 1.893에 얼어붙고 그 뒤로 한 번도 움직이지 않습니다. 기울기의 크기가 $4.2\times 10^{-9}$ 이니 float64에서도 갱신량이 사라진 것이고, 실제로 400걸음째까지 손실의 소수점 여섯 자리가 그대로입니다.

멈춘 지점을 보면 원인이 분명합니다. 그때 가장 큰 확률이 1.0입니다 — 틀린 자리 하나에 완전히 몰려 있고, 야코비안이 0이라 「그 자리가 아니다」라는 신호가 점수까지 돌아가지 못합니다. 한 번 잘못 뾰족해지면 빠져나올 길이 없다는 것이 이 실험이 보여 주는 전부입니다.

초기화·LayerNorm과 역할이 어떻게 나뉘는가

「스케일을 맞추는 장치」가 트랜스포머에 셋 있고 하는 일이 서로 다릅니다.

장치	무엇을 맞추나	언제
초기화 스케일	가중치 하나하나의 크기	학습 시작 전 한 번
LayerNorm	층에 들어가는 표현 벡터의 평균과 분산	매 순전파, 층 앞뒤
$1/\sqrt{d_k}$	두 벡터를 내적한 결과의 분산	매 순전파, softmax 직전

앞의 둘이 있어도 $1/\sqrt{d_k}$ 는 없앨 수 없습니다. LayerNorm은 벡터 하나를 평균 0·분산 1로 만들 뿐이고, 그렇게 정리된 두 벡터를 내적하면 위 계산 그대로 분산이 $d_k$ 가 되기 때문입니다. LayerNorm이 만들어 주는 조건이 오히려 이 계산의 가정을 성립시킵니다.

초기화로 대신할 수 있을까요. $W_Q, W_K$ 의 분산을 $1/\sqrt[4]{d_k}$ 배 줄여 두면 초기 점수의 분산은 1이 됩니다. 하지만 그 값은 학습이 진행되면 지켜지지 않습니다 — 가중치가 커지는 것을 막는 장치가 없기 때문입니다. $1/\sqrt{d_k}$ 는 매 순전파에서 다시 적용되는 상수이므로 그 문제가 없습니다.

정리하면 셋의 관계는 이렇습니다 — 초기화가 출발점을 잡고, LayerNorm이 층 입력을 유지하고, $1/\sqrt{d_k}$ 가 내적이라는 연산 자체가 만드는 차원 의존을 지웁니다. 앞의 둘이 아무리 잘 되어 있어도 셋째 자리는 비어 있습니다.

코드로 확인하기

import math, random

def softmax(z):
    m = max(z)
    e = [math.exp(v - m) for v in z]
    s = sum(e)
    return [v / s for v in e]

def entropy(p):    return -sum(x * math.log(x) for x in p if x > 0)
def jac_max(p):    return max(abs(p[i] * ((i == j) - p[j]))
                              for i in range(len(p)) for j in range(len(p)))

# ① Var(q·k) = d_k 인지 직접 센다
random.seed(7)
rnd = lambda d: [random.gauss(0, 1) for _ in range(d)]
for d in [8, 64, 128, 512]:
    vals = [sum(a * b for a, b in zip(rnd(d), rnd(d))) for _ in range(20000)]
    m = sum(vals) / len(vals)
    sd = math.sqrt(sum((x - m) ** 2 for x in vals) / len(vals))
    print(f"d_k={d:4d}  평균 {m:+.3f}  표준편차 {sd:7.3f}   √d_k = {math.sqrt(d):7.3f}")
# d_k=   8  평균 +0.012  표준편차   2.829   √d_k =   2.828
# d_k=  64  평균 -0.061  표준편차   7.964   √d_k =   8.000
# d_k= 128  평균 +0.029  표준편차  11.397   √d_k =  11.314
# d_k= 512  평균 +0.198  표준편차  22.669   √d_k =  22.627

# ② 같은 점수를 나눴을 때와 안 나눴을 때
raw = [1.585, 0.554, 4.287, -0.693, -8.705, -9.108]
for name, z in [("나눈 것", [v / 8 for v in raw]), ("안 나눈 것", raw)]:
    p = softmax(z)
    print(f"{name:8s} 최대 {max(p):.4f}  엔트로피 {entropy(p):.4f}  야코비안 {jac_max(p):.6f}")
# 나눈 것    최대 0.3066  엔트로피 1.6426  야코비안 0.212595
# 안 나눈 것  최대 0.9109  엔트로피 0.3710  야코비안 0.081198
print(f"균등일 때의 엔트로피 {math.log(6):.4f}")     # 1.7918

# ③ 배수를 키워 가면 야코비안이 커졌다가 죽는다
for c in [1, 2, 4, 8, 16, 32]:
    p = softmax([v * c / 8 for v in raw])
    print(f"×{c:<3d} 최대 {max(p):.6f}  엔트로피 {entropy(p):.4f}  야코비안 {jac_max(p):.3e}")
# ×1   최대 0.306596  엔트로피 1.6426  야코비안 2.126e-01
# ×2   최대 0.441677  엔트로피 1.3971  야코비안 2.466e-01
# ×4   최대 0.666979  엔트로피 0.9807  야코비안 2.221e-01
# ×8   최대 0.910855  엔트로피 0.3710  야코비안 8.120e-02
# ×16  최대 0.994908  엔트로피 0.0340  야코비안 5.066e-03
# ×32  최대 0.999979  엔트로피 0.0002  야코비안 2.057e-05

# ④ 스케일링을 빼면 학습이 멈추는 것을 400걸음으로 재현한다
def train(scaled, steps=400, lr=0.5, d=64, dv=8, n=8, seed=11, tgt=3):
    random.seed(seed)
    K = [[random.gauss(0, 1) for _ in range(d)] for _ in range(n)]
    V = [[random.gauss(0, 1) for _ in range(dv)] for _ in range(n)]
    q = [random.gauss(0, 1) for _ in range(d)]
    c = 1 / math.sqrt(d) if scaled else 1.0
    for t in range(steps):
        s = [c * sum(q[i] * K[j][i] for i in range(d)) for j in range(n)]
        p = softmax(s)
        o = [sum(p[j] * V[j][a] for j in range(n)) for a in range(dv)]
        r = [o[a] - V[tgt][a] for a in range(dv)]
        u  = [sum(r[a] * V[j][a] for a in range(dv)) for j in range(n)]   # ∂L/∂p
        pu = sum(p[j] * u[j] for j in range(n))
        ds = [p[j] * (u[j] - pu) for j in range(n)]                       # Jᵀu — 소거가 없다
        gq = [c * sum(ds[j] * K[j][i] for j in range(n)) for i in range(d)]
        if t in (0, 10, 50, 100, 399):
            gn = math.sqrt(sum(v * v for v in gq))
            print(f"  t={t:<4d} 손실 {0.5*sum(x*x for x in r):.4f}  "
                  f"p[3] {p[tgt]:.4f}  최대 p {max(p):.4f}  |∂L/∂q| {gn:.2e}")
        q = [q[i] - lr * gq[i] for i in range(d)]

print("나눈 것");   train(True)
#   t=0    손실 1.6703  p[3] 0.0868  최대 p 0.2085  |∂L/∂q| 7.75e-01
#   t=10   손실 0.7344  p[3] 0.1607  최대 p 0.4585  |∂L/∂q| 2.67e-01
#   t=50   손실 0.0250  p[3] 0.8520  최대 p 0.8520  |∂L/∂q| 4.38e-02
#   t=100  손실 0.0080  p[3] 0.9180  최대 p 0.9180  |∂L/∂q| 1.51e-02
#   t=399  손실 0.0014  p[3] 0.9661  최대 p 0.9661  |∂L/∂q| 2.84e-03
print("안 나눈 것"); train(False)
#   t=0    손실 1.5965  p[3] 0.0006  최대 p 0.6939  |∂L/∂q| 2.49e+00
#   t=10   손실 5.9063  p[3] 0.0000  최대 p 0.9508  |∂L/∂q| 6.43e+00
#   t=50   손실 1.8931  p[3] 0.0000  최대 p 1.0000  |∂L/∂q| 4.17e-09
#   t=100  손실 1.8931  p[3] 0.0000  최대 p 1.0000  |∂L/∂q| 4.17e-09
#   t=399  손실 1.8931  p[3] 0.0000  최대 p 1.0000  |∂L/∂q| 4.17e-09

④의 마지막 세 줄이 이 글의 결론입니다. 손실도, 확률도, 기울기의 크기도 소수점 넷째 자리까지 완전히 같습니다. 계산은 400걸음 내내 돌았지만 갱신은 50걸음에서 끝났습니다.

앞 글의 기울기 검사로 ④의 ds 줄이 맞는지 확인해 두었습니다 — 손실을 점수에 대해 중심차분으로 재고 $J^{\mathsf T}u$ 와 맞대면 최대 상대오차가 $1.1\times10^{-10}$ 입니다.

정리

• 성분이 평균 0·분산 1로 독립이면 $q\cdot k$ 의 평균은 0, 분산은 정확히 $d_k$ 다. 항 하나의 분산이 1이고 독립인 항이 $d_k$ 개라 그냥 더해진 것이다.
• 그래서 표준편차가 $\sqrt{d_k}$ 이고, 그 수로 나누면 차원과 무관하게 분산이 1이 된다. $\sqrt{d_k}$ 는 튜닝 상수가 아니라 이 계산의 답이다.
• softmax는 차이의 지수를 보므로 점수를 $c$ 배 하면 확률의 비가 $e^{c\Delta}$ 로 벌어진다. 같은 여섯 점수가 나눈 쪽에서는 엔트로피 1.643, 안 나눈 쪽에서는 0.371이다.
• 분포가 뾰족해지면 야코비안 $\operatorname{diag}(p) - pp^{\mathsf T}$ 의 모든 성분이 0으로 간다. 어텐션에는 교차엔트로피 같은 약분 상대가 없어서 그 0이 그대로 곱해진다.
• 야코비안이 가장 큰 자리는 최대 확률이 0.5 근처일 때의 0.25이고, 배수 16이면 0.005, 32면 $10^{-5}$ 다.
• 실험에서 스케일링을 뺀 쪽은 50걸음 만에 손실 1.893에 얼어붙었다. 최대 확률이 1.0, 기울기 크기가 $4.2\times10^{-9}$ — 틀린 자리에 몰린 채 빠져나오지 못한다.
• 초기화·LayerNorm이 $1/\sqrt{d_k}$ 를 대신하지 못한다. LayerNorm은 벡터를 정리할 뿐이고, 그렇게 정리된 두 벡터의 내적이 다시 분산 $d_k$ 를 만든다. 초기화로 한 번 맞춰도 학습 중에 지켜지지 않는다.

$\sqrt{d_k}$ 한 기호가 분산 계산 한 줄에서 나오고, 그것을 빼면 학습이 멈춘다는 것까지 왔습니다. 그런데 야코비안이 죽는 이야기를 여기서는 점수 쪽에서만 봤습니다. 다음 글에서 그 야코비안을 어텐션 전체에 통과시켜 $\partial L/\partial Q$, $\partial L/\partial K$, $\partial L/\partial V$ 를 shape까지 맞춰 끝까지 유도합니다.

---

읽어주셔서 감사합니다. 😊