실습 2 — 셀프 어텐션 손으로 계산하기

import torch
from torch import nn
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

torch.manual_seed(42)
words = ["이불", "배송", "언제", "와요"]
EMB = 8                                      # 좌표 수를 작게 잡아 계산을 눈으로 따라간다

emb = nn.Embedding(len(words), EMB)          # 낱말 4개의 임베딩 표
X = emb(torch.arange(len(words)))            # (4, 8): 낱말마다 좌표 8개
print(X.shape)

# bias=False: 곱하기만 하는 층. Q, K, V 를 만드는 가중치는 서로 다르다
Wq = nn.Linear(EMB, EMB, bias=False)
Wk = nn.Linear(EMB, EMB, bias=False)
Wv = nn.Linear(EMB, EMB, bias=False)

with torch.no_grad():                        # 학습이 아니라 계산만 하므로 기록을 끈다
    Q, K, V = Wq(X), Wk(X), Wv(X)
    scores = Q @ K.T / EMB ** 0.5            # 낱말 쌍마다 점수, 좌표 수의 제곱근으로 나눈다
print(scores.shape)
print(scores.round(decimals=2))

with torch.no_grad():
    # 마스크 없는 비율: 모든 낱말을 본다
    w_full = torch.softmax(scores, dim=-1)

    # 인과적 마스크: 아래 삼각형만 남기고 뒤쪽 칸의 점수를 -무한대로
    mask = torch.tril(torch.ones(4, 4))
    w_causal = torch.softmax(scores.masked_fill(mask == 0, float("-inf")), dim=-1)

    out = w_causal @ V                       # 비율만큼 내용을 섞은 결과
print("마스크 있는 비율:")
print(w_causal.round(decimals=2))
print("줄마다 합:", w_causal.sum(dim=-1).round(decimals=2).tolist())
print("결과 모양:", out.shape)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for ax, w, title in [(axes[0], w_full, "마스크 없음"), (axes[1], w_causal, "인과적 마스크")]:
    ax.imshow(w, vmin=0, vmax=1)                       # 비율이 클수록 밝은 색
    for i in range(4):
        for j in range(4):
            ax.text(j, i, f"{w[i, j]:.2f}", ha="center", va="center", color="white" if w[i, j] < 0.5 else "black")
    ax.set_xticks(range(4), words)
    ax.set_yticks(range(4), words)
    ax.set_xlabel("보는 대상(키)")
    ax.set_title(title)
axes[0].set_ylabel("보는 쪽(질문)")
plt.tight_layout()
plt.show()