レイヤー正規化 (layer normalization) [Transformerでよく用いるバッチ正規化層]

1. レイヤー正規化 (Layer Normalization)とは [概要]

レイヤー正規化 (Layer Normalization)とは,可変長の系列データが入力の系列モデル・系列変換モデルでも使用しやすいように,元となるバッチ正規化を「バッチ内で,レイヤー方向の正規化を行う」ようにアレンジしたものである.当初のレイヤー正規化 [Ba et al., 2016] は,バッチ正規化層が使いづらいRNNむけに提案された.

この記事では,バッチ正規化 [Ioffe and Szegedy, 2015] からの変更点・差異の視点から,レイヤー正規化の処理内容と利点などをまず説明したのち(2節),後半(3節)では,Transformer [Vaswani et al., 2017]系モデルのうち,どのモデルでよく用いられているかを,整理していきたい.2.2節が,本題の「RNNで用いるバッチ正規化の処理」についてである.

レイヤー正規化(Layer Normalization)
図1. レイヤー正規化(Layer Normalization)
[RNN, FNN, Transformerの隠れ層における特徴ベクトル向け]

レイヤー正規化 (図1, RNN, FNN, Transformerの隠れ層における特徴ベクトル向け)は,バッチ長$B$には依存していない,「サンプル単位での正規化」である(2.2節で詳しく).サンプル間をまたいでの正規化を行わないで,サンプル内で正規化が済むゆえ,「系列モデル」に向いているという性質に加えて,RNN系の自己回帰モデルで必要となりがちな「オンライン学習」にも向いている.その後,Transformer [Vaswani et al., 2017]に採用され,たことでTransformerを用いた大規模言語モデルのGPT/BERTや,画像認識向けのViT/DETRなどでも使用される展開になった.

レイヤー正規化は,トークン系列モデル向けのバッチ正規化の発展型であるゆえ,自然言語処理由来の「(単語・サブ単語の)トークンの系列モデル」において使われ,威力を発揮する.また,NLP以外も含めて「種々のTransformer系モデル(BERT/GPTやDETRなど)」においても,(画像や音のトークンも含めた)トークン系列を取り扱うことが多くなり、Transformer経由でレイヤー正規化が多用される展開になっている(3節で整理).

関連記事:

RNNなどの系列モデルむけの「可変長系列データ」に旧来のバッチ正規化を適用しようとすると,各サンプルごと(各文章ごと)に長さが異なってしまう.よって,バッチ全体に渡って,各サンプルのチャンネル方向に正規化する方針の「バッチ正規化」では,可変長データを取り扱いづらかった.

それが,各(隠れ)層の,サンプル文内に閉じた単位で正規化を行う「レイヤー正規化」が提案されたことにより,バッチサイズNが小さかったり,最悪N=1でも使用できるようになり,系列モデルや系列変換モデルでも,バッチ正規化型の層の恩恵を受けられるようになった (2.2節で詳しく).

1.1 記事の構成

2節以降は以下の構成である

  • 2. レイヤー正規化の詳細
    • 2.1 レイヤー正規化の処理 (バッチ正規化とバッチ正規化の比較)
      • 2.1.1 レイヤー正規化 (CNN特徴マップの場合)
      • 2.1.2 バッチ正規化 (CNN特徴マップの場合)
    • 2.2 本題:RNNでのレイヤー初期化層
    • 2.3 動機:系列処理・オンライン学習むけにアレンジ
    • 2.4 元論文での実験
  • 3.Transformer系モデルでの使用
    • 3.1 元祖Transformerでの使用
    • 3.2 各応用モデルので使用
      • 3.2.1 Transformer自体の広範囲での活用
      • 3.2.2 大規模言語モデル(BERT/GPTなど)でのTransformerの応用
      • 3.2.3 物体認識・物体検出バックボーンでのTransformerの応用
  • 4. まとめ

2. レイヤー正規化の詳細

2.1 レイヤー正規化とバッチ正規化の比較

2.1.1. レイヤー正規化 (CNN特徴マップの場合)

レイヤー正規化 (図2, CNNの特徴マップ向け) では,名前の通り「レイヤー単位」で,特徴値(グリッド上の各1つの特徴)の正規化を行う.ここでいう「レイヤー(方向での)正規化」とは,RNNで処理している系列中の,$l$番目の隠れ層(レイヤー)の単位で,各隠れにてデータ正規化を行う事をさす.つまりは,「バッチ内における,レイヤーごとの正規化」がレイヤー正規化である.

※ 図1では,サンプルの添え字 $n = {1,2, \ldots,n} $が,2.3節でRNNで説明する(元論文中での)隠れ層の添え字$l$に相当している.親記事のバッチ正規化層とその発展型とおなじ図かつおなじ添え字で統一して読みやすくするために,$n$による図示をこの子記事でも採用する.

