Hayami Kaito Web
~/tech

正規化層のしくみ:バッチ正規化・層正規化・RMSNorm で深いネットワークを学習させる

published

ニューラルネットワークは、層を深くするほど複雑な関係を表せるようになります。しかし、単純に層を重ねるだけでは、途中の値が層ごとに小さくなったり大きくなったりして、うまく学習できなくなります。この問題をやわらげるために、今のほとんどのネットワークに入っているのが 正規化層 です。

画像のネットワーク(ResNet など)では バッチ正規化、Transformer や LLM では 層正規化 や RMSNorm が使われています。アテンションと Transformer のしくみ では「残差接続と層正規化で深くしても学習できるようにする」とだけ書きました。この記事では、その中身を説明します。

結論を先に書くと、次のようになります。

  • 層を重ねると、値の大きさが層ごとに一定の倍率で縮んだり膨らんだりする。20 層で 1000 分の 1 にも 1000 倍にもなる
  • 正規化層は、層ごとに値を「平均 0・分散 1」にそろえ直す。重みの大きさによらず、値の大きさが保たれる。計算例では、正規化のない 20 層のネットワークはまったく学習できず、バッチ正規化を入れると正解率 0.9 を超えた
  • バッチ正規化はミニバッチの中の データ方向、層正規化は 1 件の中の 特徴方向で平均と分散を取る。バッチ正規化はバッチが小さいと崩れ、層正規化はバッチの大きさに左右されない
  • バッチ正規化は、学習のときと推論のときで計算が違う。推論では、学習中に貯めておいた平均と分散を使う

説明の途中に出てくる計算例は、Python と NumPy で実際に動かした結果です。

環境

計算例を動かした環境です。

項目バージョン
Python3.13
NumPy2.5

層を重ねると値が消える・膨らむ

20 層を通した値の大きさ

64 個のニューロンの層を 20 個重ね、それぞれで「重みを掛けて ReLU を通す」計算をします。重みの初期値の大きさを 3 通りに変えて、1・5・10・20 層目の出力の標準偏差を見ます。

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(256, 64))
for name, scale in (("小さめの初期値", 1.0), ("He の初期値", np.sqrt(2)), ("大きめの初期値", 2.0)):
    h, stds = X, []
    for layer in range(20):
        W = rng.normal(0, scale / np.sqrt(64), (64, 64))
        h = np.maximum(0, h @ W)
        stds.append(h.std())
    print(name, [f"{stds[i]:.1e}" for i in (0, 4, 9, 19)])

h, stds = X, []
for layer in range(20):
    W = rng.normal(0, 2.0 / np.sqrt(64), (64, 64))
    a = h @ W
    a = (a - a.mean(0)) / np.sqrt(a.var(0) + 1e-5)   # バッチ正規化(γ = 1, β = 0)
    h = np.maximum(0, a)
    stds.append(h.std())
print("大きめ + バッチ正規化", [f"{stds[i]:.1e}" for i in (0, 4, 9, 19)])

実行結果:

小さめの初期値 ['5.8e-01', '1.3e-01', '1.6e-02', '7.5e-04']
He の初期値 ['8.3e-01', '9.0e-01', '6.0e-01', '3.5e-01']
大きめの初期値 ['1.2e+00', '5.5e+00', '3.3e+01', '1.7e+03']
大きめ + バッチ正規化 ['5.8e-01', '5.8e-01', '5.8e-01', '5.7e-01']
  • 小さめの初期値(標準偏差 1/√64)では、ReLU が負の側を 0 にするぶん、1 層ごとに値がおよそ 1/√2 倍に縮み、20 層目では 1000 分の 1 以下になりました
  • 大きめの初期値(2/√64)では、逆に 1 層ごとに約 √2 倍に膨らみ、20 層目では 1000 倍以上です
  • He の初期値(√2/√64)は、ReLU のときに大きさがちょうど保たれるように決めた値です(He ら, 2015)。それでも、乱数のばらつきで 20 層目では 0.35 まで縮んでいます
  • バッチ正規化を入れると、大きめの初期値でも 20 層を通して値の大きさが変わらない。各層で平均 0・分散 1 にそろえ直すからです

