GPTではなぜレイヤー正規化が大切なのか?

こんにちは。ゆうせいです。

前回、位置エンコーディングを扱いました。今回はレイヤー正規化です。

Transformerの図を見ると、Add and Normという箱が何度も出てきます。説明されるときは「学習を安定させるため」の一言で済まされることが多く、なぜ必要なのかが分からないまま先へ進んでしまいます。

この記事では、レイヤー正規化が何をしているのか、なぜバッチ正規化ではだめなのか、そしてGPT-2が元のTransformerから配置を変えた理由を解説します。

結論

レイヤー正規化が重要な理由は、三つあります。

一つ目は、値の暴走を防ぐことです。Transformerは残差接続によって層ごとに値を足し込み続けるため、何もしなければ数値が際限なく大きくなります。

二つ目は、勾配を安定させることです。層を深く積むほど、勾配が消えるか発散するかの問題が深刻になります。

三つ目は、バッチに依存しないことです。文の長さが可変で、生成時には一トークンずつ処理するGPTにとって、バッチ統計に頼る正規化は使えません。

さらにGPT-2は、正規化の位置を元のTransformerから変更しました。この変更により、学習率のウォームアップなしでも安定して学習できるようになっています。

レイヤー正規化が何をしているか

計算式

入力ベクトル x = (x_1, x_2, \ldots, x_d) に対し、次の計算を行います。

まず平均を求めます。

\mu = \frac{1}{d} \sum_{i=1}^{d} x_i

次に分散を求めます。

\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2

そして正規化し、学習可能なパラメータで再調整します。

y_i = \gamma_i \cdot \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_i

\gamma はスケール、\beta はシフトを表す学習パラメータです。\epsilon はゼロ除算を防ぐための小さな値で、10^{-5} 程度が使われます。

重要なのは、平均と分散をどの方向で取るかです。レイヤー正規化は、一つのトークンの中で、768個の特徴量にわたって統計を計算します。他のトークンや他のサンプルは一切関与しません。

数値で確認する

次元を4に簡略化して計算します。

x = (2, 4, 4, 10)

平均を求めます。

\mu = \frac{2 + 4 + 4 + 10}{4} = \frac{20}{4} = 5

各要素から平均を引きます。

(2-5, 4-5, 4-5, 10-5) = (-3, -1, -1, 5)

分散を求めます。

\sigma^2 = \frac{(-3)^2 + (-1)^2 + (-1)^2 + 5^2}{4} = \frac{9 + 1 + 1 + 25}{4} = \frac{36}{4} = 9

標準偏差は次の通りです。

\sigma = \sqrt{9} = 3

正規化します。\epsilon は無視できる大きさなので省略します。

\left( \frac{-3}{3}, \frac{-1}{3}, \frac{-1}{3}, \frac{5}{3} \right) = (-1, -0.333, -0.333, 1.667)

結果を確認します。平均はゼロ、分散は1になっています。

\frac{-1 - 0.333 - 0.333 + 1.667}{4} = 0

元のベクトルの値の大小関係は保たれたまま、尺度だけが揃いました。これがレイヤー正規化の働きです。

幾何学的な意味

もう少し踏み込むと、この操作は二段階の射影として理解できます。

平均を引く操作は、すべての成分が1であるベクトルの方向の成分を取り除くことに相当します。

\mathbf{1} = (1, 1, \ldots, 1)

x - \mu \mathbf{1} \perp \mathbf{1}

標準偏差で割る操作は、長さを揃えることに相当します。正規化後のベクトルの大きさは、次のようになります。

| y | = \sqrt{\sum_{i=1}^{d} y_i^2} = \sqrt{d}

上の例で確認します。

\sqrt{(-1)^2 + (-0.333)^2 + (-0.333)^2 + (1.667)^2} = \sqrt{1 + 0.111 + 0.111 + 2.778} = \sqrt{4} = 2 = \sqrt{4}

つまり、レイヤー正規化は入力を、半径 \sqrt{d} の球面上に押し出しています。

以前の記事で、埋め込みの意味は「方向」にあると書きました。レイヤー正規化は、まさに長さの情報を捨てて方向だけを残す操作です。この対応関係を押さえておくと、なぜこの正規化が言語モデルで機能するのかが見えてきます。

なぜバッチ正規化ではないのか

画像認識ではバッチ正規化が標準です。しかし言語モデルでは、ほとんど使われません。理由は四つあります。

