[NLP] Layer Normalization & Residual Connection — Transformer를 깊게 쌓는 비결 (Pre-LN vs Post-LN)

💡

한 줄 요약 — Residual Connection은 그래디언트와 정보를 깊은 층까지 안정적으로 흘려보내는 고속도로, Layer Normalization은 각 토큰의 활성값 분포를 일정하게 유지하는 수평 보정기. 두 컴포넌트가 함께 쓰여야 깊은 Transformer의 학습이 안정된다. 그리고 순서를 어떻게 두느냐(Pre-LN vs Post-LN)가 학습 난이도를 크게 바꾼다.

이전 글(Positional Encoding)에서 위치 정보를 넣는 방법까지 정리했다. 이번 글에서는 Transformer를 수십, 수백 층까지 깊게 쌓아도 학습이 안정적으로 진행되도록 만드는 두 컴포넌트 — Residual Connection 과 Layer Normalization — 을 본다.

왜 정규화와 잔차가 필요한가

층을 깊게 쌓으면 일반적으로 두 가지 문제가 생긴다.

그래디언트 소실/폭주 — 역전파에서 미분값의 곱이 누적되면서 그래디언트가 0에 수렴하거나 발산한다.
활성값 분포 불안정 — 층마다 입력 분포가 들쭉날쭉해져서 학습률을 정하기 어렵고, 학습이 더뎌진다.

Residual Connection 은 첫 번째 문제를, Layer Normalization 은 두 번째 문제를 직접 겨냥한다.

Residual Connection

ResNet에서 처음 제안된 아이디어. 어떤 서브레이어 $F$ 의 출력을 그대로 내보내지 않고 입력을 그대로 더해서 내보낸다.

$\text{output} = x + F(x)$

별것 아닌 것처럼 보이지만 효과가 크다.

왜 깊은 층에서도 그래디언트가 살아남는가

역전파 시 출력에 대한 입력의 미분은:

$\frac{\partial\, \text{output}}{\partial x} = 1 + \frac{\partial F(x)}{\partial x}$

즉 $F$ 가 어떤 형태든 항상 1이라는 직접 경로가 추가된다. $F$ 의 그래디언트가 곱셈 누적으로 0에 가까워져도, 이 1 덕분에 최소한의 신호는 보장 된다.

직관적으로는 "변화가 어렵다면 그냥 통과시켜라" 는 의미. 모델이 항등 함수(identity) 를 학습하는 것을 매우 쉽게 해 주는 구조다.

residual.py

def residual(x, sublayer):
    return x + sublayer(x)

✅

잔차 = 변화량을 학습 — $F(x) = \text{output} - x$ 이므로 $F$ 는 입력에서 변화해야 할 양 만 학습하면 된다. 변화가 작으면 $F \approx 0$ 이 되어 자연스럽게 input을 통과시킨다. 이 변화 학습 프레임이 깊은 네트워크가 잘 작동하는 핵심 직관.

Layer Normalization

각 토큰 벡터의 feature 차원을 따라 평균과 분산을 계산해서 정규화한다.

$\text{LN}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

$\mu, \sigma^2$ : 한 토큰 벡터 내부에서 계산한 평균/분산 ( $d_{\text{model}}$ 차원에 걸쳐)
$\gamma, \beta$ : 학습 가능한 scale/shift 파라미터 (각각 $d_{\text{model}}$ 차원)
$\epsilon$ : 분모 0 방지용 작은 수 (보통 $10^{-5}$ )

layer_norm.py

def layer_norm(x, gamma, beta, eps=1e-5):
    # x: (batch, seq_len, d_model)
    mu = x.mean(dim=-1, keepdim=True)                    # 마지막 차원(feature)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    return gamma * (x - mu) / torch.sqrt(var + eps) + beta

BatchNorm과 무엇이 다른가

구분	BatchNorm	LayerNorm
정규화 축	배치 축 + 공간 축	feature 축 (한 토큰 내부)
배치 크기 의존성	강함 (작은 배치에서 불안정)	없음
가변 길이 시퀀스	불편함 (마스킹 필요)	자연스러움
학습/추론 차이	running stats 필요	동일

