RMSNorm(Root Mean Square Layer Normalization)は、ベクトルを二乗平均平方根(RMS)で割って大きさを揃える正規化です。LayerNormとの主な違いは、平均を引かずRMSで割ることです。PyTorchの標準実装では、加算するバイアスも持ちません。2019年にBiao ZhangとRico SennrichがarXiv:1910.07467(NeurIPS 2019)で提案し、2023年のLLaMA以降、Mistral・Qwen・Gemma・DeepSeekなど公開LLMの多くが採用しています。この記事では数式とLayerNormとの違いを数値で確かめ、PyTorchの torch.nn.RMSNorm の使い方と、学習済みモデルの重みを読むときに起きるずれまでを、手元で実行した結果をもとに説明します。
まとめ:RMSNormの要点
- 平均を引かずRMSで割り、ゲインを掛けるだけ。分散は計算しない(「分散のみを使う」という説明は誤り)。
- 正規化する各ベクトルの平均が0で、正規化軸・ゲイン・εが同じなら、LayerNorm(バイアスなし)と数学的に同じ出力になる。平均がずれていると結果は大きく変わる。
- PyTorchの既定設定では、学習パラメータは隠れ次元dに対してd個(LayerNormは重みとバイアスで2d個)。
- PyTorchは2.4以降で
nn.RMSNormとF.rms_normを標準搭載。2.3以前には無い。 - 論文の速度改善は7〜64%だが、CPUでは標準の
nn.RMSNorm(PyTorch 2.14.1)も自作も、LayerNormの約3倍遅かった(実測)。速さは実装しだい。 - 学習済みの重みを読むときは、epsをconfigの値で渡し、Gemmaの
(1 + weight)方式を取り違えない。 - GPT-3はLayerNormで、RMSNormではない。
RMSNormの定義と数式
数式と記号(RMS・ε・ゲイン)
次元数 d のベクトル x に対し、RMSNormは次のように計算します。
RMS(x) = √( (1/d) × Σi=1..d xi² + ε )
出力i = ( xi ÷ RMS(x) ) × gi
g は次元ごとの学習可能なゲインで、初期値は1です。εは0除算を防ぐ小さな定数で、原論文の定義式には無く、実装側が平方根の内側に足します。PyTorch 2.14のドキュメントの式も sqrt(ε + mean(x²)) の形です。
LayerNormは「平均を引く(再センタリング)」と「標準偏差で割る(再スケーリング)」の2段階でした。RMSNormは前者を外しています。RMSと標準偏差が一致するのは平均が0のときだけなので、RMSNormを「分散で割る」と説明するのは誤りです。
原論文 arXiv:1910.07467 の主張
論文の仮説は、LayerNormの効果は再スケーリング不変性から来ており、再センタリング不変性は不要(dispensable)だというものです。機械翻訳や画像・キャプション検索など複数のタスクで、RMSNormはLayerNormと同等の性能を保ちながら実行時間を7〜64%短縮したと報告しています。削減幅の開きはモデルと実装の違いによるもので、どの環境でも速くなるという意味ではありません(後述の実測を参照)。
同じ論文は、先頭の一部の要素だけでRMSを推定するpRMSNorm(partial RMSNorm)も提案しており、実験では全体の6.25%の要素でRMSNormと遜色ない性能になったとしています。ただし主要なLLM実装で使われているのは全要素で計算する通常のRMSNormです。
RMSNormとLayerNormの数値比較
入力平均によるRMSNormとLayerNormの出力差
違いを手元で確かめるため、PyTorch 2.2.2(Intel Mac・CPU)で隠れ次元4096のランダム入力を使い、LLaMA式の自作RMSNormとバイアスなしの nn.LayerNorm を比べました。結果は次のとおりです。
| 確かめたこと | 最大誤差 | 読み取れること |
|---|---|---|
| 元の入力と5倍した入力の正規化出力を比較 | 2.4×10⁻⁶ | スケールには不変 |
| 全要素に3を足した入力と比較 | 3.89 | 平均のずれには不変でない |
| 平均0の入力でLayerNormと比較 | 4.8×10⁻⁷ | 実質一致 |
| 平均3の入力でLayerNormと比較 | 3.87 | 別物になる |
1行目は今回の入力では正の定数倍に対してほぼ不変であること、2行目は定数加算に対して不変でないことを示します。εを加える実装では、再スケーリング不変性は厳密には成立しません。3行目からは、入力がもともと平均0付近に分布しているなら、平均を引く処理を省いても出力はほとんど変わらないことが分かります。ただし、この実験は平均を明示的に0にしたランダム入力での比較であり、Transformerの隠れ状態や学習性能を検証したものではありません。
パラメータ数と計算の手数
隠れ次元4096で数えると、nn.LayerNorm の学習パラメータは8,192個(重みとバイアス)、RMSNormは4,096個(重みのみ)でした。とはいえ、LLMの重み行列に比べればどちらも誤差の範囲で、パラメータ削減を採用理由にするほどの差ではありません。計算面では、平均の計算と減算が1回ずつ消え、統計量の集計が二乗和の1回で済みます。
RMSNormとLayerNormのCPU実測速度
同じMac(Core i9-9880H・4スレッド)で、形状(64, 512, 4096)のfloat32テンソルを使い、勾配計算を無効にした順伝播を3回ウォームアップした後、20回の平均時間を測りました。PyTorch 2.2.2はmacOS上で直接、2.14.1は同じMac上のDocker(Linux・CPU版)で実行しています。
| 実装 | 1回あたり |
|---|---|
| nn.LayerNorm(2.2.2) | 376 ms |
| 自作RMSNorm(2.2.2・LLaMA式) | 1,080 ms |
| 自作RMSNorm(2.2.2・昇格なし) | 1,144 ms |
| nn.LayerNorm(2.14.1) | 151 ms |
| nn.RMSNorm(2.14.1・標準) | 446 ms |
| 自作RMSNorm(2.14.1・LLaMA式) | 398 ms |
計算量の少ないRMSNormのほうが、どちらの版でも約3倍遅くなりました。2.14.1では標準の nn.RMSNorm も自作とほぼ同じ時間で、CPUでは専用の最適化カーネルを使うLayerNormに届いていません。RMSNormの計算は二乗・平均・平方根・乗算に分かれ、別々に実行されるとそのたびにテンソル全体を読み書きします。論文の7〜64%は論文が使ったモデルと実装での値で、どの実行環境でも再現する数字ではありません。
実務上の結論は、速度を理由に正規化層だけを差し替えないことです。正規化層の時間はモデル全体では行列積に比べて小さく、効くかどうかは使う実行基盤の実装で決まります。GPUで動かすなら、vLLMや学習フレームワークが持つ融合済みの実装を使い、自分の環境でLayerNorm版と並べて測ってから判断してください(GPUでの速度はこの記事では測っていません)。
PyTorchでのRMSNorm実装
torch.nn.RMSNorm と F.rms_norm(2.4以降)
PyTorchは2.4でモジュールの torch.nn.RMSNorm と関数の torch.nn.functional.rms_norm を追加しました(ソースでは v2.4.0の normalization.py から class RMSNorm が現れ、v2.3.0には存在しません)。シグネチャは nn.RMSNorm(normalized_shape, eps=None, elementwise_affine=True, device=None, dtype=None) で、バイアスの引数はありません。
import torch
import torch.nn as nn
import torch.nn.functional as F
x = torch.randn(2, 16, 4096) # (バッチ, 系列長, 隠れ次元)
norm = nn.RMSNorm(4096, eps=1e-6) # 末尾の次元で正規化する
y = norm(x)
# 関数版:重みを自分で渡す
y2 = F.rms_norm(x, (4096,), weight=norm.weight, eps=1e-6)
print(torch.allclose(y, y2)) # True
normalized_shape には末尾の次元を渡します。(4096,) なら隠れ次元だけ、(16, 4096) なら系列長と隠れ次元をまとめて1つのRMSを取るので、トークンごとに正規化したいTransformerでは隠れ次元だけを指定します。PyTorch 2.14では、eps を省略すると計算型のマシンイプシロンが使われ、fp16・bf16・fp32の入力ではいずれも torch.finfo(torch.float32).eps(約1.19×10⁻⁷)になります。2.4.0のソースでは既定値の説明が torch.finfo(x.dtype).eps(入力dtypeのeps)になっており、版によって説明が違います。版をまたいで同じ結果を得たいなら、epsは省略せず明示してください。
PyTorch 2.4未満向けの自作RMSNorm(LLaMA式)
PyTorch 2.2.2で hasattr(torch.nn, "RMSNorm") を確かめると False でした。macOS(Intel)向けの公式バイナリが2.2系で止まっている環境や、古い学習基盤を使う場合は自作します。2.14.1では、eps=1e-6・初期重み・float32の条件で標準の nn.RMSNorm と下の自作クラスの出力が最大差0.0で一致しました。標準実装へ移行するときは、eps・重み・dtype・正規化する次元を揃え、対象環境で出力を比較してください。Hugging Face transformersの LlamaRMSNorm と同じ手順で、二乗平均をfloat32で計算し、元のdtypeに戻してから重みを掛けます。
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(初期値1)
def forward(self, x):
xf = x.float() # 二乗の前にfloat32へ上げる
y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
return self.weight * y.type_as(x) # 元のdtypeへ戻してから重みを掛ける
# 2.4以降なら標準を、無ければ自作を使う
Norm = getattr(nn, "RMSNorm", RMSNorm)
norm = Norm(4096, eps=1e-6) # 標準実装と自作でepsを揃える
float32への昇格を省けないのは、fp16の上限が65,504だからです。全要素が300のfp16ベクトルで試すと、300²=90,000がfp16で inf になり、昇格なしの実装は出力がすべて0.0になりました。昇格した実装は正しく1.0を返します。混合精度で学習・推論するなら昇格は必須です。
nn.LayerNorm から置き換えるときの引数対応
| 項目 | nn.LayerNorm | nn.RMSNorm |
|---|---|---|
| 平均の減算 | あり | なし |
| 学習パラメータ | weight・bias | weightのみ |
| bias引数 | あり(既定True) | なし |
| eps既定値 | 1e-5 | None(計算型のeps) |
| 追加された版 | 初期から | 2.4 |
コード上は nn.LayerNorm(d) を nn.RMSNorm(d) に書き換えるだけで動きますが、epsの既定値が違う点は見落としやすいところです。LayerNormで学習したモデルの重みをRMSNormへ差し替えて使うことはできません。平均を引かない分、出力の分布が変わるため、置き換えは新規学習かファインチューニングのやり直しとセットで行います。
学習済みモデルの重みを読むときの落とし穴
モデルのconfigに合わせたeps設定
公開モデルはconfig.jsonの rms_norm_eps に学習時のεを持っています。Qwen2.5-7BとDeepSeek-V3は1e-06、Mistral-7B-v0.1は1e-05です。nn.RMSNorm(d) とεを省略して重みだけ読み込むと約1.19×10⁻⁷が使われ、学習時と違う計算になります。PyTorch 2.14.1で、要素の大きさが0.001程度の入力を eps=None と eps=1e-5 で正規化して比べると、出力の最大差は2.51でした。値が小さい隠れ状態ほど差が出るので、自前の推論コードへ移植するときは eps=config.rms_norm_eps を必ず渡してください。
Gemmaのゲイン表現:1 + weight
Gemma系はRMSNormの重みの持ち方が違います。transformersの GemmaRMSNorm は重みを0で初期化し、正規化後に (1.0 + weight) を掛けます。LLaMA式(重みを1で初期化してそのまま掛ける)と初期状態の出力は一致しますが、同じ重みの値を両方に読み込むと結果が変わります。手元の比較では、重みが小さい値(標準偏差0.1)でも最大差は4.32でした。Gemmaの重みをLLaMA式の実装で読むときは、保存値に1を足してから使います。
もう1つの違いは昇格の順序です。LlamaRMSNormは元のdtypeに戻してから重みを掛け、GemmaRMSNormはfloat32のまま重みを掛けてから戻します。bf16では最下位の桁が変わりうるので、既存実装と出力を突き合わせるときは同じ順序に合わせます。
fp16の二乗演算におけるオーバーフロー
前述のとおり、fp16では絶対値が約256を超える要素があると二乗が上限を超えます。LLMの隠れ状態には一部の次元だけ極端に大きな値が出ることがLLM.int8()の論文(arXiv:2208.07339)で報告されており、正規化の入力はこの外れ値をそのまま受け取ります。モデルを低精度で動かすときは、正規化層の二乗平均がfloat32で計算されているかを先に確かめてください。精度形式の違いはFP8とは?FP16・BF16・FP4との違い、量子化の手順は量子化(モデル量子化)とはで扱っています。
採用モデル:T5・LLaMAからMistral・Qwen・Gemmaへ
RMSNormを「GPT-3以降のモデルが採用」と書く解説がありますが、誤りです。GPT-3の論文はGPT-2と同じアーキテクチャを使うと明記しており、GPT-2は各サブブロックの入力側にLayerNormを置くプレノーム構成です。
| モデル | 正規化 | 確認した根拠 |
|---|---|---|
| GPT-2・GPT-3 | LayerNorm(プレノーム) | GPT-3論文 |
| T5(2019) | 平均もバイアスも持たない簡略版 | transformers T5LayerNorm |
| LLaMA(2023) | RMSNorm(プレノーム) | LLaMA論文 |
| Mistral-7B | RMSNorm(eps 1e-5) | config.json |
| Qwen2.5・DeepSeek-V3 | RMSNorm(eps 1e-6) | config.json |
| Gemma | RMSNorm(1+weight方式) | transformers GemmaRMSNorm |
T5の正規化は transformers のソースで「No bias and no subtraction of mean」と説明されており、計算の形はRMSNormと同じです。デコーダ型LLMで広めたのはLLaMA(arXiv:2302.13971)で、論文は学習の安定のため各サブ層の入力を正規化し、その関数にZhangとSennrichのRMSNormを使うと書いています。LLaMA系の日本語モデルでも同じ構成で、たとえばLlama-3-ELYZA-JP-8Bも正規化はRMSNormです。正規化層がTransformerブロックのどこに入るかはTransformerとは?自己注意の仕組みとBERT・GPTの違いで確認できます。
RMSNormとLayerNormの選び方
新しくTransformerを学習するなら、RMSNormを選んで問題ありません。主要な公開LLMが採用しており、vLLMなどの推論エンジンもRMSNorm専用の実装を持っています。精度は論文でもLayerNormと同等で、どちらかが明確に勝つわけではありません。
避けるべきなのは次の場面です。
- LayerNormで学習済みの重みを流用する:差し替えると出力分布が変わる。正規化はモデルの学習時の選択に従う。
- 速度だけを理由に差し替える:上の実測のとおり、CPUでは標準実装でもLayerNormより遅い。実行基盤で測ってから決める。
- 入力への定数加算に対する不変性が必要な設計:RMSNormは平均のずれを吸収しない。その性質が要るなら平均を引くLayerNormなどを比較候補にする。原論文はCNN(CIFAR-10)でも評価しており、CNNであることだけではRMSNormを除外できない。
社内でLLMを組み込んだシステムを作る段階では、正規化層の選び方よりモデル選定と推論基盤の構成が効きます。ISSOHでは生成AI導入支援として、モデル選定から推論環境の構築まで支援しています。
よくある質問
RMSNormは分散を使いますか?
使いません。割る値はRMS(二乗平均平方根)で、平均を引かないため分散は計算しません。RMSと標準偏差が一致するのは入力の平均が0のときだけです。
torch.nn.RMSNormはどのバージョンから使えますか?
PyTorch 2.4からです。torch.nn.functional.rms_norm も同じ版で追加されました。2.3以前では自作クラスで代用します。
RMSNormの論文はどれですか?
Biao ZhangとRico Sennrichの「Root Mean Square Layer Normalization」(arXiv:1910.07467、NeurIPS 2019)です。RMSNormと、一部の要素でRMSを推定するpRMSNormを提案しています。
RMSNormとLayerNormはどちらが精度が高いですか?
原論文の実験では同等です。差が出るのは速度と実装の簡潔さで、速度差は演算の融合状況、比較対象の実装、入力形状、実行デバイスによって変わります。
GPT-3はRMSNormを使っていますか?
いいえ。GPT-3はGPT-2と同じプレノーム構成のLayerNormです。RMSNormをデコーダ型LLMに広めたのは2023年のLLaMAです。