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

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

特に、デコーダ・ブロックの内部にあるマスクされたマルチヘッド・アテンション（Masked Multi-head Attention）とソースターゲットのアテンション（Source-Target Attention)を重点的に解説します。

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

デコーダによる自己回帰

第五弾で解説しましたが、まずはデコーダによる自己回帰の仕組みを簡単に復習します。そのほうが後の話がわかりやすくなります。

トランスフォーマーではエンコーダ・ブロックとデコーダ・ブロックが多階層に連なっており、自己アテンションのプロセスを何度も踏むことで文脈が各トークンへと吸収されていきます。

デコーダは自己回帰によって翻訳の出力文章を積み上げていきます。デコーダへの最初の入力は、文章の始まりを意味する<SOS>（Start Of Sentence）です。同時に、エンコーダからの文脈も入力として与えます。これらをもとにデコーダは最初のトークンを予測します。

デコーダへの次の入力は、<SOS>に前の出力を足したものになります。

同様の処理を続けていきます。

つまり、デコーダへの入力文章は、<SOS>から始まって徐々に積み上げられていきます。言い方を変えると、デコーダへの入力文章は、デコーダからの出力文書に<SOS>を足した（出力文章を右にシフトした）ものです。

最後はデコーダが<EOS>を予測することで翻訳が終わります。

しかし、訓練を行う際にトークンを一つ一つ積み重ねる方法は効率的ではありません。できればデコーダへの入力文章は一度で済ませたいものですが、自己回帰を行う限り無理です。もっと欲を言えば、複数の文章をまとめてバッチにして効率的に訓練したいのですが、不可能に見えます。

トランスフォーマーの訓練ではこれらの問題が解決されています。特にマスクされたマルチヘッド・アテンションが重要になります。順を追って一つ一つ見ていきましょう。

訓練中のバッチ入力

トランスフォーマーの訓練では複数の文書をバッチ入力します。つまり、複数の文章をまとめて扱うので、バッチ内の入力文章はすべて同じ固定長として扱われます。もちろん文章によって長さ（トークンの数）が異なるので、バッチ内で一番長い文章に固定の長さを合わせます。それより短い文章では余りの位置が残るので、特別なトークンである<PAD>（足りない位置を埋めるためのトークン）を指定します。

例えば、固定長が10だとして文章が４つのトークンだけならば、残りの部分は<PAD>のトークンで埋められます。

このようにエンコーダもデコーダも入力文章が使わない位置に<PAD>を指定します。こうすることで訓練中は、複数の文章をまとめたバッチを作ることが可能となります。

しかし、デコーダで自己回帰を使うなら、バッチ入力はできません。なぜならデコーダが何を出力するのかはやってみないと分からないからです。でもこれを解決する方法があります。

訓練中の教師強制

仮に自己回帰を使ってデコーダの訓練を行うとしたらどうなるかを考えましょう。

例えば、訓練中に、以下の英語の文章を日本語に翻訳しているとします。

A dog walks.

そして、モデルから期待される出力が以下の日本語の文章だとします。

犬が歩く。

英語の文章はエンコーダによって処理され文脈としてデコーダに渡されます。さらに、デコーダは<SOS>を最初の入力として受け取り翻訳を始めます。

最初に期待されるデコーダからの出力は「犬」です。しかし、訓練中のモデルは期待された値を出力しない可能性が高いので、モデルの出力をそのまま次回の入力にまわすとラベルとは全く異なる文章を生成するので訓練が安定しません。

よって、訓練では教師強制（Teacher forcing）を行います。原理は簡単で、デコーダの出力を使わずに、ラベルから期待される値を使います。

上記の例で続けると、ラベルからの値である「犬」を強制的にデコーダからの出力とします。つまり、モデルが「自転車」を予測したとしてもそれを次回の入力には使いません。次回の入力へは「犬」が追加されます。これはもはや本当の自己回帰ではなく、前もって用意した正解をデコーダが出力したという想定で訓練を行っています。

よって、<SOS>と「犬」のトークンがデコーダへと入力されることになります。

