transformers의 Self-attention과 multi-head

# 트랜스포머의 어텐션 메커니즘: 셀프 어텐션, 크로스 어텐션, 멀티헤드

어텐션 메커니즘은 AI 모델이 입력 시퀀스의 요소 중요도를 동적으로 가중치로 계산할 수 있게 합니다. 각 토큰에 대해 쿼리(Query), 키(Key), 밸류(Value) 행렬을 사용해 가중치를 계산하며, 문맥화된 표현을 생성합니다. 셀프 어텐션은 단일 시퀀스 내부 연결을 분석하고, 크로스 어텐션은 인코더와 디코더를 연결하며, 멀티헤드 어텐션은 병렬 처리로 데이터의 다양한 측면을 포착합니다.

셀프 어텐션 이론

입력 시퀀스는 행렬 $X \in \mathbb{R}^{n \times d}$로 표현되며, 여기서 $n$은 토큰 수, $d$는 임베딩 차원입니다. $X$로부터 $Q = XW^Q$, $K = XW^K$, $V = XW^V$ 행렬을 투영하며, 가중치 $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$를 사용합니다.

어텐션 스코어 행렬 $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$는 $i$번째 토큰에 대한 $j$번째 토큰의 관련성을 측정합니다. $\sqrt{d_k}$로 스케일링해 소프트맥스에서 그래디언트 소실을 방지하고 안정화합니다.

행별로 정규화된 행렬 $N = \mathrm{softmax}(A)$는 확률 분포를 제공합니다. 출력: $\mathrm{Attention}(Q, K, V) = N V$.

셀프 어텐션 예제

"Karina가 가게에 간다" 문장을 ["karina", "is", "to", "the", "store"]로 토큰화($n=5$)했다고 가정합니다. 임베딩(간소화):

| 토큰 | 임베딩 |

|--------|-----------|

| karina | [1, 0] |

| is | [0, 1] |

| to | [1, 1] |

| the | [1, 1] |

| store | [0.5, 1] |

행렬 $X_{5\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$. 단순화를 위해 $W^Q = W^K = W^V = I$, 그래서 $Q=K=V=X$.

$QK^T = \begin{bmatrix} 1 & 0 & 1 & 1 & 0.5 \\ 0 & 1 & 1 & 1 & 1 \\ 1 & 1 & 2 & 2 & 1.5 \\ 1 & 1 & 2 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.5 & 1.25 \end{bmatrix}$, $\sqrt{2} \approx 1.41$로 나눔:

$A \approx \begin{bmatrix} 0.71 & 0 & 0.71 & 0.71 & 0.35 \\ 0 & 0.71 & 0.71 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.35 & 0.71 & 1.06 & 1.06 & 0.89 \end{bmatrix}$.

행별 소프트맥스 후: $N \approx \begin{bmatrix} 0.25 & 0.15 & 0.25 & 0.25 & 0.10 \\ 0.17 & 0.23 & 0.23 & 0.23 & 0.23 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.14 & 0.23 & 0.28 & 0.28 & 0.07 \end{bmatrix}$.

출력 $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$. 새로운 임베딩은 이제 문맥을 반영: "is" 벡터는 추상 시제 대신 인간 이동을 나타냅니다.

import torch  
import torch.nn as nn  

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

크로스 어텐션: 시퀀스 연결

크로스 어텐션에서 쿼리는 디코더에서($Q = X_{dec}W^Q$), 키와 밸류는 인코더에서($K = X_{enc}W^K$, $V = X_{enc}W^V$) 가져옵니다. 계산은 동일: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$.

주요 차이:

| 구성 요소 | 셀프 어텐션 | 크로스 어텐션 |

|-----------|-------------|---------------|

| Q | X | X_dec |

| K | X | X_enc |

| V | X | X_enc |

이로 디코더가 인코더 문맥의 관련 부분에 집중할 수 있습니다.

import torch  
import torch.nn as nn  

class CrossAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x_dec, x_enc):
        Q = self.q_linear(x_dec)
        K = self.k_linear(x_enc)
        V = self.v_linear(x_enc)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

멀티헤드 어텐션: 병렬 관점

멀티헤드 어텐션은 $h$개의 독립 어텐션 계산을 병렬 실행하며, $Q^{(i)} = X W_Q^{(i)}$, $K^{(i)} = X W_K^{(i)}$, $V^{(i)} = X W_V^{(i)}$로 투영합니다. 각 헤드: $\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$.

결과를 연결하고 투영: $\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$. 헤드 크기 $d_k = d / h$, $d$는 임베딩 차원입니다.

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):  
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.q = nn.Linear(embed_dim, embed_dim)
        self.k = nn.Linear(embed_dim, embed_dim)
        self.v = nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, T, D = x.shape
        Q = self.q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, D)
        return self.out(out)

주요 요점

• 셀프 어텐션은 토큰 문맥화: 각 출력 벡터가 전체 시퀀스에 의존합니다.
• $\sqrt{d_k}$ 스케일링 필수: 소프트맥스와 그래디언트 안정화.
• seq2seq용 크로스 어텐션: 트랜스포머에서 인코더-디코더 연결.
• 멀티헤드 다양한 측면 포착: 문법, 의미, 장거리 의존성.
• PyTorch 구현은 Linear와 bmm 사용: 배치 효율성.

— Editorial Team