셀프 어텐션
Self-Attention
시퀀스 내 각 요소가 다른 모든 요소와의 관계를 계산하는 메커니즘. Transformer의 핵심.
Self-Attention
시퀀스 내 각 요소가 다른 모든 요소와의 관계를 계산하는 메커니즘. Transformer의 핵심.
셀프 어텐션(Self-Attention)은 시퀀스 내의 각 요소가 같은 시퀀스의 다른 모든 요소와의 관계를 동시에 계산하는 메커니즘입니다. 2017년 "Attention Is All You Need" 논문에서 소개된 Transformer 아키텍처의 핵심 구성 요소로, RNN의 순차적 처리 한계를 극복하고 병렬 처리를 가능하게 했습니다.
셀프 어텐션은 Query, Key, Value 세 가지 벡터를 사용합니다. 입력 시퀀스의 각 토큰은 이 세 벡터로 변환됩니다. Query와 Key의 내적으로 attention score를 계산하고, softmax를 적용한 후 Value와 가중합하여 출력을 만듭니다. 이 과정을 통해 "어떤 토큰이 어떤 토큰에 얼마나 주목해야 하는지"를 학습합니다.
수식으로 표현하면 Attention(Q,K,V) = softmax(QK^T / sqrt(d_k)) * V 입니다. sqrt(d_k)로 나누는 scaled dot-product attention은 차원이 커질수록 내적 값이 커지는 문제를 방지합니다. 실제로는 여러 개의 어텐션 헤드를 병렬로 사용하는 Multi-Head Attention을 적용하여 다양한 관점에서 관계를 학습합니다.
셀프 어텐션의 장점은 장거리 의존성(long-range dependency)을 효과적으로 모델링한다는 것입니다. RNN은 먼 거리의 토큰 정보가 희석되지만, 셀프 어텐션은 모든 토큰 쌍을 직접 연결합니다. 단점은 시퀀스 길이의 제곱에 비례하는 O(n^2) 계산 복잡도입니다. 이를 해결하기 위해 Linformer, Performer, FlashAttention 등 효율적인 변형들이 개발되었습니다.
PyTorch로 구현한 Scaled Dot-Product Self-Attention입니다.
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class SelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads=8):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
# Q, K, V 프로젝션
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
# Q, K, V 계산
Q = self.q_proj(x) # (B, S, D)
K = self.k_proj(x)
V = self.v_proj(x)
# Multi-head 분리: (B, S, D) -> (B, H, S, D/H)
Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# Scaled Dot-Product Attention
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, V) # (B, H, S, D/H)
# 헤드 결합
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_len, self.embed_dim)
return self.out_proj(attn_output), attn_weights
# 사용 예시
model = SelfAttention(embed_dim=512, num_heads=8)
x = torch.randn(2, 100, 512) # (batch, seq_len, embed_dim)
output, weights = model(x)
print(f"Output shape: {output.shape}") # (2, 100, 512)
print(f"Attention weights shape: {weights.shape}") # (2, 8, 100, 100)
| 어텐션 변형 | 계산 복잡도 | 메모리 복잡도 | 특징 |
|---|---|---|---|
| Standard Self-Attention | O(n^2 * d) | O(n^2) | 기본, 정확함 |
| FlashAttention | O(n^2 * d) | O(n) | 메모리 효율적, 빠름 |
| Linformer | O(n * k * d) | O(n * k) | K,V 프로젝션 |
| Performer | O(n * d^2) | O(n * d) | Random Features |
| Sparse Attention | O(n * sqrt(n)) | O(n * sqrt(n)) | 패턴 기반 희소화 |
* n: 시퀀스 길이, d: 임베딩 차원, k: 프로젝션 차원
"긴 시퀀스 처리 때문에 메모리 문제가 있어서 FlashAttention으로 교체했어요. 셀프 어텐션의 O(n^2) 메모리를 O(n)으로 줄여줍니다."
영어로는 "We switched to FlashAttention to reduce the O(n^2) memory footprint of self-attention"이라고 표현합니다. "scaled dot-product attention", "multi-head attention", "attention score/weight" 등의 용어도 자주 사용됩니다. 시각화할 때는 "attention heatmap"을 그려서 모델이 어디에 주목하는지 보여주세요.