トランスフォーマーの自己アテンションの理解⑦エンコーダ・ブロックの詳細

この記事では、エンコーダ・ブロックの詳細を解説します。論文の図１の左側になります。

エンコーダ・ブロックの内部にある、マルチヘッド・アテンション（Multi-head Attention）、位置ごとのフィード・フォワード（Position-wise Feed-Forward）、残差接続（Residual Connection）、レイヤー正規化（Layer Normalization）などを解説します。

では、さっそく始めましょう。

エンコーダ・ブロック

オリジナルのトランスフォーマーのエンコーダは６つのエンコーダ・ブロックが階層をなしています。上図では、Nxと書かれていますが、これは「N階の層になっている」の意味です。６つの層を展開するなら以下のようになります。

上手では、「エンコーダ・ブロック」を単に「エンコーダ」としています。これは図を複雑にしないためでもありますが、同じ仕組みを積み上げただけなので「エンコーダ」と省略して呼んでも間違いではないからです。

以下は、エンコーダの図を日本語にしたものです。

マルチヘッド（多頭）アテンションのところに３つの矢が刺さっているのは、第六弾で解説したクエリ・キー・バリューを意味します。

マルチヘッド・アテンション

これまでの解説では、クエリ・キー・バリューを一つの組み合わせとしていましたが、オリジナルのトランスフォーマーでは８つのクエリ・キー・バリューの組み合わせを使います。

つまり、８個の自己アテンションの処理が別々に行われ、様々な観点からトークン間の関係の強さを測り文脈を取り入れることが可能になっています。これをマルチヘッド（多頭）アテンションと呼びます。

まず、埋め込みベクトルを線形層（Linear Layer）を使って８個の別々なクエリ・キー・バリューのベクトルへと変換します。線形層による計算は行列による積と基本的には同じですが、積を行った後にバイアス（訓練によって決まる定数）を加算するところが異なります。しかし、目的は同じで埋め込みベクトルからクエリ・キー・バリューに関する特徴を抽出しています。

具体的には、512次元の埋め込みベクトルから64次元のクエリ・キー・バリューを生成します。つまり、512次元から64次元のクエリ・キー・バリューへと線形層で投影します。これを８組分行います。実装では、一つずつ別々に８組作るのではなく並列に生成するのですが、ここでは８個の別のクエリ・キー・バリューの組があると考えた方がわかりやすいでしょう。

この８組のクエリ・キー・バリューでそれぞれ自己アテンションの計算を全てのトークンに対して行います。最終的にはクエリとキーによる注目度に応じて加重平均されたバリューのベクトルが各トークンに対して８個出来上がります。一つのバリューが64次元なので８個結合すると全部で512次元になり、元の埋め込みベクトルと同じサイズになります。最後にもう一つの線形層を通して、埋め込みベクトルが更新されます。

このようにして、一つのトークンに対して８組のクエリ・キー・バリューによる自己アテンションの計算を行うことができ、多様な関係性に基づく文脈の抽出が可能となっています。

複雑なアテンションの見える化

マルチヘッド・アテンションがもたらす効果のイメージを持つために、論文から引用した以下の英文を使います。とりあえずは、意味を理解する必要はありません。

論文では、この英文をエンコーダに通して得た結果を下図で見える化しています。

上図では、「making」というトークンからの自己アテンションの度合いを色で示しています。一つのトークンに対して最大８個の色のついた箱が見えます。色が濃いところは関係が強く、色がついていないところは「making」とほぼ関係がないトークンになります。

よって、「making」は「more」と「difficult」と複雑な関係があるのが分かります。また、「making」と「making」自体の関係はそれなりに強くあるものの、「more」や「difficult」と比べるとより単純に見えます。

以下に日本語に訳した文章で見てみます。太字のところが「making」と関係が強い部分です。

「より（more)」「困難に(difficult)」「している(making)」の部分が密接に関わっていると判断しており、マルチヘッド・アテンションが効いているのが分かります。

ちなみに上図の文章の最後にある<EOS>とは文章の終わりを意味する特別なトークンです。<pad>は空埋めのトークンで無視されます。

