마스크드 어텐션
정의
각각 쿼리, 키, 밸류라 불리는 세 행렬 $\mathbf{Q} \in M^{d_{k} \times n}$, $\mathbf{K} \in M^{d_{k} \times m}$, $\mathbf{V} \in M^{d_{v} \times m}$가 주어졌다고 하자. 성분이 $0$ 또는 $-\infty$인 행렬 $\mathbf{M} \in M^{m \times n}$을 마스크mask라 하고, 다음과 같이 정의되는 함수를 마스크드 어텐션masked attention이라 한다.
$$ \operatorname{MaskedAttention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) := \mathbf{V} \operatorname{Softmax} \left( \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} + \mathbf{M} \right) $$
여기서 $\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} $$
설명
$[\mathbf{M}]_{ij} = 0$이라는 것은 $\operatorname{score}(\mathbf{k}_{i}, \mathbf{q}_{j})$를 그대로 둔다는 것이고, $[\mathbf{M}]_{ij} = -\infty$라는 것은 $\operatorname{score}(\mathbf{k}_{i}, \mathbf{q}_{j})$의 값을 보정하여 $-\infty$로 강제한다는 뜻이다. $e^{-\infty} = 0$이므로 소프트맥스의 정의에 따라 마스킹된 성분은 소프트맥스를 거치면서 어텐션 확률이 정확히 $0$이 된다. 즉 마스킹은 쿼리가 특정 키/밸류를 보지 못하게 강제로 가리는 장치다.
트랜스포머 논문에서는 디코더의 첫번째 서브층에서 쿼리가 자기 자신과 그 이전의 키만 볼 수 있도록 하는 인과 마스크causal mask를 사용한다. 즉 $j$번째 쿼리가 $i$번째 키를 참조할 수 있는지 여부가 $i \gt j$인지에 따라 결정된다. 이를 함수로써 표현하면 아래와 같다.
$$ \begin{bmatrix} \operatorname{Mask}(\mathbf{K}^{\mathsf{T}} \mathbf{Q}/ \sqrt{d_{k}}) \end{bmatrix}_{ij} = \begin{cases} -\infty & i \gt j \\ \mathbf{k}_{i}^{\mathsf{T}} \mathbf{q}_{j} / \sqrt{d_{k}} & i \le j \end{cases} $$
$$ \operatorname{MaskedAttention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \circ \operatorname{Mask} \left( \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} \right) $$
한편 실제 구현에서는 $-\infty$ 대신 $-10^{9}$ 정도의 아주 큰 음수를 더하는 것이 보통이다.
