"Scaled" dot product

의문

transformer에서 사용되는 Scaled dot-product attention에서 score는

$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$

여기서 왜 $\sqrt{d_k}$로 나눌까?

이유

$S = QK^T$라 했을 때, S의 원소는 $S_{ij} = \sum^{d_k}_{n=1}Q_{in}K_{jn}$이라 할 수 있다.(내적이므로)

Q와 K의 모든 값이 독립이면서 평균이 0, 분산이 1이라 할 때
$\mathbb{E}[X_{n}] = \mathbb{E}[Q_{in} K_{jn}] = \mathbb{E}[Q_{in}]\mathbb{E}[K_{jn}] = 0$

$Var(Q_{1, i} K_{i, 1}) = (Var(Q_{1, i}) + \mathbb{E}[Q_{1, i}]^2)(Var(K_{i, 1}) + \mathbb{E}[K_{i, 1}]^2) - \mathbb{E}[Q_{1, i}]^2\mathbb{E}[K_{i, 1}]^2 = 1$

$X_n$들이 독립이면 분산은 더해져
$Var(S_{ij}) = \sum_{n=1}^{d_k}Var(X_n) = \sum_{n=1}^{d_k}1 = d_k$

따라서 $QK^T$의 각 원소는 평균이 0, 분산이 $d_k$의 분포를 갖게 된다.

$S_{ij}$에 $\sqrt{d_k}$를 나눠주면서 차원 $d_k$가 커져도 분산이 1 근처로 유지되게 만든다.

softmax는 $p_i = \frac{e^{s_i}}{\sum_je^{s_j}}$이기에 score 벡터인 s의 값이 커지면 차이가 커지게 된다. 이렇게 되면 softmax가 포화되어 one-hot에 가깝게 형성된다.

이 경우, softmax 미분 시에도 확률들의 곱으로 표현되기 때문에 확률값이 작아지면 gradient vanishing 문제가 발생할 수 있다. 따라서 scailing을 통해 score 값들이 과도하게 커지는 것을 막는 것이다.