Peri-LN: Revisiting Normalization Layer in the Transformer Architecture [1] は、各サブ層の入力と出力の両方に正規化を置く配置を Peri-LN と名付け、Post-LN・Pre-LN と比べて中間出力の分散の増大がゆるやかで、勾配が安定し、学習が発散しにくいことを解析と実験で示した論文です。NAVER Cloud・KAIST の Jeonghoon Kim らによって著され、ICML 2025 に採択されました。
参考文献
- Jeonghoon Kim, Byeongchan Lee, Cheonbok Park, Yeontaek Oh, Beomjun Kim, Taehwan Yoo, Seongjin Shin, Dongyoon Han, Jinwoo Shin, Kang Min Yoo. Peri-LN: Revisiting Normalization Layer in the Transformer Architecture. Proceedings of the 42nd International Conference on Machine Learning (ICML 2025), vol. 267, pp. 30400−30436, 2025.
https://proceedings.mlr.press/v267/kim25u.html
https://arxiv.org/abs/2502.02732
https://arxiv.org/html/2502.02732
以下の内容は arXiv 版 (v3) に基づきます。
文献概要
Peri-LN の定義
$l$ 番目のサブ層 (アテンションまたは FFN) を $\mathrm{Module}$ として、Post-LN は $x_{l+1} = \mathrm{Norm}(x_l + \mathrm{Module}(x_l))$、Pre-LN は $x_{l+1} = x_l + \mathrm{Module}(\mathrm{Norm}(x_l))$ です。Peri-LN は、Pre-LN のサブ層の出力側にも正規化を足して
x_{l+1} = x_l + \mathrm{Norm}\Bigl(\mathrm{Module}\bigl(\mathrm{Norm}(x_l)\bigr)\Bigr)
とします。Pre-LN と同じく最終層の出力も正規化し、さらに、最初の層に入力するトークンの埋め込みも正規化してよいとしています。
この配置そのものは著者らの考案ではなく、Gemma 系・OLMo 系のモデルですでに採用されていましたが、統一された呼び名がなく、その性質もほとんど調べられていませんでした。そこで著者らは、サブ層を前後から正規化で囲むかたちであることから、「周囲」を意味する接頭辞 peri- (peripheral の peri-) をとってこの配置を Peri-LN と名付け、その性質を調べています。
中間出力の分散と勾配
著者らは、3 つの配置について、中間出力の分散が深さとともにどう増えるかを解析し、以下のように説明しています。特に、Pre-LN は安定に学習できるというのが通説ですが、著者らの Pre-LN についての説明はこれと異なり、学習が進むと中間出力が増大しうるとするものです。
- Post-LN: 残差接続の後で正規化するので中間出力の分散はほぼ一定に保たれるが、深いモデルでは勾配が届きにくくなりうる。
- Pre-LN: 残差パス上に正規化がないので、初期化時点で中間出力の分散が深さに対して線形に増え、学習が進んでサブ層の出力が大きくなると指数的に増えうる。
- Peri-LN: 残差パスに足されるのは正規化後の出力なので、その分散をほぼ一定の $\beta_0$ とすると $\mathrm{Var}(x_{l+1}) \approx \mathrm{Var}(x_l) + \beta_0$ となり、分散の増大は線形程度にとどまる。
勾配についても著者らは解析しており、FFN の 2 層目の重みの勾配ノルムが、Pre-LN では 1 層目の出力 $h$ の大きさに比例するため、中間出力が巨大になると勾配も爆発しうるのに対し、Peri-LN では出力側の正規化によって $h$ が大きくても有界に保たれることを示しています (Proposition 3.1)。
実験
400M・1.5B・3.2B パラメータ (埋め込みを除く) の Transformer を、それぞれ 30B トークン・5 シードで事前学習しています (正規化は主に RMSNorm)。評価損失1 (C4) とベンチマーク 5 種 (ARC-Easy・HellaSwag・PIQA・SIQA・Winogrande) の平均スコア2は以下の通りで、どのサイズでも Peri-LN が最良でした。
| サイズ | Post-LN | Pre-LN | Peri-LN | |
|---|---|---|---|---|
| 400M | 評価損失 | 7.46 | 3.43 | 3.34 |
| 平均スコア | 42.45 | 49.69 | 51.57 | |
| 1.5B | 評価損失 | 5.38 | 3.29 | 3.18 |
| 平均スコア | 45.49 | 53.71 | 56.55 | |
| 3.2B | 評価損失 | 4.43 | 3.20 | 3.11 |
| 平均スコア | 46.45 | 56.69 | 58.56 |
このほか、Pre-LN では学習初期の発散がシードによらず頻繁に起きたこと (発散したチェックポイントは上の表から除外されています)、Peri-LN ではシード間のスコアの標準偏差が小さいこと、3.2B の Pre-LN では中間出力が学習のごく初期 (0.5B トークン時点) に FP16 の最大値を超えることを報告しています。