レイヤー正規化 (CNNの場合)
図2. レイヤー正規化 (CNNの場合)

(親記事にも配置している) 図2は,「n 個の特徴マップから構成される,CNN隠れ層での,特徴マップのミニバッチ」のうち,レイヤー単位で平均・分散を計算して正規化する単位を,色でグループ分けして見やすく可視化した図である.2.1.3節では,元論文中のRNNでの例を見る前に,先に2.1.2節(親記事)での「CNN特徴マップのバッチ正規化」と比較するため,レイヤー正規化も特徴マップベースでみておきたい.

CNNRNNが,隠れ層(の活性化値) において,学習用のミニバッチに「(H x W x C) サイズの「活性化済み特徴マップ $\bm{x}_i^l$ 」× N個 に対して,レイヤー正規化を行う各範囲を,4色で分けて表示している.

レイヤー正規化では,各層$l (=n)$番目の特徴マップ$\bm{x}^l$ 内で,$H =(X \times \times H \times C)$ 個の特徴$x_i^l$ に対する平均と分散 $(\mu^l, \sigma^l)$を計算し,レイヤーごと(= サンプル内ごと)に特徴値$x_i^l$の正規化を行う.

レイヤー正規化での平均・分散は,以下のように計算する:

レイヤー正規化での,レイヤーごとの平均・分散の算出

\begin{align} \mu^{l} = \frac{1}{H} \sum_{i=1}^H x_i^l \\
\sigma^{l} = \sqrt{ \frac{1}{H} \sum_{i=1}^H (x_i^l - \mu^{l})^2 } \end{align}

2.1.2 バッチ正規化の復習 (CNN特徴マップの場合)

バッチ正規化 (Batch Normalization)
図2. バッチ正規化 (Batch Normalization)

一方,図3はバッチ正規化での,CNN特徴マップ中における正規化グループ単位を,色分けして可視化した図である.バッチ正規化は,図2のように,サンプル$\bm{x}_i$間をまたいで,バッチ内全体で,各チャンネル c ごとに,各特徴値の正規化(とスケーリング・シフト)を行う.

※ つまりは「バッチ内でのチャンネルごと正規化」とでも呼んだ方がわかりやすい.

レイヤー正規化(図2)と比較しやすくするために,バッチ正規化(図3)でのバッチ平均・バッチ分散の算出手順も,以下に示す(より詳しくは,親記事を参照):

バッチ正規化での,バッチ平均・バッチ分散の算出:

\begin{align} \mu^{\beta} = \frac{1}{N} \sum_{i=1}^N x_i^l \\ \sigma^{\beta} = \sqrt{ \frac{1}{N} \sum_{i=1}^N (x_i^l - \mu^{\beta})^2 } \end{align}

2.2 本題: RNNでのレイヤー正規化

レイヤー正規化(RNNとFNNの場合)
図1 [再掲]. レイヤー正規化 (Layer Normalization)
[RNN, FNN, Transformerの隠れ層における特徴ベクトル向け]

※ 元論文では 「3.1 Layer normalized recurrent neural networks 」で「レイヤー正規化を挿入したRNN」の例を紹介している.

さて,2.1.1 と2.1.2は,元のCNNでのバッチ正規化の「3次元特徴マップ B個のミニバッチ」の例であった.よって,ここでは本来のレイヤー正規化の動機(2.2節)である,この記事の本題の「RNNRNNLMなどの系列モデル」の場合での「隠れ層の特徴ベクトル」をレイヤー正規化する処理を,図とともにみていく.

図1 [再掲]の,隠れ層が特徴ベクトルであるRNN, FNN, Transformerなどで,正規化を行う範囲をレイヤー正規化 v.s. バッチ正規化で比較したものを再度みてもらいたい.既に概要と2.1でも見てきたように,レイヤー正規化 (図1-a)では,RNNの毎タイムフレーム$l$における自己回帰において,$\bm{x}^l$の中のみで,各特徴$x_i$の平均・分散をとって,「レイヤー方向の正規化」を行う.この処理は,バッチ内のサンプル数Bに依存しないので,自己回帰モデル・RNNTransformer系の系列モデルの学習に向いている.

一方,特徴ベクトル群のミニバッチで,バッチ正規化(図1-b)を行おうとすると,2.2節のCNN都特徴マップのときのように,各特徴次元のバッチ方向で,それぞれ正規化を行う.

以上の「RNNでのレイヤー正規化の処理」を,最後に数式でもまとめておきたい.RNNの入力の $l$層における (活性化結果の) 特徴ベクトル $\bm{x}^l$ のバッチに対して,レイヤー正規化を行った出力 $\bm{h}$とすると,処理手順は以下のように書ける.

