RMSNorm

Root mean square layer normalization, RMSNorm; RMS 정규화

레이어 정규화에서 평균을 빼는 단계를 없애고, 벡터를 제곱평균제곱근(RMS)으로만 나눠 크기를 맞추는 정규화 방식

RMSNorm(Root Mean Square Layer Normalization)은 레이어 정규화(LayerNorm)를 단순화한 정규화 층이다. LayerNorm이 평균을 빼고(재중심화) 표준편차로 나누는(재스케일링) 두 단계를 하는 데 비해, RMSNorm은 평균을 빼지 않고 RMS로만 나눈다. 2019년 에든버러 대학의 장뱌오(Biao Zhang)와 리코 센리히(Rico Sennrich, 에든버러·취리히 대학)가 NeurIPS 2019에서 발표했다. 논문은 LayerNorm의 효과가 재중심화가 아니라 재스케일링에서 온다고 가정했고, 여러 과제에서 LayerNorm과 비슷한 성능에 실행 시간을 7~64% 줄였다고 보고했다.

계산이 단순하고 편향 항도 없어, 메타 LLaMA 계열, 구글 젬마, 딥시크 등 2020년대 대형 언어 모델 상당수가 LayerNorm 대신 RMSNorm을 쓴다. 이 문서는 RMSNorm 자체와 함께, 트랜스포머에서 정규화를 어디에 두느냐(Post-LN, Pre-LN과 그 변형)도 다룬다. 정규화 층의 종류보다 위치가 학습 안정성에 더 큰 영향을 주기 때문이다.

수식

입력 벡터 $a\in {ℝ}^n$에 대해 LayerNorm은

$\mu =\frac{1}{n}\sum _{i=1}^n{a}_i,\sigma =\sqrt{\frac{1}{n}\sum _{i=1}^n(a_i-\mu {)}^2},{\overline{a}}_i=\frac{a_i-\mu }{\sigma }{g}_i+b_i$

이고, RMSNorm은

$\mathrm{RMS}(a)=\sqrt{\frac{1}{n}\sum _{i=1}^n{a}_i^2},{\overline{a}}_i=\frac{a_i}{\mathrm{RMS}(a)}{g}_i$

이다. 구현에서는 0으로 나누지 않도록 $\sqrt{\frac{1}{n}\sum a_i^2+ϵ}$을 쓴다. 학습 매개변수는 이득(gain) $g$ 하나이고 편향 $b$는 보통 두지 않는다. 입력의 평균이 0이면 RMS는 표준편차와 같아지므로 RMSNorm은 LayerNorm과 정확히 같아진다.

RMSNorm을 거친 벡터는 길이가 $\sqrt{n}$인 구 위에 놓인다(이득을 곱하기 전). 방향 정보는 그대로 두고 크기만 고정하는 셈이다.

불변성

논문은 정규화 방식별로 어떤 변환에 출력이 변하지 않는지를 정리했다.

방식	가중치 행렬 재스케일링	가중치 행렬 재중심화	가중치 벡터 재스케일링	데이터셋 재스케일링	데이터셋 재중심화	개별 표본 재스케일링
배치 정규화	불변	아님	불변	불변	불변	아님
LayerNorm	불변	불변	아님	불변	아님	불변
RMSNorm	불변	아님	아님	불변	아님	불변

RMSNorm은 $\mathrm{RMS}(\alpha x)=\alpha \mathrm{RMS}(x)$라는 선형성 덕분에 크기 변화에는 불변이지만, 모든 원소에 같은 값을 더하는 이동에는 불변이 아니다. 논문은 이 재중심화 불변성이 은닉 상태나 기울기의 분산을 줄이지 않으므로 없어도 된다고 주장했다. 또 가중치 크기에 반비례하는 기울기 덕분에 암묵적으로 학습률이 조절되는 효과가 있다고 분석했다.

속도와 부분 RMSNorm