値が縮むと、逆伝播で戻ってくる勾配も同じように縮み、前の方の層がほとんど学習されません(ニューラルネットワークの学習のしくみ の勾配消失)。初期値を工夫すれば最初は大きさを保てますが、学習が進んで重みが変わると、またずれていきます。正規化層は、学習のあいだずっと、値の大きさを保ち続けます。

バッチ正規化と層正規化

平均と分散を取る向きが違う

正規化層は、層の出力 a(正規化の前の値)を次のように変換します。

x̂ = (a − 平均) / √(分散 + ε)
出力 = γ × x̂ + β

ε は 0 で割らないための小さな数です。γ と β は学習する値で、正規化でそろえた大きさと位置を、学習で必要なら変えられるようにするためのものです(γ = 標準偏差、β = 平均にすれば、元に戻すこともできます)。

バッチ正規化と層正規化の違いは、平均と分散を、どの値の集まりで計算するか だけです。ミニバッチを「データの数 × 特徴の数」の表と考えると、次のようになります。

種類平均と分散を取る向き意味
バッチ正規化(Ioffe と Szegedy, 2015)表の縦(データ方向)特徴ごとに、ミニバッチの中のデータで平均 0・分散 1 にする
層正規化(Ba ら, 2016)表の横(特徴方向)データ 1 件ごとに、その件の特徴の値で平均 0・分散 1 にする

画像の場合は、バッチ正規化は「チャンネルごとに、バッチの全画像・全画素」で平均を取ります。ほかにも、画像 1 枚のチャンネルをいくつかのグループに分けて平均を取る グループ正規化(Wu と He, 2018)などがあります。

実装

バッチ正規化と層正規化の順伝播と逆伝播を書きます。ax を 0 にするか 1 にするかだけで、両方を同じコードで計算できます。逆伝播は、数値微分と一致することを確かめてあります。

バッチ正規化では、学習中にミニバッチの平均と分散の移動平均(rm・rv)を貯めておき、推論(train=False)ではそれを使います。理由は後で説明します。

import numpy as np

def norm_forward(a, kind, L, train):
    """バッチ正規化はデータ方向(axis=0)、層正規化は特徴方向(axis=1)で平均と分散を取る"""
    if kind == "none":
        return a, None
    if kind == "bn" and not train:
        xhat = (a - L["rm"]) / np.sqrt(L["rv"] + 1e-5)   # 推論では学習中に貯めた平均と分散を使う
        return L["g"] * xhat + L["be"], None
    ax = 0 if kind == "bn" else 1
    mu, var = a.mean(ax, keepdims=True), a.var(ax, keepdims=True)
    if kind == "bn":
        L["rm"] = 0.9 * L["rm"] + 0.1 * mu[0]; L["rv"] = 0.9 * L["rv"] + 0.1 * var[0]
    xhat = (a - mu) / np.sqrt(var + 1e-5)
    return L["g"] * xhat + L["be"], (xhat, var, ax)

def norm_backward(d, cache, kind, L):
    if kind == "none":
        return d, None, None
    xhat, var, ax = cache
    dxhat = d * L["g"]
    da = (dxhat - dxhat.mean(ax, keepdims=True) - xhat * (dxhat * xhat).mean(ax, keepdims=True)) / np.sqrt(var + 1e-5)
    return da, (d * xhat).sum(0), d.sum(0)

逆伝播の式の dxhat.mean(...) と xhat * (dxhat * xhat).mean(...) の 2 つの引き算は、「平均を引いた」「分散で割った」ことの微分です。平均と分散そのものが入力に依存しているので、その分の勾配を差し引いています。

深いネットワークを学習させる

