정규화 레이어: 배치 정규화에서 RMSNorm까지

지난 글에서 가중치 초기화가 층을 지나는 신호의 분산을 지켜 준다는 것을 봤다. 그런데 초기화는 학습이 시작되기 직전 한 번만 하는 일이다. 파라미터가 갱신되기 시작하면 각 층이 내보내는 값의 분포가 다시 흔들린다. 이 흔들림을 학습 도중에도 계속 붙잡아 두려고 층 사이에 끼워 넣는 것이 정규화 레이어(Normalization Layer)다.

이 글은 배치 정규화(BatchNorm), 레이어 정규화(LayerNorm), 그룹 정규화(GroupNorm), RMSNorm을 한자리에서 다룬다. 넷을 따로 외울 필요가 없기 때문이다. 네 방법이 쓰는 수식은 사실상 하나이고, 평균과 분산을 어느 축으로 재느냐만 다르다. 축이 정해지면 배치 크기에 의존하는지, 훈련과 추론이 갈리는지, 시퀀스 데이터에 쓸 수 있는지가 전부 따라 나온다. 그래서 순서도 축을 기준으로 잡는다.

정규화 레이어

내부 공변량 이동

신경망을 학습하면 모든 층의 가중치가 동시에 갱신된다. 3번 층 입장에서 보면 자기 가중치만 바뀌는 것이 아니라 입력으로 들어오는 값의 분포 자체가 매 스텝 달라진다. 2번 층이 갱신되었으니 그 출력이 달라지는 것은 당연하다. 이렇게 학습 도중 각 층의 입력 분포가 계속 이동하는 현상을 내부 공변량 이동(Internal Covariate Shift)이라 부른다. 2015년 Ioffe와 Szegedy가 배치 정규화를 제안하면서 붙인 이름이다.

간단한 MLP에 랜덤 입력을 흘려 보기만 해도 층마다 활성값의 통계가 다르다는 것이 보인다. ReLU를 세 번 지난 값의 평균과 표준편차가 층마다 어떻게 달라지는지 재 보자.

import torch
import torch.nn as nn

torch.manual_seed(0)                  # 값을 재현하려면 시드를 고정한다
model = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(),
    nn.Linear(256, 256), nn.ReLU(),
    nn.Linear(256, 256), nn.ReLU(),
)

h = torch.randn(64, 784)
for i, layer in enumerate(model):
    h = layer(h)
    if isinstance(layer, nn.ReLU):
        print(f"ReLU {i//2+1}: mean={h.mean():.2f}, std={h.std():.2f}")
# ReLU 1: mean=0.23, std=0.34
# ReLU 2: mean=0.09, std=0.14
# ReLU 3: mean=0.04, std=0.06

층을 지날수록 평균과 표준편차가 함께 줄어든다. 줄어드는 속도가 만만치 않다 — 표준편차가 한 층을 지날 때마다 대략 2.4배씩 작아져서, 세 층 만에 0.34가 0.06이 된다. 층을 열 개만 더 쌓으면 활성값이 사실상 0으로 수렴한다는 뜻이다. 시드를 바꾸면 소수점 아래가 조금씩 달라지지만 줄어드는 방향과 배율은 그대로다.

원인은 두 군데에 나뉘어 있다. 하나는 ReLU 자체다. 평균 0 근처의 값에서 음수 절반을 0으로 눌러 버리므로 통과한 값의 제곱평균이 절반으로 준다. 다른 하나는 PyTorch nn.Linear의 기본 초기화다. 입력 차원이 $n$ 일 때 가중치를 $U(-1/\sqrt{n},\ 1/\sqrt{n})$ 에서 뽑는데, 이 분산은 ReLU를 전제한 He 초기화보다 작아서 신호가 층마다 조금씩 깎인다. 둘을 합치면 위와 같은 감쇠가 나온다.

여기서 중요한 것은 숫자 자체가 아니라 이 값이 학습 전 초기 상태의 숫자라는 점이다. 가중치가 갱신되기 시작하면 층마다 깎이는 정도가 스텝마다 달라지고, 어떤 층은 반대로 커지기도 한다. 뒤쪽 층은 매번 스케일이 다른 입력을 받으며 학습하는 셈이다.

