正規化レイヤー全史 — BatchNormからRMSNormまで

深層学習を「実際に深くできるもの」に変えた正規化レイヤーを前提知識ゼロから辿る。BatchNormが効く理由をめぐる内部共変量シフト論争、LayerNormがバッチ軸を捨てた理由、そしてLLMが一様にRMSNormへ行き着いた理由まで。

Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

積み上げるほど学習が止まる

2015年ごろまで、ニューラルネットには奇妙な壁がありました。理屈では層を深くするほど表現力が上がるはずなのに、実際に10段20段と積むと学習が進まない。学習率を大きくすれば発散し、小さくすれば損失が下がらないのです。

原因のひとつは、層のあいだを流れる数値のスケールでした。組立ラインで前工程の部品寸法が日によって2倍や10分の1になれば、後工程は毎朝、道具の当て方から調整し直すはめになります。学習中のネットでも同じで、前の層の重みが更新されるたび、後ろの層は「入ってくる値の大きさ」から学び直すことになる。正規化レイヤーは、層の出口に寸法を揃える治具を挟んでこれを断ち切る道具です。この記事では2015年のBatchNormから、いまLLMがほぼ一様に使うRMSNormまでを、なぜ乗り換えが起きたのかという理由込みで辿ります。

値がはみ出すと、勾配が消える

なぜスケールがそこまで問題になるのか。活性化関数を動かしてみるのが早い。

平らな領域では傾きがほぼ0です。学習は勾配の向きに重みを動かすことでしか進まないので、勾配が0になった層はそこで学ぶのをやめます。しかも誤差逆伝播は勾配を掛け算で後ろへ運ぶため、途中の1層が0に近い値を返せば手前の層すべてが道連れです（誤差逆伝播を1から理解する）。正規化の狙いは、層に入る値をこの平らな崖の手前へ押し戻すことです。

平均0・分散1にして、また戻す

正規化レイヤーの中身は、どの手法でも次の2式に尽きます。

$x$ は正規化したい値、$\mu$（ミュー）はその値が属する集まりの平均、$\sigma^2$ は分散、$\varepsilon$ はゼロ割りを防ぐ小さな定数です。$\gamma$（ガンマ）と $\beta$（ベータ）は学習で決まり、「もう一度どれだけ広げるか／ずらすか」を担います。式(1)は要するに、いったん平均0・広がり1に揃えてから、必要なぶんだけ学習で広げ直すと言っています。

揃えたものを戻すのは表現力を守るためです。全層の出力を強制的に平均0・分散1にすると、値の大きさそのもので情報を運べなくなる。$\gamma$ と $\beta$ があれば、モデルは必要なとき $\gamma=\sigma$、$\beta=\mu$ を学んで正規化を打ち消すこともできます。

ここまでは全手法に共通で、3手法を分けるのはたった1点——$\mu$ と $\sigma$ を、どの範囲の値から計算するかだけです。

BatchNorm: バッチ全体で平均を取る

BatchNorm（Ioffe & Szegedy, 2015）は、ミニバッチのサンプル方向に統計を取ります。畳み込み層なら、チャネルごとに「バッチ内の全画像・全ピクセル」を1つの平均と1つの分散にまとめる。つまり、ある1枚がどう変換されるかがたまたま同じバッチに居合わせた他の画像に依存する。ここがBatchNormの効きどころであり、同時にあらゆる厄介ごとの根っこです。

推論時はバッチが1枚かもしれず、バッチ統計が使えません。そこで学習中に平均と分散の移動平均を貯めておき、推論ではそちらを使います。学習時と推論時で計算式そのものが違う——PyTorchで model.train() と model.eval() を切り替えねばならない理由はここにあります。

「内部共変量シフト」論争

原論文はBatchNormが効く理由を、内部共変量シフト（Internal Covariate Shift, ICS）——前の層が更新されるたびに後ろの層へ入る値の分布が動く現象——の低減だと説明しました。論文タイトルにも "Reducing Internal Covariate Shift" と入っています。直感的で覚えやすく、長く教科書的な理解として定着しました。

ところが2018年、Santurkarらの "How Does Batch Normalization Help Optimization?" がこれを正面から否定します。彼らはBatchNormの直後にわざとノイズを注入しました。平均が0でなく、分散も1でなく、しかもステップごとに変わるノイズです。ICSの定義からすれば最悪の条件のはずが、このネットワークは通常のBatchNorm付きとほぼ同じ速さで学習しました。ICSを意図的に増やしても性能が落ちないなら、ICSの低減は効果の説明になりません。