Dot-product Attention과 Scaling 필요성

Intro

Attention Is All You Need 논문에선 $Q$와 $K$ 벡터의 유사도를 0 ~ 1 사이의 확률 값으로 정규화하기 위해서 Dot-product Attention를 사용한다. 이때, Dot-product Attention을 그대로 사용하지 않고 $K$벡터의 차원 값인 $d_k$를 활용하여 scaling된 형태로 사용한다. 

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

 

scaling의 근거는 논문에서 언급되며 다음과 같다.

1. $d_k$가 커지면 $QK^{T}$ 값이 커진다.
2. $QK^{T}$ 값이 커지면 softmax 함수의 gradient vanishing이 발생한다.
3. scaling을 통해서 gradient vanishing 줄일 수 있다.
4. scaling은 $\sqrt d_k$의 값으로 진행한다.

이를 확인해보면 다음과 같다.

1. $d_k$가 커지면 $QK^{T}$ 값이 커진다. (+ 4. scaling은 $\sqrt d_k$의 값으로 진행한다.)

$Q, K$의 각 성분을 $q, k$라 하면 다음과 같다.

$$QK^{T} = \sum_{i = 1}^{d_k}q_ik_i$$

이때, $q_i, k_i$가 서로 독립이라고 가정하면 다음 두 식이 성립한다.

참고

$$
\mathrm{E}(q_i k_i)=\mathrm{E}(q_i) \cdot \mathrm{E}(k_i)
$$

$$
\operatorname{Var}(q_i k_i)=\left(\sigma_{q_i}^{2}+\mu_{q_i}^{2}\right)\left(\sigma_{k_i}^{2}+\mu_{k_i}^{2}\right)-\mu_{q_i}^{2} \mu_{k_i}^{2}
$$

계산을 간단하게 하기 위해 $q_i, k_i$가 gaussian distribution을 따른다고 가정하면 $q_ik_i$의 기댓값과 분산은 다음과 같다.

$$
\mathrm{E}(q_i k_i)=0
$$

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

이를 통해 다음을 얻을 수 있다.

$$E(QK^{T})=E\left(\sum_{i=1}^{d_{k}} q_{i} k_{i}\right)=\sum_{i=1}^{d_{k}} E\left(q_{i} k_{i}\right)=0
$$

$$
\operatorname{Var}(QK^{T})=\operatorname{Var}\left(\sum_{i=1}^{d_{i}} q_{i} k_{i}\right)=\sum_{i=1}^{d_{i}} \operatorname{Var}\left(q_{i} k_{i}\right)=d_{k}
$$

수식을 통해 $d_k$가 커짐에 따라 분산이 커짐을 확인했다. 여기서 $\sqrt{d_k}$를 QK^{T}에 곱해주면 $d_k$의 분산을 다시 1로 줄일 수 있게 되며 이것이 scaling 값으로 $\sqrt{d_k}$를 사용한 근거가 된다. 4에서 언급한 scaling은 $\sqrt d_k$ 값으로 한다. 는 여기서 확인할 수 있다.

 

분산이 크다는 것은 결국 $QK^{T}$가 큰 값을 가질 확률이 커졌다는 것이다. 실제로 표준 편차가 각각 1일 때, $\sqrt {512}$일때의 분포는 다음과 같다.

mu, sigma = 0, 1

mu, sigma = 0, sqrt {512}

2. $QK^{T}$ 값이 커지면 softmax 함수의 gradient vanishing이 발생한다. (+ 3. scaling을 통해서 gradient vanishing 줄일 수 있다.)

softmax 함수는 logit 값을 0 ~ 1 사이 확률 값으로 정규화 시킨다.

 

$$\sigma(z_i) = \frac{e^{z_{i}}}{\sum_{j=1}^K e^{z_{j}}} \ \ \ for\ i=1,2,\dots,K$$

 

위 식에서 볼 수 있듯이 각 $z_i$의 크기에 따라서 다른 확률 값으로 계산되기 때문에

softmax는 아래와 같이 scale variant한 특성을 가지고 있다. 

https://towardsdatascience.com/transformer-networks-a-mathematical-explanation-why-scaling-the-dot-products-leads-to-more-stable-414f87391500

softmax는 $\mathbb{R^{n}}\rightarrow\mathbb{R^{n}}$인 함수 이기 때문에 gradient vanishing을 확인하기 위해선

jacobian 행렬을 확인해야 한다.  $n = 4$라고 가정했을 때 softmax의 jacobian은 다음과 같다.

위 식에서 $s = (s1, s2, s3, s4)$라고 할 때 $s$가 다음과 같은 형태를 가질 경우 jacobian 행렬이 0이 돼 gradient vanishing이 발생하게 된다. 

 

이때, 1 에서 확인한 $d_k$가 커짐에 따라 $QK^{T}$의 분산이 커진다는 성질이 만나면 위와 같은 softmax의 결과로 위와 같은 $s$를 얻게 돼 gradient vanishing이 발생하게 된다. 정리하면 다음과 같다.

• $d_k$가 커짐에 따라 $QK^{T}$의 분산이 커진다.
• 분산이 커짐에 따라 $QK^{T} = (q_1k_1, q_2k_2, q_3k_3, q_4k_4)$의 성분 사이 편차도 커진다.
• 성분 사이 큰 편차는 scale variant한 softmax 함수와 만나 $s$가 위와 같이 gradient vanishing을 일으키는 형태가 된다.

 

참고

Attention is All you need

Distribution of the product of two random variables - Wikipedia

Transformer Networks: A mathematical explanation why scaling the dot products leads to more stable…