셀프 어텐션 직접 구현(Self-Attention from Scratch)

어텐션(Attention)은 모든 단어가 "나에게 중요한 것은 누구인가?"라고 묻고 그 답을 학습하는 조회 테이블(lookup table)입니다.

유형: Build 언어: Python 선수 지식: Phase 3 (Deep Learning Core), Phase 5 Lesson 10 (Sequence-to-Sequence) 예상 시간: 약 90분

학습 목표

• NumPy만 사용해 쿼리(Query), 키(Key), 값(Value) 사영(projection)과 소프트맥스 가중합(softmax-weighted sum)을 포함한 스케일드 점곱 셀프 어텐션(scaled dot-product self-attention)을 직접 구현합니다.
• 헤드(head)를 나누고, 병렬 어텐션(parallel attention)을 계산한 뒤, 결과를 연결(concatenate)하는 멀티헤드 어텐션 계층(multi-head attention layer)을 만듭니다.
• 어텐션 행렬(attention matrix)이 토큰(token) 관계를 어떻게 포착하는지 추적하고, sqrt(d_k)로 스케일링(scaling)하는 이유가 소프트맥스 포화(softmax saturation) 방지임을 설명합니다.
• 양방향 어텐션(bidirectional attention)을 자기회귀적(autoregressive), 디코더 스타일(decoder-style) 어텐션으로 바꾸기 위해 인과 마스킹(causal masking)을 적용합니다.

문제

RNN은 시퀀스(sequence)를 한 토큰(token)씩 처리합니다. 50번째 토큰에 도달할 때쯤 첫 번째 토큰의 정보는 50번의 압축 단계(compression step)를 통과한 뒤입니다. 장거리 의존성(long-range dependency)은 고정 크기 은닉 상태(fixed-size hidden state)라는 병목(bottleneck) 안에 눌려 들어갑니다. LSTM 게이팅(gating)을 아무리 잘 써도 이 병목을 완전히 제거하지는 못합니다.

2014년 바다나우 어텐션(Bahdanau attention) 논문은 해결책을 보여주었습니다. 디코더(decoder)가 모든 인코더 위치(encoder position)를 다시 보고, 현재 단계(step)에 중요한 위치를 고르게 한 것입니다. 하지만 여전히 RNN 위에 붙은 구조였습니다. 2017년 "Attention Is All You Need"는 더 날카로운 질문을 던졌습니다. 어텐션만 있으면 어떨까? 순환(recurrence)도 합성곱(convolution)도 없이 어텐션만 쓰면 어떨까?

셀프 어텐션(Self-attention)은 시퀀스의 모든 위치(position)가 다른 모든 위치를 한 번의 병렬 단계(parallel step)에서 참조(attend)하게 합니다. 이것이 트랜스포머(transformer)를 빠르고 확장 가능(scalable)하며 지배적인 아키텍처(architecture)로 만든 핵심입니다.

개념

데이터베이스 조회(Database Lookup) 비유

어텐션은 부드러운 데이터베이스 조회(soft database lookup)처럼 생각할 수 있습니다.

전통적인 데이터베이스:
  쿼리(Query): "capital of France"  -->  정확히 일치(exact match)  -->  "Paris"

어텐션(Attention):
  쿼리(Query): "capital of France"  -->  모든 키(Key)와의 유사도(similarity)  -->  모든 값(Value)의 가중 결합(weighted blend)

모든 토큰은 세 벡터(vector)를 만듭니다.

• 쿼리(Query, Q): "나는 무엇을 찾고 있는가?"
• 키(Key, K): "나는 무엇을 담고 있는가?"
• 값(Value, V): "선택되면 어떤 정보를 제공하는가?"

쿼리와 모든 키의 내적(dot product)이 어텐션 점수(attention score)를 만듭니다. 높은 점수는 "이 키가 내 쿼리와 잘 맞는다"는 뜻입니다. 이 점수가 값에 가중치(weight)를 줍니다. 출력(output)은 값들의 가중합(weighted sum)입니다.

Q, K, V 계산

각 토큰 임베딩(token embedding)은 세 개의 학습된 가중치 행렬(learned weight matrix)을 통과합니다.

입력 임베딩 (n개 토큰의 시퀀스, 각 토큰은 d차원):

  X = [x1, x2, x3, ..., xn]       shape: (n, d)

세 개의 가중치 행렬:

  Wq  shape: (d, dk)
  Wk  shape: (d, dk)
  Wv  shape: (d, dv)

사영(Projection):

  Q = X @ Wq    shape: (n, dk)      각 토큰의 쿼리(query)
  K = X @ Wk    shape: (n, dk)      각 토큰의 키(key)
  V = X @ Wv    shape: (n, dv)      각 토큰의 값(value)