초기화의 한계

Xavier나 He 초기화는 이 문제의 절반만 해결한다. 초기화 공식은 「층을 통과할 때 분산이 유지되도록 초기 가중치의 스케일을 맞춘다」는 것인데, 그 계산은 가중치가 아직 랜덤이라는 가정 위에 서 있다. 한 번 학습이 시작되면 가중치는 더 이상 랜덤이 아니고, 분산 보존 조건도 함께 깨진다.

깨진 결과가 실무에서 드러나는 방식은 두 가지다. 첫째는 학습률을 크게 잡을 수 없다는 것이다. 어느 층의 활성값이 유난히 커져 있으면 그 층의 기울기도 함께 커지고, 큰 학습률을 곱한 갱신이 그 층을 발산시킨다. 그래서 모든 층이 견딜 만한 작은 학습률로 낮춰 잡게 되고, 학습이 느려진다. 둘째는 초기화에 지나치게 민감해진다는 것이다. 같은 구조를 시드만 바꿔 돌렸는데 어떤 시드는 수렴하고 어떤 시드는 손실이 발산한다면, 그 모델은 초기 분산의 좁은 구간에 겨우 얹혀 있는 상태다.

정규화 레이어의 발상은 여기서 나온다. 초기값을 잘 고르는 대신, 층마다 값을 강제로 다시 맞춰 놓는다. 앞 층이 무엇을 내보내든 다음 층이 받는 값은 평균 0, 분산 1 근처로 고정되므로 뒤쪽 층은 안정적인 입력을 전제로 학습할 수 있다.

정규화의 공통 식

네 방법이 공유하는 식은 네 줄이다. 정규화할 값들의 집합을 하나 정하고, 그 집합의 평균과 분산을 구해 표준화한 뒤, 학습 가능한 두 파라미터로 다시 늘리고 옮긴다.

$$\begin{aligned} \mu &= \frac{1}{m}\sum_{i=1}^{m} x_i \\ \sigma^2 &= \frac{1}{m}\sum_{i=1}^{m}(x_i - \mu)^2 \\ \hat x_i &= \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \\ y_i &= \gamma\,\hat x_i + \beta \end{aligned}$$

앞의 세 줄은 통계 시간에 배우는 표준화 그대로다. $\epsilon$ 은 분모가 0이 되는 것을 막는 아주 작은 값이고 보통 $10^{-5}$ 를 쓴다. 눈여겨볼 것은 네 번째 줄이다.

$\gamma$ (감마)와 $\beta$ (베타)는 학습되는 파라미터다. 세 번째 줄에서 애써 평균 0, 분산 1로 맞춰 놓고 왜 다시 늘리고 옮기느냐는 물음이 자연스럽게 나오는데, 이유는 표현력이다. 시그모이드를 예로 들면 입력이 0 근처인 구간은 거의 직선이라, 모든 활성값을 그 구간에 몰아넣으면 비선형성이 사라진다. 정규화가 오히려 신경망을 얕게 만드는 셈이다. $\gamma$ 와 $\beta$ 는 「이 층에는 표준정규분포가 최적이 아니다」라고 판단할 여지를 모델에 되돌려 준다. 초기값은 $\gamma = 1$, $\beta = 0$ 이므로 학습 시작 시점에는 순수한 표준화로 출발하고, 필요한 만큼만 데이터가 그 값을 밀어낸다.

이 네 줄에서 아직 정해지지 않은 것은 「$m$ 개가 무엇이냐」 하나뿐이다. 그 집합을 배치 방향으로 잡으면 배치 정규화가 되고, 특성 방향으로 잡으면 레이어 정규화가 된다. 다음 절이 그 이야기다.

정규화의 축

배치 방향과 특성 방향

크기 $(N, C)$ 의 텐서를 떠올려 보자. $N$ 은 배치 안의 샘플 수, $C$ 는 특성 수다. 이 표에서 평균을 낼 방법은 크게 둘이다.

배치 정규화는 세로로 잰다. 특성 하나를 고정하고 배치 안의 모든 샘플에 대해 평균과 분산을 구한다. 「이 배치에서 3번 특성의 평균은 얼마인가」를 묻는 것이므로, 통계는 특성마다 하나씩 총 $C$ 개가 나오고 $\gamma$ 와 $\beta$ 도 각각 $C$ 개다. 각 특성이 서로 다른 스케일을 가지고 있어도(픽셀 밝기와 나이처럼) 특성별로 따로 맞춰지므로 서로 간섭하지 않는다.