なお、この注目度のビジュアル化の図は、５番目のエンコーダ・ブロックからの値を使っています。ここまで高階層に来るとかなり文章の意味をとらえているようですね。

マルチヘッド・アテンションで得たの文脈を含んだ埋め込みベクトルからさらに特徴量を引き出すのが位置ごとのフィード・フォワードになります。

位置ごとのフィードフォワード

位置ごとのフィード・フォワード（Position-wise Feed-Forward）では、各トークンの埋め込みベクトルを線形層、ReLU、線形層に通すことでさらに特徴量を引き出します。

線形層は、全結合層（Fully-Connected Layer）とも呼ばれ、全ての入力値が計算に含まれます。上図の線形層でも全結合が行われますが、トークンごとに入力が独立しており、トークン間のデータのやり取りがないので「位置ごとの」（Position-wise）と修飾されています。つまり、トークンの埋め込みのベクトルはこれらの層によって別々に処理されます。

ただし、同じニューラルネットワークが全てのトークンに対して使われます。また、順番に一つずつ処理するわけではなく、並列に処理がなされます。

つまり、すべての埋め込みベクトルを同時に同じニューラルネットワークで処理するのですが、よくある線形層の直前でデータをフラット（平坦化）にするようなことはありません。埋め込みベクトルを全て一列にまとめて処理するわけではないということです。

線形層と非線形層によって構成されたこれらの層がデータから学ぶことで複雑な関数となり、おのおのの埋め込みベクトルから有益な情報を抽出するようになります。端的に言えば、各トークンの特徴量をさらに抽出します。

また、複数のエンコーダ・ブロックが階層をなすことで、マルチヘッド・アテンションによる文脈の抽出と位置ごとのフィード・フォワードによる各トークンの特徴量の抽出を繰り返します。

このようにして、文章の意味に関するより抽象的な情報が形成されていくことになります。

続いて「加算」と「正規化」をそれぞれ見ていきましょう。

残差接続

図の中にある「加算」とは残差接続（Residual Connection）のことです。MicrosoftのResNetでよく知られるようになった残差接続は複数の層を通ることによって生じる勾配消失を防ぐ効果があります。実際、ResNetでは最大152の層を持つモデルの訓練が可能となりました。スキップ結合（Skip connection）とも呼ばれます。

エンコーダでは、下図の赤線で記した部分が残差接続になっています。

勾配消失が起こると学習中にネットワークのパラメータを効率よく更新できなくなり、訓練が非常に長くなったり、あるいは全然進まなくなってしまいます。よって残差接続はニューラルネットワークへ入力した値を出力された値に加算します。

つまり、ニューラルネットワークが入力値$${x}$$から $${H(x)}$$という関数を学習すると考えると、残差接続がある場合では、入力値$${x}$$に$${F(x)}$$という関数を加えたものを学習することになるので、以下の関係が成り立ちます。

$$
H(x) = F(x) + x
$$

よって、入力値$${x}$$に対して残差$${F(x)}$$を学習するようになります。

このため、残差接続があると、入力されたもともとの埋め込みベクトルからの情報（位置情報を含む）を維持することができます。ハイレベルな意味や抽象的な概念だけでなく、もともとのトークンの情報を失わないようにする意図があります。

残差接続の後に、次に解説する「正規化」が行われます。

レイヤー正規化

図の中の「正規化」とはレイヤー正規化（Layer Normalization）のことです。

画像処理などでよく行われるバッチ正規化（Batch Normalization）と同等の役割を持ち、訓練中にデータの平均や標準偏差が大きく変わることを防ぎ訓練が安定することを目指しています。

なので、まずはバッチ正規化から解説します。ここはニューラルネットワークの知識がないとちょっとわかりづらいかもしれません。正規化をすることで訓練が安定することだけ抑えてあとは飛ばしても全体像の理解には差し支えないありません。

バッチ正規化ではニューロンの活性化の値に対し、バッチ毎にデータの平均と標準偏差を計算して正規化（平均が０で標準偏差が１になるようにする）を行います。