20 個の入力から、sin と掛け算で決まる 2 値の答えを当てる問題を、隠れ層 20 層(64 ニューロン)のネットワークで解きます。重みの初期値は、上の「小さめの初期値」です。最適化は、ミニバッチ 64 件のふつうの確率的勾配降下法で、1500 回更新します。

def make_data(rng, n):
    X = rng.normal(size=(n, 20))
    y = (np.sin(X[:, 0] * 2) + X[:, 1] * X[:, 2] > 0).astype(int)
    return X, y

def init(rng, depth, H=64):
    dims = [20] + [H] * depth + [2]
    return [{"W": rng.normal(0, 1 / np.sqrt(dims[i]), (dims[i], dims[i + 1])), "b": np.zeros(dims[i + 1]),
             "g": np.ones(dims[i + 1]), "be": np.zeros(dims[i + 1]),
             "rm": np.zeros(dims[i + 1]), "rv": np.ones(dims[i + 1])} for i in range(len(dims) - 1)]

def forward(layers, X, kind, train):
    caches, h = [], X
    for L in layers[:-1]:
        an, nc = norm_forward(h @ L["W"] + L["b"], kind, L, train)
        out = np.maximum(0, an)
        caches.append((h, nc, out)); h = out
    caches.append((h, None, None))
    return h @ layers[-1]["W"] + layers[-1]["b"], caches

def step(layers, X, y, kind, lr):
    z, caches = forward(layers, X, kind, train=True)
    z = z - z.max(1, keepdims=True); P = np.exp(z); P /= P.sum(1, keepdims=True)
    loss = -np.mean(np.log(P[np.arange(len(y)), y] + 1e-12))
    d = P; d[np.arange(len(y)), y] -= 1; d /= len(y)
    for i in reversed(range(len(layers))):
        L = layers[i]; hin, nc, out = caches[i]
        if i < len(layers) - 1:
            d, dg, dbe = norm_backward(d * (out > 0), nc, kind, L)
            if dg is not None:
                L["g"] -= lr * dg; L["be"] -= lr * dbe
        dW, db = hin.T @ d, d.sum(0)
        d = d @ L["W"].T
        L["W"] -= lr * dW; L["b"] -= lr * db
    return loss

def run(kind, lr, depth=20, batch=64, steps=1500):
    rng = np.random.default_rng(0)
    X, y = make_data(rng, 4000)
    Xte, yte = make_data(rng, 2000)
    layers = init(rng, depth)
    for _ in range(steps):
        idx = rng.choice(4000, batch, replace=False)
        loss = step(layers, X[idx], y[idx], kind, lr)
    if not np.isfinite(loss):
        return float("nan")
    return np.mean(forward(layers, Xte, kind, train=False)[0].argmax(1) == yte)

正規化なし・バッチ正規化・層正規化のそれぞれで、学習率を 0.01・0.1・0.5 に変えて、評価用のデータでの正解率を比べます。

np.seterr(all="ignore")
for kind in ("none", "bn", "ln"):
    print(kind, {lr: round(float(run(kind, lr)), 3) for lr in (0.01, 0.1, 0.5)})

実行結果:

none {0.01: 0.501, 0.1: 0.501, 0.5: 0.501}
bn {0.01: 0.583, 0.1: 0.904, 0.5: 0.916}
ln {0.01: 0.736, 0.1: 0.818, 0.5: 0.501}
  • 正規化なしでは、どの学習率でもまったく学習できない(0.501 は当て推量と同じ)。20 層を通るうちに値と勾配が 1000 分の 1 以下に縮み、学習率 0.5 でも重みがほとんど動きません
  • バッチ正規化を入れると、正解率は 0.9 を超えた。しかも、学習率を 0.5 まで大きくしても学習が壊れず、いちばん良い成績になりました。大きな学習率が使えるのは、バッチ正規化の大事な長所です。重みが大きくなっても、正規化で出力の大きさがそろえられるので、更新しすぎの影響が出にくいのです
  • 層正規化も学習はできるが、この問題ではバッチ正規化に届かず、学習率 0.5 では壊れた。全結合のネットワークや CNN では、バッチ正規化の方が良い成績になりやすいことが知られています。層正規化が力を発揮するのは、次に説明する RNN や Transformer です