레이어 정규화는 가로로 잰다. 샘플 하나를 고정하고 그 샘플이 가진 $C$ 개 특성에 대해 평균과 분산을 구한다. 「이 토큰의 512차원 임베딩의 평균은 얼마인가」를 묻는 것이므로 통계는 샘플마다 하나씩 나오고, 옆에 어떤 샘플이 있든 결과가 달라지지 않는다. 2016년 Ba 등이 제안한 이 방식이 지금 GPT·BERT를 포함한 사실상 모든 대규모 언어 모델의 기본값이다. LLaMA 계열이 쓰는 RMSNorm도 축은 이것과 같다 — 같은 가로 방향에서 평균 빼기 한 단계만 덜어 낸 것이고, 마지막 절에서 따로 본다.

x = torch.randn(4, 128, 512)          # (Batch, Seq_len, d_model)

ln = nn.LayerNorm(512)                # 마지막 차원만 정규화한다
out = ln(x)                           # (4, 128, 512), 모양은 그대로

print(out[0, 0].mean().item())        # ≈ 0.0  토큰 하나 안에서 평균 0
print(out[0, 0].std().item())         # ≈ 1.0

배치 의존성

축이 정해지면 나머지는 전부 따라 나온다. 그 첫 번째가 배치 의존성이다.

배치 정규화가 쓰는 통계는 배치를 가로질러 계산되므로, 같은 샘플이라도 어떤 샘플들과 함께 배치에 들어갔느냐에 따라 출력이 달라진다. 이 사실은 두 가지를 강제한다. 하나는 배치 크기가 충분해야 통계가 믿을 만하다는 것이고(다음다음 절), 다른 하나는 추론 시점에 쓸 통계를 따로 마련해 두어야 한다는 것이다. 서비스에서는 샘플 하나만 들어오는 일이 흔한데, 그때는 배치 분산이 0이라 정규화 자체가 성립하지 않는다.

레이어 정규화는 이 문제가 아예 없다. 통계가 샘플 하나 안에서 닫혀 있으므로 배치 크기가 1이든 512든 같은 출력이 나오고, 훈련과 추론이 같은 식을 쓴다. 저장해 둘 이동 평균도, 모드 전환도 없다. 정규화 레이어를 고를 때 실제로 갈리는 지점의 대부분이 이 한 줄에서 나온다.

정규화 변형

축을 바꾸는 방식으로 만들어진 변형이 몇 개 더 있다. 이미지 텐서 $(N, C, H, W)$ 에서 보면 차이가 분명하다 — $H$ 와 $W$ 는 세로·가로 픽셀 위치다.

방법	통계를 재는 축	주요 용도	배치 의존성
BatchNorm	N·H·W (채널마다)	CNN, MLP	있다
LayerNorm	C (샘플마다)	트랜스포머, RNN	없다
InstanceNorm	H·W (샘플·채널마다)	스타일 트랜스퍼	없다
GroupNorm	채널 G그룹 × H·W	소배치 CNN, 객체 탐지	없다
RMSNorm	C (평균 빼기 없이)	LLM (LLaMA 등)	없다

인스턴스 정규화(InstanceNorm)는 이미지 한 장의 채널 하나만 놓고 정규화한다. 스타일 트랜스퍼에서 쓰는 이유가 있다 — 그 채널의 평균과 분산이 곧 「이 이미지의 색조와 대비」에 해당하므로, 그것을 지우면 내용만 남고 스타일이 씻겨 나간다. 그룹 정규화(GroupNorm)는 그 표의 두 줄 사이를 잇는 손잡이다. 그룹을 1개로 두면 레이어 정규화가 되고 채널 수만큼 두면 인스턴스 정규화가 되며, 그 사이 어디쯤을 고르는 것이 실제 쓰임이다(자세한 것은 뒤의 「그룹 정규화」에서 본다).

