AI

RMSNormとは?LayerNormとの違い・数式・PyTorch実装をわかりやすく解説

RMSNorm(Root Mean Square Layer Normalization)は、入力ベクトルを二乗平均平方根(RMS)で割って揃える正規化手法です。ポイントは「平均を引かない」こと。LayerNormが行う2つの処理――平均を引く再センタリングと、ばらつきで割る再スケーリング――のうち、RMSNormは再スケーリングだけを残します。2019年にBiao ZhangとRico Sennrichが提案し(arXiv:1910.07467、NeurIPS 2019)、いまではLLaMA・Mistral・Gemma・Qwenなど主要な大規模言語モデルの標準的な正規化層になっています。

まとめ:RMSNormの要点

  • 平均を引かず、RMSで割るだけ。分散を計算する処理は無い(後述のとおり「分散のみを使う」という説明は誤り)。
  • LayerNormから再センタリングを省いた分、計算が軽い。論文では実行時間を7〜64%削減(モデル依存)と報告。
  • 正規化後は学習可能なゲイン g で各次元をスケールし直す。
  • PyTorchは2.4以降で torch.nn.RMSNormtorch.nn.functional.rms_norm を標準搭載。自作しなくてよい。
  • 採用例はLLaMA系・Mistral・Gemma・Qwen・DeepSeekなど。GPT-3はLayerNormで、RMSNormではない。

RMSNormとは:正規化で「平均を引かない」という選択

ニューラルネットワークの各層では、入力のスケールが層ごとにばらつくと学習が不安定になります。これを層ごとに揃えるのが正規化で、Transformer系ではLayerNormが長く標準でした。RMSNormは、そのLayerNormから「平均を引く」ステップを丸ごと外し、大きさ(RMS)だけで割り戻します。省いても精度がほとんど落ちないことを実験で示したのが元論文の貢献です。

数式(公式)と各記号の意味

次元数 d のベクトル x に対し、RMSNormは次のように定義されます。

RMS(x) = √( (1/d) × Σi=1..d xi² + ε )

出力i = ( xi ÷ RMS(x) ) × gi

ここで ε は0除算を防ぐ小さな定数(PyTorchでは既定で計算精度型のマシンイプシロン)、g は次元ごとの学習可能なスケール(ゲイン)です。なお ε は原論文の定義式には現れず、数値安定化のために実装側で付与するものです。LayerNormと違って平均 μ を引く項が式に現れないこと、バイアス項を持たないことが特徴です。ε は平方根の内側に足す実装が標準で、外側に足すと数値挙動が変わるため注意します。

論文の核心:再センタリング不変性は捨てられる

元論文(Zhang & Sennrich, 2019)の主張は明快です。LayerNormが持つ性質のうち、学習の安定に効いているのは入力を「同じ大きさに揃える」再スケーリング不変性であって、「平均をゼロに揃える」再センタリング不変性は実質的に不要(dispensable)だ、というものです。だから平均を引く計算をやめても性能は保てる――これが根拠であり、単なる高速化のための近似ではありません。論文はさらに、和の一部(p%)だけからRMSを推定するpRMSNormも提案しています。

LayerNormとの違い:再センタリングを省くと何が変わるか

RMSNormとLayerNormはどちらもバッチサイズに依存せず、1サンプル(1トークン)ごとに特徴量方向へ正規化します。ここは同じです。違いは、平均を引くかどうかの一点に集約されます。旧来「RMSNormは分散のみを使う」と説明されることがありますが、これは不正確です。分散は平均を引いてから二乗平均を取る量であり、平均を引かないRMSNormは分散を計算していません。RMSと標準偏差が一致するのは、たまたま平均が0のときだけです。

手法 平均を引く 割る量 バッチ依存 バイアス項
BatchNorm あり バッチ統計の標準偏差 あり あり
LayerNorm あり 標準偏差(μを引いた後) なし あり(任意)
RMSNorm なし RMS(二乗平均平方根) なし なし

計算コスト削減の仕組み

LayerNormは1回のフォワードで平均・分散の2つの統計量を求め、平均を各要素から引く処理が入ります。RMSNormはこの平均計算と減算を省くため、要素あたりの演算とメモリアクセスが減ります。1層あたりの差はわずかでも、数十〜百層を数千トークン分回すLLMでは累積し、論文の実測で7〜64%の実行時間短縮につながっています。精度がほぼ変わらないまま速くなるので、新規にTransformerを設計するならRMSNormを既定に選んで問題ありません。

PyTorchでのRMSNorm実装:標準モジュールと自作

「rmsnorm pytorch」「torch.nn.rmsnorm」「F.rms_norm」で調べる人が多いので、まず結論を書きます。PyTorch 2.4以降なら自作は不要で、torch.nn.RMSNorm(モジュール)か torch.nn.functional.rms_norm(関数)を使えます。対応バージョンは公式ドキュメントで確認してください。

標準API:torch.nn.RMSNorm と F.rms_norm