RNNでのレイヤー正規化

\begin{align}
\mu^{l} &= \frac{1}{N} \sum_{i=1}^N x_i^l \\
\sigma^{l} &= \sqrt{ \frac{1}{N} \sum_{i=1}^N (x_i^l - \mu^{l})^2 }\\
\bm{h}^{l}_{out} &= f\left[ \frac{ \bm{\gamma}}{\sigma^{l} + \eta} \odot (\bm{x}^l - \mu^{l}) + \bm{\beta} \right]
\end{align}

ここで,$\gamma_i$と$\beta_i$は,バッチ正規化と同じく,各次元$i$を正規化したのち(リ)スケール,(リ)センターを行うためのパラメータであり,ネットワーク学習中に,一緒に推定される.また,$f$は,RNNの隠れ状態の遷移の際に用いられる全結合層の変換を表わした関数である

※ 元論文[Ba et al.,2016] の3.1節の式(4)に対応している.

このように「RNNと(トークン)系列データに対応できるようにした」点が,レイヤー正規化の改良点である.上記の処理により,レイヤー正規化層前のRNNの活性化済み出力 $\bm{x_i}^l$が,上記の処理で $\bm{h}^{l}_{out}$ へと正規化され系列長が異なる,学習データをまとめたミニバッチでも,「バッチサンプル数$B$に依存していない」ので正規化が(各タイムステップで)可能となる.

2.3 動機:系列データ処理・オンライン学習むけにアレンジ

さて,先に2.2で処理内容を紹介したが,この2.3節では [Ba et al., 2016]の「発案の動機」にもどって,「バッチ正規化をもとにレイヤー正規化をアレンジした流れ」をおさえておきたい.

※ 元論文[Ba et al.,2016] の「2 Background」の内容に相当する話を行う

学習データ中の「各系列サンプルの系列長」が異なる「可変長の入力$\bm{x}_i$」およびそれらか構成される「異なる系列長から構成されるバッチ」と,バッチ内合計数が一定でないので,「バッチ正規化」流の正規化では,「ミニバッチ内の平均・分散を同一の分母で計算するこができなくなってしまう.これだとバッチ正規化は可変長の系列データには,そのまま使用できない

それに対して,レイヤー正規化 [Ba et al., 2016] は,層単位で,サンプル内の全チャンネルに対する正規化を行い,バッチ内の1サンプル内で正規化範囲を留める正規化を提案した(図1).このアレンジおかげで,「バッチサイズに依存せず,各ステップ$t$ ごとに正規化」でき,可変長入力にも問題なく適用できるようになった.

またRNNのBPTT(Back Propagation Through-Time)ような,(自己回帰モデルなどでの)オンライン学習タスクにも向いているバッチ正規化型の層となった.これにより「バッチサイズ1だけの短いバッチ長」でも,正規化が可能となった.Transfomer時代に入ると,入力の系列長N(=トークン数)も数百トークンの長いものになっていくので,省メモリで「GPUフレンドリー」なレイヤー正規化が活躍することにもなる.

以上のように,レイヤー正規化を用いると,系列モデルのRNN(LSTMGRU)やTransfomerに相性が良い「層単位での正規化 + スケーリング」が行えるようになり,(バッチ)正規化の「学習高速化」や「収束安定化」「汎化性能向上」の恩恵を,系列モデルにおいても得やすくなった.

※ ただしバッチ正規化の挿入でも,LSTMの学習は,ある程度高速になる.

2.4 元論文での実験

元論文 [Ba et al., 2016] の6節では,レイヤー正規化をあらたに用いて,以下の6タスクにおける実験が行われ,レイヤー正規化の有効性(ある時vs ない時) が確かめられた:(各実験は内容が古く,詳細も省略するので,論文を参照のこと):

  • 6.1節 Order-embedding [Vendrov et al., 2016]でCOCO captionのモーダル間検索.
  • 6.2節 機械 (machine comprehension) [Cooijmans et al., 2016]
  • 6.3節 Skip-thought べクトルを,文予測の下流タスクに適応
  • 6.4節 2値化MNISTから自己回帰生成モデルのDRAW[Gregor et al., 2015]で画像生成.
  • 6.5節 LSTMによる手書きシーケンス画像の生成[Graves et al., 2013]
  • 6.6節 (RNNだけでなく)FFNでも,MNISTクラス識別を実験.

いずれの実験でも,主に「学習速度の高速化」と「性能向上」が示されており,NLPやVision-Language向けの埋め込み空間の学習,自己回帰画像生成モデル(DRAW)などで,レイヤー正規化が(RNNとコンビを組ませることで)力を発揮することが示された.

