現代の主要な大規模言語モデル (LLaMA / Mistral / Gemma / Qwen など) の正規化は、ほぼ共通してPre-LN の位置に RMSNorm を置く構成 (Pre-RMSNorm) になっています。例えば LLaMA は「訓練の安定性を高めるため、(出力ではなく) 各ブロック (MHA・FFN) の入力を正規化し (= Pre-LN)、正規化関数には RMSNorm を用いる」と述べています [5]。この構成は独立な 2 つの選択の組み合わせであり、それぞれ別の論文で提案されました。
- 軸1 (正規化する位置) : Post-LN vs Pre-LN … 正規化を残差パスの外に置くか中に置くか → 学習の安定性を決める。On Layer Normalization in the Transformer Architecture (ICML 2020) による [1]。
- 軸2 (どんな正規化をするか) : LayerNorm vs RMSNorm … 平均を残すか捨てるか → 計算量を決める。Root Mean Square Layer Normalization (NeurIPS 2019) による [3]。
本記事では、それぞれの文献概要とトイコードを記します。
参考文献
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. On Layer Normalization in the Transformer Architecture. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), 2020.
Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019.
- PMLR: On Layer Normalization in the Transformer Architecture
- arXiv: On Layer Normalization in the Transformer Architecture
- arXiv: Root Mean Square Layer Normalization
- GitHub: bzhangGo/rmsnorm
- arXiv: LLaMA: Open and Efficient Foundation Language Models
- arXiv: Layer Normalization
- arXiv: DeepNet: Scaling Transformers to 1000 Layers
- arXiv: CogView: Mastering Text-to-Image Generation via Transformers
- arXiv: Peri-LN: Revisiting Normalization Layer in the Transformer Architecture
- arXiv: Spike No More: Stabilizing the Pre-training of Large Language Models
文献概要
軸1 : Post-LN vs Pre-LN (正規化する位置)
- オリジナルの Transformer は (原論文の記述上は) Post-LN です。各ブロック (MHA ブロック、FFN ブロック) の後、残差加算の後に LayerNorm を置きます。
- $x \leftarrow \mathrm{LN}(x + \mathrm{Block}(x))$
- Post-LN の訓練には学習率の warmup が経験的に必要で、外すと発散しやすいことが知られていました (T-Fixup の記事 参照)。
- Xiong らは、この不安定性の原因が初期化時の勾配にあることを、解析的に示しました。
- Post-LN では、初期化時に出力層に近いパラメータの勾配が大きく、warmup なしだと訓練初期の大きな更新で発散しやすい。
- Pre-LN (LN をブロックの前に置き、残差合流後には LN を挟まず素通しする) では、初期化時の勾配が、深さを増やしても爆発も消失もしない。
- $x \leftarrow x + \mathrm{Block}(\mathrm{LN}(x))$
- したがって Pre-LN は warmup を外すことができ、学習率にも頑健になる。
- ポイントは、Pre-LN では残差合流後に正規化を挟まないので、勾配がそのまま浅い層まで届きます。ただし、そのままだと出力の大きさが深さとともに増えるため、最終段には正規化を 1 つ置きます。
- ただし、Post-LN は上手く訓練できれば最終性能がやや高いという報告もあり、正規化する位置の議論は完全には決着していません (後述)。
軸2 : LayerNorm vs RMSNorm (どんな正規化をするか)
- LayerNorm はもともと Transformer 専用の手法ではなく、Ba らが 2016 年に提案した汎用の正規化手法です [6] (元は RNN などの学習安定化が動機)。具体的に、特徴次元にわたって「中心化 (平均を引く) + スケール (標準偏差で割る)」を行い、gain と bias を掛け足すものでした。
- Zhang と Sennrich は、LayerNorm の効きの主因はスケール不変性にあり、中心化の寄与は小さいと仮説を立てました。
- そこで、中心化を捨て、RMS (二乗平均平方根) で割るだけの RMSNorm を提案しました ($a$ は入力ベクトル、$g$ は gain)。
- $\bar{a}_i = \dfrac{a_i}{\mathrm{RMS}(a)} g_i, \quad \mathrm{RMS}(a) = \sqrt{\dfrac{1}{d}\sum_{j=1}^{d} a_j^2}$
- 平均を計算しないので統計量が 1 つで済み、bias も持たないため計算が軽い。
- 検証実験では、機械翻訳などで LayerNorm と同等の精度を保ちつつ、実行時間を短縮できました [3]。
- RMSNorm は入力を定数倍しても出力が変わらないスケール不変性は保ちます。入力に定数を足す不変性 (中心化不変性) は失いますが、実用上は問題にならないことが多い、というのが要旨です。
組合せ : Pre-RMSNorm
- 現代の LLM は、この 2 軸の組合せである Pre-RMSNorm (Pre-LN の位置に RMSNorm を置く) を採用しています。
- 採用理由は、「安定化」と「高速化」です。2 つの軸は独立なので、両方を同時に得られます。
- Pre-LN → 安定化 : warmup なしでも訓練が安定し、深く・大規模にスケールしやすい。LLaMA も「訓練の安定性を高めるため」に Pre 化したと明記しています [5]。
- RMSNorm → 高速化 : 中心化を省く分計算が軽く、原論文では実行時間を 7〜64% 削減 (モデルや実装で幅があり、大きいのは RNN、Transformer では差は小さめ)[3]。精度は LayerNorm と同等を維持します。
- 精度について : Pre-RMSNorm は精度を上げるための選択ではありません。むしろ最終精度だけを見れば、Post-LN + LayerNorm の方が有利な場合もあります (前述)。それでも、安定して大規模化でき、かつ軽いことから広く使われています。後続研究でも、Pre-LN は「近年の Transformer のデファクトスタンダード」と述べられています [10]。
発展
- 正規化する位置の議論は今も続いており、Pre / Post 以外の配置も提案されています。
- DeepNorm [7] : Post-LN の残差を定数倍してから LayerNorm し、理論的に導いた初期化と組み合わせる。「Post-LN の性能」と「Pre-LN の安定性」を両立し、1000 層まで学習可能にした (Wang ら, 2022)。
- Sandwich-LN [8] : Pre-LN に加え、各ブロックの出力側にも LayerNorm を足して前後で挟む。深い Transformer の学習安定化のため CogView (Ding ら, 2021) が提案。ただし GLM-130B の学習では崩壊した例も報告されている。
- Peri-LN [9] : 各ブロックの入力側・出力側の両方に LayerNorm を置く配置 (Sandwich 系を整理・分析したもの)。分散の増大や勾配を安定させ、近年の Gemma・OLMo 系で採用されている (Kim ら, ICML 2025)。
トイコード
以下のトイコードは、$N$ 層の簡易モデルを warmup なしの Adam で数ステップ学習させ、正規化の位置 (Post-LN / Pre-LN) と種類 (LayerNorm / RMSNorm) の 4 通りで損失の下がり方がどう違うかを見るものです。モデルの構成は次のとおりです。
- 各ブロック = FFN (ReLU を挟む 2 層の線形変換) を残差でくるんだもの。正規化をブロックの後に置けば Post-LN ($x \leftarrow \mathrm{LN}(x + \mathrm{Block}(x))$)、前に置けば Pre-LN ($x \leftarrow x + \mathrm{Block}(\mathrm{LN}(x))$) です。
- このブロックを $N=32$ 層積み、最後に正規化と分類ヘッドを置いて 10 クラスのスコアを出します。
学習させるのはダミーデータです。ランダムなベクトル 64 個 ($d=32$) に、ランダムな 10 クラスのラベルを付けたものを交差エントロピーで丸暗記させ、この 64 サンプルで損失がどこまで下がるか (= 学習が進むか) を見ます。
T-Fixup の記事 は初期化を工夫して warmup を不要にする話でした。ここでは逆に、warmup を外すと Post-LN は損失が下がらず、Pre-LN は下がることを、実際に学習して確かめます。
import torch
class RMSNorm(torch.nn.Module):
def __init__(self, d, eps=1e-5):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(d))
self.eps = eps
def forward(self, x):
return self.weight * x / x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
def make_norm(kind, d):
return torch.nn.LayerNorm(d) if kind == 'LN' else RMSNorm(d)
def xavier(lin):
torch.nn.init.xavier_uniform_(lin.weight)
torch.nn.init.zeros_(lin.bias)
return lin
class Net(torch.nn.Module):
# FFN ブロックを N 層積み、最後に正規化 + 分類ヘッドを置いたモデル
def __init__(self, N, d, placement, kind='LN', n_class=10):
super().__init__()
self.placement = placement # 'Post' or 'Pre'
self.norms = torch.nn.ModuleList([make_norm(kind, d) for _ in range(N)])
self.fc1 = torch.nn.ModuleList([xavier(torch.nn.Linear(d, 4 * d)) for _ in range(N)])
self.fc2 = torch.nn.ModuleList([xavier(torch.nn.Linear(4 * d, d)) for _ in range(N)])
self.final_norm = make_norm(kind, d)
self.head = xavier(torch.nn.Linear(d, n_class))
def forward(self, x):
for norm, fc1, fc2 in zip(self.norms, self.fc1, self.fc2):
if self.placement == 'Pre': # Pre-LN: 正規化 → ブロック、残差は素通し
x = x + fc2(torch.relu(fc1(norm(x))))
else: # Post-LN: ブロック → 残差 → 正規化
x = norm(x + fc2(torch.relu(fc1(x))))
return self.head(self.final_norm(x))
def train_losses(placement, kind, N, d, steps, lr):
# warmup なしの Adam で数ステップ学習し、各ステップの損失を返す
torch.manual_seed(0)
model = Net(N, d, placement, kind)
X = torch.randn(64, d)
y = torch.randint(0, 10, (64,))
opt = torch.optim.Adam(model.parameters(), lr=lr)
losses = []
for _ in range(steps + 1):
loss = torch.nn.functional.cross_entropy(model(X), y)
losses.append(loss.item())
opt.zero_grad()
loss.backward()
opt.step()
return losses
def main():
d, N, steps, lr = 32, 32, 30, 0.01
runs = {f'{p}-{k}': train_losses(p, k, N, d, steps, lr)
for p in ['Post', 'Pre'] for k in ['LN', 'RMS']}
print(f'warmup なし・N={N} 層・Adam(lr={lr}) の損失推移')
print(f'{"step":>4} | ' + ' | '.join(f'{name:>8}' for name in runs))
for t in [0, 1, 2, 3, 5, 10, 20, 30]:
print(f'{t:>4} | ' + ' | '.join(f'{v[t]:>8.4f}' for v in runs.values()))
if __name__ == '__main__':
main()
実行すると以下のようになります。
warmup なし・N=32 層・Adam(lr=0.01) の損失推移
step | Post-LN | Post-RMS | Pre-LN | Pre-RMS
0 | 2.7399 | 2.7438 | 2.6015 | 2.6482
1 | 3.2239 | 3.4183 | 2.9216 | 3.0025
2 | 2.8667 | 2.9987 | 2.5541 | 2.4276
3 | 2.5200 | 2.6105 | 2.2131 | 2.1588
5 | 2.2596 | 2.2676 | 1.8619 | 1.8424
10 | 2.2218 | 2.2416 | 0.9175 | 0.8868
20 | 2.2087 | 2.2087 | 0.0920 | 0.0637
30 | 2.2052 | 2.2034 | 0.0201 | 0.0121
読み取れることは以下です。
- Post-LN / Post-RMS : どちらも損失が 2.2 前後で頭打ちになり、ほとんど学習が進みません (乱数当ての $\ln 10 \approx 2.30$ とほぼ同水準)。初期化時の大きな Adam 更新が深い残差を通って増幅され、訓練を不安定にするためです。
- Pre-LN / Pre-RMS : どちらも損失は下がり続け、30 step で 0.01〜0.02 まで下がります (64 個のダミーをほぼ丸暗記できている状態)1。残差合流後に正規化を挟まず初期化時の勾配が爆発も消失もしないので、warmup なしでも安定です。
- 正規化を LayerNorm から RMSNorm に変えてもほぼ同じ挙動です。学習の成否を決めるのは正規化の位置 (Pre / Post) であり、種類 (LayerNorm / RMSNorm) は無関係です。
-
交差エントロピーは 1 サンプルあたり $-\ln p$ ($p$ は正解クラスに与えた確率) で、その平均が損失です。0.02 は平均で $p \approx e^{-0.02} \approx 0.98$、つまりどのサンプルも約 98% の確率で正解している状態です。仮に 64 個のうち 1 個でも当てずっぽう ($p = 0.1$) だと、それだけで平均に $2.30 / 64 \approx 0.036$ を足してしまい 0.02 を超えます。よって 0.02 は「ひどく外したサンプルが 1 つもない (= 64 個をほぼ丸暗記できている)」ことを意味します。 ↩