실습 4 — Add & Norm 과 FFN 으로 Encoder 블록 완성하기

import numpy as np

np.set_printoptions(precision=2, suppress=True)

def layer_norm(x, eps=1e-5):
    # 낱말(행)마다 평균과 표준편차를 구한다
    mean = x.mean(axis=-1, keepdims=True)
    std = x.std(axis=-1, keepdims=True)
    # eps: 표준편차가 0 일 때 0 으로 나누지 않게 더하는 아주 작은 수
    return (x - mean) / (std + eps)

x = np.array([[1.0, 2.0, 3.0, 6.0],
              [10.0, 20.0, 30.0, 60.0]])   # 둘째 줄은 첫째 줄의 10배
y = layer_norm(x)
print(y)
# 결과: 두 줄이 거의 같다 → 크기(눈금)가 달라도 같은 꼴로 맞춘다
print("행 평균:", y.mean(axis=1).round(2), " 행 표준편차:", y.std(axis=1).round(2))
# 결과: 행 평균: [0. 0.]  행 표준편차: [1. 1.]

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

def multi_head_attention(x, W_Q, W_K, W_V, W_O, h):
    n, d_model = x.shape
    d_head = d_model // h
    # 실습 3과 같은 계산: 나누기 → head 별 Attention → 이어 붙이기 → W_O
    split = lambda M: M.reshape(n, h, d_head).transpose(1, 0, 2)
    Qh, Kh, Vh = split(x @ W_Q), split(x @ W_K), split(x @ W_V)
    weights = softmax(Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d_head))
    concat = (weights @ Vh).transpose(1, 0, 2).reshape(n, d_model)
    return concat @ W_O

np.random.seed(2)
n, d_model, h = 4, 8, 2
x = np.random.randn(n, d_model)                     # 입력 (4, 8)
attn_W = [np.random.randn(d_model, d_model) * 0.5 for _ in range(4)]   # W_Q, W_K, W_V, W_O

attn_out = multi_head_attention(x, *attn_W, h)      # (4, 8)
x1 = layer_norm(x + attn_out)                       # Add(잔차) 뒤 Norm
print("attn_out:", attn_out.shape, " x1:", x1.shape)
# 결과: attn_out: (4, 8)  x1: (4, 8)
print("x1 첫 낱말:", x1[0])

d_ff = 4 * d_model                     # 넓힐 폭 = 32 (원 논문 512 → 2048 과 같은 4배)
W1 = np.random.randn(d_model, d_ff) * 0.3
b1 = np.zeros(d_ff)                    # 편향(bias): 더하는 값. 처음엔 0
W2 = np.random.randn(d_ff, d_model) * 0.3
b2 = np.zeros(d_model)

def ffn(x):
    # 8 → 32 로 넓히고, ReLU(음수는 0) 로 비선형을 준 뒤, 32 → 8 로 줄인다
    hidden = np.maximum(0, x @ W1 + b1)
    return hidden @ W2 + b2

x2 = layer_norm(x1 + ffn(x1))          # Add(잔차) 뒤 Norm
print("x2:", x2.shape)
# 결과: x2: (4, 8)

# FFN 은 낱말마다 따로: 한 줄만 넣은 결과와 네 줄을 한꺼번에 넣은 결과의 그 줄이 같다
print("낱말마다 따로 계산되는가:", np.allclose(ffn(x1[2:3]), ffn(x1)[2:3]))
# 결과: True

def encoder_block(x, attn_W, W1, b1, W2, b2, h):
    # 1) Multi-Head Attention → Add & Norm
    x = layer_norm(x + multi_head_attention(x, *attn_W, h))
    # 2) FFN → Add & Norm
    hidden = np.maximum(0, x @ W1 + b1)
    x = layer_norm(x + hidden @ W2 + b2)
    return x

# 층마다 가중치를 따로 만든다(실제 모델도 층마다 다른 가중치)
np.random.seed(3)
layers = []
for _ in range(2):
    layers.append(([np.random.randn(d_model, d_model) * 0.5 for _ in range(4)],
                   np.random.randn(d_model, d_ff) * 0.3, np.zeros(d_ff),
                   np.random.randn(d_ff, d_model) * 0.3, np.zeros(d_model)))

out = x
for i, params in enumerate(layers):
    out = encoder_block(out, *params, h)
    print(f"{i + 1}층 통과 뒤 shape:", out.shape)
# 결과: 1층, 2층 모두 (4, 8) → 모양이 같아 몇 층이든 쌓을 수 있다

# 블록 하나의 파라미터 수
attn_params = 4 * d_model * d_model                       # 256
ffn_params = d_model * d_ff + d_ff + d_ff * d_model + d_model   # 552
print("Attention:", attn_params, " FFN:", ffn_params, " FFN 비율:", round(ffn_params / (attn_params + ffn_params), 2))
# 결과: Attention: 256  FFN: 552  FFN 비율: 0.68 → 블록의 약 2/3