한 토큰에 대해 시각화하면 다음과 같습니다.

             Wq
  x_i ------[*]------> q_i    "나는 무엇을 찾는가?"
       |
       |     Wk
       +----[*]------> k_i    "나는 무엇을 담고 있는가?"
       |
       |     Wv
       +----[*]------> v_i    "나는 무엇을 제공하는가?"

어텐션 행렬(Attention Matrix)

모든 토큰의 Q, K, V가 준비되면 어텐션 점수는 행렬(matrix)을 이룹니다.

Scores = Q @ K^T    shape: (n, n)

              k1    k2    k3    k4    k5
        +-----+-----+-----+-----+-----+
   q1   | 2.1 | 0.3 | 0.1 | 0.8 | 0.2 |   <- q1이 각 키에 두는 어텐션 정도
        +-----+-----+-----+-----+-----+
   q2   | 0.4 | 1.9 | 0.7 | 0.1 | 0.3 |
        +-----+-----+-----+-----+-----+
   q3   | 0.2 | 0.6 | 2.3 | 0.5 | 0.1 |
        +-----+-----+-----+-----+-----+
   q4   | 0.9 | 0.1 | 0.4 | 1.7 | 0.6 |
        +-----+-----+-----+-----+-----+
   q5   | 0.1 | 0.3 | 0.2 | 0.5 | 2.0 |
        +-----+-----+-----+-----+-----+

각 행(row): 한 토큰이 전체 시퀀스에 걸쳐 두는 어텐션

왜 스케일링(scaling)하는가?

내적은 차원(dimension) dk가 커질수록 커집니다. dk = 64라면 내적이 수십 단위가 될 수 있고, 소프트맥스(softmax)가 그래디언트(gradient)가 사라지는 영역으로 밀려날 수 있습니다. 해결책은 sqrt(dk)로 나누는 것입니다.

Scaled scores = (Q @ K^T) / sqrt(dk)

이 스케일링은 소프트맥스가 유용한 그래디언트를 만들 수 있는 범위에 값을 유지합니다.

소프트맥스(Softmax)가 점수(score)를 가중치(weight)로 바꾼다

소프트맥스는 원시 점수(raw score)를 각 행(row)의 확률 분포(probability distribution)로 바꿉니다.

q1의 원시 점수(raw score):  [2.1, 0.3, 0.1, 0.8, 0.2]
                                   |
                                softmax
                                   |
어텐션 가중치(attention weight):  [0.52, 0.09, 0.07, 0.14, 0.08]   (합계 약 1.0)

이제 각 토큰은 다른 모든 토큰에 얼마나 어텐션을 둘지 나타내는 가중치 집합(weight set)을 갖습니다.

값(Value)의 가중합(Weighted Sum)

각 토큰의 최종 출력은 모든 값 벡터(value vector)의 가중합입니다.

output_i = sum( attention_weight[i][j] * v_j  for all j )

토큰 1에 대해:
  output_1 = 0.52 * v1 + 0.09 * v2 + 0.07 * v3 + 0.14 * v4 + 0.08 * v5

전체 파이프라인(Full Pipeline)

                    +-------+
  X (input)  ----->|  @ Wq  |-----> Q
                    +-------+
                    +-------+
  X (input)  ----->|  @ Wk  |-----> K
                    +-------+                     +----------+
                    +-------+                     |          |
  X (input)  ----->|  @ Wv  |-----> V ---------->| weighted |----> output
                    +-------+          ^          |   sum    |
                                       |          +----------+
                              +--------+--------+
                              |    softmax      |
                              +---------+-------+
                                        ^
                              +---------+-------+
                              | Q @ K^T / sqrt  |
                              +-----------------+

한 줄 공식은 다음과 같습니다.

Attention(Q, K, V) = softmax( Q @ K^T / sqrt(dk) ) @ V

직접 만들기

Step 1: Softmax 직접 구현

소프트맥스(softmax)는 원시 로짓(raw logit)을 확률(probability)로 바꿉니다. 수치적 안정성(numerical stability)을 위해 최댓값(max)을 빼 줍니다.

import numpy as np

def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

logits = np.array([2.0, 1.0, 0.1])
print(f"logit:   {logits}")
print(f"softmax: {softmax(logits)}")
print(f"합계:    {softmax(logits).sum():.4f}")

Step 2: 스케일드 점곱 어텐션(Scaled Dot-Product Attention)

핵심 함수입니다. Q, K, V 행렬(matrix)을 받아 어텐션 출력(attention output)과 가중치 행렬(weight matrix)을 반환합니다.

def scaled_dot_product_attention(Q, K, V):
    dk = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(dk)
    weights = softmax(scores)
    output = weights @ V
    return output, weights

Step 3: 학습된 사영(Learned Projection)을 가진 셀프 어텐션(Self-Attention) class

Xavier와 비슷한 스케일링으로 초기화된 Wq, Wk, Wv 가중치 행렬을 가진 셀프 어텐션 모듈(module)입니다.

