🤖 AI/ML

Layer Normalization

Layer Normalization

레이어 출력을 정규화. Transformer에서 Batch Norm 대신 사용.

📖 상세 설명

Layer Normalization(LayerNorm)은 각 샘플의 특성(feature) 차원을 따라 정규화하는 기법입니다. 2016년 Ba 등이 제안했으며, Transformer의 핵심 구성 요소로 GPT, BERT, LLaMA 등 모든 현대 언어 모델에 사용됩니다.

Batch Normalization과의 차이: BatchNorm은 배치 차원을 따라 정규화하여 배치 크기에 의존하고 추론 시 running statistics가 필요합니다. LayerNorm은 각 샘플 독립적으로 정규화하여 배치 크기와 무관하고, 시퀀스 길이가 가변인 NLP에 적합합니다.

수식: y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta. mean과 var는 hidden dimension을 따라 계산. gamma(scale)와 beta(shift)는 학습 가능한 파라미터로 모델이 정규화 정도를 조절합니다.

변형: Pre-LN(레이어 전에 정규화, GPT-2+), Post-LN(레이어 후 정규화, 원래 Transformer), RMSNorm(mean 제거로 더 빠름, LLaMA). Pre-LN이 학습 안정성이 높아 현재 표준입니다.

💻 코드 예제

import torch
import torch.nn as nn

# 1. PyTorch LayerNorm 사용
layer_norm = nn.LayerNorm(
    normalized_shape=768,  # hidden dimension
    eps=1e-5,             # 수치 안정성
    elementwise_affine=True  # gamma, beta 학습
)

x = torch.randn(32, 100, 768)  # (batch, seq_len, hidden)
output = layer_norm(x)  # 각 (batch, seq) 위치에서 768차원 정규화

# 2. LayerNorm 직접 구현 (이해용)
class ManualLayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(hidden_size))
        self.beta = nn.Parameter(torch.zeros(hidden_size))
        self.eps = eps

    def forward(self, x):
        # x: (batch, seq, hidden)
        mean = x.mean(dim=-1, keepdim=True)  # hidden dim 평균
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta

# 3. RMSNorm (LLaMA 스타일, 더 빠름)
class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        # mean 계산 생략 -> 더 빠름
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

# 4. Pre-LN vs Post-LN Transformer Block
class PreLNBlock(nn.Module):
    """GPT-2+ 스타일: 더 안정적인 학습"""
    def __init__(self, hidden_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(hidden_size)
        self.attn = nn.MultiheadAttention(hidden_size, num_heads=8)
        self.ln2 = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )

    def forward(self, x):
        # Pre-LN: LayerNorm을 먼저 적용
        x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]
        x = x + self.ffn(self.ln2(x))
        return x

🗣️ 실무에서 이렇게 말하세요

💬 회의에서
"학습이 불안정해서 원인 분석 중인데, Post-LN 구조라 gradient가 깊은 레이어에서 폭발하는 것 같습니다. Pre-LN으로 바꾸고 learning rate warmup 추가하면 해결될 거예요."
💬 면접에서
"LayerNorm과 BatchNorm의 차이는 정규화 축입니다. BatchNorm은 배치 차원, LayerNorm은 feature 차원을 정규화해요. NLP에서 시퀀스 길이가 가변이고 배치 크기가 작아서 LayerNorm이 표준입니다."
💬 기술 토론에서
"LLaMA가 RMSNorm 쓰는 이유는 속도에요. mean 계산을 생략해서 15% 정도 빨라지고, 성능 차이는 거의 없습니다. 대규모 모델에서 이런 최적화가 학습 비용에 큰 영향을 줘요."

⚠️ 흔한 실수 & 주의사항

CNN에 LayerNorm 무조건 적용

이미지는 채널 간 통계가 중요해서 BatchNorm이 더 효과적. LayerNorm은 NLP/Transformer에 최적화됨.

eps 값을 너무 작게 설정

float16 학습 시 eps=1e-12는 수치 불안정 유발. float16은 1e-5, bfloat16은 1e-6 권장.

올바른 방법

Transformer는 Pre-LN 구조 권장. 큰 모델은 RMSNorm으로 속도 최적화. normalized_shape은 마지막 차원(들)과 일치시키세요.

🔗 관련 용어

📚 더 배우기