NLP는 시퀀스 길이가 가변적이고 배치 크기도 시나리오에 따라 다르다. 그래서 배치 통계에 의존하지 않는 LayerNorm이 자연스러운 선택이다.

✅

왜 feature 차원만 정규화? 한 토큰의 표현 벡터 내부에서만 평균/분산을 잡기 때문에, 다른 토큰이나 다른 배치 샘플의 영향을 전혀 받지 않는다. 즉 한 토큰을 그 자체로 표준화 한다.

Transformer 블록의 전체 구조

Residual과 LayerNorm을 어떤 순서로 적용하느냐에 따라 두 가지 변형이 생긴다. 이 차이가 학습 난이도에 큰 영향을 준다.

Post-LN (원 논문 방식)

원 Transformer 논문이 채택한 방식.

$y = \text{LN}\bigl(x + \text{SubLayer}(x)\bigr)$

post_ln_block.py

def post_ln_block(x, attn, ff, ln1, ln2):
    x = ln1(x + attn(x))   # attention 후 residual → LN
    x = ln2(x + ff(x))     # FFN 후 residual → LN
    return x

Pre-LN (현대 추세)

GPT-2 이후 대부분의 대형 LLM이 채택한 방식.

$y = x + \text{SubLayer}\bigl(\text{LN}(x)\bigr)$

pre_ln_block.py

def pre_ln_block(x, attn, ff, ln1, ln2):
    x = x + attn(ln1(x))   # LN 먼저, 그 결과를 attention → residual
    x = x + ff(ln2(x))     # LN 먼저, FFN → residual
    return x

두 방식 비교

항목	Post-LN	Pre-LN
학습 안정성	깊어질수록 불안정	안정적
Warmup 스케줄	필수 (없으면 발산)	거의 불필요
최종 성능	잘 튜닝하면 약간 우수	비슷하거나 약간 낮음
깊은 모델 (50층+)	학습 매우 어려움	무리 없이 학습 가능
채택 모델	원 Transformer, BERT	GPT-2 이후 대부분의 LLM

Pre-LN의 핵심 이점은 잔차 경로가 LayerNorm을 거치지 않는 깨끗한 identity 경로가 된다는 점이다. Post-LN은 매 블록마다 LN이 잔차 합산 결과에 적용되므로, 깊어질수록 그 효과가 누적되면서 초반 학습이 불안정해지기 쉽다.

💡

원 논문이 왜 Post-LN을 썼나? 원 논문의 모델은 6층밖에 안 됐고, learning rate warmup으로 충분히 학습이 됐기 때문. Pre-LN의 학습 안정성 우위는 모델이 깊어지면서 드러난 사실이고, 후속 분석 논문들(Xiong et al., 2020 등)이 이를 이론·실험적으로 정리했다.

정리

Residual Connection은 그래디언트와 정보의 고속도로. $y = x + F(x)$ 라는 단순한 식이 깊은 학습을 가능하게 한다.
Layer Normalization은 각 토큰 벡터를 feature 차원으로 표준화. 배치 크기와 시퀀스 길이에 영향받지 않는다.
Post-LN vs Pre-LN — 잔차와 정규화의 순서 가 학습 안정성을 좌우. 현대 대형 LLM은 거의 Pre-LN을 쓴다.
이 두 컴포넌트가 합쳐져서 Transformer를 수십~수백 층까지 깊게 쌓는 게 비로소 가능해진다.

Attention 시리즈에서 메커니즘(Part 1·2), 위치 정보(Part 3), 학습 안정화(Part 4)까지 정리했다. 다음 글에서는 Transformer 블록의 마지막 퍼즐인 Feed-Forward Network를 다룰 예정.

다음 글 예고

Feed-Forward Network — 사실상의 표현력 담당 (왜 Attention 헤드보다 파라미터가 더 많을까?)
더 멀리 가면: 학습 트릭 모음 (Adam + warmup, label smoothing, dropout 위치)