정의

RMS 정규화(Root Mean Square Layer Normalization, RMSNorm)는 입력의 평균을 빼지 않고, 제곱 평균의 제곱근인 RMS를 이용해 입력의 크기만 조절하는 정규화 방법이다.

하나의 token에 대응하는 입력 vector를 x∈Rd\mathbf{x} \in \mathbb{R}^{d}라고 하자. 입력의 RMS는 다음과 같다.

RMS⁡(x)=1d∑i=1dxi2\operatorname{RMS}(\mathbf{x}) = \sqrt{ \frac{1}{d} \sum_{i=1}^{d}x_i^2 }

실제 정규화에서는 0으로 나누는 문제를 막기 위해 분모에 ϵ\epsilon을 더한다. RMSNorm은 각 원소를 같은 RMS로 나눈 뒤, 학습 가능한 gain γi\gamma_i를 곱한다.

x^i=xiRMS⁡(x)2+ϵ\hat{x}_i = \frac{x_i}{ \sqrt{ \operatorname{RMS}(\mathbf{x})^2+\epsilon } } yi=γix^iy_i = \gamma_i\hat{x}_i

각 기호의 의미는 다음과 같다.

  • dd: 정규화할 feature dimension이다.
  • ϵ\epsilon: RMS가 0에 가까울 때 0으로 나누는 문제를 막기 위한 작은 값이다.
  • γi\gamma_i: 정규화 이후 각 feature의 크기를 다시 조절하는 학습 가능한 parameter이다.

학습 가능한 gain을 적용하기 전에는 정규화된 vector의 RMS가 거의 1이 된다.

1d∑i=1dx^i2=RMS⁡(x)2RMS⁡(x)2+ϵ≈1\frac{1}{d} \sum_{i=1}^{d} \hat{x}_i^2 = \frac{ \operatorname{RMS}(\mathbf{x})^2 }{ \operatorname{RMS}(\mathbf{x})^2+\epsilon } \approx 1

즉 RMSNorm은 각 feature를 따로 정규화하는 것이 아니라, 하나의 token vector 전체에 공통 scale을 적용한다.

LayerNorm과의 차이

LayerNorm은 하나의 입력 vector 안에서 평균과 분산을 계산한다.

μ=1d∑i=1dxi\mu = \frac{1}{d} \sum_{i=1}^{d}x_i σ2=1d∑i=1d(xi−μ)2\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i-\mu)^2

이후 평균을 빼 vector의 중심을 0으로 이동하고, 표준편차로 나누어 scale을 조절한다.

yiLN⁡=γixi−μσ2+ϵ+βiy_i^{\operatorname{LN}} = \gamma_i \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}} +\beta_i

반면 RMSNorm은 평균을 빼지 않고 RMS로 나누는 과정만 수행한다.

yiRMS⁡=γixi1d∑j=1dxj2+ϵy_i^{\operatorname{RMS}} = \gamma_i \frac{x_i}{ \sqrt{ \frac{1}{d}\sum_{j=1}^{d}x_j^2+\epsilon } }

분산과 RMS 사이에는 다음 관계가 있다.

RMS⁡(x)2=σ2+μ2\operatorname{RMS}(\mathbf{x})^2 = \sigma^2+\mu^2

따라서 입력 평균이 0이라면 RMS와 표준편차가 같아진다. 하지만 평균이 0이 아니라면 LayerNorm은 평균을 제거한 값의 scale을 조절하고, RMSNorm은 원래 중심을 유지한 채 vector 전체의 크기만 조절한다.

구분LayerNormRMSNorm
중심 이동평균을 빼 중심을 0으로 이동한다.평균을 빼지 않는다.
scale 조절 기준표준편차를 사용한다.RMS를 사용한다.
학습 parameter보통 gain γ\gamma와 shift β\beta를 사용한다.보통 gain γ\gamma만 사용한다.
batch 의존성다른 sample의 통계에 의존하지 않는다.다른 sample의 통계에 의존하지 않는다.

두 방법 모두 sample 또는 token마다 feature 축의 통계를 계산하므로 BatchNorm과 달리 running mean과 running variance가 필요하지 않다. 학습과 추론에서도 같은 방식으로 계산한다.

Scale 변화에 대한 성질

RMSNorm에서 중요한 성질은 입력 전체의 scale이 변해도 정규화 결과의 scale은 유지된다는 점이다. 양의 상수 cc에 대해 RMS는 다음과 같이 변한다.