なぜ正規化で学習しやすくなるのか

バッチ正規化の最初の論文は、「前の層の重みが変わるたびに、次の層への入力の分布が変わってしまう問題(内部共変量シフト)を減らすため」と説明していました。しかし、その後の研究(Santurkar ら, 2018)で、効果の主な理由は別のところにあり、損失の地形をなめらかにして、勾配の向きと大きさを安定させること だと示されました。勾配が急に変わらないので、大きな学習率で進んでも行き過ぎにくくなります。

バッチ正規化の 2 つの弱点

学習と推論で計算が違う

推論のときは、データが 1 件ずつ来ることがあります。1 件だけで「データ方向の平均と分散」を取ると、平均はその値そのもの、分散は 0 になり、出力がすべて β になってしまいます。また、同じ画像でも、一緒にバッチに入った画像によって答えが変わるのは困ります。

そこでバッチ正規化は、学習中に各ミニバッチの平均と分散の 移動平均 を貯めておき、推論ではそれを使います。上の norm_forward の rm・rv がそれです。この切り替えには注意が必要です。

よくある失敗起きること
推論なのに学習のモードのままバッチの中身によって答えが変わる。1 件だと答えがおかしくなる
微調整のデータが少ないのに、移動平均を更新し続ける少ないデータの偏った統計に引っ張られる。転移学習 では統計を固定することもある
学習と推論でデータの前処理が違う貯めた平均と分散が、推論のデータに合わなくなる

PyTorch では、model.train() と model.eval() でこのモードを切り替えます。

バッチが小さいと崩れる

バッチ正規化は、ミニバッチの中の数件から平均と分散を見積もります。件数が少ないと見積もりのばらつきが大きくなり、うまく学習できなくなります。高解像度の画像の物体検出やセグメンテーションでは、メモリーの都合でバッチが 1〜2 枚になることがあり、これが実際の問題になります。

隠れ層 10 層のネットワークで、ミニバッチの大きさを 64・8・2 にして比べます。学習で見るデータの総数がそろうように更新の回数を増やし、学習率はバッチの大きさに比例させています。

np.seterr(all="ignore")
for batch in (64, 8, 2):
    steps = 1500 * 64 // batch   # 見るデータの数をそろえる
    print(batch, {kind: round(float(run(kind, 0.1 * batch / 64, depth=10, batch=batch, steps=steps)), 3) for kind in ("bn", "ln")})

実行結果:

64 {'bn': 0.855, 'ln': 0.866}
8 {'bn': 0.853, 'ln': 0.872}
2 {'bn': 0.506, 'ln': 0.873}

バッチ正規化は、バッチが 8 件までは成績を保ちましたが、2 件では 0.506 と、まったく役に立たないモデルになりました。2 件の平均と分散は、データ全体の平均と分散とは大きく違うので、学習中の計算と、推論で使う移動平均の計算が食い違ってしまいます。層正規化は、1 件の中で計算するので、バッチの大きさによらず同じ成績です。

Transformer の正規化:層正規化と RMSNorm

系列の長さがまちまちでも使える

文章を扱うモデルでは、文の長さがまちまちで、同じ位置に意味のある単語があるとは限りません。データ方向に平均を取るバッチ正規化は、こうしたデータと相性が悪く、RNN や Transformer では層正規化が使われてきました。1 つのトークンのベクトルの中で正規化するので、バッチの大きさにも系列の長さにも左右されず、生成のときに 1 トークンずつ計算するのにも向いています。

RMSNorm:平均を引かない

最近の LLM の多くは、層正規化から 平均を引く処理と β を省いた RMSNorm(Zhang と Sennrich, 2019)を使っています。

層正規化:出力 = γ × (a − 平均) / √(分散 + ε) + β
RMSNorm:出力 = γ × a / √(a の 2 乗の平均 + ε)

