Self-Attention 직접 구현

import numpy as np

np.random.seed(42)

# Step 0: 입력 — 3개 토큰, 4차원 임베딩
tokens = ["나는", "고양이를", "좋아한다"]
X = np.array([
    [1.0, 0.0, 1.0, 0.0],   # "나는"
    [0.0, 1.0, 0.0, 1.0],   # "고양이를"
    [1.0, 1.0, 0.0, 0.0],   # "좋아한다"
])
d_k = 4

print("=== Step 0: 입력 임베딩 ===")
for token, vec in zip(tokens, X):
    print(f"  {token:6s} → {vec}")

# Step 1: 가중치 행렬
W_Q = np.random.randn(4, 4) * 0.5
W_K = np.random.randn(4, 4) * 0.5
W_V = np.random.randn(4, 4) * 0.5

# Step 2: Q, K, V 계산
Q = X @ W_Q   # (3,4) × (4,4) = (3,4)
K = X @ W_K
V = X @ W_V
print(f"\n=== Step 2: Q, K, V 크기 = {Q.shape} ===")

# Step 3: 어텐션 스코어 = Q × K^T
scores = Q @ K.T   # (3,4) × (4,3) = (3,3)
print(f"\n=== Step 3: 스코어 크기 = {scores.shape} ===")

# Step 4: 스케일링
scores_scaled = scores / np.sqrt(d_k)

# Step 5: Softmax → 확률 가중치
def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

attention_weights = softmax(scores_scaled)

print(f"\n=== Step 5: 어텐션 가중치 ===")
print("         나는   고양이를 좋아한다")
for i, token in enumerate(tokens):
    w_str = "  ".join(f"{w:.3f}" for w in attention_weights[i])
    print(f"  {token:6s} [{w_str}]")

# Step 6: 가중 합
output = attention_weights @ V   # (3,3) × (3,4) = (3,4)
print(f"\n=== Step 6: 최종 출력 크기 = {output.shape} ===")
print("→ 각 토큰에 다른 토큰들의 정보가 섞여 들어갔습니다!")