📂머신러닝

멀티헤드 어텐션

도입¹

각각 쿼리/키/밸류라 불리는 세 행렬 $\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} $$

이렇게 계산된 어텐션의 함숫값은 쿼리와 키의 적합도를 가중치로 삼아 밸류를 가중평균한 행렬이다(자세한 설명은 어텐션 문서를 참고하라).

$$ \operatorname{Attention} (\mathbf{Q},\mathbf{K},\mathbf{V}) \in M^{d_{v} \times n} $$

여기서 잠시 이미지 데이터와 합성곱 신경망(CNN)을 생각해보자. 크기가 $(1, N, N)$인 입력 이미지는 CNN을 거치며 크기가 $(C, M, M)$인 배열로 가공된다. 이때 $C$를 채널의 차원이라 부르며, 각 채널에는 입력 데이터에 대한 추상 정보가 담겨 있다. 이 추상 정보가 색상이라면 컬러 사진은 RGB 세 채널을 가져 $(3, N, N)$과 같이 표현되고, 흑백 사진은 채널이 하나뿐이므로 $(1, N, N)$과 같이 표현된다. 물론 CNN을 통과하면서 얻어지는 $(C, M, M)$ 배열의 채널에 담긴 정보는 이처럼 인간이 직관적으로 알 수 있는 것이 아니지만, 어쨌든 입력 데이터의 정보가 $C$가지 종류로 분해되었다고 이해할 수 있다. 마찬가지로 어텐션을 여러 개 써서 $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$의 정보를 여러 채널로 표현할 수 있을 것이다. 어텐션에서는 이러한 채널을 헤드^head라 부른다.

그런데 $(1)$을 보면 $\operatorname{Attention}$ 함수 자체는 학습가능한 파라미터를 포함하지 않는다. 그저 소프트맥스와 행렬곱만으로 이루어져 있다(물론 $f$를 신경망으로 구성하면 학습가능한 파라미터가 포함되지만, 여기서는 그렇지 않다고 하자). 그래서 쿼리, 키, 밸류를 $\operatorname{Attention}$에 보내기 전에 학습되는 선형변환행렬로 각각을 미리 여러 공간으로 사상한다.

쿼리/키/밸류 벡터의 차원이 $d_{\text{model}}$이라고 하자. 다시 말해 $\mathbf{Q} \in M^{d_{\text{model}} \times n}$, $\mathbf{K} \in M^{d_{\text{model}} \times m}$, $\mathbf{V} \in M^{d_{\text{model}} \times m}$이다. 각각의 $h = 1, \cdots, H$에 대해서 쿼리/키를 $d_{k}$차원으로, 밸류를 $d_{v}$차원으로 압축하는 선형변환이 다음과 같이 주어졌다고 하자.

$$ \mathbf{W}_{h}^{Q} \in M^{d_{k} \times d_{\text{model}}}, \qquad \mathbf{W}_{h}^{K} \in M^{d_{k} \times d_{\text{model}}}, \qquad \mathbf{W}_{h}^{V} \in M^{d_{v} \times d_{\text{model}}} $$

그러면 $(\mathbf{Q}, \mathbf{K}, \mathbf{V})$는 헤드마다 서로 다른 공간으로 사상되어, $H$가지 종류의 정보로, 즉 $H$개의 헤드로 분해된다.

$$ \operatorname{head}_{h} := \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \in M^{d_{v} \times n} \tag{2} $$

$(1)$에서는 $\mathbf{Q} \in M^{d_{k} \times n}$인 것에 반해, $(2)$에서는 $\mathbf{Q} \in M^{d_{\text{model}} \times n}$이고 $\mathbf{W}_{h}^{Q}\mathbf{Q} \in M^{d_{k} \times n}$임에 주의하라.

$$ \begin{align*} \text{in } (1) &: \quad \mathbf{Q} \in M^{d_{k} \times n} \\ \text{in } (2) &: \quad \mathbf{Q} \in M^{d_{\text{model}} \times n}, \quad \mathbf{W}_{h}^{Q}\mathbf{Q} \in M^{d_{k} \times n} \end{align*} $$

이를 모두 모으면 아래와 같은 꼴로 표현할 수 있다.

$$ \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} = \begin{bmatrix} \operatorname{Attention}(\mathbf{W}_{1}^{Q}\mathbf{Q}, \mathbf{W}_{1}^{K}\mathbf{K}, \mathbf{W}_{1}^{V}\mathbf{V}) \\ \operatorname{Attention}(\mathbf{W}_{2}^{Q}\mathbf{Q}, \mathbf{W}_{2}^{K}\mathbf{K}, \mathbf{W}_{2}^{V}\mathbf{V}) \\ \vdots \\ \operatorname{Attention}(\mathbf{W}_{H}^{Q}\mathbf{Q}, \mathbf{W}_{H}^{K}\mathbf{K}, \mathbf{W}_{H}^{V}\mathbf{V}) \end{bmatrix} \in M^{H d_{v} \times n} $$

그리고 다음과 같은 블록행렬 $\mathbf{W}^{O} \in M^{d_{\text{model}} \times H d_{v}}$를 곱해주면 헤드들의 선형결합을 얻을 수 있다.

$$ \mathbf{W}^{O} = \begin{bmatrix} \mathbf{W}_{1}^{O} & \mathbf{W}_{2}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix}, \quad \mathbf{W}_{h}^{O} \in M^{d_{\text{model}} \times d_{v}} $$