計算が少し減り、成績はほとんど変わらないことが報告されています。値の位置(平均)より、大きさをそろえることが大事だということです。

正規化を置く場所:Post-LN と Pre-LN

Transformer のブロックでは、正規化を残差接続のどこに置くかで、学習のしやすさが変わります。

置き方計算特徴
Post-LN(元の Transformer)x ← 正規化(x + 層(x))残差の道にも正規化が入る。深くすると学習が不安定になりやすく、学習率をゆっくり上げるウォームアップが欠かせない
Pre-LN(今の LLM の主流)x ← x + 層(正規化(x))残差の道が正規化を通らないので、勾配がそのまま流れる。深くしても安定して学習できる(Xiong ら, 2020)

残差接続の「足し算の道」を何も通さずに残しておくのは、RNN と LSTM のしくみ で説明した LSTM のセルと同じ考え方です。

正規化層の手法の流れ

年手法・出来事ポイント
2010Xavier の初期値(Glorot と Bengio)層を通しても値の分散が変わらないように、重みの初期値の大きさを決める
2015He の初期値ReLU のときに分散が保たれる初期値
2015バッチ正規化ミニバッチの統計で特徴ごとに正規化する。大きな学習率で速く学習できるようになった
2016層正規化・重み正規化1 件ごとに正規化する。RNN でも使える
2016インスタンス正規化画像 1 枚・チャンネルごとに正規化する。画風の変換で使われる
2017Transformer残差接続のあとに層正規化を置く(Post-LN)
2018グループ正規化チャンネルをグループに分けて正規化する。小さなバッチでも使える
2018バッチ正規化の効果の再検討(Santurkar ら)効果の主な理由は、損失の地形をなめらかにすること
2019RMSNorm平均を引かずに大きさだけをそろえる。今の LLM で広く使われる
2020Pre-LN の分析(Xiong ら)正規化を層の前に置くと、深い Transformer が安定して学習できる

まとめ

  • 層を重ねると、値と勾配の大きさが層ごとに一定の倍率で縮んだり膨らんだりする。初期値の工夫だけでは、学習中のずれまでは防げない
  • 正規化層は、層ごとに値を平均 0・分散 1 にそろえ、学習する γ と β で大きさと位置を調整する。正規化のない 20 層のネットワークは学習できなかったが、バッチ正規化を入れると正解率 0.9 を超え、大きな学習率も使えた
  • バッチ正規化はデータ方向、層正規化は特徴方向で平均と分散を取る。バッチ正規化は推論で学習中の移動平均を使い、バッチが小さいと崩れる
  • Transformer や LLM では、バッチの大きさや系列の長さに左右されない層正規化や RMSNorm を、層の前に置く(Pre-LN)のが主流

参考

  • X. Glorot and Y. Bengio, "Understanding the difficulty of training deep feedforward neural networks", AISTATS 2010
  • K. He et al., "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification", ICCV 2015(arXiv:1502.01852)
  • S. Ioffe and C. Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift", ICML 2015(arXiv:1502.03167)
  • J. L. Ba, J. R. Kiros and G. E. Hinton, "Layer Normalization", 2016(arXiv:1607.06450)
  • T. Salimans and D. P. Kingma, "Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks", NIPS 2016(arXiv:1602.07868)
  • D. Ulyanov, A. Vedaldi and V. Lempitsky, "Instance Normalization: The Missing Ingredient for Fast Stylization", 2016(arXiv:1607.08022)
  • Y. Wu and K. He, "Group Normalization", ECCV 2018(arXiv:1803.08494)
  • S. Santurkar et al., "How Does Batch Normalization Help Optimization?", NeurIPS 2018(arXiv:1805.11604)
  • B. Zhang and R. Sennrich, "Root Mean Square Layer Normalization", NeurIPS 2019(arXiv:1910.07467)
  • R. Xiong et al., "On Layer Normalization in the Transformer Architecture", ICML 2020(arXiv:2002.04745)