인지야공

인지야공/딥러닝 기초 정리/9번째 글

Transformer 각 층 상세 해부

딥러닝 기초 정리 9/56읽는 데 18분

실행: python 04_transformer_layers.py 직접 만든 인코더/디코더 층이 PyTorch 공식 nn.TransformerEncoderLayer / nn.TransformerDecoderLayer 와 오차 정확히 0.0 으로 일치함을 확인했습니다.

0. 전체 지도

Transformer 전체 구조 원문은 인코더 6층을, 번역문은 디코더 6층을 지난다. 인코더의 출력(memory)이 디코더의 크로스 어텐션에 키와 값으로 들어가고, 디코더 출력이 Linear와 Softmax를 거쳐 단어 확률이 된다. 원문 "I am a student" 번역문 "나는 학생 이다" (한 칸 밀린 정답) ① 토큰 임베딩 ① 토큰 임베딩 + ② 위치 인코딩 + ② 위치 인코딩 인코더 층 × N (=6) 디코더 층 × N (=6) memory ⑦ Linear → Softmax → 단어 확률 ③ Multi-Head 셀프어텐션 ④ Add & LayerNorm ⑤ FFN (확장 → 축소) ④ Add & LayerNorm ③' Masked 셀프어텐션 ④ Add & LayerNorm ⑥ Cross Attention Q = 디코더, K/V = 인코더 ④ Add & LayerNorm ⑤ FFN ④ Add & LayerNorm

모든 서브층은 예외 없이 이 형태로 감싸져 있다:

out = LayerNorm( x + Sublayer(x) )

① 토큰 임베딩 (Token Embedding)

단어 번호(정수)를 벡터로 바꾸는 거대한 조회 테이블이다.

emb = nn.Embedding(vocab_size, d_model)
emb(torch.tensor([[5, 17, 42]]))     # (1, 3) → (1, 3, d_model)
  • 파라미터 수 = vocab_size × d_model (예: 32000 × 512 = 1638만개)
  • 학습 가능. 비슷한 의미의 단어가 비슷한 벡터로 수렴한다.

왜 √d_model 을 곱하나?

원논문은 임베딩 출력에 √d_model 을 곱한다 (d_model=512면 22.6배).

이유: nn.Embedding 초기값은 표준편차가 작은데, 바로 뒤에 더할 위치 인코딩은 진폭이 -1~1로 고정되어 있다. 스케일을 안 맞추면 단어 정보가 위치 정보에 묻힌다. √d_model 을 곱해 두 신호의 크기를 비슷하게 만든다.

참고: 최근 모델(GPT 계열)은 이 스케일링을 생략하고 초기화 표준편차(0.02 등)로 대신 조절한다.

② 위치 인코딩

→ 별도 문서 위치 인코딩 — 어텐션이 잃어버린 순서를 되돌려 놓는 법 참조

③ Multi-Head Self-Attention

→ 별도 문서 Attention 은 어떻게 계산되는가 — Q·K·V 와 병렬 계산의 원리 참조

파라미터 수 = 4 × (d_model² + d_model) (Q, K, V, O 4개의 정사각 행렬) d_model=512 기준 1,050,624개.

④ Add & LayerNorm — 두 가지가 한 세트

(a) 잔차 연결 (Residual Connection) = 기울기 고속도로

out = x + Sublayer(x)

x 를 그대로 더하므로, 역전파 때 미분값 안에 항상 1이 남는다.

d(out)/dx = 1 + d(Sublayer)/dx

실측 (f’ = 0.5 가정, 50층 통과 후 기울기 크기):

기울기
잔차 없음 (0.5를 50번 곱함)8.882e-16 ← 소실
잔차 있음 (1.5를 50번 곱함)6.376e+08 ← 살아남음

이 덕분에 층을 6개, 24개, 100개까지 쌓아도 학습이 된다. Transformer가 깊게 쌓일 수 있는 이유의 절반은 잔차 연결이다.

(b) LayerNorm — 왜 BatchNorm이 아닌가

y = γ · (x - μ) / √(σ² + ε) + β

여기서 μ, σ 는 “토큰 1개의 d_model 차원 안에서만” 계산한다. (BatchNorm은 배치 전체에서 같은 채널끼리 계산)

BatchNorm을 못 쓰는 이유 3가지:

  1. 문장 길이가 제각각이고 패딩이 섞여 배치 통계가 오염된다.
  2. 추론할 때 배치 크기가 1이면 배치 통계가 무의미하다.
  3. 시퀀스는 위치마다 분포가 달라 “같은 위치끼리” 정규화하는 게 부자연스럽다.