class SelfAttention:
    def __init__(self, d_model, dk, dv, seed=42):
        rng = np.random.default_rng(seed)
        scale = np.sqrt(2.0 / (d_model + dk))
        self.Wq = rng.normal(0, scale, (d_model, dk))
        self.Wk = rng.normal(0, scale, (d_model, dk))
        scale_v = np.sqrt(2.0 / (d_model + dv))
        self.Wv = rng.normal(0, scale_v, (d_model, dv))
        self.dk = dk

    def forward(self, X):
        Q = X @ self.Wq
        K = X @ self.Wk
        V = X @ self.Wv
        output, weights = scaled_dot_product_attention(Q, K, V)
        return output, weights

Step 4: 문장에 실행하기

문장의 가짜 임베딩(fake embedding)을 만들고 어텐션 가중치(attention weight)를 봅니다.

sentence = ["The", "cat", "sat", "on", "the", "mat"]
n_tokens = len(sentence)
d_model = 8
dk = 4
dv = 4

rng = np.random.default_rng(42)
X = rng.normal(0, 1, (n_tokens, d_model))

attn = SelfAttention(d_model, dk, dv, seed=42)
output, weights = attn.forward(X)

print("어텐션 가중치(각 row는 해당 token이 어디를 보는지 나타냄):\n")
print(f"{'':>6}", end="")
for token in sentence:
    print(f"{token:>6}", end="")
print()

for i, token in enumerate(sentence):
    print(f"{token:>6}", end="")
    for j in range(n_tokens):
        w = weights[i][j]
        print(f"{w:6.3f}", end="")
    print()

Step 5: ASCII 히트맵(heatmap)으로 어텐션 시각화

어텐션 가중치를 문자에 매핑(mapping)해서 빠르게 시각화합니다.

def ascii_heatmap(weights, tokens, chars=" ░▒▓█"):
    n = len(tokens)
    print(f"\n{'':>6}", end="")
    for t in tokens:
        print(f"{t:>6}", end="")
    print()

    for i in range(n):
        print(f"{tokens[i]:>6}", end="")
        for j in range(n):
            level = int(weights[i][j] * (len(chars) - 1) / weights.max())
            level = min(level, len(chars) - 1)
            print(f"{'  ' + chars[level] + '   '}", end="")
        print()

ascii_heatmap(weights, sentence)

사용해보기

PyTorch의 nn.MultiheadAttention은 우리가 만든 것에 멀티헤드 분할(multi-head splitting)과 출력 사영(output projection)을 더한 것입니다.

import torch
import torch.nn as nn

d_model = 8
n_heads = 2
seq_len = 6

mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads, batch_first=True)

X_torch = torch.randn(1, seq_len, d_model)

output, attn_weights = mha(X_torch, X_torch, X_torch)

print(f"입력 shape:             {X_torch.shape}")
print(f"출력 shape:             {output.shape}")
print(f"어텐션 가중치 shape:     {attn_weights.shape}")
print(f"\n헤드 평균 어텐션 가중치:")
print(attn_weights[0].detach().numpy().round(3))

핵심 차이는 멀티헤드 어텐션(multi-head attention)이 여러 어텐션 함수(attention function)를 병렬로 실행한다는 점입니다. 각 헤드는 자기만의 Q, K, V 사영(projection)을 가지며, 크기는 dk = d_model / n_heads입니다. 이후 결과를 연결(concatenate)합니다. 덕분에 모델(model)은 여러 관계 유형을 동시에 참조(attend)할 수 있습니다.

산출물 만들기

이 lesson은 다음 산출물을 만듭니다.

• outputs/prompt-attention-explainer.md - 데이터베이스 조회(database lookup) 비유로 어텐션을 설명하는 프롬프트(prompt)

연습문제

1. scaled_dot_product_attention이 선택적 마스크 행렬(optional mask matrix)을 받도록 수정합니다. 특정 위치를 소프트맥스 전에 음의 무한대(negative infinity)로 만드는 방식입니다. 이것이 인과/디코더 마스킹(causal/decoder masking)의 원리입니다.
2. 멀티헤드 어텐션을 직접 구현합니다. Q, K, V를 n_heads개의 청크(chunk)로 나누고 각 헤드에서 어텐션을 실행한 뒤 연결하고 최종 가중치 행렬(final weight matrix) Wo로 사영합니다.
3. 길이가 같은 서로 다른 두 문장을 같은 SelfAttention 인스턴스(instance)에 넣고 어텐션 패턴(attention pattern)을 비교합니다. 무엇이 바뀌고 무엇이 유지됩니까?

핵심 용어

더 읽을거리

• Attention Is All You Need (Vaswani et al., 2017) - 트랜스포머(Transformer)를 처음 제안한 원 논문입니다.
• The Illustrated Transformer (Jay Alammar) - 트랜스포머 전체 아키텍처(architecture)를 시각적으로 따라가기 좋습니다.
• The Annotated Transformer (Harvard NLP) - PyTorch 구현을 한 줄씩 설명합니다.