Q, K, V — 질문하고, 검색하고, 읽다

import math

# === Q, K, V 어텐션 직접 계산 ===

# 3개 토큰, 차원 d_k = 4
Q = [[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]]
K = [[1, 1, 0, 0], [0, 0, 1, 1], [1, 0, 0, 1]]
V = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]

d_k = len(Q[0])  # 차원 수

# 1단계: Q × Kᵀ (유사도 점수)
scores = []
for i in range(len(Q)):
    row = []
    for j in range(len(K)):
        dot = sum(Q[i][k] * K[j][k] for k in range(d_k))
        row.append(dot / math.sqrt(d_k))  # 스케일링
    scores.append(row)

print("스케일된 점수 (Q×Kᵀ/√d_k):")
for i, row in enumerate(scores):
    print(f"  토큰{i}: {[round(v, 3) for v in row]}")

# 2단계: Softmax (확률 분포)
def softmax(xs):
    mx = max(xs)
    exps = [math.exp(x - mx) for x in xs]
    s = sum(exps)
    return [e / s for e in exps]

weights = [softmax(row) for row in scores]
print("\n어텐션 가중치 (softmax):")
for i, w in enumerate(weights):
    print(f"  토큰{i}: {[round(v, 3) for v in w]}")

# 3단계: Weights × V (가중 합)
output = []
for i in range(len(weights)):
    out = [0.0] * len(V[0])
    for j in range(len(V)):
        for k in range(len(V[0])):
            out[k] += weights[i][j] * V[j][k]
    output.append([round(v, 4) for v in out])

print("\n어텐션 출력:")
for i, o in enumerate(output):
    print(f"  토큰{i}: {o}")