Transformer Block 완성: FFN·LayerNorm·Residual을 빼면 깊은 모델이 왜 흔들리나
Attention만으로는 Transformer가 완성되지 않는다. 오늘은 Attention 뒤에 붙는 FFN, 값의 크기를 맞춰주는 LayerNorm, 입력을 그대로 더해주는 Residual을 하나씩 만들어 Transformer Block 하나를 완성하고 N개로 쌓는다. 그리고 Residual과 LayerNorm을 빼면 깊은 모델의 학습이 어떻게 불안정해지
By 송민성
오늘의 목표
Multi-Head Attention(MHA)에 FFN, LayerNorm, Residual을 붙여 Transformer Block 하나를 직접 구현한다. 그 블록을 N개 쌓아 작은 모델을 만들고, Residual과 LayerNorm을 빼면 학습이 얼마나 불안정해지는지 실험으로 확인한다.
개념과 직관
지금까지 만든 MHA는 "토큰끼리 정보를 주고받는" 부품이다. 하지만 Transformer Block에는 부품이 세 종류 더 필요하다. 하나씩 보자.
1) FFN (Feed Forward Network, 토큰마다 독립적으로 적용하는 작은 신경망)
Attention이 "누구를 참고할지"를 정한다면, FFN은 참고해서 모은 정보를 "각 토큰이 혼자 곱씹는" 단계다. 구조는 단순하다.
- 확장: 차원을 d_model에서 4*d_model 정도로 키운다 (4배는 관례이며 절대 규칙은 아니다).
- 활성(activation, 비선형 함수를 통과시켜 단순한 직선 이상의 표현을 가능하게 하는 것): GELU 또는 ReLU를 적용한다.
- 축소: 다시 d_model로 줄인다.
FFN(x) = W2 * GELU(W1 * x + b1) + b2비유하면 좁은 책상(d_model)에서 메모를 큰 작업대(4*d_model)에 펼쳐 놓고 정리한 뒤, 다시 책상 크기로 요약해 가져오는 것이다. 활성 함수가 없으면 W1과 W2가 합쳐져 선형 변환 하나와 같아지므로, 비선형 함수가 꼭 필요하다.
2) LayerNorm (Layer Normalization, 한 토큰 벡터의 값 분포를 평균 0·분산 1로 맞추는 것)
층을 지날수록 값이 점점 커지거나 작아지면 학습이 흔들린다. LayerNorm은 토큰 하나의 벡터(길이 d_model)만 보고 그 안에서 정규화한다.
mean = (x1 + x2 + ... + xd) / d
var = ((x1-mean)^2 + ... + (xd-mean)^2) / d
y_i = gamma_i * (x_i - mean) / sqrt(var + eps) + beta_i- eps는 0으로 나누는 것을 막는 아주 작은 수(예: 1e-5)다.
- gamma와 beta는 학습되는 파라미터다. 정규화 후에도 모델이 원하는 크기와 위치로 되돌릴 수 있게 해준다.
- 배치 방향이 아니라 "토큰 하나의 특징 방향"으로 계산한다는 점이 핵심이다. 그래서 배치 크기나 문장 길이에 영향을 받지 않는다.
3) Residual connection (잔차 연결, 층의 입력을 출력에 그대로 더해주는 지름길)
y = x + F(x)F는 Attention이나 FFN이다. 비유하면 문서를 고쳐 쓸 때 원본을 버리고 새로 쓰는 대신 "원본 + 수정 사항"만 기록하는 방식이다. 층이 별로 할 일이 없으면 F(x)가 0 근처가 되고, 입력이 그대로 통과한다.
gradient(기울기, 파라미터를 바꿨을 때 손실이 얼마나 변하는지 알려주는 값) 관점에서도 중요하다. y = x + F(x)를 x로 미분하면 1 + F'(x)다. 앞의 "1"이 있어서 F'(x)가 작아도 gradient가 깊은 층까지 끊기지 않고 전달된다. 이 "1"이 없으면 층을 지날 때마다 작은 수가 곱해져 앞쪽 층의 gradient가 사라지기 쉽다. (gradient가 실제로 어떻게 계산되는지는 PyTorch가 자동으로 해주므로, 오늘은 "지름길이 있다"는 직관만 잡으면 된다.)
4) Pre-norm vs Post-norm
LayerNorm을 어디에 놓느냐의 차이다.
Post-norm (원래 Transformer 논문 방식):
x = LayerNorm(x + Attention(x))
x = LayerNorm(x + FFN(x))
Pre-norm (요즘 LLM에서 흔한 방식):
x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))Pre-norm은 Residual 경로(x가 그대로 흐르는 길)에 LayerNorm이 끼지 않아서, 깊이 쌓을 때 학습이 더 안정적인 경향이 있다. 오늘은 pre-norm으로 구현한다. 현대 LLM 다수가 pre-norm 계열을 쓰는 것으로 알려져 있지만, 모델마다 세부 구현은 다르다.
5) Causal mask (인과 마스크, 미래 토큰을 못 보게 가리는 장치)
GPT 같은 언어 모델은 "다음 토큰 맞히기"를 학습하므로, 위치 t는 t보다 뒤의 토큰을 보면 안 된다. Attention 점수 행렬에서 대각선 윗부분을 -inf로 채운 뒤 softmax를 하면, 그 위치의 확률이 0이 된다.
허용(1) / 차단(0)
t=0: 1 0 0 0
t=1: 1 1 0 0
t=2: 1 1 1 0
t=3: 1 1 1 1코드로 직접 해보기
PyTorch가 설치되어 있다고 가정한다. 파일 하나(block.py)에 순서대로 이어 붙이면 된다.
1단계: LayerNorm을 수식 그대로 만들고 검증
import torch
import torch.nn as nn
import torch.nn.functional as F
class MyLayerNorm(nn.Module):
def __init__(self, d, eps=1e-5):
super().__init__()
self.gamma = nn.Parameter(torch.ones(d))
self.beta = nn.Parameter(torch.zeros(d))
self.eps = eps
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
return self.gamma * (x - mean) / torch.sqrt(var + self.eps) + self.beta
x = torch.randn(2, 5, 8) * 10 + 3
mine = MyLayerNorm(8)(x)
ref = nn.LayerNorm(8)(x)
print(torch.allclose(mine, ref, atol=1e-5)) # True가 나와야 함
print(mine.mean(-1)[0, :2], mine.std(-1)[0, :2]) # 평균≈0, 표준편차≈12단계: Causal Self-Attention (멀티 헤드)
앞선 Day에서 만든 MHA와 같은 원리다. 여기서는 한 클래스로 묶고 causal mask를 넣는다.
class CausalSelfAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.h = n_heads
self.dk = d_model // n_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.out = nn.Linear(d_model, d_model)
def forward(self, x):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
# (B, T, D) -> (B, h, T, dk)
q, k, v = [t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v)]
scores = q @ k.transpose(-2, -1) / (self.dk ** 0.5) # (B, h, T, T)
mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
scores = scores.masked_fill(~mask, float("-inf"))
w = F.softmax(scores, dim=-1)
y = (w @ v).transpose(1, 2).contiguous().view(B, T, D)
return self.out(y)3단계: FFN
class FFN(nn.Module):
def __init__(self, d_model, expansion=4):
super().__init__()
self.w1 = nn.Linear(d_model, expansion * d_model) # 확장
self.w2 = nn.Linear(expansion * d_model, d_model) # 축소
def forward(self, x):
return self.w2(F.gelu(self.w1(x))) # 확장 -> 활성 -> 축소4단계: Transformer Block (pre-norm, residual 스위치 포함)
실험을 위해 use_residual, use_norm 옵션을 넣어둔다.
class Block(nn.Module):
def __init__(self, d_model, n_heads, use_residual=True, use_norm=True):
super().__init__()
self.attn = CausalSelfAttention(d_model, n_heads)
self.ffn = FFN(d_model)
self.ln1 = MyLayerNorm(d_model) if use_norm else nn.Identity()
self.ln2 = MyLayerNorm(d_model) if use_norm else nn.Identity()
self.use_residual = use_residual
def forward(self, x):
a = self.attn(self.ln1(x))
x = x + a if self.use_residual else a
f = self.ffn(self.ln2(x))
x = x + f if self.use_residual else f
return x5단계: 블록을 N개 쌓아 작은 언어 모델 만들기
class TinyLM(nn.Module):
def __init__(self, vocab, d_model=64, n_heads=4, n_layers=8, max_len=64,
use_residual=True, use_norm=True):
super().__init__()
self.tok = nn.Embedding(vocab, d_model)
self.pos = nn.Embedding(max_len, d_model)
self.blocks = nn.ModuleList(
[Block(d_model, n_heads, use_residual, use_norm) for _ in range(n_layers)]
)
self.ln_f = MyLayerNorm(d_model) if use_norm else nn.Identity()
self.head = nn.Linear(d_model, vocab)
def forward(self, idx):
B, T = idx.shape
x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
for blk in self.blocks:
x = blk(x)
return self.head(self.ln_f(x))
model = TinyLM(vocab=50)
idx = torch.randint(0, 50, (2, 16))
print(model(idx).shape) # torch.Size([2, 16, 50])6단계: causal mask 확인
미래 토큰을 바꿔도 앞쪽 출력이 변하지 않아야 한다.
model.eval()
a = torch.randint(0, 50, (1, 16))
b = a.clone(); b[0, -1] = (b[0, -1] + 1) % 50 # 마지막 토큰만 변경
with torch.no_grad():
print(torch.allclose(model(a)[0, :-1], model(b)[0, :-1], atol=1e-5)) # True오늘의 실험
질문: Residual과 LayerNorm을 빼면 깊은 모델 학습이 정말 불안정해지는가?
설정: 같은 작은 데이터(예: 길이 32의 반복 패턴이나 간단한 규칙 시퀀스에서 다음 토큰 맞히기)로 아래 4가지 조합을 학습한다. 깊이는 n_layers=12처럼 충분히 깊게 하고, seed·학습률(예: 3e-4)·스텝 수(예: 300)는 모두 동일하게 둔다.
| 조합 | use_residual | use_norm |
|---|---|---|
| A (기준) | True | True |
| B | True | False |
| C | False | True |
| D | False | False |
import random
def make_batch(B=32, T=32, vocab=50):
# 간단한 규칙: 다음 토큰 = (현재 토큰 + 1) % vocab
start = torch.randint(0, vocab, (B, 1))
seq = (start + torch.arange(T + 1)) % vocab
return seq[:, :-1], seq[:, 1:]
def run(use_residual, use_norm, n_layers=12, steps=300, seed=0):
torch.manual_seed(seed)
m = TinyLM(50, n_layers=n_layers, use_residual=use_residual, use_norm=use_norm)
opt = torch.optim.Adam(m.parameters(), lr=3e-4)
losses, gnorms = [], []
for s in range(steps):
x, y = make_batch()
loss = F.cross_entropy(m(x).view(-1, 50), y.reshape(-1))
opt.zero_grad()
loss.backward()
g = sum(p.grad.norm() ** 2 for p in m.parameters() if p.grad is not None) ** 0.5
opt.step()
losses.append(loss.item()); gnorms.append(g.item())
return losses, gnorms
for ur, un in [(True, True), (True, False), (False, True), (False, False)]:
l, g = run(ur, un)
print(f"residual={ur}, norm={un} | 첫 loss {l[0]:.3f} -> 마지막 loss {l[-1]:.3f} | "
f"grad norm 최대 {max(g):.2f}, 최소 {min(g):.4f}")관찰할 것:
- loss가 내려가는 속도와 최종 값 (내려가지 않고 멈추는가, 튀는가, NaN이 되는가)
- gradient norm(기울기 전체 크기)의 범위: 너무 작아서 거의 0인가, 너무 커서 튀는가
- 깊이를 4, 12, 24로 바꾸며 A와 D의 차이가 벌어지는지
기록 방법: 표로 정리한다. 열은 조합, 깊이, 첫 loss, 마지막 loss, grad norm 최대/최소, 특이사항(NaN 여부 등). 가능하면 matplotlib으로 loss 곡선을 한 그래프에 겹쳐 그리고, 마지막에 "내 결과에서는 어떤 조합이 가장 불안정했고, 그 이유를 한 문장으로 어떻게 설명할 수 있나"를 적는다. 결과가 이론과 다르게 나오면(예: 얕은 모델에서는 D도 학습됨) 그 자체가 중요한 기록이다. 작은 모델·짧은 학습에서는 차이가 약하게 보일 수 있으니, 깊이를 늘려가며 비교한다.
왜 이걸 배우나 (LLM / Post-training 연결)
GPT 계열 같은 decoder-only LLM은 사실상 오늘 만든 블록을 수십 개 쌓은 구조이고, 모델 크기를 키운다는 것은 d_model, 층 수, 헤드 수를 늘린다는 뜻이다. Post-training(사전학습된 모델을 SFT·정렬 등으로 추가 학습하는 단계)에서 학습이 발산하거나 loss가 튈 때, 어느 층에서 값과 gradient가 어떻게 흐르는지 읽어낼 수 있어야 하므로 블록 내부를 손으로 짜 본 경험이 디버깅의 기반이 된다. 이후 LoRA 같은 기법을 적용할 때도 "어느 Linear가 Attention이고 어느 것이 FFN인지"를 아는 것이 출발점이다.
초보자가 막히는 지점
- shape 에러 (`view`/`transpose`): 각 단계에서
print(t.shape)를 찍어 (B, h, T, dk) 순서를 확인한다.transpose뒤view가 실패하면.contiguous()를 붙였는지 본다. - LayerNorm을 배치 방향으로 계산하는 실수: 평균·분산은 반드시 마지막 차원(
dim=-1)에서 구한다. 1단계의 검증 코드에서allclose가 False면 여기부터 의심한다. - `var`의 `unbiased=False`: LayerNorm은 N으로 나누는 분산을 쓴다. 기본값(N-1로 나눔)을 쓰면
nn.LayerNorm과 값이 미세하게 달라진다. - mask를 softmax 뒤에 적용하는 실수: -inf는 반드시 softmax 이전에 채운다. 이후에 하면 확률 합이 1이 아니게 된다.
- post-norm과 pre-norm 혼동: 오늘 코드는 "서브레이어 입력 쪽에서 정규화하고, 결과를 x에 더한다"는 pre-norm이다. 마지막 출력 전에
ln_f를 한 번 더 두는 이유도 pre-norm에서는 Residual 경로에 정규화가 없기 때문이다. - FFN이 토큰끼리 섞는다는 오해: FFN은 각 토큰에 독립적으로 적용된다. 토큰 간 정보 교환은 Attention만 한다.
- 실험에서 NaN이 나왔을 때: 버그가 아니라 관찰 대상일 수 있다. 다만 A(기준)에서도 NaN이 나오면 코드 버그이니, 학습률을 낮추거나 shape·mask부터 점검한다.
알고리즘 (병행 1시간)
LeetCode Medium 트리 문제 2문제를 푼다. 시간은 문제당 약 30분으로 잡고, 막히면 20~25분쯤에서 힌트를 본 뒤 직접 다시 구현한다. 예시 후보(다른 트리 Medium도 무방):
- Binary Tree Level Order Traversal (BFS로 레벨별 순회)
- Validate Binary Search Tree (재귀 + 값의 범위 전달)
풀이 후 각 문제마다 아래를 기록한다.
- 접근: 어떤 순회(DFS/BFS)를 왜 골랐는가, 재귀의 종료 조건은 무엇인가
- 시간복잡도 / 공간복잡도: 노드 수 n 기준으로 쓰고, 재귀 깊이나 큐 크기가 공간에 어떻게 영향을 주는지 적는다
- 틀린 이유: 틀린 제출이나 막힌 지점이 있다면 원인(예: None 처리 누락, 범위 경계에서 등호 처리 실수)을 한 줄로 남긴다