3.1 Self-Attention Mathematics

2017년 “Attention Is All You Need” (Vaswani et al.) 논문에서 소개된 트랜스포머 아키텍처는 순환 구조(RNN 등)를 Self-Attention이라는 매커니즘으로 대체하여 AI 분야에 혁명을 일으켰습니다. 이를 통해 모델은 시퀀스의 모든 토큰을 동시에 처리하고, 거리에 관계없이 토큰 간의 의존성을 모델링할 수 있게 되었습니다.

비하인드 스토리: 논문의 제목인 “Attention Is All You Need(필요한 것은 어텐션뿐이다)“는 당시로서는 매우 도발적인 주장이었습니다. 그동안 시퀀스 데이터 처리의 왕좌를 지키고 있던 RNN이나 LSTM을 완전히 배제하고, 오직 어텐션 메커니즘만으로 최고 성능을 낼 수 있다는 뜻이었기 때문입니다. 구글의 저자들은 번역 작업에서 RNN의 느린 속도에 답답함을 느껴 이 모델을 개발했고, 이 도발적인 제목은 결국 AI 역사를 바꾸는 예언이 되었습니다.

이 핵심 매커니즘을 친숙한 비유를 통해 이해해 봅시다.

비유: 파일 서랍장 시스템

당신이 매우 효율적인 파일 시스템을 갖춘 도서관에서 연구를 하고 있다고 가정해 봅시다.

• Query ($Q$): 이것은 당신이 찾고 있는 것입니다. 머릿속에 있는 주제입니다 (예: “새는 어떻게 나는가?”).
• Key ($K$): 이것들은 파일 폴더에 붙어 있는 라벨이나 태그입니다. 각 폴더에는 안에 무엇이 들어 있는지에 대한 요약이 있습니다.
• Value ($V$): 이것은 폴더 안에 들어 있는 실제 내용입니다.

필요한 정보를 찾기 위해:

1. 당신의 Query를 모든 Key와 비교하여 어떤 폴더가 관련이 있는지 확인합니다.
2. 각 폴더에 대한 **관련성 점수 (Attention weight)**를 계산합니다.
3. 폴더에서 Value를 추출하되, 관련성 점수가 높은 폴더의 값에 더 많은 가중치를 부여합니다.

Self-Attention에서는 문장 속의 모든 단어가 Query, Key, Value로 작용하여 다른 모든 단어와 상호작용합니다.

임베딩에서 Q, K, V로: 선형 투영 (Linear Projections)

실제로는 입력 단어 임베딩을 Queries, Keys, Values로 직접 사용하지 않습니다. 대신, 학습 가능한 가중치 행렬을 사용하여 이들을 서로 다른 공간으로 투영(project)합니다.

$T$를 시퀀스 길이, $d_{model}$을 임베딩 차원이라고 할 때, 행렬 $X \in \mathbb{R}^{T \times d_{model}}$로 표현되는 입력 시퀀스가 주어지면 다음과 같이 계산합니다:

$Q = XW_Q$ $K = XW_K$ $V = XW_V$

여기서 학습 가능한 가중치 행렬은 다음과 같습니다:

• $W_Q \in \mathbb{R}^{d_{model} \times d_k}$
• $W_K \in \mathbb{R}^{d_{model} \times d_k}$
• $W_V \in \mathbb{R}^{d_{model} \times d_v}$

왜 이렇게 할까요? 만약 세 가지 역할 모두에 $X$를 직접 사용한다면, self-attention 연산은 정적인 임베딩에만 전적으로 의존하게 될 것입니다. 선형 투영을 사용함으로써 모델은 동일한 단어에서 역할에 따라 다른 측면을 추출하는 방법을 배울 수 있습니다. 예를 들어, 어떤 단어가 다른 단어를 찾는 Query 역할을 할 때의 표현과, 검색 대상이 되는 Key 역할을 할 때의 표현을 다르게 가질 수 있습니다. 이는 모델의 표현력(expressiveness)을 크게 향상시킵니다.

Scaled Dot-Product Attention의 수학적 원리

Self-Attention의 핵심 연산은 Scaled Dot-Product Attention입니다. 행렬 $Q, K, V$로 묶인 Query, Key, Value 세트가 주어지면, 연산은 다음과 같이 정의됩니다:

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

여기서:

• $Q \in \mathbb{R}^{T \times d_k}$ (Query 행렬)
• $K \in \mathbb{R}^{T \times d_k}$ (Key 행렬)
• $V \in \mathbb{R}^{T \times d_v}$ (Value 행렬)
• $T$는 시퀀스 길이입니다.
• $d_k$는 Key(및 Query)의 차원입니다.
• $\sqrt{d_k}$는 스케일링 요소입니다.

단계별 분석

1. Dot Product ($QK^T$): 각 Query와 모든 Key 사이의 원시 유사도를 측정합니다.
2. Scaling ($\frac{1}{\sqrt{d_k}}$): 내적(dot product)의 크기가 너무 커지는 것을 방지합니다. 크기가 커지면 softmax 함수가 매우 작은 기울기를 갖는 영역으로 밀려나기 때문입니다.
3. Softmax: 스케일링된 점수를 합이 1이 되는 확률(attention weights)로 변환합니다.
4. Weighted Sum ($V$): Attention weight에 Value를 곱하여 최종 출력을 얻습니다.

PyTorch 구현

이 연산을 PyTorch로 처음부터 구현해 보겠습니다. 이것이 트랜스포머의 핵심 빌딩 블록입니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Scaled Dot-Product Attention 계산.
    """
    d_k = query.size(-1)
    
    # Step 1 & 2: 내적(Dot product) 및 스케일링
    # query shape: (batch, heads, seq_len, d_k)
    # key.transpose(-2, -1) shape: (batch, heads, d_k, seq_len)
    # scores shape: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5)
    
    # 선택 사항: 마스크 적용 (예: 인과적 autoregressive 디코딩용)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Step 3: Softmax를 통한 attention weights 획득
    attention_weights = F.softmax(scores, dim=-1)
    
    # Step 4: Value의 가중 합
    # output shape: (batch, heads, seq_len, d_v)
    output = torch.matmul(attention_weights, value)
    
    return output, attention_weights

# 예제 사용법
batch_size = 1
seq_len = 4
d_k = 8
d_v = 8

# 랜덤 텐서
Q = torch.randn(batch_size, 1, seq_len, d_k)
K = torch.randn(batch_size, 1, seq_len, d_k)
V = torch.randn(batch_size, 1, seq_len, d_v)

output, weights = scaled_dot_product_attention(Q, K, V)
print("Output Shape:", output.shape)
print("Attention Weights Shape:", weights.shape)

예제: 어텐션 가중치 행렬

단어들이 서로에게 어떻게 주의(attention)를 기울이는지 시각화해 보세요. 문장 속의 단어를 클릭하면 해당 단어가 다른 모든 단어에 부여하는 Attention weight를 볼 수 있습니다. (시뮬레이션된 데이터입니다).

Quizzes

References

1. Vaswani, A., et al. (2017). Attention is all you need. In Advances in neural information processing systems (pp. 5998-6008). arXiv:1706.03762.