論点バッチ正規化レイヤー正規化
統計を取る方向バッチ内の同じ特徴量にわたって一つのサンプルの特徴量にわたって
バッチサイズへの依存小さいと不安定依存しない
系列長が可変の場合パディング部分が統計を歪める影響しない
推論時訓練中の移動平均を使う訓練時と同じ計算

系列長の問題

文の長さは一定ではありません。バッチの中に10トークンの文と500トークンの文が混在します。短い文はパディングで埋められますが、この無意味な部分がバッチ統計に混入します。位置ごとに有効なサンプル数が異なるため、統計の信頼性が位置によって変わってしまいます。

生成時の問題

GPTは、一トークンずつ生成します。このとき、バッチサイズは実質1です。バッチ正規化は、バッチ内の複数サンプルから統計を取る仕組みですから、この状況では機能しません。

訓練時の移動平均を使う回避策はありますが、訓練時と推論時で計算が変わるという不整合が生じます。

レイヤー正規化には、この問題がありません。一つのトークンの中だけで完結するため、バッチサイズが1でも、系列長が何であっても、まったく同じ計算になります。

残差接続との関係

値が増え続ける問題

Transformerの各層は、次の形をしています。

x_{l+1} = x_l + F(x_l)

前の層の出力に、新しい計算結果を足し込みます。この足し込みが12層、24層と続きます。

何も制御しなければ、値は層を追うごとに大きくなります。分散に注目すると、足し合わせるベクトルが独立に近い場合、次の関係が成り立ちます。

\text{Var}(x_l + F(x_l)) \approx \text{Var}(x_l) + \text{Var}(F(x_l))

層を重ねるほど分散が積み上がります。層数を L とすると、最終層の大きさはおおよそ \sqrt{L} に比例して増加します。

24層であれば、\sqrt{24} \approx 4.9 倍です。この程度なら扱えますが、層数が増え、しかも各層の出力が独立でない場合には、もっと急激に増大します。

レイヤー正規化は、この積み上がりを各層でリセットします。

残差の流れを壊さない配置

ただし、単純に足した後で正規化すると、別の問題が生じます。ここが次の話題につながります。

GPT-2が変えたこと

二つの配置

正規化の置き場所には、二つの方式があります。

元のTransformer論文の方式は、次の形です。

x_{l+1} = \text{LN}(x_l + F(x_l))

足してから正規化するため、後置正規化(Post-LN)と呼ばれます。

GPT-2が採用したのは、次の形です。

x_{l+1} = x_l + F(\text{LN}(x_l))

正規化してから計算に入れるため、前置正規化(Pre-LN)と呼ばれます。GPT-2の論文には、レイヤー正規化を各サブブロックの入力へ移動し、さらに最後の自己注意ブロックの後に正規化を一つ追加した、と記載されています。

なぜ移動したのか

差は、勾配の流れに現れます。

後置正規化では、入力から出力までのすべての経路が正規化を通過します。逆伝播の際、勾配は各層の正規化を通るたびに変形されます。層が深いほど、この変形が累積します。

前置正規化では、x_l から x_{l+1} への直接の経路に、正規化が入りません。加算だけです。この経路は、勾配をそのまま通します。

\frac{\partial x_{l+1}}{\partial x_l} = I + \frac{\partial F(\text{LN}(x_l))}{\partial x_l}

第一項が単位行列であることが要点です。層を何段重ねても、勾配が減衰せずに入力側まで届きます。

この経路は、残差ストリームと呼ばれます。層をまたいで情報と勾配を運ぶ、幹線道路のようなものです。前置正規化は、この幹線に信号機を置かない設計だといえます。

ウォームアップが不要になる

実務上の最大の違いがこれです。

後置正規化のモデルは、学習の初期に勾配が不安定になります。そのため、学習率を最初は極小さくし、徐々に上げていくウォームアップが必須でした。

前置正規化では、初期から勾配が安定するため、ウォームアップなしでも学習が進みます。学習率の設定にも寛容になります。

大規模モデルの学習では、一度の失敗の代償が極めて大きくなります。安定性は、それ自体が価値を持ちます。

最後の正規化

前置正規化には、一つ副作用があります。最終層の出力が正規化を通らないため、値の大きさが揃いません。

そのためGPT-2では、全ブロックの後、出力層に入る前に、もう一度レイヤー正規化を置いています。これが論文にある追加の正規化です。

位置目的
各サブブロックの入力計算に入る値の尺度を揃える
全ブロックの後出力層に渡す値の尺度を揃える

