トランスフォーマーの自己アテンションの理解④位置エンコーディング
この記事では、本題であるトランスフォーマーの自己アテンションの仕組みに触れ始めます。再帰の構造を使わずにトークンの位置関係を知る手法である位置エンコーディングを解説します。
では、さっそく始めましょう。
「必要なのはアテンションだけ」の意味
2017年に発表されたトランスフォーマーの論文のタイトルは「Attention Is All You Need」（必要なのはアテンションだけ）ですが、これは再帰は必要がないという意味でもあります。さらに、再帰がないだけでなく、アテンションの計算は全てのトークンに対して同時（並列）に行うので、トークンを順番に処理することもありません。
しかし、文章の理解においてトークンの順番を知ることは欠かせません。以下の二つの文章を比べれば全く異なる意味になっているのが分かるでしょう。
「犬が私を噛んだ」
「私が犬を噛んだ」
両方とも同じトークンを使っています。第一弾で登場したBoW（Bag of Words）だったら同じ文章として判断されてしまいます。
第二弾で解説したRNNを使ったエンコーダ・デコーダでは再帰の構造がトークンを順番に処理することで文脈を伝播するのでトークンの前後関係が情報として組み込まれています。
単純なRNNでは、再帰による勾配消失の問題がありました。特に長い文章では離れたトークン間の関係を反映したモデルの学習が困難でした。LSTMによって改善されましたが、再帰の構造上の問題である勾配消失は完全に無くすことはできませんでした。
トランスフォーマーでもエンコーダ・デコーダという枠組みが使われていますが、その中で再帰の構造は使われていません。アテンションの仕組みだけで文脈を取り込んでいきます。
再帰がなくトークンを順番に処理しないトランスフォーマーには、再帰の構造に起因する勾配消失の問題はありません。しかし、どのようにトークンの位置関係を認識しているのでしょうか。
実はトランスフォーマーでは埋め込みのベクトルにトークンの位置情報を追加します。なので、まずはトークンの埋め込みついて見ていきましょう。
埋め込みを学習する
他の言語モデルと同じく、トランスフォーマーへの入力文章はまずトークン化されます。これによって各トークンは整数のIDを与えられます。
このIDを使って、埋め込みテーブルから埋め込み（分散表現）のベクトルを引き出します。
この埋め込みテーブルは、以下の項目をあらかじめ設定しておいたものです。
テーブルに収めるトークンの総数
分散表現ベクトルの次元数
収めるトークンの数はある程度大きくしておく必要があります。しかし、予期されないトークンが登場した場合は、特別に知らないトークン用のIDを使います。例えば、ID=0は知らないトークンとしておきます。英語ではよく<UNK>トークンと呼びます。UNKはUnknown（知られていない）の略です。
また、分散表現ベクトルの次元数は、モデルが使用する次元に合わせます。トランスフォーマーのベースのモデルでは512次元の埋め込みベクトルを使っているので埋め込みテーブルでも512を指定します。
埋め込みテーブルは、TensorFlowやPyTorchなどのライブラリに標準で付属するので、そう言ったものを使えば、処理は数行で完了します。
ちなみに、訓練が行われる前の埋め込みベクトルはランダムに初期化されたものです。よって、訓練中の埋め込みベクトルの要素の数値は固定されていません。翻訳のタスクを通しての学習を行いながら埋め込みベクトルの数値も更新されていきます。つまり、トランスフォーマーでは翻訳タスクにモデルを最適化しながら同時に埋め込みの分散表現をデータから学習します。
さて、ここまでの話で、トランスフォーマーへの入力文章がトークン化され、トークンIDに対応した分散表現ベクトルが埋め込みテーブルから取り出されました。
この段階では、まだ埋め込みベクトルにはトークンの位置情報は含まれていません。
次に、トランスフォーマーはトークンの位置情報を埋め込みに加えます。位置情報はトークンのシーケンス（順番に並んだもの）における位置を表現します。
しかし、ちょっと遠回りになりますが、まずは三角関数を簡単に復習します。その方が位置情報の話が簡単になります。
三角関数と時計の針
まず、以下のような半径１の円（単位円）を考えます。青色の弧（円周の一部）に対応する角度を$${\theta}$$と呼ぶことにします。$${\theta}$$はギリシャ文字ですが、角度を表すときによく使われるので慣用に従っています。なお、$${\theta}$$を日本語の発音にするとシータ、セータ、セイタ、テータなど人によって異なるので表記は日本語にせず$${\theta}$$とし、発音はお任せします。
また、角度は時計回りに増加するとします。下図の角度は45度よりは少し小さいので、40度ぐらいでしょうか。ただし、三角関数を扱う場合には、我々が日常使う角度の度数とは異なる単位のラジアン（radian）を使います。radianを日本語の発音にしたときラジアンが一般的だと思われるのでそう表記します。
この$${\theta}$$の値をラジアンで表すと上記の単位円の青い弧の長さと同じになるようにラジアンは定義されています。単位円の円周は$${2\pi}$$なので、弧が一周して360度をカバーするならば、ラジアンによる角度は$${2\pi}$$になります。 180度ならば$${\pi}$$です。
よって、日常使う角度とラジアンには以下の関係があります。
$$
\\
\theta = \frac{日常の角度}{180}\pi
$$
なぜ、ラジアンを使うのかというと三角関数であるサイン（sine)とコサイン(cosine)を使った公式などが単純になるからです。そのような公式はここでは一切登場しないので、これ以上は深掘りはしません。
ただし、サインとコサインの意味は知る必要があります。
$${\theta}$$に対するサインの値を$${\sin \theta}$$と表記します。これは上図の単位円の中に描かれた三角形の緑の辺の長さになります。また、$${\theta}$$に対するコサインの値を$${\cos \theta}$$と表記し、その値は赤の辺の長さになります。
よって、$${\left[\ \sin \theta,\ \cos \theta\ \right]}$$は、下図の矢印が指す単位円上の点に相当します。
$${\theta}$$が大きくなると、矢印も時計回りに動きます。そこで$${\theta}$$を時間tの関数だとします。まず、単純な関数としてtをそのまま$${\theta}$$として設定します。
$$
\theta = t
$$
tが0から$${2\pi}$$へと増えると$${\theta}$$も同様に増え、矢印も時計回りに円を描きます。次に、２分の１の速度で時計の針が回るとします。
$$
\theta = \frac{t}{2}
$$
ここでは、tが0から$${4\pi}$$へと増えると$${\theta}$$が0から$${2\pi}$$になって時計の針が一周します。つまり、一周するのに2倍の時間がかかっています。
よって、分母の数が大きくなると時計はよりゆっくりと進むことになります。上記二つの例を比べると短針と長針の違いと似ています。
ここまで理解できれば、トークンの位置情報エンコーディングもすんなりと頭に入るはずです。
トークンの位置情報
トランスフォーマーのベースモデルは埋め込みに512次元を使っていますが、直感的に考えるのが難しいので、まずは2次元で考えましょう。
まずは、トークンの位置情報をサインとコサインの値で指定することを考えます。その際の$${\theta}$$を単純にトークンの位置$${\text{pos}}$$をそのまま指定してみましょう。時間tを直接$${\theta}$$に設定するのと同じ考えです。
$$
\text{PE}(\text{pos}) = \left[ \ \sin(\text{pos}), \ \cos(\text{pos}) \ \right]
$$
ここで、トークンの位置情報を出力する関数を$${\text{PE}}$$と呼んでいます。これは、トランスフォーマーの論文の中で位置情報を追加することをPositional Encoding（位置エンコーディング）と呼んでいるからです。また、トークンの位置$${\text{pos}}$$は0から始まる整数になります。
$$
\text{PE}(\text{pos}=0) = \left[ \ \sin(0), \cos(0) \ \right] = \left[ \ 0\ , \ 1 \ \right]
$$
$$
\text{PE}(\text{pos}=1) = \left[ \ \sin(1), \cos(1) \ \right] = \left[ \ 0.84, \ 0.54 \ \right]
$$
$$
\text{PE}(\text{pos}=2) = \left[ \ \sin(2), \cos(2) \ \right] = \left[ \ 0.91, \ -0.42 \ \right]\\
\\
…
$$
ただし、このままだと、$${\text{pos}}$$が大きくなると針が一周以上回って他のトークンの位置とかぶるかもしれません。時計に例えると、針の動きが速すぎます。
そこで$${\text{pos}}$$を10000で割ることにします。つまり、時計の針がゆっくり動くようにするのです。
$$
\text{PE}(\text{pos}) = \left[ \ \sin ( \frac{\text{pos}}{10000} ), \ \cos (\frac{\text{pos}}{10000}) \ \right]
$$
これだとトークンの位置が１万でもラジアンで考えると１なので約57.3度にしかなりません。トークンの位置が$${2\pi}$$の１万倍（約62,831）にならないかぎり針は一周しません。かなりの数のトークンの位置を重複することなく表現できます。
$$
\text{PE}(\text{pos}=0) = \left[ \ \sin(\frac{0}{10000}), \cos(\frac{0}{10000}) \ \right]
$$
$$
\text{PE}(\text{pos}=1) = \left[ \ \sin(\frac{1}{10000}), \cos(\frac{1}{10000}) \ \right] \\
$$
$$
\text{PE}(\text{pos}=2) = \left[ \ \sin(\frac{2}{10000}), \cos(\frac{2}{10000}) \ \right] \\
\\
…
$$
では、埋め込みの次元数を4次元に拡大しましょう。それに合わせてサインとコサインのペアもふたつに増やします。針がふたつある時計と同等になりました。
$$
\text{PE}(\text{pos}) = \left[ \ \sin ( \text{pos} ), \ \cos ( \text{pos} ) , \ \sin ( \frac{\text{pos}}{10000} ), \ \cos (\frac{\text{pos}}{10000}) \ \right]
$$
最初のサインとコサインのペアは速く動く針で、最後のサインとコサインのペアはゆっくり動く針です。例えると時計の秒針と時針の違いになります。この4次元の位置エンコーディングでも遅い針が一周しない限りは値が被ることはありません。
$$
\text{PE}(\text{pos}=0) = \left[ \ \sin ( 0 ), \ \cos ( 0 ) , \ \sin ( \frac{0}{10000} ), \ \cos (\frac{0}{10000}) \ \right]
$$
$$
\text{PE}(\text{pos}=1) = \left[ \ \sin ( 1 ), \ \cos ( 1 ) , \ \sin ( \frac{1}{10000} ), \ \cos (\frac{1}{10000}) \ \right]
$$
$$
\text{PE}(\text{pos}=2) = \left[ \ \sin ( 2 ), \ \cos ( 2 ) , \ \sin ( \frac{2}{10000} ), \ \cos (\frac{2}{10000}) \ \right]
\\
…
$$
では、6次元の埋め込みに拡張してみます。長くなるので見やすいようにペア（針）ごとに表示します。
$$
\text{PE}(\text{pos}, i=0) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{0}{6}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{0}{6}}}) \ \right] \\
\\
\text{PE}(\text{pos}, i=1) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{2}{6}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{2}{6}}}) \ \right] \\
\\
\text{PE}(\text{pos}, i=2) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{4}{6}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{4}{6}}}) \ \right]
$$
上式では$${i = 0,1,2}$$で針を区別しています。各ペアの分母に対し$${\frac{2i}{6}}$$、つまり、$${\frac{0}{6}}$$、$${\frac{2}{6}}$$、$${\frac{4}{6}}$$を分母の10000のべき乗数として針の速さがだんだん遅くなるように調節しています。
こうして6次元のユニークな位置を表現しています。あるいは、３つの速さの異なる針（例えるならば、秒針、分針、時針）によるユニークな時間とも考えられます。いずれにせよ、6次元のユニークな値であることに変わりはありません。
この手法を使えば、512次元の埋め込みベクトルも256個の異なるスピードで動く針を持つ時計の時間として表現できます。
$$
\text{PE}(\text{pos}, i) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{2i}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{2i}{512}}}) \ \right]
$$
$${i}$$は0から255までで、合計256本の針に相当します。
$$
\text{PE}(\text{pos}, i=0) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{0}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{0}{512}}}) \ \right] \\
\\
\text{PE}(\text{pos}, i=1) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{2}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{2}{512}}}) \ \right] \\
\\
… \\
\\
\text{PE}(\text{pos}, i=254) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{508}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{508}{512}}}) \ \right] \\
\\
\text{PE}(\text{pos}, i=255) = \left[ \ \sin ( \frac{\text{pos}}{10000^{\frac{510}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{510}{512}}}) \ \right]
$$
ここで、$${i}$$を1から256にした方が最終的な分母が10000になるので良さそうですが、論文や実装の例などを見ていると0から255なのでそれに従っています。何事も0から始めるのが通常なのでここだけ１から始めるのも唐突で不自然ではあります。また、512次元だと最後のべき乗数が$${\frac{510}{512} \approx 0.996}$$になり、1と大差はありません。そんなに長い文章を扱うことがなければ問題なしです。いずれにせよ、位置エンコーディングの理解に変わりはありません。
上記の位置エンコーディングを一つのベクトルにまとめると以下になります（横スクロールすると終わりまで見れます）。
$$
\text{PE}(\text{pos}) = \left[ \sin ( \frac{\text{pos}}{10000^{\frac{0}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{0}{512}}}) \ \sin ( \frac{\text{pos}}{10000^{\frac{2}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{2}{512}}}),
…, \ \sin ( \frac{\text{pos}}{10000^{\frac{510}{512}}} ), \ \cos (\frac{\text{pos}}{10000^{\frac{510}{512}}}) \right]
$$
位置$${\text{pos}}$$は0から始まる整数なので、各トークンの位置は以下のように表現されます。
$$
\text{PE}(\text{pos=0}) = \left[ \sin( \frac{0}{10000^{\frac{0}{512}}} ), \cos(\frac{0}{10000^{\frac{0}{512}}}) \ \sin ( \frac{0}{10000^{\frac{2}{512}}} ), \ \cos (\frac{0}{10000^{\frac{2}{512}}}),
…, \ \sin ( \frac{0}{10000^{\frac{510}{512}}} ), \ \cos (\frac{0}{10000^{\frac{510}{512}}}) \right] \\
$$
$$
\text{PE}(\text{pos=1}) = \left[ \sin( \frac{1}{10000^{\frac{0}{512}}} ), \cos(\frac{1}{10000^{\frac{0}{512}}}) \ \sin ( \frac{1}{10000^{\frac{2}{512}}} ), \ \cos (\frac{1}{10000^{\frac{2}{512}}}),
…, \ \sin ( \frac{1}{10000^{\frac{510}{512}}} ), \ \cos (\frac{1}{10000^{\frac{510}{512}}}) \right] \\
$$
$$
\text{PE}(\text{pos=2}) = \left[ \sin( \frac{2}{10000^{\frac{0}{512}}} ), \cos(\frac{2}{10000^{\frac{0}{512}}}) \ \sin ( \frac{2}{10000^{\frac{2}{512}}} ), \ \cos (\frac{2}{10000^{\frac{2}{512}}}),
…, \ \sin ( \frac{2}{10000^{\frac{510}{512}}} ), \ \cos (\frac{2}{10000^{\frac{510}{512}}}) \right] \\
$$
つまり、各トークンの位置がユニークなベクトルで表現されているのがわかります。
位置情報のビジュアル化
下記の図ではトークンの位置情報ベクトル内の値をカラーで表現しています。
縦軸がトークンの位置でここでは0から127までになっています。これだと最大128の位置しかカバーできません。もっと長いトークンのシーケンス（順番に並んだもの）を扱う場合は、もっと多くの位置を計算しておく必要がありますが、ここではビジュアル化が目的なのでこれでOKです。
横軸が埋め込みベクトルの次元になります。512次元あるのでトランスフォーマーの埋め込みと同じサイズになります。埋め込みベクトル内の数値が１に近いほど黄色になり、-1に近いほどダークブルーに近づきます。0だと少し暗めの緑です。
下図は、位置情報を16まで、次元数を32までにして、より色の違いを分かりやすくしたものです。各行のベクトルが異なる値を含んでいるのが見えます。
$${pos=0}$$の１次元目は$${\sin( \frac{0}{10000^{\frac{0}{512}}}) = \sin 0 = 0}$$なので濃いめの緑になっています。また$${pos=0}$$の2次元目は$${\cos( \frac{0}{10000^{\frac{0}{512}}}) = \cos 0 = 1}$$なので黄色になっています。
埋め込みに位置情報を追加する
埋め込みテーブルから取り出した分散ベクトルにあらかじめ計算しておいたトークンの位置情報を追加します。最初のトークンの分散ベクトルにはPE(0)を足します。次のトークンの埋め込みベクトルにはPE(1)を足します。そうして、最後のトークンまで位置情報を追加します。これが位置情報エンコーディングです。
やっていること単純なベクトルの足し算です。つまり、ベクトルの要素ごとに値を足し合わせるだけです。
例えば、３次元のベクトルの足し算は以下の例のようになります。
$$
v_1 = [1, 2, 3] \\
v_2 = [2, 4, 5] \\
\\
v_1 + v_2 = [1+2, 2+4, 3+5] = [3, 6, 8]
$$
これが512次元であっても同じことです。要素同士を足し合わせるだけです。これで位置情報が埋め込みベクトルに取り入れられました。
モデルはどう位置を理解するのか
それにしても、モデルがどうこの位置を理解するのでしょうか。ここからは論文の中であくまでも推測として語られていることを解説します。
まず、彼らがサインやコサインの三角関数を選択したのは、ある位置から別の位置への相対的な位置は角度の増減、つまり回転によって得られるからです。回転は線形代数では行列による転換として捉えられるのでニューラルネットワークが位置情報を上手く扱えるようになるだろうと想定したそうです。ニューラルネットワークは行列の重みを調節して様々な特徴量を抽出するのでこれは無理のない想定だと言えます。
また、彼らは次のような仮定もしています。訓練中にあまり出ないような長い文章でも、サインやコサイン関数からの値ならば、簡単に訓練中に学んだ経験から延長してトークンの相対位置を考えることが可能だろう、と。これも納得な理由です。
これらの仮定や推測を直接確かめるのは難しいですが、結果として彼らは成功しました。ニューラルネットワークを訓練して翻訳のタスクに最適化させることで位置情報を利用できるようになったことで、再帰の構造なしでトークン同士の位置関係を捉えられるようになり、トランスフォーマーは「犬が私を噛んだ」と「私が犬を噛んだ」を混同することはありません。
これで自己アテンションへの入力の準備が整い、文脈を取り込んでいくプロセスが可能となりました。
次回は、それを見ていきましょう。
（続く）