さらに、あたかも「犬」と「が」が出力されたのかのように次の入力に回していきます。このように、すべての入力に対してデコーダが期待通りの予測をしたかのように処理していきます。

もちろん実際の出力は損失関数に渡され訓練中のパラメータの更新で使われます。ここで話しているのはデコーダへの入力に対して教師強制を行うということです。

上記の例にあるように、教師強制による入力と出力の関係は、<SOS>を初めにつけた入力文章と文章の終わりを意味する<EOS>を最後につけた正解文章（ラベル）になります。また、すべて必要なものを前もって準備することができるので、固定長になるように<PAD>を追加すればバッチ入力をすることが可能になります。

例えば、入力データと期待される出力データ（ラベル）を以下のように準備します。ここではバッチ内の固定長が10であると仮定しています。

ただし、このまま入力データ（左）をデコーダへ入力したら正解がわかってしまいます。また、<PAD>の部分は文章と関係ないので無視する必要があります。

ここでマスクの概念が登場します。

PADを無視する

<PAD>の部分は本来の文章の一部ではなく文脈情報もないのでアテンションの計算から除外する必要があります。そこで訓練を行う際にマスクを準備します。マスクでは<PAD>がある位置に０を、それ以外は１に設定します。

これはエンコーダとデコーダで共通の仕組みですが、ここではエンコーダの入力文章を使ってマスクの例を見てみます。

このマスクを使って<PAD>の位置を無視します。まず、第六弾でクエリとキーを使ってアテンションの計算を行ったのを思い起こしてください。クエリとキーで関係の強さを計算した後にソフトマックスで各トークンへ注目する割合を決めます。そこで、マスクがかかった位置に対しては、クエリとキーによって計算した値を使わずに、非常に大きな負の値（例えば、-10億）を設定します。この値は、ソフトマックスによってほぼゼロに変換されるので、その位置に対する注目度がなくなり無視されます。

よって、訓練を行う際に入力データとともにマスクの情報もエンコーダに渡します。これもバッチ入力が可能です。

第七弾でマルチヘッド・アテンションを見える化した図を紹介しました。そこでも<PAD>の部分はまったく注目されていないのが分かります。

もちろん、デコーダでも同様のマスクを使用して<PAD>を無視します。つまり、<PAD>を無視する仕組みはエンコーダもデコーダも同じです。

デコーダからの出力は最後は線形層などを経たのちに、損失関数（Loss Function）へと与えられますが、ここでも<PAD>の位置は無視するようになっています。なぜなら、<PAD>の位置は本来の文章には関係がないので、それに対する予測は意味がないからです。

なお、実用としてモデルを使う場合は<PAD>は必要ありません。一文ずつ入力するのでバッチを作る必要がなく、固定長にする必要もないからです。よってマスクも使われません。

次に、デコーダへの入力において正解を隠して必要な部分だけを使う方法を解説します。

マスクされたマルチヘッド・アテンション

もう一度、デコーダへの入力データ（右）とラベル（左）の例を見ましょう。

このままデコーダへ入力するとすべてのトークン間をアテンションを計算しまいます。本来なら、<SOS>は次に続く「犬」や「が」などのトークンに対するアテンションを計算するべきではありません。

よって、ここでもマスクを使います。<SOS>からのアテンションの計算で処理されたくない部分にマスクで０を設定しておけばアテンションの計算で無視されるのでデコーダの予測に影響がありません。つまり、入力データにとともと無かったのと同じことになります。

よって以下のようにすべてのトークンに対しマスクを用意しておきます。

このマスクを使うことで、<SOS>は<SOS>に対するアテンションだけを処理します。「犬」は<SOS>と「犬」へのアテンションを計算します。「が」は<SOS>、「犬」、「が」へのアテンションを計算します。このように各トークンはまるで自己回帰を行っているかのように見える範囲が限定されます。

さらに、同じマスクで<PAD>も無視するようにしておけば、両方のマスクを同時に処理できます。

上記のように固定長で10のトークンがある文章では 10 x 10でトークン同士の関係の強さの計算が生じます。その内、マスクで０に相当するアテンションの計算は無視しされます。しかも、これらの処理はすべてのトークンに対して同時に行われます。よって、自己回帰のように一つずつ入力を積み上げる必要がなくバッチ処理が行えます。

