스케일드 닷프로덕트 어텐션
도입
각각 쿼리/키/밸류라 불리는 세 행렬 $\mathbf{Q} \in M^{d_{k} \times n}$, $\mathbf{K} \in M^{d_{k} \times m}$, $\mathbf{V} \in M^{d_{v} \times m}$과 쿼리와 키에 대한 스코어 함수 $f : M^{d_{k} \times n} \times M^{d_{k} \times m} \to M^{m \times n}$이 주어졌다고 하자. 어텐션이란 다음과 같이 정의되는 함수를 말한다.
$$ \begin{align*} \operatorname{Attention}: M^{d_{k} \times n} \times M^{d_{k} \times m} \times M^{d_{v} \times m} &\to M^{d_{v} \times n} \\ (\mathbf{Q},\mathbf{K},\mathbf{V}) &\mapsto \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} \end{align*} $$
여기서 $\operatorname{Softmax}$는 주어진 행렬 $\mathbf{X} = \begin{bmatrix} \mathbf{x}_{1} & \cdots & \mathbf{x}_{N}\end{bmatrix}$에 대해서 각 열벡터에 소프트맥스 $\operatorname{softmax}$를 취해서 각 열의 합이 $1$이 되게 하는 함수를 말한다.
$$ \operatorname{Softmax}(\mathbf{X}) := \begin{bmatrix} \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{1})}} & \cdots & \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{N})}} \end{bmatrix} $$
$f$를 어떻게 정의하느냐에 따라 어텐션의 종류가 구분되는데, Ashish Vaswani 외 7명이 트랜스포머 구조를 최초로 제안한 논문 『Attention Is All You Need』1에서 제안한 형태의 어텐션은 아래와 같다.
정의
아래 형태의 어텐션 함수를 스케일드 닷프로덕트 어텐션scaled dot-product attention이라 한다.
$$ \operatorname{Attention} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \left( \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} \right) $$
설명
이름 그대로 쿼리 벡터 $\mathbf{q}$와 키 벡터 $\mathbf{k}$의 스코어 함수를 스케일링 팩터가 포함된 내적으로 둔 것이다.
$$ \operatorname{score}(\mathbf{q}, \mathbf{k}) = \dfrac{1}{\sqrt{d_{k}}} \mathbf{k}^{\mathsf{T}} \mathbf{q} $$
$$ f(\mathbf{Q}, \mathbf{K}) = \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} = \frac{1}{\sqrt{d_{k}}} \begin{bmatrix} \mathbf{k}_{1} \cdot \mathbf{q}_{1} & \mathbf{k}_{1} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{1} \cdot \mathbf{q}_{n} \\ \mathbf{k}_{2} \cdot \mathbf{q}_{1} & \mathbf{k}_{2} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{2} \cdot \mathbf{q}_{n} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{k}_{m} \cdot \mathbf{q}_{1} & \mathbf{k}_{m} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{m} \cdot \mathbf{q}_{n} \end{bmatrix} $$
왜 $\sqrt{d_{k}}$로 나누는가
내적을 그대로 스코어로 쓰지 않고 굳이 $\sqrt{d_{k}}$로 나누는 이유는, 차원 $d_{k}$가 커질수록 내적의 크기가 $\sqrt{d_{k}}$ 스케일로 커져서 소프트맥스를 포화시키기 때문이다1. 쿼리 $\mathbf{q}$와 키 $\mathbf{k}$의 각 성분 $q_{i}$, $k_{i}$가 모두 평균이 $0$, 분산이 $1$이고 서로 독립인 확률변수라고 가정하자. 그러면 기댓값의 선형성과 독립성에 의해 내적 $\mathbf{k}^{\mathsf{T}} \mathbf{q} = \sum_{i=1}^{d_{k}} k_{i} q_{i}$의 평균은 다음과 같이 $0$이다.
$$ \mathbb{E} \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = \mathbb{E} \left[ \sum_{i=1}^{d_{k}} k_{i} q_{i} \right] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} ] \mathbb{E} [ q_{i} ] = 0 $$
성분들이 모두 서로 독립이므로 곱 $k_{i} q_{i}$들도 서로 독립이고, 독립인 확률변수들의 합의 분산은 각 분산의 합이다. 또한 평균이 $0$이라 $\mathbb{E} [ k_{i}^{2} ] = \Var [ k_{i} ] = 1$이므로 내적의 분산은 다음과 같이 정확히 $d_{k}$가 된다.
$$ \begin{align*} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] &= \sum_{i=1}^{d_{k}} \Var [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \left( \mathbb{E} [ k_{i}^{2} q_{i}^{2} ] - \left( \mathbb{E} [ k_{i} q_{i} ] \right)^{2} \right) \\ &= \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i}^{2} ] \mathbb{E} [ q_{i}^{2} ] = \sum_{i=1}^{d_{k}} \Var [ k_{i} ] \Var [ q_{i} ] = d_{k} \end{align*} $$
다시 말해 스코어 $\mathbf{k}^{\mathsf{T}} \mathbf{q}$의 표준편차가 $\sqrt{d_{k}}$라는 뜻이고, 차원이 $d_{k} = 512$라면 표준편차가 $\sqrt{512} \approx 22.6$이니 절댓값이 수십에 달하는 스코어가 나오게 된다. 절댓값이 큰 값들이 소프트맥스에 들어가면 지수함수가 최댓값 성분을 압도적으로 키워서 출력이 원-핫 벡터에 가깝게 포화되는데, 포화 영역에서는 입력이 웬만큼 변해도 출력이 거의 변하지 않으므로 그래디언트가 극히 작아져 기울기 소실로 학습이 저해된다. 반면 스코어를 $\sqrt{d_{k}}$로 나누면 분산이 다시 $1$로 정규화되므로, $d_{k}$가 아무리 커져도 스코어의 크기가 일정하게 유지된다.
$$ \Var \left[ \frac{\mathbf{k}^{\mathsf{T}} \mathbf{q}}{\sqrt{d_{k}}} \right] = \frac{1}{d_{k}} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = 1 $$
아래 그림이 이 차이를 보여준다. 표준정규분포에서 성분을 샘플링해 $d_{k} = 512$짜리 쿼리 $\mathbf{q}$ 한 개와 키 $\mathbf{k}_{1}, \dots, \mathbf{k}_{10}$ 열 개를 만들고, 스코어 $\mathbf{k}_{j}^{\mathsf{T}} \mathbf{q}$ 열 개를 소프트맥스에 통과시킨 결과다. 왼쪽처럼 스케일링 없이 그대로 통과시키면 확률이 사실상 키 하나에 쏠려 원-핫으로 포화된다. 반면 같은 스코어를 $\sqrt{d_{k}} = \sqrt{512}$로 나눈 뒤 통과시키면 오른쪽처럼 여러 키에 확률이 고루 분산되어, 어텐션이 특정 키 하나만 보는 대신 여러 키를 두루 참조할 수 있게 된다.