ここでニューロンと呼んでいるのは、下図にある単純なニューラルネットワークでいうと、隠れ層にある紫の丸の一つに相当します。

このニューロンからの出力である活性化の値に対してバッチ正規化を行う場合、仮にバッチサイズが$${N}$$ならば、同じニューロンに対する$${N}$$個の活性化の値から平均と標準偏差を計算して正規化を行います。なお、標準偏差は分散の平方根です。

ここで、この$${N}$$個の値を$${x = (x_1, x_2, \dots, x_N)}$$と呼ぶことにします。平均$${E[x]}$$と分散$${Var(x)}$$は、下記のように表現できます。

$$
E[x] = \frac{1}{N} \sum\limits_{i=1}^N x_i \\
 \\
Var(x) = \frac{1}{N} \sum\limits_{i=1}^N (x_i - E[x])^2
$$

これで正規化を行えます。

$$
(正規化された x) = \frac{x - E[x]}{\sqrt{Var(x) + \epsilon}}
$$

$${\epsilon}$$（エプシロン）がついているのは分散がゼロの場合（全ての$${x_i}$$が同じ値を持つ）にゼロの割り算となってエラーが生じないようにするためです。

さらに、正規化された値に調整（スケールとシフト）を加えます。

$$
(スケール) \times (正規化された x) + (シフト)
$$

このスケールとシフトは学習によって決まる値です。数式では以下のようにスケールを$${\gamma}$$（ガンマ）、シフトを$${\beta}$$（ベータ）と表記することが多いです。

$$
y = \boldsymbol{\gamma} \frac{x - E[x]}{\sqrt{Var(x) + \epsilon}} + \boldsymbol{\beta}
$$

上記の式を直感的に説明すると、バッチごとに正規化をすることで平均０で標準偏差１にしてから、値をスケールしたり、シフトすることで後の層から見るとちょうど良い範囲から常に似たような値の入力値がやってくるようになり、訓練が安定化します。これによって訓練が収束する（損失値が小さくなる）スピードも速くなります。

このバッチ正規化は、RNNではうまく機能しません。なぜなら、RNNには再帰の構造があり、シーケンス（時系列など順番に並んだデータ）の長さも毎回異なるためにバッチ毎に平均や分散を計算するのが困難だからです。

そこでRNNでは平均や分散の計算をRNNのステップ毎に計算し正規化を行います。RNNの層（レイヤー）からの値（隠れ状態）で平均と分散を計算し正規化を行うのでレイヤー正規化と呼びます。よってバッチの中にある他のデータに依存することがなく正規化を行うことができます。

よって文章によってステップの数が異なっていても問題がありません。

オリジナルのトランスフォーマーは言語モデルとして登場したので、このレイヤー正規化を採用しています。埋め込みベクトル内で平均と分散を計算して正規化を行うので、バッチ正規化より単純で計算も速くなります。

なお、後の論文で、レイヤー正規化を多頭アテンションや位置ごとのフィード・フォワードの前に行う手法の方がより訓練が安定することがわかりました。これを事前レイヤー正規化（Pre-LN）と呼びます。よって、多くの実装はPre-LNに基づいています。

逆に、オリジナルのトランスフォーマーの図では、後にレイヤー正規化を行うように描かれており、これを事後レイヤー正規化（Post-LN）と呼びます。実装されたコードを読む際にPre-LNがあることを知っておくと驚かなくて済むでしょう。

まとめ

エンコーダ・ブロックも一つ一つ順番に見ていけばそれほど複雑ではありません。

また、デコーダ・ブロックでもマルチヘッド・アテンション、位置ごとのフィード・フォワード、残差接続、レイヤー正規化が使われているのでデコーダ・ブロックもほぼ理解できていることになります。

あとは、ソース・ターゲットのアテンション（Source Target Attention）やマクスされたマルチヘッド・アテンション（Masked Multi-head Attention）などを押さえていけば、デコーダ・ブロックの仕組みをカバーできます。

次回は、デコーダ・ブロックの詳細を解説します。

（続く）