표를 세로로 읽으면 규칙이 하나 보인다. 통계를 재는 축에 $N$ 이 들어간 것은 배치 정규화 하나뿐이고, 그래서 배치 의존성 칸이 「있다」인 것도 그 한 줄뿐이다. 나머지 넷은 축을 샘플 하나 안에 가둬 두었기 때문에 배치 크기와 무관하다. 넷을 각각 외우는 대신 「이 방법은 샘플 바깥을 보는가」 하나만 물으면 표의 오른쪽 절반이 저절로 채워진다.

훈련 모드와 추론 모드

배치 통계

배치 정규화 층은 훈련 중에 두 가지 일을 동시에 한다. 하나는 지금 들어온 배치의 평균 $\mu_B$ 와 분산 $\sigma^2_B$ 로 그 배치를 정규화하는 것이고, 다른 하나는 나중에 추론에서 쓸 통계를 몰래 쌓아 두는 것이다. 이 두 번째가 자주 잊힌다.

쌓는 방식은 지수 이동 평균이다. 매 스텝 아래 갱신이 일어난다.

$$\mu_{\text{run}} \leftarrow (1 - \text{momentum}) \cdot \mu_{\text{run}} + \text{momentum} \cdot \mu_B$$

PyTorch의 기본 momentum은 0.1이므로 새 배치의 통계가 10%만큼 섞인다. 대략 최근 수십 배치의 평균을 들고 있는 값이라고 보면 된다. 이 값들은 running_mean과 running_var라는 이름의 버퍼로 저장되며, 기울기가 흐르지 않는다 — 학습되는 것이 아니라 관측되어 누적되는 값이다. 그래서 state_dict에는 들어가지만 parameters()에는 안 나온다.

model = nn.Sequential(
    nn.Linear(784, 256), nn.BatchNorm1d(256), nn.ReLU(),
    nn.Linear(256, 10),
)
bn = model[1]

print(bn.weight.shape)        # torch.Size([256])  γ, 특성마다 하나
print(bn.bias.shape)          # torch.Size([256])  β
print(bn.running_mean[:3])    # 학습 중 누적되는 버퍼
print(bn.momentum)            # 0.1

FC 층 뒤에는 BatchNorm1d, 합성곱 층 뒤에는 BatchNorm2d를 쓴다. BatchNorm2d는 채널마다 $N \times H \times W$ 개 값을 모아 통계를 내므로, 배치가 32여도 28×28 특성 맵이면 채널당 25,088개 값에서 평균을 구하는 셈이다. CNN에서 배치 정규화가 유난히 잘 통하는 이유 중 하나다.

이동 평균

model.eval()을 호출하면 배치 정규화 층은 배치 통계를 무시하고 누적해 둔 running_mean·running_var를 쓴다. 이 전환이 필요한 이유는 추론의 성질에 있다. 추론은 샘플 하나에 대해 결정론적인 답을 내야 하는데, 배치 통계를 쓰면 같은 입력이라도 함께 들어온 이웃에 따라 답이 달라진다. 열 명의 요청을 묶어 처리하느냐 하나씩 처리하느냐로 결과가 바뀌는 서비스는 디버깅이 불가능하다.

전환은 재귀적이라 모델 최상단에서 한 번 부르면 안쪽 모든 층에 적용된다. 드롭아웃도 같은 스위치를 공유하므로, 두 층을 함께 쓰는 모델에서는 이 한 줄이 두 가지를 동시에 바꾼다.

model.train()                 # 배치 통계 사용 + running 통계 갱신
out = model(x_batch)

model.eval()                  # running 통계 사용, 갱신 없음
with torch.no_grad():
    out = model(x_single)     # 샘플 하나여도 안전하다

eval 전환 누락

배치 정규화에서 가장 흔한 버그가 이 전환을 빠뜨리는 것이고, 증상이 원인과 멀어서 찾기 어렵다.

가장 극적인 경우는 배치 크기 1로 추론할 때다. 샘플 하나의 배치 분산은 정의상 0이므로 정규화 식의 분모가 $\sqrt{\epsilon}$ 만 남고, 분자도 0이라 출력이 통째로 $\beta$ 가 된다. 입력이 무엇이든 같은 값이 나오는 것이다. BatchNorm1d는 이 경우를 알고 있어서 훈련 모드에서 배치 크기 1이 들어오면 아예 에러를 던지지만, BatchNorm2d는 $H \times W$ 덕분에 값이 여러 개라 조용히 넘어간다.