LayerNorm은 샘플 하나 안에서만 정규화하므로 길이/배치와 완전히 무관하다.

실측: LayerNorm 후 각 토큰의 평균 = -0.000000, 표준편차 = 0.999995 ✔ 파라미터 = 2 × d_model (γ, β) — d_model=512면 1024개로 아주 작다.

최근 모델은 평균 빼기를 생략한 RMSNorm 을 많이 쓴다. y = γ · x / √(mean(x²) + ε). 계산이 더 싸고 성능은 비슷하다. (LLaMA 등)

⑤ FFN (Position-wise Feed-Forward Network)

FFN(x) = Linear2( ReLU( Linear1(x) ) )
#         d_model → d_ff(보통 4배) → d_model

“Position-wise”의 뜻: 토큰마다 똑같은 가중치를 독립적으로 적용한다. 토큰끼리 섞는 일은 어텐션이 다 했고, FFN은 각 토큰의 정보를 혼자서 가공한다.

왜 4배로 늘렸다 줄이나? 좁은 공간에서는 표현할 수 없는 비선형 변환을, 넓은 공간으로 펼쳐서 수행한 뒤 다시 압축해 돌아오는 것이다. (오토인코더의 반대 모양)

파라미터 비중이 압도적이다.

구성d_model=512, d_ff=2048 기준비중
Multi-Head Attention1,050,62433.3%
FFN2,099,71266.6%
LayerNorm × 22,0480.1%
합계 (층 1개)3,152,384100%

FFN이 파라미터의 2/3를 차지한다. 그래서 최근 연구는 FFN을 “모델이 사실 지식을 저장하는 곳(key-value memory)” 으로 해석한다. Mixture-of-Experts(MoE)가 바로 이 FFN만 여러 개로 늘린 구조다.

활성함수 변천: ReLU(원논문) → GELU(BERT/GPT) → SwiGLU(LLaMA/PaLM) GELU(1.0) = 0.8413, ReLU(1.0) = 1.0 — GELU가 부드럽게 깎아 기울기가 안정적이다.

⑥ Cross Attention (디코더 전용)

디코더 층에만 있는 두 번째 어텐션이다.

c, _ = self.cross_attn(query=tgt,       # Q = 디코더의 현재 상태
                       key=memory,      # K = 인코더 출력
                       value=memory)    # V = 인코더 출력
셀프 어텐션크로스 어텐션
Q자기 자신디코더
K, V자기 자신인코더 출력(memory)
의미“내 문장 안에서 뭘 볼까”“지금 만들 단어 기준으로 원문 어디를 볼까”
마스크디코더는 causal 필요불필요 (원문은 다 봐도 됨)

번역기의 핵심이 바로 여기다. “학생”이라는 단어를 생성할 때 원문의 “student”에 높은 가중치를 주는 것 — 이게 예전 SMT의 “정렬(alignment)“을 학습으로 대체한 것이다.

⑦ 인코더 층 전체 코드 (Post-LN, 원논문 방식)

class MyEncoderLayer(nn.Module):
    def __init__(self, d_model, n_head, d_ff):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x, mask=None, key_padding_mask=None):
        # 서브층 1: 셀프 어텐션 + 잔차 + 정규화
        attn_out, _ = self.self_attn(x, x, x, attn_mask=mask,
                                     key_padding_mask=key_padding_mask,
                                     need_weights=False)
        x = self.norm1(x + attn_out)

        # 서브층 2: FFN + 잔차 + 정규화
        ff_out = self.linear2(F.relu(self.linear1(x)))
        x = self.norm2(x + ff_out)
        return x

✔ nn.TransformerEncoderLayer(dropout=0.0, norm_first=False) 와 최대 오차 0.0

디코더 층 전체 코드

def forward(self, tgt, memory, tgt_mask=None):
    # (1) 마스크드 셀프어텐션 — 자기 자신과 '과거'만
    a, _ = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, need_weights=False)
    tgt = self.norm1(tgt + a)

    # (2) 크로스 어텐션 — Q=디코더, K/V=인코더 출력
    c, _ = self.cross_attn(tgt, memory, memory, need_weights=False)
    tgt = self.norm2(tgt + c)

    # (3) FFN
    f = self.linear2(F.relu(self.linear1(tgt)))
    tgt = self.norm3(tgt + f)
    return tgt

✔ nn.TransformerDecoderLayer 와 최대 오차 0.0

causal 마스크는 nn.Transformer.generate_square_subsequent_mask(N) 로 만든다 (0과 -inf로 채워진 float 마스크):

