トランスフォーマーの論文を読む④全体
今回から「モデルの構造」のセクションを読み解いていきます。これによって論文の図１を掘り下げながら理解していきます。
まずはトランスフォーマーの全体像をざっくり理解することを目指します。
そもそもオリジナルのトランスフォーマーは、機械翻訳モデルです。つまり、入力文章（例えば英文）を受け取って出力文章（例えばフランス語）を生成します。
よって、トランスフォーマーを鳥瞰してみると、以下のようになります。
ここでは、英語の文章「Hello world!」がフランス語の「Bonjour le monde!」と翻訳されています。
でも、このままだと上述の論文の図１とは全然似ていませんね。そこで、この記事では、論文を読み解きながらトランスフォーマーの鳥瞰図から論文の図１へ徐々に近づけるように解説していきます。
エンコーダ・デコーダの全体像
論文のセクション３では、図１に引き続き、次のように書かれています。
トランスフォーマーは、全体的に図１にあるようなアーキテクチャに従います。そこでは、スタックされ（積み重ねられ）た自己アテンションと位置ごとの全結合層がエンコーダ（図の左）とデコーダ（図の右）の両方で使われています。
The Transformer follows this overall architecture using stacked self-attention and point-wise, fully connected layers for both the encoder and decoder, shown in the left and right halves of Figure 1, respectively.
下の図は、論文図１に赤と青の箱をつけたものです。
左側の赤い箱がエンコーダで、右側の青い箱がデコーダに相当する部分です。そして、その両方の中に自己アテンションや位置ごとの全結合層（フィードフォワード）がスタックされ（積み重ねられ）ています。
また、エンコーダの出力は、デコーダへと受け渡されています。
中身を無視して、エンコーダ・デコーダという大枠の構造だけで考えると、このような構造は、トランスフォーマーだけでなく、それ以前の機械翻訳でもよく見られたモデル構造です。エンコーダが入力文章から文脈を抽出して、デコーダがそれを受け取って出力文章を生成します。
これを単純化した図に当てはめると次のようになります。
つまり、エンコーダ・デコーダ構造を使った機械翻訳では、エンコーダが入力文章の情報（特徴量）を文脈（何らかの中間言語できなもの）へと変換し、その情報を元にデコーダが出力文章を生成することになります。
繰り返しますが、このエンコーダ・デコーダという大枠の仕組みとしては、トランスフォーマーもそれ以前の機械翻訳と同じアプローチをとっています。ただし、エンコーダとデコーダの中身は大きく異なっています。そこには、回帰や畳み込みは使われていません。その代わりに、エンコーダとデコーダそれぞれの中で自己アテンションや位置ごとのフィードフォワードが使われています。これを論文では説明しているわけです。
エンコーダとデコーダのスタック
論文は次にセクション3.1へと進みます。そこでは、エンコーダとデコーダそれぞれについての概要が述べられています。
エンコーダ：エンコーダは、N = 6の同一のレイヤー（層、ブロック）のスタック（積み重ね）で構成されています。
...
デコーダ：デコーダも、N = 6の同一のレイヤー（層、ブロック）のスタック（積み重ね）で構成されています。
…
Encoder: The encoder is composed of a stack of N = 6 identical layers.
…
Decoder: The decoder is also composed of a stack of N = 6 identical layers.
…
つまり、エンコーダとデコーダの中には6個の層がスタックされ（積み重ねられ）ているということです。しかも、同じ構造（ブロック）が繰り返し使われていることになります。
エンコーダのブロックが6つ、デコーダのブロックも6つとスタックされて使われています。これは、画像分類モデルであるResNetなどでブロックを繰り返して特徴量を抽出するやり方と似ています。ただし、画像分類とは異なり画像のサイズを小さくしていくようなことは行いません。
ちなみに、トランスフォーマー構造を応用したBERTのベースモデルは12階層のエンコーダを使っていまし、初代のGPTは12階層のデコーダを使っています。
上図のエンコーダ・デコーダを縦に描き直すと次のようになります。
エンコーダからの文脈は、デコーダのそれぞれの文脈に渡されるので次のように描けます。
論文の図１では、Nx という表現を使ってスタック（積み重ね）を表現しています。つまり、N個のブロックがスタック（積み重ね）になっていることを次のように描いています。
同様の表現を使えば、単純化した図も次のようになります。
これでトランスフォーマーのエンコーダ・デコーダの全体構造が見えてきました。要するに、入力文章（例えば英文）をスタックされ（積み重ねられ）たエンコーダに通すことで文脈（特徴量、中間言語的なもの）を抽出し、それをデコーダに渡すことで出力言語を使った文章（例えば、仏文）を生成します。
全体像がわかると論文の著者が主張しているようにシンプルな構造に見えてきます。
我々は、新しいシンプルなネットワークアーキテクチャであるトランスフォーマーを提案します。これは、再帰や畳み込みを一切使用せず、完全にアテンション機構のみに基づいています。
We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely.
もちろん、詳細はもっと複雑です。しかし、全体像を知っているとそうでないとでは、理解のしやすさも変わってくるでしょう。
ここからは、エンコーダとデコーダの内部構造をもう少し掘り下げていきます。
エンコーダ・ブロックとサブレイヤー
セクション3.1の続きを読んでいきます。
エンコーダ：エンコーダは、N = 6の同一のレイヤー（層、ブロック）のスタック（積み重ね）で構成されています。各レイヤーにはさらに2つのサブレイヤーに分かれています。最初のサブレイヤーはマルチヘッド（多頭）自己アテンション機構であり、2番目のサブレイヤーはシンプルな位置ごとの全結合型フィードフォワードネットワークです。
Encoder: The encoder is composed of a stack of N = 6 identical layers. Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed-forward network.
ここでマルチヘッド（多頭）自己アテンションや位置ごとのフィードフォワードといった言葉に対する説明はまだありません。よって、この段階では、そういう名前のものがあるのだなくらいの捉え方に留めておきます。前回にこの論文を深掘りせずに読んだときに、これらの言葉の解説は後の方にあることはわかっているので、「この二つのサブレイヤーが文脈抽出に役立っているのだな」的な理解で十分です。
下図は、エンコーダの図を日本語にしたものです。エンコーダの中に、多頭アテンションと位置ごとのフィードフォワードがあるのがわかります。
なお、上図にある位置エンコーディングなども論文の後の方で登場するので、入力文章に対する前処理ということ以外、今は深掘りしません。
論文は次のように続きます。
我々は、2つのサブレイヤーのそれぞれに対して、残差接続（スキップ結合）[11]を用いて、その後にレイヤー正規化（Layer Normalization）[1]を適用しています。つまり、各サブレイヤーの出力は
LayerNorm( x + Sublayer(x) )
であり、ここでの Sublayer(x) とは、サブレイヤー（多頭アテンション、あるいは位置ごとにフィードフォワード）によって実装される関数です。
We employ a residual connection [11] around each of the two sub-layers, followed by layer normalization [1]. That is, the output of each sub-layer is LayerNorm(x + Sublayer(x)), where Sublayer(x) is the function implemented by the sub-layer itself.
つまり、二つのサブレイヤー（多頭アテンションと位置ごとのフィードフォワード）では、ResNetにあるような残差接続（スキップ結合）を使っています。これはたくさんの層があるニューラルネットワークでは常套手段となっています。スキップ結合があるとたくさんの層を保つことで生じやすい勾配消失の問題が軽減できるからです。
論文の中にある番号で[11]とあるのは、巻末にある参照のところで定義されている番号のことで、ResNetの論文を参照しています。
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770–778, 2016.
次に、レイヤー正規化は、バッチ正規化のように活性化関数への入力値の分布が大きく乱れないように値を調節して学習を安定化させるものです。ただし、バッチ正規化とは異なり、バッチごとに平均と分散を計算するのではなく、入力シーケンスごとで正規化を行います。
ここで入力シーケンスと呼んでいるのは、入力文章がトークン化され、さらに各トークンが埋め込みになり、前処理などを施されたデータの並びを意味します。それが、エンコーダの中で繰り返し処理されているイメージです。サブレイヤーに対する入力データですが、処理され続ける埋め込みベクトルが並んだものになります。
レイヤー正規化により、シーケンス内の全トークンが一貫したスケールで扱われ、モデルの学習プロセスが安定しやすくなります。これは、後で登場するアテンションの仕組みにも大きく関わってきます。
なお、[1]とは Layer Normalization の論文を参照しています。
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
ちなみに、後の研究によって、レイヤー正規化を多頭アテンションや位置ごとのフィード・フォワードの前に行う手法の方がより訓練が安定することがわかりました。これを事前レイヤー正規化（Pre-LN）と呼びます。よって、多くの実装はPre-LNに基づいています。
ただし、トランスフォーマーの論文では、事後レイヤー正規化（Post-LN）が行われています。他の論文や実装されたコードを読む際にPre-LNがあることを知っておくと驚かなくて済むでしょう。
こういった情報は継続してさまざまな論文を読んだり、まとめ記事などに目を通していないと見逃すことがあります。ただし、人工知能関連の論文の数は多すぎるので専門家でない限り見逃しを避けるのは相当難しくはなりますが。
セクション３.１のエンコーダの説明は次のように結ばれています。
これらの残差接続（スキップ結合）を容易にするために、モデル内の全てのサブレイヤーおよび埋め込み層は、出力の次元が $${d_\text{model} = 512}$$ となるように設計されています。
To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension $${d_\text{model} = 512}$$.
つまり、画像処理などと違って、次元数を拡大・縮小するようなことはなく、全てのトークンの埋め込みの次元は512となっています。
よって、エンコーダを通過し出力される文脈は、入力シーケンスの埋め込みに対して繰り返し多頭アテンションと位置ごとのフィードフォワードによる更新を施したものになりますが、処理が容易になるように次元数が固定されているわけです。
このようにして処理されるシーケンスは、データの形式（シーケンス内の埋め込みの数やその次元）は、元々の入力言語と同じなのですが、中身は異なるものになります。イメージとしては、元々の入力言語が中間言語によるシーケンスへと変換されている感じになります。これは、出力言語とも異なるのですが、文章を生成するのに必要な情報（文脈、トークン間の関係など）が含まれています。しかし、これはあくまでもイメージとしての説明であり、実際に中間言語の文法などがわかるわけではありません。
なお、言語モデルにおけるトークンや埋め込みの考え方の解説はこちらにあります。必要に応じて参照してください。
デコーダ・ブロックとサブレイヤー
セクション3.1は、次にデコーダの説明に移ります。
デコーダ：デコーダも、N = 6の同一のレイヤーのスタックで構成されています。デコーダでは、2つのサブレイヤーに加えて、エンコーダの出力に対して多頭（マルチヘッド）アテンションを実行する第3のサブレイヤーを挿入します。エンコーダと同様に、各サブレイヤーの周りに残差接続（スキップ結合）を使用し、その後にレイヤー正規化を行います。
Decoder: The decoder is also composed of a stack of N = 6 identical layers. In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization.
ここでは、デコーダ・ブロック内のサブレイヤーについて述べられています。基本的には、エンコーダと同じように多頭アテンションと位置ごとのフィードフォワードがあり、残差接続（スキップ結合）やレイヤー正規化が行われています。ただし、デコーダには第３のサブレイヤーがあります。
下図を見てください。左側のエンコーダから右側にデコーダに対して文脈が渡された先に第３のサブレイヤーがあります。
これは、多頭アテンションですが、入力としてエンコーダからの文脈が混じっています。つまり、ここでは入力シーケンスからの情報を取り入れた上での出力シーケンスの処理が行われています。これは、機械翻訳が原文の情報を必要とすることを考えれば不思議なことではありません。しかし、これが具体的に何を意味するのかはこの段階では解説されていません。よって、ここでは、入力シーケンスからの文脈がここで繋がっているとだけ理解しておきます。
論文は次のように続きます。もし、ちょっと読みにくいと感じたら、最初は解説に目を通した後で読み返してみてください。
デコーダスタック内の自己アテンションのサブレイヤーも修正し、シーケンスにおける後続の位置にアテンションを適用することを防ぎます。このマスキングは、出力埋め込みが一つの位置によってオフセットされる事実と組み合わせることで、位置 j における予測が j より前の位置における既知の出力にのみ依存することを保証します。
We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i.
これは、デコーダの最下部にある多頭アテンションのことを言及しています。
名前に「Masked Multi-Head Attention」とあるように、単なる多頭アテンションとは異なっています。この理由はデコーダが出力文章を生成する仕組みを考えるとよくわかります。
トランスフォーマーに限らず、言語を生成するデコーダはには、まず文章の始まりを意味するトークン（の埋め込み）が与えられます。この文章の始まりを意味するトークンを <SOS>（Start Of Sentence）などと表記します。
デコーダは、まず<SOS>を最初の多頭アテンション（マスクされた多頭アテンション）で処理します。これで、文章の始まりという情報が受け入れられます。そして、その次の多頭アテンション（第３のサブレイヤー）に渡します。そこでは、エンコーダからの入力文章の文脈（それは中間言語的なものによるシーケンス）が入ってくるので、<SOS>と合わせて考えると、翻訳文の最初のトークン（埋め込み）を生み出すための情報を埋め込みに織り込むようになります。
つまり、中間言語的なもののシーケンスから最初の翻訳の言葉を生み出すための情報を持った埋め込みへと変換していくわけです。その処理には、後続する位置ごとのフィードフォワードも加わっています。また、デコーダのスタックを通して繰り返し処理が行われることになります。そして、デコーダの出力として最終的には、出力言語のトークンへと変換されるわけです。
このようにして出力されたトークン（埋め込み）は、再び、デコーダの最下部にある入力へと回されます。よって、デコーダは自分が予測したトークンに対する次のトークンを生成します。この繰り返しによって翻訳された文章が生成されるわけです。
ただし、訓練中は正解文章がデコーダに入力として与えられるために、正解を先読みしないようにマスクする必要があります。なぜなら、デコーダが生成するトークン（の埋め込み）は、それ以前のトークンからの情報でのみ生成されるべきだからです。よって、ある位置 j のトークンに関する予測は、それ以前の位置（< j）にある情報からのみに依存するべきです。よって、将来的な情報が影響しないよう後続位置（> j）にあるの情報をマスクする必要があります。それが「マスクされた多頭アテンション」となっています。
このように、マスクされた多頭アテンションは、訓練中のデコーダが未来のトークンの情報を先読みしないようにします。各予測位置でモデルが利用できる情報は、その位置より前にあるトークンに限定されるため、自然言語の生成プロセスを模倣することができます。
ここまでセクション3. 1を読んできました。エンコーダ・デコーダの大まかな枠組みとサブレイヤーの配置などがわかりました。しかし、アテンションについての詳細はまだまだ不明です。よって多頭アテンションもマスクされた多頭アテンションも具体的には何をやっているのかはわかっていません。
次回予告
次回は、セクション３.２のアテンションを読みながら解説を続けます。
お楽しみに！