더 흔한 것은 검증 정확도가 훈련 정확도보다 이상하게 낮거나 에폭마다 크게 튀는 경우다. 검증 루프에서 model.eval()을 안 불렀다면, 검증 배치의 통계로 정규화하는 것에 더해 검증 데이터가 running_mean을 오염시킨다. 검증셋의 통계가 훈련 통계에 섞여 들어가므로 다음 에폭의 훈련까지 영향을 받는다. 반대로 훈련 루프로 돌아오면서 model.train()을 다시 안 부르면 이번에는 훈련이 고정된 통계로 진행되어 손실이 잘 안 내려간다. 두 줄이 짝이라는 것을 기억해 두면 절반은 막힌다.

배치 크기의 한계

소배치의 흔들림

배치 정규화의 통계는 표본 통계다. 표본 평균의 표준오차가 $\sigma / \sqrt{N}$ 이므로, 배치 크기가 4분의 1이 되면 통계의 흔들림은 2배가 된다. 배치 2에서 구한 「평균」은 두 값의 중점일 뿐이고, 그 값으로 정규화한 결과는 이웃 샘플이 무엇이냐에 따라 크게 달라진다.

이것을 코드로 확인할 때 한 번 헛짚기 쉬운 자리가 있다. 정규화된 출력의 표준편차를 재면 안 된다. 훈련 모드의 배치 정규화는 특성마다 평균 0, 분산 1로 맞춰 내보내도록 정의되어 있으므로, 배치가 2든 128이든 out.std()는 언제나 1.0 근처가 나온다. 정의를 확인하는 것일 뿐 아무것도 재지 못한다. 흔들리는 것은 출력이 아니라 그 출력을 만든 통계 쪽이다. 그러니 같은 분포에서 배치를 여러 번 새로 뽑아, 특성별 배치 평균이 시행마다 얼마나 달라지는지를 본다.

torch.manual_seed(0)
for B in [2, 8, 32, 128]:
    means = torch.stack([torch.randn(B, 256).mean(0) for _ in range(200)])
    print(f"B={B:3d}: 배치 평균의 흔들림 {means.std(0).mean():.2f}")
# B=  2: 배치 평균의 흔들림 0.71
# B=  8: 배치 평균의 흔들림 0.35
# B= 32: 배치 평균의 흔들림 0.18
# B=128: 배치 평균의 흔들림 0.09

입력이 표준정규분포이므로 $\sigma = 1$ 이고, 나온 값은 $1/\sqrt{B}$ 그대로다 — 0.71은 $1/\sqrt{2}$, 0.09는 $1/\sqrt{128}$ 이다. 배치 2에서는 정규화의 기준점이 매 스텝 ±0.7씩 제자리를 옮기는 셈이고, 배치 128에서는 그 폭이 ±0.09로 준다. 같은 샘플을 같은 가중치에 통과시켜도 배치 2에서는 앞뒤 스텝의 출력이 눈에 띄게 다르다는 뜻이다. 그리고 이 흔들림은 훈련에서 끝나지 않는다 — running_mean은 그 흔들리는 값들의 이동 평균이므로, 배치가 작으면 추론에서 쓸 기준점도 함께 부정확해진다.

이 흔들림이 전부 나쁘기만 한 것은 아니다. 매 스텝 조금씩 다른 통계로 정규화되므로 활성값에 잡음이 섞이고, 그 잡음이 약한 정규화(regularization) 효과를 낸다. 배치 정규화를 넣으면 드롭아웃 비율을 낮춰도 되는 경우가 많은 이유다. 하지만 배치가 8보다 작아지면 잡음이 신호를 덮기 시작하고, 훈련 통계와 추론 통계의 차이도 벌어져 검증 성능이 무너진다. 경험적인 하한이 32이고, 그 아래로 내려가야 하는 상황이면 다른 정규화를 쓴다.