3. Transformer系モデルでの使用

3.1 元祖Transformerでの使用

TransformerのEncoder構成
図3. Transformer Encoder
TransformerのDecoder構造
図4. Transformer-Decoder

Transformer [Vaswani et al., 2017] では,各ブロックにおいて「マルチヘッドアテンション後」に,毎回レイヤー正規化が配置されて使用されている.

下記の本,Transformers for Natural Language Processing (2nd Edition)では,このマルチヘッドアテンション後のレイヤー正規化のことを「Post-layer normalization」と独自に呼びながら,しっかり構成を説明しており,わかりやすい.また,その次の層であるFFNのあとにも,レイヤー正規化をもう一度行う(図3,4).

こうして,オリジナルTransformerにおいて,ブロック内での正規化にレイヤー正規化が採用されたことにより,「Transformerを使用したモデル(3.2系)」がのちにたくさん登場するなかで,レイヤー正規化もその部品として頻繁に使用される展開となった.

3.2 各応用モデルので使用

3.2.1 Transformer自体の広範囲での活用

まずTransformer自体は,旧来seq2seq with attentionで解かれていた多数の系列変換モデル問題を中心に使用されていく.ただし,ソース系列だけがトークン入力なだけでなく,変換後のターゲット系列側も,トークン入力かつsoftmaxクラス出力(one-hotベクトル)である問題でないと,(トークン)系列変換モデルをそのまま使用できない.よって,Generator(Decoder)側で「高次元の綺麗な生データ」を生成する問題では,Transformer-Encoder(図3)のみが使用される展開ともなる(例:Text-to-SpeechやText-to-Imageなど).

\[
\text{LayerNormalizaiton( x + \text{Sublayer(x)} ) }
\]

3.2.2. 大規模言語モデル(BERT/GPTなど)でのTransformerの応用

自己教師有り事前学習を用いた,BERT/GPTでTransformerが主部品に使われたことで,事前学習済みのBERTやGPTをタスクにチューニングする解き方でも,Transformer(の一部)が使用されていくことで,レイヤー正規化も部品として同時に使用されていった.

ちなみに,Googleの BERT は,Transformer-Encoderを主に使用したLLMである.一方,OpenAI のGPTは Transformer-Deocderを使用したLLMである.

3.2.3 物体認識・物体検出バックボーンでのTransformerの応用

画像認識でも,Transformerを応用した以下2つが登場して以降,その中でレイヤー正規化がよく使用されるようになっていった:

  1. 物体認識向けのViT (Vision Transformer) (オリジナルのViTや,Swin-Tranformerなど)
  2. 物体検出向けのDETR (DEtection TRansformer) (DETRやdeformable DETRなど)

ただし,Transformerは畳込み層が排除されたまま当初は画像認識でも使用されていたものの,2022年後半から,CNNバックボーン畳み込み層も復権の兆しがあることには注意.音声系でも,畳み込み層も使用したConformer系ネットワークが強い.

また,レイヤー正規化は(2次元データで3次元特徴マップをつくる)画像データ相手には不自然なので,Transformer的全体構造でも,バッチ正規化を正規化層に用いたケースも,あとから主流になってきていった点には注意(例 YOLOFなど).レイヤー正規化は,1次元トークン系列が相手の場合に使用すると良い仕組みである.

4. まとめ

この記事では,RNNやTransformer向けの「バッチ正規化型の層」であるレイヤー正規化,バッチ正規化と比較しながら紹介した(2節).

また,Transformer系のどのようなネットワークで,レイヤー正規化が用いられているか,もしくは用いられなくなったかを整理した.

関連書籍

References

  • [Ba et al., 2016] J. Ba, R. Kiros, and G. E. Hinton. Layer normalization. ArXiv, abs/1607.06450, 2016.
  • [Cooijmans et al., 2016] Tim Cooijmans, Nicolas Ballas, Ce ́sar Laurent, and Aaron Courville. Recurrent batch normalization. arXiv preprint arXiv:1603.09025, 2016.
  • [Gregor et al., 2015] K. Gregor, I. Danihelka, A. Graves, and D. Wierstra. DRAW: a recurrent neural network for image generation. arXiv:1502.04623, 2015.
  • [Ioffe and Szegedy, 2015] S. Ioffe and C. Szegedy, “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” In ICML , 2015.
  • [Vaswani et al., 2017] Ashish Vaswani, et al. Attention is all you need. In NIPS, 2017.
  • [Vendrov et al., 2016] Ivan Vendrov, Ryan Kiros, Sanja Fidler, and Raquel Urtasun. Order-embeddings of images and language. ICLR, 2016.

参照外部リンク