논문 실험(텐서플로·파이토치·테아노 구현, 번역용 GRU 기반 RNNSearch와 트랜스포머, CNN/Daily Mail 독해, 이미지-캡션 검색, CIFAR-10 분류)에서 RMSNorm은 LayerNorm과 비슷한 성능을 내며 실행 시간을 모델에 따라 7~64% 줄였다. 원소 일부($p$%)만으로 RMS를 추정하는 부분 RMSNorm(pRMSNorm)도 제안했는데, 6.25%만 써도 비슷한 결과를 냈지만 비율이 작으면 기울기가 불안정해지는 문제가 있었다. 오늘날 GPU에서 정규화 층은 보통 다른 연산과 융합(fusion)되므로 절감폭은 이보다 작을 수 있고, 대형 모델이 RMSNorm을 고르는 데는 편향이 없고 구현이 단순하다는 점도 작용한다.

구현

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))   # 이득 g, 편향 없음

    def forward(self, x):
        x32 = x.float()                                # bf16/fp16 입력도 통계는 float32로
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)

x = torch.randn(4, 16, 512)
print((RMSNorm(512)(x) - nn.RMSNorm(512, eps=1e-6)(x)).abs().max())   # 0 (PyTorch 2.4+)

구현마다 다른 세부 사항이 있어 체크포인트를 옮길 때 주의해야 한다.

구현	특징
PyTorch torch.nn.RMSNorm	2.4부터 제공. eps를 주지 않으면 입력 자료형의 머신 엡실론을 쓴다.
메타 LLaMA 공식 코드	float32로 정규화한 뒤 원래 자료형으로 바꾸고 이득을 곱한다. 라마 3 기본 norm_eps는 10-5.
허깅 페이스 LlamaRMSNorm	위와 같은 순서. 코드 주석에 "T5LayerNorm과 같다"고 적혀 있다.
허깅 페이스 GemmaRMSNorm	이득을 0으로 초기화하고 $(1+g)$를 곱한다. 자료형 변환 전에 이득을 곱해 라마와 반올림 순서가 다르다.

구글 T5(2019)는 RMSNorm이라는 이름을 쓰지 않았지만 "활성값을 재스케일링만 하고 편향을 더하지 않는 단순화된 레이어 정규화"를 썼는데, 이것이 사실상 RMSNorm과 같은 연산이다.

정규화의 위치: Post-LN과 Pre-LN

트랜스포머 블록에서 정규화를 잔차 연결의 어디에 두느냐에 따라 학습 특성이 크게 달라진다. $F$를 어텐션 또는 FFN 부분층, $x_l$을 $l$번째 층 입력이라 하자.

• Post-LN (원 트랜스포머, BERT): $x_{l+1}=\mathrm{Norm}(x_l+F(x_l))$
• Pre-LN (GPT-2 이후 대부분의 LLM): $x_{l+1}=x_l+F(\mathrm{Norm}(x_l))$, 마지막 층 뒤에 최종 정규화를 한 번 더 둔다.

GPT-2는 레이어 정규화를 각 부분 블록의 입력 쪽으로 옮기고 마지막 셀프 어텐션 블록 뒤에 정규화를 하나 더했다. 시옹(Ruibin Xiong) 등은 평균장 이론으로, Post-LN은 초기화 직후 출력층 근처 매개변수의 기울기 기댓값이 커서 큰 학습률을 쓰면 불안정하고 그래서 학습률 워밍업이 꼭 필요한 반면, Pre-LN은 기울기가 잘 정돈되어 워밍업 없이도 학습된다는 것을 보였다. LLaMA는 GPT-3를 따라 각 부분층 입력을 정규화하는 Pre-LN에 RMSNorm을 썼다.