$$ \begin{bmatrix} \mathbf{W}_{1}^{O} & \mathbf{W}_{2}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix} \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} = \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{head}_{h} \in M^{d_{\text{model}} \times n} $$

이 행렬이 멀티헤드 어텐션의 출력이다. 이제 함수 $\operatorname{MultiHead}$를 아래와 같이 정의하자.

정의

차원이 $d_{\text{model}}$인 벡터들의 행렬 $\mathbf{Q} \in M^{d_{\text{model}} \times n}$, $\mathbf{K} \in M^{d_{\text{model}} \times m}$, $\mathbf{V} \in M^{d_{\text{model}} \times m}$를 각각 쿼리, 키, 밸류 행렬이라 하자. 각각의 $h = 1, \cdots, H$에 대해서, 이 행렬들의 열벡터를 각각 $d_{k}$, $d_{k}$, $d_{v}$차원으로 사상하는 선형변환들을 $\mathbf{W}_{h}^{Q} \in M^{d_{k} \times d_{\text{model}}}$, $\mathbf{W}_{h}^{K} \in M^{d_{k} \times d_{\text{model}}}$, $\mathbf{W}_{h}^{V} \in M^{d_{v} \times d_{\text{model}}}$라 하자. $\operatorname{head}_{h}$를 아래와 같이 정의하자.

$$ \operatorname{head}_{h} := \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \in M^{d_{v} \times n} $$

함수 멀티헤드 어텐션^{multi-head attention} $\operatorname{MultiHead} : M^{d_{\text{model}} \times n} \times M^{d_{\text{model}} \times m} \times M^{d_{\text{model}} \times m} \to M^{d_{\text{model}} \times n}$를 아래와 같이 정의한다.

$$ \begin{align*} \operatorname{MultiHead} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) &:= \mathbf{W}^{O} \begin{bmatrix} \operatorname{head}_{1} \\ \operatorname{head}_{2} \\ \vdots \\ \operatorname{head}_{H} \end{bmatrix} \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{head}_{h} \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \operatorname{Attention} (\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}, \mathbf{W}_{h}^{V}\mathbf{V}) \\ &= \sum_{h=1}^{H} \mathbf{W}_{h}^{O} \left( \mathbf{W}_{h}^{V}\mathbf{V} \right) \operatorname{Softmax} \left( f(\mathbf{W}_{h}^{Q}\mathbf{Q}, \mathbf{W}_{h}^{K}\mathbf{K}) \right) \end{align*} $$

여기서 $\mathbf{W}^{O} = \begin{bmatrix} \mathbf{W}_{1}^{O} & \cdots & \mathbf{W}_{H}^{O} \end{bmatrix} \in M^{d_{\text{model}} \times H d_{v}}$는 학습되는 가중치 행렬이고, 각 블록의 크기는 $\mathbf{W}_{h}^{O} \in M^{d_{\text{model}} \times d_{v}}$이다.

수학 논문에서^{2 3}

수학 논문에서 주로 쓰이는 정의는 위보다 훨씬 단순하다. 셀프 어텐션의 경우를 생각하면, 하나의 데이터 $\mathbf{X}$로부터 다음과 같이 쿼리, 키, 밸류를 얻는다.

$$ \mathbf{Q} = \mathbf{W}^{QX}\mathbf{X},\qquad \mathbf{K} = \mathbf{W}^{KX}\mathbf{X},\qquad \mathbf{V} = \mathbf{W}^{VX}\mathbf{X} $$

그러면 어텐션의 입력이 $\mathbf{W}^{X}\mathbf{W}^{QX}\mathbf{X}$와 같이 지저분한 형태로 표현되는데, 선형변환의 합성도 선형변환이므로 $\mathbf{W}^{X}\mathbf{X}$으로 표현하는 것과 본질적으로 같다. 또한 수학적으로는 어텐션 함수가 멀티헤드어텐션에서 $H=1$인 특수한 경우에 불과하므로, 잔차층을 포함하고 스코어함수를 내적으로 둔 다음의 함수를 어텐션이라 정의한다.

$$ \operatorname{Attention}(\mathbf{X}) = \mathbf{X} + \sum_{h=1}^{H} \left( \mathbf{W}_{h}^{V}\mathbf{X} \right) \operatorname{Softmax} \left[ (\mathbf{W}_{h}^{K}\mathbf{X})^{\mathsf{T}} (\mathbf{W}_{h}^{Q}\mathbf{X}) \right] $$

설명

트랜스포머 논문에서는 행벡터 기준으로 기술하고 있어서 행렬들의 차원이 반대이고, 노테이션도 아래와 같이 되어있다.

$$ \operatorname{head}_{h} = \operatorname{Attention} (\mathbf{Q}\mathbf{W}_{h}^{Q}, \mathbf{K}\mathbf{W}_{h}^{K}, \mathbf{V}\mathbf{W}_{h}^{V}) \in M^{n \times d_{v}} $$

$$ \operatorname{MultiHead} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \begin{bmatrix} \operatorname{head}_{1} & \operatorname{head}_{2} & \cdots & \operatorname{head}_{H} \end{bmatrix} \begin{bmatrix} \mathbf{W}_{1}^{O} \\ \mathbf{W}_{2}^{O} \\ \vdots \\ \mathbf{W}_{H}^{O} \end{bmatrix} $$