[[  0., -inf, -inf, -inf],
 [  0.,   0., -inf, -inf],
 [  0.,   0.,   0., -inf],
 [  0.,   0.,   0.,   0.]]

⑧ Post-LN vs Pre-LN — 요즘 모델이 바꾼 것

Post-LN (원논문)Pre-LN (GPT, LLaMA, ViT)
식x = LayerNorm(x + Sublayer(x))x = x + Sublayer(LayerNorm(x))
잔차 경로LayerNorm을 통과함뚫려 있음 (정규화를 안 거침)
학습 안정성불안정. warmup 필수안정. warmup 없이도 됨
깊이6~12층 정도가 한계수십~수백 층 가능
활성값매 층 std ≈ 1로 고정층이 깊어질수록 커짐

실측 (12층 통과, 활성값 표준편차):

층Post-LNPre-LN
11.0010.984
61.0011.196
121.0011.390

Pre-LN은 잔차가 계속 누적되어 값이 커진다. → 그래서 마지막에 final LayerNorm 을 한 번 더 둔다.

class PreLNEncoderLayer(nn.Module):
    def forward(self, x):
        h = self.n1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]   # 잔차가 정규화를 안 거침
        x = x + self.ff(self.n2(x))
        return x

PyTorch에서는 norm_first=True 한 줄로 바꿀 수 있다.

nn.TransformerEncoderLayer(d_model, nhead, norm_first=True)

⑨ 파라미터 예산 (원논문 base 모델)

d_model=512, d_ff=2048, h=8, N=6

계산개수
MHA 1개4 × (512² + 512)1,050,624
FFN 1개512×2048+2048 + 2048×512+5122,099,712
LayerNorm 2개2 × 2 × 5122,048
인코더 층 1개3,152,384
인코더 6층 + 디코더 6층디코더는 크로스어텐션이 추가약 4,414만
+ 임베딩/출력층 포함약 6,500만 (논문값)

⑩ 실제 동작 검증 — 문장 뒤집기 과제

nn.Transformer 로 인코더-디코더를 만들어 “입력 수열을 뒤집어 출력하라” 를 학습시켰다. (어휘 12개, 길이 8, 2층, d=64, 20 epoch)

epoch  5 | loss 1.0726 | 토큰 정확도  61.1%
epoch 10 | loss 0.3663 | 토큰 정확도  95.0%
epoch 15 | loss 0.1234 | 토큰 정확도  99.6%
epoch 20 | loss 0.0460 | 토큰 정확도 100.0%

입력       : [7, 8, 4, 6, 4, 6, 9, 5]
정답(뒤집기) : [5, 9, 6, 4, 6, 4, 8, 7]
모델 생성   : [5, 9, 6, 4, 6, 4, 8, 7]   ← 성공

여기서 꼭 봐야 할 학습 패턴 (Teacher Forcing)

logits = net(src, tgt[:, :-1])                        # 입력은 마지막 토큰 제외
loss = lf(logits.reshape(-1, V), tgt[:, 1:].reshape(-1))   # 정답은 한 칸 밀어서
  • 디코더 입력: <BOS> 5 9 6 4 6 4 8
  • 정답 라벨: 5 9 6 4 6 4 8 7

한 칸 밀린 정답을 맞히도록 학습하고, causal 마스크로 미래를 가린다. 이 덕분에 학습은 전체 문장을 한 번에 병렬 처리할 수 있다.

추론(생성)은 반대로 순차다:

gen = torch.tensor([[BOS]])
for _ in range(max_len):
    nxt = net(src, gen)[:, -1].argmax(-1, keepdim=True)
    gen = torch.cat([gen, nxt], dim=1)

한 장 요약

층역할파라미터핵심 포인트
토큰 임베딩단어 → 벡터vocab × d√d_model 곱해 스케일 맞춤
위치 인코딩순서 정보 주입0 (고정형)더한다 (concat 아님)
Multi-Head Attn토큰 사이 정보 교환4d²head를 배치축으로 → 병렬
Add & Norm기울기 보존 + 안정화2d잔차가 핵심, LayerNorm은 토큰별
FFN토큰 내부 정보 가공8d²파라미터의 2/3, 지식 저장소
Cross Attn원문 참조4d²Q=디코더, K/V=인코더
Linear+Softmax벡터 → 단어 확률d × vocab임베딩과 가중치 공유하기도 함

한 문장 요약:

Transformer 층 = 어텐션(토큰끼리 섞기) + FFN(토큰 내부 가공) 을 잔차 + LayerNorm 으로 감싼 것. 이것을 N번 반복한 게 전부다.

표시는 이 브라우저에만 남는다. 서버로 가는 것은 없다.