シグネチャは torch.nn.RMSNorm(normalized_shape, eps=None, elementwise_affine=True)normalized_shape に正規化したい末尾次元(通常は隠れ次元)を渡します。eps=None のときは入力dtype(fp16/bf16/fp32いずれもfloat32基準)のマシンイプシロンが自動採用され、elementwise_affine=True なら学習可能なゲイン g が付きます。LayerNormと違いバイアスは持ちません。

import torch
import torch.nn as nn
import torch.nn.functional as F

x = torch.randn(2, 8, 512)          # (batch, seq, hidden)

# モジュール版(学習可能なゲインを持つ)
norm = nn.RMSNorm(512)
y = norm(x)

# 関数版(重みは任意、eps未指定でdtype既定値)
y2 = F.rms_norm(x, (512,))
# ゲインがones初期化のモジュール版と関数版は一致する
print(y.shape, torch.allclose(y, y2))

自作クラス:LLaMA式の実装と精度の注意点

2.4より前の環境や、挙動を厳密に合わせたいときは自作します。Metaのllama実装が実質の標準で、要点は二乗平均をfloat32に上げて計算し、元のdtypeへ戻してからゲインを掛けることです。fp16のまま二乗平均を取ると大きな値で精度が崩れるため、この昇格が実務上重要になります。

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 _norm(self, x):
        # 末尾次元でRMSを取り、rsqrtで割る(epsは平方根の内側)
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        # float32で正規化してから入力dtypeへ戻す(fp16の精度崩れ対策)
        out = self._norm(x.float()).type_as(x)
        return out * self.weight

NumPyだけで確かめるなら x / np.sqrt(np.mean(x**2) + eps) の1行で本質は表せます。ただし学習ではゲイン g とdtype昇格が要るので、実運用は上記のクラスか標準APIを使ってください。

どのモデルが採用しているか:LLaMAが広めた標準正規化

RMSNormを「GPT-3以降のモデルが採用」と説明する記事が多いですが、これは誤りです。GPT-3(2020)はプレノーム構成のLayerNormを使っています。RMSNormを最初期に取り入れたのは2019年のT5系(平均もバイアスも持たない簡素な正規化=実質RMSNorm)で、デコーダ型LLMで一気に広めたのが2023年のLLaMA(Meta)です。その後Mistral・Mixtral・Gemma・Qwen・DeepSeek・Phiなど、2023年以降に登場した主要モデルの多くがRMSNormを標準採用しています。

配置はいずれもプレノーム(各サブ層の入力側で正規化)です。Transformer全体の中で正規化層がどこに入るかは、Transformerとは?自己注意の仕組みとBERT・GPTの違いを合わせて読むと位置づけがつかみやすくなります。RMSNormが動いている具体例としては、日本語特化のLlama-3-ELYZA-JP-8BのようなLLaMA系モデルが分かりやすい対象です。

導入時の実装設定:eps・dtype・配置

実装で詰まりやすいのは細部です。順に押さえます。

  • ε の位置:平方根の内側(mean(x**2) + eps の形)に足す。外側に足すと小さな値で挙動が変わる。
  • dtype:fp16/bf16学習では二乗平均をfloat32で計算してから戻す。標準の nn.RMSNorm は内部でこの精度確保を行う。
  • バイアスは足さない:RMSNormはゲイン g のみ。LayerNormからの移植でバイアスを残すと定義がずれる。
  • 正規化する次元normalized_shape は隠れ次元(末尾)に合わせる。系列長やバッチ次元を含めない。
  • 置き換えの是非:LayerNorm前提で学習済みの重みにそのまま差し込むと崩れる。RMSNormは新規学習・ファインチューニング時に採用する。

大規模言語モデル全体の中でこうした正規化がどんな役割を果たすかは、大規模言語モデル(LLM)とはで基礎から確認できます。

よくある質問(FAQ)

RMSNormとLayerNormはどちらが精度が高いですか?

元論文の実験では、RMSNormはLayerNormと同等の精度を保ちつつ実行時間を7〜64%短縮します。精度はほぼ互角で、速度と実装の簡潔さでRMSNormが有利、というのが実務的な結論です。

RMSNormは分散を使いますか?

使いません。RMSNormが割るのはRMS(二乗平均平方根)で、分散ではありません。分散は平均を引いてから二乗平均を取る量ですが、RMSNormは平均を引かないため分散を計算しません。RMSと標準偏差が一致するのは平均が0のときだけです。

torch.nn.RMSNormはどのバージョンから使えますか?

PyTorch 2.4以降で torch.nn.RMSNormtorch.nn.functional.rms_norm が標準提供されています。それ以前のバージョンでは自作クラスで代用します。最新の対応状況は公式ドキュメントで確認してください。

RMSNormの論文はどれですか?

Biao ZhangとRico Sennrichによる「Root Mean Square Layer Normalization」(arXiv:1910.07467、NeurIPS 2019)です。再センタリング不変性が不要であることを示し、RMSNormとその部分版pRMSNormを提案しています。

GPT-3はRMSNormを使っていますか?

いいえ。GPT-3はプレノーム構成のLayerNormを使っています。RMSNormを広めたのは2023年のLLaMAで、以降のMistral・Gemma・Qwen・DeepSeekなどが標準採用しています。

関連記事

資料請求

RELATED POSTS 関連記事