작은 배치는 취향이 아니라 제약인 경우가 많다. 고해상도 세그멘테이션이나 3D 의료 영상처럼 샘플 하나가 GPU 메모리를 크게 먹는 과제에서는 배치 2~4가 한계다. 기울기 누적으로 유효 배치를 키워도 소용이 없다 — 누적은 기울기를 더할 뿐이고, 정규화 통계는 순전파 시점의 실제 배치에서 계산되기 때문이다.

시퀀스 데이터

배치 크기와 별개로, 자연어 처리에서는 축 자체가 안 맞는다.

$(N, T, C)$ 모양의 텍스트 배치에 배치 정규화를 걸면 통계를 $N$ 과 $T$ 를 가로질러 계산하게 된다. 즉 서로 다른 문장의 서로 다른 위치에 있는 토큰들을 한 통에 넣고 평균을 낸다. 5번째 토큰이 어떤 문장에서는 명사이고 다른 문장에서는 조사인데, 그 자리들의 임베딩을 함께 평균 내는 것에는 별 의미가 없다.

여기에 길이 문제가 겹친다. 문장 길이가 제각각이라 짧은 문장은 패딩으로 채우는데, 패딩 토큰까지 통계에 들어가면 배치에 짧은 문장이 몇 개 섞였느냐가 정규화 결과를 바꾼다. 추론 때 길이 200짜리 문장이 들어오면 훈련에서 본 적 없는 통계 분포가 되고, 누적해 둔 running_mean은 그 상황을 대표하지 못한다. 레이어 정규화가 언어 모델의 기본값이 된 것은 유행이 아니라 축의 문제다.

그룹 정규화

CNN에서 배치를 키울 수 없을 때 자리를 대신하는 것이 그룹 정규화다. 채널 $C$ 개를 $G$ 개 그룹으로 나누고 각 그룹 안에서 $C/G \times H \times W$ 개 값의 통계를 낸다. 배치를 가로지르지 않으므로 배치 크기와 무관하고, 채널을 전부 한 통에 넣지도 않으므로 채널마다 스케일이 크게 다른 합성곱 특성 맵에 레이어 정규화보다 잘 맞는다.

gn = nn.GroupNorm(num_groups=32, num_channels=256)   # 그룹당 8채널
x = torch.randn(2, 256, 56, 56)                      # 배치가 2뿐이어도 무방하다
out = gn(x)

그룹 수는 32를 기본으로 두고 채널 수에 맞춰 조정한다. 원 논문의 실험에서 배치 32에서는 배치 정규화가 조금 앞서지만, 배치 2에서는 그룹 정규화가 오차율 기준으로 10%포인트 넘게 앞선다. 배치가 충분하면 배치 정규화, 부족하면 그룹 정규화가 CNN 쪽의 실용적인 갈림길이다.

레이어 정규화와 트랜스포머

토큰 단위 통계

레이어 정규화를 트랜스포머에 쓸 때 실제로 일어나는 일은 이렇다. $(4, 128, 512)$ 배치라면 정규화 단위가 512차원 벡터 하나이고, 그런 벡터가 $4 \times 128 = 512$ 개 있으므로 평균과 분산을 512쌍 따로 구한다. 각 토큰은 자기 임베딩 안에서만 표준화되고 옆 토큰도 다른 문장도 쳐다보지 않는다.

이 성질이 자기회귀 생성에서 특히 중요하다. GPT 계열은 토큰을 하나씩 만들어 내면서 앞의 결과를 다시 입력으로 넣는데, 정규화가 시퀀스를 가로지르면 「아직 만들지 않은 뒤쪽 토큰」의 통계가 앞쪽 계산에 새어 들어간다. 어텐션에는 마스크를 씌워 미래를 못 보게 막아 놓고 정규화가 그 옆으로 미래를 흘려보내는 꼴이라, 학습 때의 계산과 한 토큰씩 생성할 때의 계산이 달라진다. 레이어 정규화는 통계가 토큰 하나 안에서 닫혀 있으므로 그런 누출이 구조적으로 불가능하다.

$\gamma$ 와 $\beta$ 는 여전히 특성 차원 크기, 즉 512개다. 통계는 토큰마다 따로 내지만 되돌리는 손잡이는 모든 토큰이 공유한다. 「이 층에서 37번 차원은 조금 크게 유지하는 편이 낫다」 같은 지식이 위치와 무관하게 배워진다.

