AI・機械学習
中級

Layer Normalization(LayerNorm)(レイヤーノーマライゼーション)

Layer Normalization(LayerNorm)は、ニューラルネットワークの各層で特徴量をサンプル単位に正規化する手法であり、Transformer ベースの大規模言語モデル(LLM)で事実上の標準正規化として広く採用されている。Batch Normalization がバッチ方向に統計量を計算するのに対し、LayerNorm は単一サンプル内の全特徴次元にわたって平均・分散を計算するため、バッチサイズに依存せず系列長が変動する自然言語処理タスクに適している。

0 回閲覧
0 いいね

Layer Normalization(LayerNorm)は、2016年に Jimmy Lei Ba らが提案した正規化手法であり、Transformer アーキテクチャにおける最も基本的な正規化レイヤーである。バッチ方向ではなく特徴次元方向に平均と分散を計算することで、ミニバッチサイズや系列長に依存しない安定した学習を実現する。GPT・BERT・T5 をはじめとする主要 LLM のほぼ全てが LayerNorm またはその変種を採用しており、LLM の学習安定性を支える中核技術である。

概要

LayerNorm の核心は「各サンプルの各層出力を、そのサンプル自身の統計量で正規化する」点にある。Batch Normalization(BatchNorm)は 2015年に Ioffe & Szegedy が提案し画像認識で革命的な成果を上げたが、バッチ内の全サンプルから平均・分散を計算するため、バッチサイズが小さい場合や系列長が可変の場合に統計量が不安定になるという問題があった。LayerNorm はこの問題を解決するために設計された。

自然言語処理では入力系列長がサンプルごとに異なり、パディングの影響でバッチ統計量が歪むため BatchNorm は不適切である。LayerNorm は個々のサンプル・個々のタイムステップについて独立に正規化を行うため、系列長の変動やバッチサイズの変更に対してロバストである。

2017年の Transformer 論文「Attention Is All You Need」で採用されて以降、LayerNorm は自然言語処理における事実上の標準正規化手法となった。オリジナル Transformer では Multi-Head Attention および Feed-Forward Network の後に LayerNorm を適用する Post-Norm 配置が採用されていたが、後続研究で Pre-Norm 配置(LayerNorm を先に適用)の方が学習が安定することが示され、GPT-2 以降の多くのモデルは Pre-Norm を採用している。

技術的詳細と数式

LayerNorm の計算は以下の手順で行われる。入力ベクトル x = (x₁, x₂, ..., x_H) に対して、H は隠れ層の次元数である。

ステップ1: 平均の計算

μ = (1/H) × Σᵢ xᵢ

ステップ2: 分散の計算

σ² = (1/H) × Σᵢ (xᵢ − μ)²

ステップ3: 正規化

x̂ᵢ = (xᵢ − μ) / √(σ² + ε)

ステップ4: アフィン変換

yᵢ = γ × x̂ᵢ + β

ここで γ(スケール)と β(シフト)は学習可能なパラメータ、ε は数値安定性のための微小定数(通常 1e-5 または 1e-6)である。

計算コストの分析

LayerNorm の計算量は O(H) であり、隠れ層の次元数に線形比例する。LLM の典型的な隠れ次元(H = 4096〜12288)では、1トークン・1層あたり約 3H 回の浮動小数点演算が必要である。Transformer 全体の計算量に対する LayerNorm の割合は通常 1〜3% 程度だが、推論時のレイテンシではメモリアクセスがボトルネックとなるため、実測では 5〜10% の時間を占めることがある。

パラメータ説明典型的な値
H(隠れ次元)正規化対象の特徴次元数768〜12288
ε(イプシロン)ゼロ除算防止の微小値1e-5〜1e-6
γ(ガンマ)学習可能スケールパラメータ初期値 1.0
β(ベータ)学習可能シフトパラメータ初期値 0.0
追加パラメータ数γ + β で 2H 個1536〜24576

実装上の注意点

GPU 上での LayerNorm 実装には、平均・分散の計算でリダクション操作が必要であり、ナイーブな実装ではグローバルメモリへの複数回アクセスが発生する。NVIDIA の APEX ライブラリや PyTorch の fused LayerNorm 実装では、平均と分散の計算を1回のカーネル起動で行う融合カーネル(Fused Kernel)が提供されており、メモリ帯域の消費を 50% 以上削減できる。FlashAttention の開発チームが公開した Triton ベースの LayerNorm カーネルは、さらに高速な実装を提供している。

他の正規化手法との比較

手法正規化方向バッチ依存系列長依存学習パラメータ計算コスト主な用途
BatchNormバッチ方向ありなし2HO(BH)CNN(画像)
LayerNorm特徴方向なしなし2HO(H)Transformer/LLM
InstanceNormチャネル内空間方向なしなし2CO(HW)スタイル変換
GroupNormグループ内特徴方向なしなし2CO(H)小バッチ CNN
RMSNorm特徴方向(平均なし)なしなしHO(H)最新 LLM

LayerNorm のメリット

  • バッチサイズ非依存: バッチサイズ 1 でも正常に動作するため、推論時やオンライン学習に適している
  • 系列長非依存: 可変長入力に自然に対応でき、パディングの影響を受けない
  • 安定した勾配: 逆伝播時の勾配のスケールが正規化により安定し、深いネットワークでも学習が収束しやすい
  • 並列化との相性: サンプル間の依存がないため、データ並列・テンソル並列と相性が良い

