A-3：正規化と活性化関数
A-2では、Self-Attentionがトークン間から文脈を集め、各トークンの表現を更新する仕組みを見ました。ただし、文脈を集められるだけでは、何十層にも積み重なるTransformerを安定して学習させることはできません。各層の値のスケールが不安定になれば、最適化は難しくなり、表現を十分に変換できなければモデルの力も伸びません。
そこで働くのが、正規化と活性化関数です。正規化は各層の表現を扱いやすいスケールに保ち、活性化関数は線形変換だけでは表せない複雑な関係を表現へ加えます。どちらも目立つ部品ではありませんが、現代のLLMの学習安定性、性能、計算効率を左右する重要な設計です。
今回は、Transformerで広く使われるLayerNormとRMSNorm、FFNで使われるGELUとSwiGLUを中心に扱います。数式の細部を暗記するよりも、「何を安定させるのか？」「どの位置で、何を変換するのか？」を説明できることを目標にします。
1. 正規化と活性化関数は何を担うのか？
Transformerの1ブロックには、Self-AttentionとFFNという大きく異なる処理があります。前者はトークン間で情報を集め、後者は各トークンの表現を非線形に変換します。正規化と活性化関数は、この二つの処理を支える役割を分担します。
| 部品 | 主な役割 | Transformer内での位置 |
|---|---|---|
| 正規化 | 表現のスケールを整え、最適化を安定させる | AttentionやFFNの前後 |
| 活性化関数 | 非線形性を加え、表現できる関数の幅を広げる | 主にFFNの中間層 |
| 残差接続 | 入力を保ちながら更新分を重ねる | AttentionとFFNの出力 |
線形層だけを何層重ねても、全体としては一つの線形変換にまとめられます。活性化関数が入ることで、入力に対して曲線的で条件依存の変換を表せるようになります。一方、深いモデルでは層を通るごとに表現の大きさが偏りやすいため、正規化と残差接続を組み合わせて、学習を続けやすい経路を作ります。
2. TransformerでBatchNormよりLayerNormが使われるのはなぜか？
BatchNormは、ミニバッチ内の平均と分散を使って特徴を正規化する手法です。画像の畳み込みネットワークでは大きな効果を示してきましたが、Transformerの言語モデルでは標準的な選択ではありません。
言語モデルでは、系列長が異なる入力を扱い、学習時のバッチサイズが小さくなることもあります。また、自己回帰生成では1トークンずつ処理する場面があり、バッチ統計への依存が扱いにくくなります。LayerNormは各トークンの隠れ次元に対して正規化するため、同じサンプルの他の要素やバッチサイズに依存せず、訓練時と推論時で同じ計算を使えます。
| 手法 | 主な正規化の範囲 | バッチサイズへの依存 | Transformerでの扱いやすさ |
|---|---|---|---|
| BatchNorm | バッチ内の各特徴 | ある | 一般に低い |
| LayerNorm | 各トークンの隠れ次元 | ない | 高い |
| RMSNorm | 各トークンの隠れ次元のRMS | ない | 高い |
BatchNormが常に使えないわけではありません。しかし、Transformerで採用される理由は、LayerNormやRMSNormのほうが優れているという単純な優劣ではなく、系列処理と自己回帰生成の条件に合いやすいからです。
3. LayerNormは何を正規化するのか？
LayerNormは、1つのトークンの隠れベクトル $x$ に含まれる特徴次元の平均と分散を計算し、正規化します。その後、学習可能なスケール $\gamma$ とバイアス $\beta$ を適用します。
ここで $\mu$ と $\sigma^2$ は隠れ次元内の平均と分散、$\epsilon$ はゼロ除算や数値的不安定さを避けるための小さな値です。LayerNormは各トークンを個別に扱うため、入力の系列長やバッチサイズが変わっても計算の意味が変わりにくいという特徴があります。
LayerNormがあるから学習が必ず安定するわけではありません。学習率、初期化、データ、残差接続、精度形式なども関わります。正規化は、それらを不要にする仕組みではなく、深いネットワークを最適化しやすくする重要な条件の一つです。
4. RMSNormはLayerNormと何が違うのか？
RMSNormは、平均を引かず、二乗平均平方根（RMS）でスケールを整える正規化です。一般的な形は次のように書けます。
$$ \operatorname{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2 + \epsilon}} $$
LayerNormとの違いは、平均を0へ中心化する処理と、通常は加算バイアスを持たない点です。計算が少し簡潔で、大規模なDecoder-only LLMで広く採用されています。一方で、どのモデルでもRMSNormが常に最適だと断定はできません。既存実装との互換性、学習の再現性、実際の計測結果を踏まえて選ぶ必要があります。
| 観点 | LayerNorm | RMSNorm |
|---|---|---|
| 平均の減算 | 行う | 行わない |
| 分散またはRMSの利用 | 分散を利用 | RMSを利用 |
| 学習可能なバイアス | 一般に持つ | 実装により省くことが多い |
| 採用例 | 原論文のTransformer、BERT系 | LLaMA系などのDecoder-only LLM |
RMSNormが平均を計算しないことは、単に高速化のためだけではありません。中心化を省いても十分な性能が得られる場合がある、という経験的な知見に基づく設計です。モデル規模、実装、訓練条件が異なると結果も変わり得ます。
5. Pre-NormとPost-Normはどこに正規化を置くのか？
正規化は、置く位置によっても意味が変わります。原論文のTransformerで使われたPost-Normは、サブレイヤーの出力へ残差接続を足したあとに正規化します。一方、現在の多くのDecoder-only LLMで使われるPre-Normは、AttentionやFFNへ入力する前に正規化します。
| 配置 | 概念的な流れ | 特徴 |
|---|---|---|
| Post-Norm | $x \leftarrow \operatorname{Norm}(x + \operatorname{Sublayer}(x))$ | 原論文の標準形 |
| Pre-Norm | $x \leftarrow x + \operatorname{Sublayer}(\operatorname{Norm}(x))$ | 深いモデルで最適化しやすいことが多い |
Pre-Normでは残差経路に正規化されていない入力が直接流れるため、非常に深いネットワークでも勾配の経路を保ちやすいと説明されます。ただし、Pre-NormとPost-Normの選択だけで安定性が決まるわけではありません。学習率のウォームアップ、初期化、残差のスケーリングなどと合わせて設計します。
6. 活性化関数はFFNで何を変えるのか？
FFNは、各トークンに同じ変換を適用する2層程度のニューラルネットワークです。典型的には隠れ次元を一度広げ、活性化関数を通し、元の次元へ戻します。
$$ \operatorname{FFN}(x) = W_2\,\phi(W_1x+b_1)+b_2 $$
ここで $\phi$ が活性化関数です。活性化関数がなければ、$W_1$ と $W_2$ を重ねても一つの線形変換にしかなりません。$\phi$ があることで、入力の値に応じて通す情報を変える非線形な表現が可能になります。
| 関数 | 特徴 | Transformerでの位置付け |
|---|---|---|
| ReLU | 単純で高速だが、負の入力を0にする | 初期の深層学習で広く利用 |
| GELU | なめらかに入力を通す | BERTやGPT-2系で採用例が多い |
| SiLU / Swish | なめらかなゲートに使える | SwiGLUの構成要素 |
| SwiGLU | ゲート付きのFFNを作る | LLaMA系などで採用例が多い |
活性化関数は「どれが最強か？」という一問で選べるものではありません。FFNの中間次元、パラメータ予算、学習データ、実装効率と組み合わせて比較する必要があります。
7. GELUとSwiGLUはどのように使い分けるのか？
GELUは、入力をなめらかに通す活性化関数です。厳密には標準正規分布の累積分布関数 $\Phi$ を用いて、次のように定義されます。
$$ \operatorname{GELU}(x) = x\Phi(x) $$
SwiGLUは、値の経路とゲートの経路を作り、ゲート側へSiLUを適用してから要素ごとに掛け合わせる構造です。概念的には次のように表せます。
SwiGLUでは、どの情報を通すかをゲートが調整します。ゲート用の射影が増えるため、同じ中間次元ならパラメータ数と計算量は増えます。そのため実際のモデルでは、中間次元を調整して、通常のFFNと総計算量やパラメータ数を比較しやすくすることがあります。
既存モデルの再現が目的なら、そのモデルの活性化関数と中間次元を踏襲することが基本です。新規設計なら、GELUの標準FFNとSwiGLUを、同程度の計算予算で比較してから決めるのが安全です。
原論文の基本的な思想は保ちつつ、超大規模化・超深層化に耐えうる形へ、各コンポーネントが洗練・最適化された集大成。
8. 最小コードでRMSNormとSwiGLUを確認する
以下は、RMSNormとSwiGLUを最小限にしたPyTorch実装です。学習済みLLMを再現するためではなく、RMSNormが最後の次元を正規化し、SwiGLUが2本の経路を掛け合わせる構造を確認するための例です。
import torch
import torch.nn as nn
import torch.nn.functional as F
class RMSNorm(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
return x * rms * self.weight
class SwiGLU(nn.Module):
def __init__(self, hidden_size: int, intermediate_size: int):
super().__init__()
self.value = nn.Linear(hidden_size, intermediate_size, bias=False)
self.gate = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down = nn.Linear(intermediate_size, hidden_size, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down(self.value(x) * F.silu(self.gate(x)))実装でよくある失敗は、RMSNormで平均を引いてLayerNormと同じ形にしてしまうこと、またはSwiGLUのvalueとgateを掛ける前に次元を合わせないことです。テンソル形状を(batch, sequence, hidden_size)のまま追い、最後の次元だけに操作が適用されているかを確認しましょう。
9. 学習が不安定なとき、どこを確認すればよいのか？
損失がNaNになる、急に発散する、学習が進まないといった現象を、正規化か活性化関数だけの問題と決めつけるのは危険です。まずは観測できる値を分けて確認します。
| 症状 | 最初に確認すること | 正規化・活性化との関係 |
|---|---|---|
| 損失がNaNになる | 入力、学習率、混合精度、勾配の値 | 除算時の$\epsilon$、活性値の過大化も確認する |
| 損失が急に増える | 学習率、ウォームアップ、勾配ノルム | Pre/Post-Normや残差の実装を確認する |
| 学習が進まない | データ、損失、マスク、初期化 | 活性化の実装や次元設定を確認する |
| 再現結果が揺れる | 乱数、データ順序、精度、並列設定 | 正規化層のtrain/eval差も確認する |
実務では、最初から複数の設計を変えないことが重要です。基準となる構成を一つ決め、学習率、正規化、活性化関数を一度に一つずつ変更し、損失曲線と評価指標を比べます。理由を特定できないまま改善したように見える状態を避けられます。
10. まとめ
今回の要点を整理します。
- 正規化は各層の表現のスケールを整え、深いTransformerを最適化しやすくする
- LayerNormとRMSNormはバッチ統計に依存せず、系列処理に適している
- RMSNormは平均の減算を省き、大規模なDecoder-only LLMで広く使われる
- Pre-Normは、正規化をAttentionやFFNの入力側に置く設計である
- 活性化関数はFFNに非線形性を与え、GELUやSwiGLUが代表的な選択肢である
- 不安定な学習は単一原因と決めつけず、学習率や数値精度を含めて切り分ける
正規化と活性化関数は、Transformerの中で派手に見える部品ではありません。それでも、Self-Attentionが集めた情報を意味のある表現へ変換し、深い層まで安定して届けるための基盤です。
11. 今回のブログの考察：A-4への橋渡し
A-3では、Transformerの内部で表現を安定させ、非線形に変換する仕組みを見ました。しかし、どれほどAttention、FFN、正規化を丁寧に設計しても、モデルが最初に受け取る単位が適切でなければ、学習の効率や表現できる範囲には限界があります。
次に考えるべきは、文章をどの単位でモデルへ渡すのか？という問いです。日本語のように単語境界が明確でない言語も含め、文字列をトークンへ分割し、語彙をどのように設計するかは、入力長、埋め込み層、未知語への対応に直接関わります。
次回A-4では、BPEとSentencePieceを中心に、トークン化と語彙構築を扱います。Transformerの内部で起きている変換から一歩戻り、その入力がどのように作られるのか？を見ていきましょう。
参考文献
- Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). “Layer Normalization.” https://arxiv.org/abs/1607.06450
- Zhang, B., & Sennrich, R. (2019). “Root Mean Square Layer Normalization.” https://arxiv.org/abs/1910.07467
- Xiong, R., Yang, Y., He, D., et al. (2020). “On Layer Normalization in the Transformer Architecture.” https://arxiv.org/abs/2002.04745
- Hendrycks, D., & Gimpel, K. (2016). “Gaussian Error Linear Units (GELUs).” https://arxiv.org/abs/1606.08415
- Shazeer, N. (2020). “GLU Variants Improve Transformer.” https://arxiv.org/abs/2002.05202
- Touvron, H., Martin, L., Stone, K., et al. (2023). “Llama 2: Open Foundation and Fine-Tuned Chat Models.” https://arxiv.org/abs/2307.09288
このシリーズの案内
ブログA「Transformer構成」（全7回）では、現代のLLMを支えるTransformerの構成要素を、全体像から実装への接続まで段階的に扱います。
| # | テーマ | 内容 |
|---|---|---|
| A-1 | Transformer構成の基本概念と重要性 | 全体像、部品の役割、データフロー |
| A-2 | Self-Attentionの詳細 | Q・K・V、Causal Mask、マルチヘッド、計算量 |
| A-3 | 正規化と活性化関数 | LayerNorm、RMSNorm、GELU、SwiGLU（今回） |
| A-4 | トークン化と語彙構築 | BPE、SentencePiece、語彙サイズの設計 |
| A-5 | 位置エンコーディングの進化 | 絶対位置、相対位置、RoPE、ALiBi |
| A-6 | 構成要素の統合と設計 | Pre-Norm、層構成、スケーリングとの関係 |
| A-7 | 理論の総まとめ | 概念の統合と実装編への橋渡し |
コメント