Post-LN과 Pre-LN

정규화 레이어를 넣는 것만큼 어디에 넣느냐도 결과를 바꾼다. 2017년 원래 트랜스포머 논문은 서브레이어를 지나고 잔차를 더한 뒤에 정규화했다. 이것이 Post-LN이다.

$$\text{Post-LN}: \quad x \leftarrow \mathrm{LN}(x + \text{Sublayer}(x))$$ $$\text{Pre-LN}: \quad x \leftarrow x + \text{Sublayer}(\mathrm{LN}(x))$$

이후 연구에서 Pre-LN이 훨씬 안정적으로 학습된다는 것이 밝혀졌고, GPT-2 이후 사실상 모든 대형 모델이 Pre-LN을 쓴다. CNN 쪽에서 정규화를 활성화 함수 앞에 두는 원래 배치(Conv → BN → ReLU)가 여전히 잘 통하는 것과 대비된다 — 잔차 연결이 깊게 쌓이는 구조에서만 이 차이가 크게 벌어진다.

class PreNormBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        h = self.norm1(x)                                  # 정규화가 먼저다
        x = x + self.drop(self.attn(h, h, h, attn_mask=mask)[0])
        x = x + self.drop(self.ff(self.norm2(x)))
        return x

잔차 경로와 워밍업

Pre-LN이 안정적인 이유는 잔차 경로에 정규화가 끼지 않는다는 한 가지로 설명된다.

Post-LN에서는 입력이 출력에 닿기까지 모든 층의 정규화를 통과해야 한다. 역전파에서 기울기도 그 정규화들을 전부 거꾸로 지나야 하고, 층마다 스케일이 조금씩 곱해지면서 수십 층이 쌓이면 그 곱이 크게 어긋난다. Pre-LN에서는 $x + \text{Sublayer}(\mathrm{LN}(x))$ 의 $x$ 항이 아무것도 안 거치고 그대로 흐르므로, 아래쪽 층까지 기울기가 감쇠 없이 도착한다.

실무에서 이 차이가 나타나는 자리는 워밍업이다. Post-LN 트랜스포머는 학습 초반에 학습률을 0에서 서서히 올리는 워밍업 없이는 대개 발산한다. 초기의 큰 갱신이 정규화를 통과하며 증폭되기 때문이다. Pre-LN은 워밍업 없이도 학습이 시작되고, 있어도 훨씬 짧게 두면 된다. 학습률을 몇 배 크게 잡을 수 있다는 뜻이기도 하다. 다만 공짜는 아니다 — Pre-LN은 깊이가 늘어도 잔차 경로로 값이 계속 더해져 출력의 크기가 층수에 비례해 커지는 경향이 있어, 마지막 블록 뒤에 정규화를 한 번 더 두는 것이 관례다.

RMSNorm과 선택 기준

제곱평균제곱근

RMSNorm(Root Mean Square Normalization)은 2019년 Zhang과 Sennrich가 제안한 레이어 정규화의 축소판이다. 네 줄짜리 뼈대에서 평균 빼기와 $\beta$ 를 통째로 지웠다.

$$\mathrm{RMSNorm}(x) = \frac{x}{\mathrm{RMS}(x)} \cdot \gamma, \qquad \mathrm{RMS}(x) = \sqrt{\frac{1}{C}\sum_{i=1}^{C} x_i^2}$$

$C$ 는 앞에서와 같은 특성 차원 크기, 트랜스포머라면 히든 차원이다(원 논문은 이 자리를 $H$ 로 적는다). 제곱평균제곱근(RMS)은 값들을 제곱해 평균 낸 뒤 다시 제곱근을 씌운 것으로, 부호를 지운 「평균 크기」에 해당한다. 평균을 빼지 않았으므로 분산이 아니라 원점에서 잰 크기라는 점만 다르다.

남은 것은 크기 조정뿐이다. 벡터를 제곱평균제곱근으로 나누면 방향은 그대로 두고 길이만 1 근처로 맞춰진다. 원 논문의 주장은 레이어 정규화가 주는 이득의 대부분이 중심 이동이 아니라 스케일 고정에서 온다는 것이었고, 실험에서 여러 과제의 성능이 사실상 같게 나왔다. 지금은 LLaMA·Gemma·Mistral·Qwen이 모두 이 방식을 쓴다.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))   # γ만 있고 β는 없다

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
        return x / rms * self.weight

