アテンション

アテンション関数の定義

アテンション関数は、1つのクエリと、キーと値のペアの集合を受け取り、1つの出力へ写す関数として説明できる。クエリ・キー・値・出力は、いずれもベクトルである。

出力は値の重み付き和として計算される。各値に割り当てる重みは、クエリと、その値に対応するキーとの適合度関数(compatibility function)で計算される。クエリとよく適合するキーを持つ値ほど大きな重みを受け、出力に強く反映される。

Scaled Dot-Product Attention

原論文で提案されたアテンションは、Scaled Dot-Product Attention と呼ばれる。入力は、次元 d_k のクエリとキー、次元 d_v の値である。計算は次の順に進む。

1. クエリとすべてのキーとの内積を計算する
2. それぞれの内積を sqrt(d_k) で割る
3. softmax 関数を適用し、値に対する重みを得る

実際には、複数のクエリをまとめて行列 Q とし、キーと値もそれぞれ行列 K と V にまとめて、出力の行列を一度に計算する。

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

QK^T がクエリとキーの内積をまとめて求める部分、sqrt(d_k) で割る部分がスケーリング、softmax で得た重みを V に掛ける部分が値の重み付き和にあたる。

加法的アテンションと内積アテンション

よく使われるアテンション関数は、加法的アテンション(additive attention)と内積アテンション(dot-product attention、乗法的アテンションとも呼ぶ)の2つである。

• 内積アテンションは、スケーリング係数 1/sqrt(d_k) が無いことを除けば、Scaled Dot-Product Attention と同じである
• 加法的アテンションは、隠れ層を1つ持つフィードフォワードネットワークで適合度関数を計算する

両者の理論上の計算量は同程度である。しかし内積アテンションは、高度に最適化された行列の掛け算のコードで実装できるため、実務上ははるかに高速で、メモリ効率もよい。

なぜ sqrt(d_k) で割るのか

d_k が小さいときは、両者の性能は同程度である。ところが d_k が大きくなると、スケーリングをしない内積アテンションは、加法的アテンションに性能で劣る。

原論文の著者らは、その原因を次のように推測している。d_k が大きいと内積の値の大きさが大きくなり、softmax 関数が、勾配が極めて小さくなる領域へ押しやられてしまう。

この影響を打ち消すため、内積に 1/sqrt(d_k) を掛けてスケーリングする。名前の Scaled は、このスケーリングを指している。内積アテンションの速さとメモリ効率を保ったまま、d_k が大きいときの弱点を補う工夫である。