なお、<PAD>から各トークンへのアテンションが含まれるように見えますが、損失関数で<PAD>の位置は無視されるので問題はありません。

以上により、デコーダへの入力は固定長のデータをそのまま使うのですが、各トークンからのアテンションの計算はマスクされていない位置に関してのみ有効となるので、結果としては教師強制を使った自己回帰と同等になります。

なお、実用で使う場合には本当の自己回帰を使うので教師強制もマスクも不要になります。

デコーダの最初のマルチヘッド・アテンションではこのデコーダ特有のマスクを使うのでマスクされたマルチヘッド・アテンションと呼ばれます。

最後にエンコーダとデコーダをつなげるソース・ターゲットのアテンションを解説します。

ソース・ターゲットのアテンション

デコーダはエンコーダから出力された最終的な埋め込みベクトルを入力文章の文脈として取り入れます。

ここで、エンコーダ側をソース（Source)と呼び、デコーダ側をターゲット（Target）と呼びます。この方が両者の関係を説明しやすくなります。

上記を言い直すと、ソースの文脈はエンコーダからの最終出力の埋め込みベクトルで、デコーダの各デコーダ・ブロックへと渡されます。これによって、各デコーダ・ブロックはソースからの文脈をターゲットへ取り込むことが可能となっています。

ここまでは何度か繰り返し話しましたが、ここからその詳細を見ていきましょう。

各デコーダ・ブロックはソースの文脈からどれほどの情報をターゲットへと取り込むのかを計算する必要があります。ここでもアテンションが使われますが、計算するのはソースとターゲットのトークン間の関係の強さになります。これをソース・ターゲットのアテンション（Source-Target Attention）と呼びます。

ソース・ターゲットのアテンションでも、自己アテンションと同様に、クエリ・キー・バリューの手法が使われます。下図において、赤い線はソースからのバリューとキーを意味します。青い線はターゲットからのクエリです。

つまり、ターゲットからの質問をソースの文脈へ投げかけています。ターゲットのクエリとソースのキーによるターゲットとソースのトークン間の関係の強さを知れば、どのようにソースからの文脈をターゲットの各トークンに対して取り入れるかを決めることができます。

例えば、英語を日本語に訳しているとすると、英語の文章の文脈（ソース）と翻訳されていく日本語の文章の文脈（ターゲット）のトークン間の関係の強さを計算することでソースからターゲットへ文脈情報を取り込めます。

より具体的に言うと、例えば、ターゲットのトークンの一つである「犬」からソースの「A dog walks.」の文脈の各トークン（埋め込みベクトル）への関係の強さを計算して英文（ソース）からの文脈を日本語へ翻訳（ターゲット）の文脈へと取り込んでいきます。なので、クエリがターゲット（翻訳文）からきているのが自然なのが理解できます。また、ソース（原文）からキーとバリューが来ているのも、ターゲット（翻訳文）からソース（原文）に対して質問をしていると考えると自然です。

一つだけ不思議に思われるかもしれないのは、異なる言語間でのアテンションが計算できるということです。このことから、翻訳タスクによる訓練を通じてソースとターゲットの埋め込みベクトルは言語を超えた共通な情報を含むものになっているためであるのがわかります。

以上、ソース・ターゲットのアテンションによって、翻訳文として生成される文章の各トークンが入力文書のどのトークンに注目すべきかを計算することが可能になっています。結果として、ソースからの文脈がバリューとして抽出され、ターゲットの文脈へと取り込まれていきます。これが翻訳を可能にしているわけです。

まとめ

トランスフォーマーのアテンションの仕組みに関する全ての解説が終わりました。下は、論文からのお馴染みの図です。その整然さに美しささえ感じます。

ここまで読まれた方なら、なぜGoogle Brainの研究者であるAshish Vaswaniらが論文の冒頭でトランスフォーマーを「シンプル」と呼んだのか納得できるでしょう。もう一度、ここに引用します。

今後もトランスフォーマーに関する記事は書く予定ですが、自己アテンションのシリーズはここで一旦、完結といたします。

（完）