RMS Norm 정규화 쉽게 이해하기 — 레이어 정규화보다 빠른 트랜스포머 정규화

핵심 메시지

• RMS Norm은 루트 평균 제곱 레이어 정규화(root mean square layer normalization)로, 레이어 정규화를 단순화한 기법이다.
• 미스트랄·라마 같은 대형 언어 모델에서 어텐션 모듈 직후 등 여러 위치에 RMS Norm이 사용된다.
• 레이어 정규화와 달리 RMS Norm은 평균을 빼지 않고 분산도 계산하지 않으며, 루트 평균 제곱만으로 원소별 조정을 한다.
• 평균·분산 계산을 생략해 연산량이 크게 줄어, 차원 500 기준 레이어 정규화의 약 4,000회 연산보다 약 37.5% 적은 연산으로 처리한다.
• 안정성은 다소 떨어지지만 보정이 가능하며, 깊고 큰 트랜스포머에서 오히려 더 잘 작동한다.

쉽게 이해하기

RMS Norm은 정식 명칭이 “루트 평균 제곱 레이어 정규화(root mean square layer normalization)”로, 기존 레이어 정규화(layer norm)를 단순화한 버전이다. 미스트랄(Mistral) 같은 모델 코드를 보면 어텐션 모듈 바로 뒤를 비롯한 여러 위치에서 RMS Norm이 쓰인다. 핵심 아이디어는 계산 복잡도를 줄이는 것이다.

복잡도를 줄이는 이유는 분포를 평균과 분산으로 조정하는 대신, 루트 평균 제곱(RMS)만으로 원소별 조정을 하기 때문이다. 레이어 정규화는 각 원소에서 평균을 빼고 분산으로 나눠 분포를 다시 맞추지만, RMS Norm에는 평균 빼기가 없고 분산 계산도 없다. 그만큼 단순하다.

연산량 차이는 수치로 두드러진다. 차원이 500인 벡터를 예로 들면, 레이어 정규화는 평균 계산, 평균에서의 차이, 제곱, 제곱합, 차원으로 나누기, 엡실론 더하기, 제곱근, 표준편차로 나누기, 그리고 감마·베타 적용까지 거쳐 약 4,000회의 부동소수점 연산이 필요하다. 반면 RMS Norm은 입력 제곱, 제곱합, 나눗셈만으로 끝나 약 2,000~2,500회, 즉 약 37.5% 적은 연산으로 처리된다.

안정성 측면에서는 RMS Norm이 레이어 정규화보다 다소 덜 안정적이지만, 몇 가지 조정으로 개선할 수 있다. 특히 길고 깊은 어텐션 네트워크에서는 RMS Norm이 더 잘 작동하며, 미스트랄이나 라마처럼 거대한 언어 모델일수록 그 이점이 분명해진다. 더 빠르고 부동소수점 연산이 적다는 점이 핵심 장점이다.

구현 측면에서는 미스트랄 코드의 트랜스폼 레이어 아래에서 RMS Norm을 찾을 수 있다. 엡실론과 학습 파라미터를 정의한 뒤, 벡터의 제곱에 대한 역제곱근을 PyTorch의 torch.rsqrt(제곱근의 역수)로 계산한다. 평균은 마지막 차원, 즉 특징(feature) 기준으로 구하고 텐서의 형태는 그대로 유지하며, 순전파에서 정규화한 결과에 가중치를 곱해 마무리한다.

주요 인사이트

• RMS Norm이 빠른 본질적 이유는 “평균 빼기와 분산 계산을 통째로 생략”했기 때문으로, 정규화를 제곱·제곱합·나눗셈만의 연산으로 압축한다.
• 레이어 정규화 대비 약 37.5% 적은 연산은 거대 모델에서 누적되면 큰 효율 이득이 되며, 이것이 라마·미스트랄이 RMS Norm을 채택한 배경이다.
• torch.rsqrt로 1/RMS를 한 번에 계산하고 마지막 차원(특징 축) 기준으로 정규화한다는 구현 디테일은, 수식의 단순함이 코드의 단순함으로 이어짐을 보여준다.

자주 묻는 질문

관련 AI 소식