AI・機械学習
上級

パラメータ初期化の分散制御(パラメータショキカノブンサンセイギョ)

LLMの学習安定性を確保するために、各層の重みパラメータの初期分散を理論的・経験的に制御する技術。層の深さ、幅、活性化関数の種類に応じて最適な分散を計算し、勾配消失・爆発を防止する。μP(Maximal Update Parameterization)はこの分野の最新手法で、スケーリング則に基づく分散制御を実現する。

0 回閲覧
0 いいね

パラメータ初期化の分散制御は、LLMの数十〜数百層にわたるTransformerブロック全体で信号と勾配の大きさを適切に維持するための設計原理である。各層のパラメータ初期分散が大きすぎると勾配爆発(loss が NaN に発散)、小さすぎると勾配消失(学習が停滞)を引き起こす。GPT-4(推定1.8兆パラメータ、推定120層)やLlama 3 405B(126層)のように層数が増えるほど、この制御の重要性は指数関数的に高まる。2026年現在、μP(Maximal Update Parameterization)や SP(Standard Parameterization)の理論に基づく体系的な分散制御が主流となっている。

分散制御の基本原理

多層ネットワークにおいて、入力 x が L 層を通過した後の出力分散は:

  • Var(output) = Var(x) × ∏(i=1 to L) [n_i × Var(W_i)]

各層で n_i × Var(W_i) = 1 を満たせば分散は保存される(n_i は層 i の入力次元数)。これがXavier/Kaiming初期化の理論的根拠だが、実際のLLMでは以下の要因で単純な理論が当てはまらない:

  • 残差接続: x + f(x) の形で信号が加算されるため、層を経るごとに分散が増大
  • レイヤー正規化: 各層出力の分散を1に正規化するが、勾配の挙動は変わる
  • Attention機構: Q, K, V の内積が次元数に比例して増大(√d_k で除算して補正)
  • FFN層の活性化: GELU/SwiGLU の非線形性による分散変化

残差接続と分散増大

Transformerの残差接続では、L層後の出力は:

  • y_L = x + f_1(x) + f_2(...) + ... + f_L(...)

各 f_i の出力分散が σ² のとき、y_L の分散は約 L × σ² に増大する。これを防ぐための戦略:

手法スケーリング採用モデル
GPT-2スタイル出力射影を 1/√(2L) でスケールGPT-2, GPT-3
Pre-Norm + 小分散σ = 0.02 固定 + Pre-LNGPT-NeoX, Pythia
DeepNorm残差を α で、初期化を β でスケールGLM-130B
μP幅に依存する学習率×初期化スケールCerebras-GPT

GPT-2 の実装では、残差ブロックの最終射影層の初期分散を N(0, 0.02/√(2N)) とし、N はブロック数。GPT-3 175B(96層)では 0.02/√192 ≈ 0.00144 となり、極めて小さい初期値から学習を開始する。

μP(Maximal Update Parameterization)

Microsoft Research が2022年に提案したμPは、モデル幅(Hidden Size)のスケーリングに対して、初期化分散と学習率を理論的に調整する枠組みである:

  • 入力埋め込み: σ = 1(幅に依存しない)
  • 隠れ層の重み: σ = 1/√d(dは層の幅)
  • 出力層の重み: σ = 1/d
  • 学習率: 隠れ層は 1/d に比例、出力層は 1/d² に比例

この設定により、小規模モデル(128M)で探索した最適ハイパーパラメータを、大規模モデル(数十億〜数千億パラメータ)にそのまま転移できる。Cerebras-GPT(111M〜13B)で実証され、従来のハイパーパラメータ探索コストを大幅に削減した。

実践的な分散設定

2026年現在の主要モデルで採用されている初期化分散の具体値:

モデルHidden Size層数初期化分散 σ残差スケーリング
GPT-2768120.020.02/√24 ≈ 0.004
GPT-3 175B12,288960.020.02/√192 ≈ 0.0014
Llama 2 70B8,192800.021/√(80) ≈ 0.0112
Llama 3 405B16,3841261/√dモデル固有
PaLM 540B18,432118切断正規分布1/√(3×118)

デバッグと診断

初期化分散が不適切な場合の典型的な症状:

  • loss が NaN: 分散が大きすぎ → σ を半分にして再試行
  • loss が減少しない: 分散が小さすぎ → σ を2倍にして再試行
  • 学習初期に loss スパイク: 残差スケーリング不足 → DeepNorm や GPT-2 スタイルのスケーリングを適用
  • 特定の層だけ勾配が爆発: その層の初期化を個別に調整(Attention の Q,K は特に注意)

PyTorch での診断コード:

for name, param in model.named_parameters():
    print(f'{name}: mean={param.data.mean():.6f}, std={param.data.std():.6f}')

よくある質問(FAQ)

Q1: 初期化の分散はどうやって決めればよいですか? A: まず Xavier または σ=0.02 で開始し、学習初期の loss カーブを観察する。loss が発散なら σ を小さく、学習が遅すぎるなら σ を大きくする。大規模モデルでは μP を採用し、小規模実験で最適値を探索するのが効率的。

Q2: Pre-Norm と Post-Norm で初期化分散は変えるべきですか? A: Pre-Norm(Llama、GPT-NeoX方式)では残差接続の分散増大が緩和されるため、固定 σ=0.02 で安定しやすい。Post-Norm(元のTransformer方式)では DeepNorm のような追加スケーリングが必要になることが多い。

Q3: μPは2026年時点でどの程度普及していますか? A: Cerebras-GPT で実証され、学術研究では広く参照されている。しかし Llama 3 や GPT-4 が μP を採用しているかは公開情報では不明。商用モデルでは独自のスケーリング手法を使用している可能性が高い。

Q4: Attention層のQ,K,Vの初期化で特別な配慮は必要ですか? A: Q と K の内積は次元数 d_k に比例して増大するため、Attention スコアを √d_k で除算する。初期化時にも Q, K の分散を 1/d_k に設定するモデルがある(T5など)。V の初期化は通常の Xavier/カスタム分散で問題ない。

まとめ

  • 初期化分散の制御はLLM学習の安定性を決定する最重要設計要素の一つ
  • 残差接続による分散増大を補正するスケーリング手法(GPT-2式/DeepNorm/μP)が必須
  • σ=0.02 が経験的な出発点、モデル規模に応じて残差スケーリングを追加
  • μP により小規模実験→大規模モデルへのハイパーパラメータ転移が理論的に可能に
  • 初期化の不具合は loss の発散・停滞として即座に表面化するため、診断は比較的容易