RMS⁡(cx)=cRMS⁡(x)\operatorname{RMS}(c\mathbf{x}) = c\operatorname{RMS}(\mathbf{x})

따라서 ϵ\epsilon을 생략하면 다음 관계가 성립한다.

cxRMS⁡(cx)=xRMS⁡(x)\frac{c\mathbf{x}}{ \operatorname{RMS}(c\mathbf{x}) } = \frac{\mathbf{x}}{ \operatorname{RMS}(\mathbf{x}) }

앞 layer의 출력 크기가 커지더라도 RMSNorm을 통과한 vector의 전체 크기는 일정하게 조절된다. RMSNorm 논문에서는 이 성질을 re-scaling invariance라고 설명한다.

반대로 입력의 모든 원소에 같은 값을 더하는 경우에는 그 값이 제거되지 않는다. RMSNorm은 LayerNorm이 제공하는 re-centering invariance를 포기하고 scale 조절에만 집중한 방법이다.

Transformer에서 자주 사용하는 이유

RMSNorm은 LayerNorm에서 평균 계산과 평균을 빼는 연산을 제거한다. 정규화 식이 단순하면서도 activation의 전체 크기를 안정적으로 조절하고, batch 구성에 의존하지 않는다는 장점은 그대로 유지한다.

RMSNorm 원 논문은 여러 RNN, CNN, Transformer 실험에서 LayerNorm과 비슷한 성능을 보이면서, 당시 실험한 구현과 model 기준으로 실행 시간을 7%에서 64%까지 줄였다고 보고했다. 이는 평균을 제거하는 re-centering이 모든 경우에 필수는 아닐 수 있다는 실험적 근거가 된다.

이후 LLaMA는 각 Transformer sublayer의 입력을 RMSNorm으로 정규화하는 pre-norm 구조를 사용했고, Qwen2도 RMSNorm과 pre-normalization을 사용한다. 즉 최근 LLM에서는 다음 이유로 RMSNorm을 자주 선택한다.

  • 깊은 Transformer에서 activation scale을 안정적으로 유지할 수 있다.
  • LayerNorm보다 계산 과정이 단순하다.
  • batch 통계 없이 token마다 독립적으로 계산할 수 있다.
  • 대규모 model에서 실제로 검증된 architecture 구성을 그대로 활용할 수 있다.

다만 수식의 연산 수가 적다는 사실이 모든 환경에서 실제 실행 시간이 더 짧다는 뜻은 아니다. OpenELM 기술 보고서에서는 단순한 PyTorch RMSNorm 구현이 최적화된 LayerNorm보다 느린 사례도 보고했다. 실제 속도는 fused kernel, hardware, tensor shape과 같은 구현 조건에 따라 달라진다.

PyTorch 구현

입력이 (B,S,d)(B,S,d) shape이라면 마지막 feature 축 dd를 따라 제곱 평균을 계산한다. keepdim=True를 사용하면 RMS의 shape이 (B,S,1)(B,S,1)로 유지되어 원래 입력의 feature 축으로 broadcasting된다.

import torch
from torch import nn


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        inv_rms = torch.rsqrt(
            x.pow(2).mean(dim=-1, keepdim=True) + self.eps
        )
        return x * inv_rms * self.weight

shape의 흐름은 다음과 같다.

(B,S,d)→(B,S,1)→(B,S,d)(B,S,d) \rightarrow (B,S,1) \rightarrow (B,S,d)

첫 번째 화살표는 feature 축을 따라 RMS의 역수를 구하는 과정이고, 두 번째 화살표는 그 값을 입력과 곱할 때 broadcasting되는 과정이다.

중요 개념 정리

  • LayerNorm은 평균을 빼는 중심 이동과 표준편차를 이용한 scale 조절을 모두 수행한다.
  • RMSNorm은 평균을 빼지 않고 RMS를 이용한 scale 조절만 수행한다.
  • RMSNorm은 입력 vector 전체를 하나의 값으로 나누므로 activation의 전체 크기를 안정적으로 유지한다.
  • LayerNorm과 RMSNorm은 모두 batch 통계에 의존하지 않는다.
  • RMSNorm의 계산식은 더 단순하지만 실제 속도 차이는 구현과 hardware에 따라 달라진다.

참고 자료