메모리 대역폭 이득

식이 짧아진 만큼 무엇이 실제로 절약되는지는 따져 볼 만하다. 곱셈·덧셈 횟수로는 큰 차이가 아니다 — 평균을 구하는 한 번의 합산과 원소별 뺄셈 한 번이 빠질 뿐이다. 이득이 나오는 곳은 산술이 아니라 메모리 접근이다.

레이어 정규화는 평균을 구하려 벡터를 한 번 훑고, 분산을 구하려 다시 훑는다. RMSNorm은 제곱합을 위해 한 번만 훑는다. 정규화 층은 계산량 자체가 작아 GPU 연산기가 아니라 메모리 대역폭에 묶여 있으므로, 훑는 횟수가 곧 시간이다. 원 논문이 과제와 구현에 따라 7%에서 64%까지 학습 시간이 줄었다고 보고했는데, 그 폭이 이렇게 넓은 이유도 여기 있다 — 히든 차원이 크고 커널이 잘 융합된 구현일수록 정규화 층이 차지하는 몫이 작아 이득이 묻히고, 층 수가 많고 배치가 작을수록 그 몫이 커져 이득이 그대로 드러난다. 그래서 「RMSNorm으로 바꾸면 몇 퍼센트 빨라진다」는 값은 그대로 옮겨 쓸 수 있는 숫자가 아니다. 자기 모델에서 정규화 층이 전체 시간의 몇 퍼센트인지를 먼저 재 보면 기대치의 상한이 바로 나온다.

$\beta$ 가 없어진 것도 부수적으로 도움이 된다. 파라미터가 절반으로 줄고, 「정규화 층의 편향은 가중치 감쇠에서 빼야 한다」는 잔손질도 하나 줄어든다. 다만 RMSNorm은 레이어 정규화의 상위 호환이 아니다. 입력의 평균이 0에서 크게 벗어나 있는 데이터에서는 중심 이동이 실제로 일하고, 그 경우에는 평균 빼기를 지운 만큼 손해가 난다. 트랜스포머 안쪽 활성값이 대체로 0 근처에 몰려 있기 때문에 통하는 절충이다.

상황별 선택

앞의 이야기를 고르는 순서로 접으면 물음 두 개로 줄어든다. 첫째, 시퀀스인가. 그렇다면 레이어 정규화이고, 처음부터 크게 학습할 계획이면 RMSNorm에 Pre-LN 배치다. 둘째, 배치를 32 이상 잡을 수 있는가. CNN에서 그렇다면 배치 정규화가 여전히 가장 강하고, 아니라면 그룹 정규화다.

상황	고르는 것	배치할 자리
CNN, 배치 32 이상	BatchNorm	Conv → BN → ReLU
CNN, 배치 8 이하	GroupNorm	같은 자리
트랜스포머·RNN	LayerNorm	서브레이어 앞 (Pre-LN)
대규모 LLM	RMSNorm	서브레이어 앞 + 마지막에 한 번
스타일 트랜스퍼	InstanceNorm	생성기 각 블록

한 가지 덧붙일 것은 정규화 층 바로 앞의 편향은 지워도 된다는 점이다. 정규화가 평균을 빼는 순간 그 편향이 사라지므로 nn.Linear(..., bias=False)로 두면 파라미터가 조금 줄고 결과는 같다. RMSNorm처럼 평균을 안 빼는 경우는 예외다.

정규화 레이어는 활성값의 분포를 고르게 만들어 학습을 안정시키는 것이 본업이고, 배치 통계의 잡음 덕에 과적합을 조금 눌러 주는 것은 부수 효과다. 그 부수 효과를 본업으로 삼는 층이 따로 있다 — 학습 중에 뉴런 일부를 무작위로 꺼서 매 스텝 다른 서브네트워크를 훈련시키고, 결과적으로 수많은 모델의 앙상블처럼 동작하게 만드는 장치다. 다음 글에서 그 원리와, 왜 그것 역시 훈련과 추론에서 다르게 동작해야 하는지를 본다.

---

읽어주셔서 감사합니다. 😊