Pre-LN의 대가는 잔차 스트림 $x_l$이 정규화 없이 층마다 더해지기만 한다는 점이다. 깊어질수록 잔차의 크기가 커져 뒤쪽 층의 기여가 상대적으로 작아지는 문제가 지적되어 왔고, 이를 보완하는 변형이 여럿 나왔다.

방식	식(요약)	특징	사용 예
Post-LN	$\mathrm{Norm}(x+F(x))$	잘 학습되면 성능이 좋지만 워밍업이 필수, 깊으면 불안정	원 트랜스포머, BERT
Pre-LN	$x+F(\mathrm{Norm}(x))$ + 최종 Norm	안정적, 워밍업 부담 적음. 잔차 크기 증가	GPT-2, GPT-3, LLaMA
DeepNorm	$\mathrm{LN}(\alpha x+F(x))$, $\alpha >1$와 초기화 조정	Post-LN의 성능과 Pre-LN의 안정성을 함께 노림. 1,000층까지 학습	MiniMax-01
앞뒤 모두(샌드위치)	$x+\mathrm{Norm}(F(\mathrm{Norm}(x)))$	부분층 입력과 출력을 모두 정규화	젬마 2
출력 정규화 + QK-Norm	$x+\mathrm{Norm}(F(x))$, 쿼리·키에도 RMSNorm	기울기 노름의 증가와 급등을 줄임	OLMo 2

OLMo 2 보고서는 부분층 출력 쪽 정규화와 QK-Norm이 따로는 효과가 없었지만 함께 쓰면 기울기 급등 지표(spike score)를 0.108에서 0.069로 낮췄다고 보고했다. QK-Norm은 원래 구글의 ViT-22B 연구에서 나온 기법이다. 반면 멀티 헤드 잠재 어텐션처럼 추론 때 키를 만들지 않는 구조에는 QK-Norm을 쓸 수 없어, 키미 K2는 대신 가중치를 잘라 내는 QK-Clip을 썼다.

블록 코드

import torch.nn as nn

class PostLNBlock(nn.Module):          # 원 트랜스포머: x = Norm(x + F(x))
    def __init__(self, d, attn, ffn):
        super().__init__()
        self.attn, self.ffn = attn, ffn
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
    def forward(self, x):
        x = self.n1(x + self.attn(x))
        return self.n2(x + self.ffn(x))

class PreNormBlock(nn.Module):         # LLaMA 계열: x = x + F(Norm(x))
    def __init__(self, d, attn, ffn):
        super().__init__()
        self.attn, self.ffn = attn, ffn
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)   # 위에서 정의한 RMSNorm
    def forward(self, x):
        x = x + self.attn(self.n1(x))
        return x + self.ffn(self.n2(x))  # 전체 모델 끝에 최종 RMSNorm 하나를 더 둔다

실무 선택 기준

• 새 디코더 LLM이면 Pre-LN + RMSNorm이 기본 조합이다. 학습 불안정(손실 급등)이 문제면 QK-Norm, 부분층 출력 정규화, z-loss 같은 보강을 검토한다.
• 기존 모델을 미세 조정할 때는 정규화 구현(이득 초기화, $(1+g)$ 여부, $ϵ$, 자료형 변환 순서)을 원 구현과 똑같이 맞춘다. 수치가 조금만 달라도 BF16에서는 출력이 눈에 띄게 달라질 수 있다.
• RMSNorm의 통계는 float32로 계산한다. BF16은 가수 비트가 적어 제곱 평균에서 오차가 커진다(bfloat16 참고).
• 배치 크기가 작거나 시퀀스 길이가 들쑥날쑥한 언어 모델에는 배치 정규화보다 LayerNorm·RMSNorm 계열이 맞다. 배치 통계에 의존하지 않기 때문이다.

같이 보기

• 레이어 정규화
• 배치 정규화
• 트랜스포머 블록
• 트랜스포머 (인공지능)
• SwiGLU
• RoPE
• 기울기 소실
• 학습률 스케줄링
• 혼합 정밀도 학습