現在の主流はRMSNorm

近年の大規模モデルでは、簡略版が使われることが増えています。RMSNormです。

y_i = \gamma_i \cdot \frac{x_i}{\sqrt{\frac{1}{d}\sum_{j=1}^{d} x_j^2 + \epsilon}}

平均を引く操作を省き、シフトのパラメータ \beta も持ちません。二乗平均平方根で割るだけです。

先ほどの数値例で計算してみます。

x = (2, 4, 4, 10)

\frac{1}{4}(2^2 + 4^2 + 4^2 + 10^2) = \frac{4 + 16 + 16 + 100}{4} = \frac{136}{4} = 34

\sqrt{34} \approx 5.831

\left( \frac{2}{5.831}, \frac{4}{5.831}, \frac{4}{5.831}, \frac{10}{5.831} \right) \approx (0.343, 0.686, 0.686, 1.715)

平均を引いていないため、結果は先ほどと異なります。

省略しても性能がほとんど変わらないことが実験的に示されており、計算量が減る分だけ有利です。長さを揃えるという本質的な機能は保たれています。方向の情報が主役であるという先ほどの見方と、整合的な簡略化だといえます。

解釈可能性への影響

以前の記事で、埋め込みの一部の次元が桁違いに大きな値を持つという現象に触れました。これらの外れ値次元は、レイヤー正規化と関係があると指摘されています。

正規化によって全体の尺度が揃えられる環境では、特定の次元だけを極端に大きくすることが、他の次元を実質的に抑制する手段になります。モデルがこの性質を利用している可能性が議論されています。

また、レイヤー正規化は非線形な操作です。内部の計算を線形の回路として読み解こうとする解釈可能性の研究では、この非線形性が障害になります。正規化のパラメータを後続の重み行列に畳み込んで扱う手法が用いられますが、分散で割る部分は完全には線形化できません。

安定性のために導入された仕組みが、理解の妨げにもなっている。この点は、興味深い緊張関係です。

研修での教え方

このテーマを扱う場合、次の順序を推奨します。

  1. 残差接続の式を示し、値が積み上がることを確認させる
  2. 上の数値例を手計算させる。平均ゼロ、分散1になることを自分で確かめさせる
  3. バッチ正規化との違いを、統計を取る方向の図で示す
  4. 生成時にバッチサイズが1になることを指摘する
  5. 後置と前置の式を並べ、勾配の経路の違いを説明する
  6. GPT-2が前置を採用した理由に触れる

所要時間は50分程度です。2番の手計算が理解の要になります。

受講者からよく出る質問と回答例を挙げます。

質問。正規化すると、せっかく計算した情報が失われませんか。 回答。長さの情報は失われます。しかし方向の情報は保たれます。言語モデルにおいて意味を担っているのは主に方向であるため、実用上の問題になりません。また、学習可能な \gamma \beta により、必要な尺度は復元できます。

質問。全ての層に必要ですか。 回答。GPT-2では、各ブロック内で自己注意の前とMLPの前の二箇所、そして最後に一箇所です。層数が12であれば、合計で25箇所になります。

理解度を確認する問い

  • レイヤー正規化は、どの方向で平均と分散を計算しますか
  • 言語モデルでバッチ正規化が使いにくい理由を、三つ挙げてください
  • 後置正規化と前置正規化の式を書き、勾配の経路の違いを説明してください
  • GPT-2が最後にもう一度正規化を置いている理由は何ですか
  • RMSNormは、レイヤー正規化から何を省いていますか

まとめと次の学習ステップ

レイヤー正規化は、残差接続によって積み上がる値を各層でリセットし、勾配を安定させる仕組みです。一つのトークンの中で完結する計算であるため、系列長が可変でも、生成時にバッチサイズが1でも、同じように動作します。

GPT-2は、正規化を加算の後から前へ移しました。これにより残差の経路が正規化を通らなくなり、勾配が入力側まで素直に届くようになります。学習率のウォームアップなしでも安定して学習できるのは、この配置のおかげです。

次に学ぶなら、残差ストリームという考え方に進むとよいでしょう。各層が共通の情報経路に読み書きしていくという見方で、Transformerの内部を理解する上で有力な枠組みです。この記事で扱った前置正規化の利点も、その視点から見ると自然に理解できます。

あわせて、重みの初期化と学習率スケジュールも押さえておくと、学習の安定性という主題が立体的になります。正規化、初期化、学習率の三つは、いずれも同じ問題への異なる角度からの対処です。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。