LayerNorm のデメリット

  • 計算オーバーヘッド: 平均と分散の2つの統計量を計算するため、RMSNorm に比べて約 30% 多い計算量が必要
  • リダクション操作: GPU 上で平均・分散計算のリダクションがメモリ帯域のボトルネックになりうる
  • β パラメータの冗長性: 研究により、シフトパラメータ β はほとんどのケースで学習後にゼロ付近に収束し、実質的に不要であることが示されている

採用モデルと実装例

LayerNorm を採用している主要な LLM とその配置:

モデルパラメータ数LayerNorm 配置備考
BERT(2018)110M-340MPost-Normオリジナル Transformer 配置
GPT-2(2019)117M-1.5BPre-NormPre-Norm の先駆的採用
T5(2019)60M-11BPre-Normエンコーダ・デコーダ両方
GPT-3(2020)175BPre-Norm大規模モデルでの安定学習に貢献
PaLM(2022)540BPre-Norm並列配置と組み合わせ
BLOOM(2022)176BPre-Norm多言語対応 LLM

PyTorch での実装例

PyTorch では torch.nn.LayerNorm として標準提供されており、以下のように使用する:

import torch
import torch.nn as nn

# 隠れ次元 768 の LayerNorm
layer_norm = nn.LayerNorm(768, eps=1e-5)

# 入力: (batch_size, seq_len, hidden_dim)
x = torch.randn(32, 512, 768)
output = layer_norm(x)

Transformer ブロック内での典型的な配置(Pre-Norm):

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForwardNetwork(d_model)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

最新の研究動向

2024年以降、LayerNorm に関する研究は主に以下の方向で進展している:

1. パラメータフリー正規化: γ と β を排除し、計算コストとパラメータ数を削減する試みが増えている。Meta の LLaMA シリーズは RMSNorm を採用し、β パラメータの不要性を実証した。

2. 融合カーネルの進化: NVIDIA Hopper アーキテクチャ(H100)では、Transformer Engine が FP8 量子化と LayerNorm を融合した専用カーネルを提供しており、スループットが従来比 2〜3倍向上している。

3. 条件付き正規化: Mixture of Experts(MoE)モデルでは、エキスパートごとに異なる正規化パラメータを持つ Conditional LayerNorm が研究されている。

4. 量子化との相互作用: INT8/INT4 量子化環境で LayerNorm の精度を維持するため、正規化前に外れ値を分離する手法(SmoothQuant など)が実用化されている。

FAQ

Q: LayerNorm と BatchNorm はどのように使い分けるべきか?

A: 基本的なガイドラインとして、画像認識(CNN)には BatchNorm、自然言語処理(Transformer/RNN)には LayerNorm を使用する。BatchNorm はバッチ統計量に依存するため、バッチサイズが小さい場合やテスト時にバッチサイズが変わる場合に不安定になる。LayerNorm はサンプル内の統計量のみを使用するため、バッチサイズに依存せず、可変長入力にも対応できる。ただし、近年は画像認識でも Vision Transformer(ViT)の普及に伴い LayerNorm の使用が増加している。

Q: LayerNorm の ε(イプシロン)はどの値に設定すべきか?

A: 多くのフレームワークのデフォルト値は 1e-5 であり、ほとんどのケースで適切である。ただし、混合精度学習(FP16/BF16)を使用する場合は、数値精度の制約から 1e-6 ではアンダーフロー問題が発生することがある。BF16 では指数部が FP32 と同じため 1e-5 で安定するが、FP16 では場合によっては 1e-4 に緩和する必要がある。量子化推論(INT8)では、ε の値が量子化誤差に影響を与えるため、キャリブレーション時に確認が必要である。

Q: Pre-Norm と Post-Norm のどちらが優れているか?

A: 現在の LLM 開発では Pre-Norm が主流であり、学習の安定性という観点では Pre-Norm が優れている。Pre-Norm は残差接続の入力側に LayerNorm を配置するため、勾配が残差パスを通じて直接伝播し、深いネットワークでも勾配消失が起きにくい。一方、Post-Norm は理論的にはより高い表現力を持つとされ、一部の研究では Post-Norm の方が最終的な精度が高いケースも報告されている。DeepNorm のような手法は、Post-Norm の精度優位性を維持しつつ学習安定性を改善する試みである。詳細は用語「Pre-Norm vs Post-Norm アーキテクチャ」を参照されたい。

Q: なぜ最近の LLM は LayerNorm ではなく RMSNorm を使うのか?

A: RMSNorm は LayerNorm から平均の減算を省略し、二乗平均平方根(RMS)のみで正規化する手法である。LLaMA・Mistral・Gemma などの最新 LLM が RMSNorm を採用している理由は、(1)平均の計算を省くことで計算コストが約 10〜30% 削減される、(2)シフトパラメータ β が不要になりパラメータ数が半減する、(3)実験的に LayerNorm と同等以上の性能が確認されている、ためである。ただし、BERT・GPT-3・PaLM など多くの既存モデルは LayerNorm を使用しており、LayerNorm が劣っているわけではない。