トランスフォーマーの論文を読む⑤注意
今回は「モデルの構造」からセクション3. 2の「アテンション」を読み進めます。アテンションはよく「注意」と訳されますが、あまりしっくりこない気がします。まだ「注目」の方がいいかなとも思いますが、この記事では英語のAttentionをカタカナ読みした「アテンション」で通します。
このセクションでは次の二つの図が登場します。
（左）スケールされた内積アテンション
（右）多頭アテンション
多頭アテンションも「マルチヘッド・アテンション」の方が響きが良いかもしれないですが、この記事では「多頭アテンション」としています。そもそも英語名「Multi-Head Attention」をそのまま使う方が良いのかもしれませんが。
上図を見るとわかるように、左の「スケールされた内積アテンション」は、右の「多頭アテンション」の中で使われています。つまり、「スケールされた内積アテンション」が複数あるのが「多頭アテンション」ということになります。
では、そもそも「スケールされた内積アテンション」はどのような仕組みを持ち、何を目指しているのでしょうか。
RNNにおけるアテンション
トランスフォーマーではRNNやCNNは使われていません。ただし、前提知識としてRNNにおけるアテンションについて知っていた方が理解が深まるので、ここで簡単に解説します。
RNNを使った機械翻訳のモデルで代表的なものは、エンコーダ・デコーダ構造を持ちます。RNNを使ったエンコーダ・デコーダでは回帰を使うことで文章（埋め込みのシーケンス）から文脈を抽出します。
RNNエンコーダは入力シーケンスから各トークンの特徴量（隠れ状態）のシーケンス$${h_1, h_2, \dots, h_{n-1}, h_n}$$として抽出しています。下図では、RNNエンコーダの処理の様子を表現しています。
このエンコーダからの出力シーケンス$${h_1, h_2, \dots, h_{n-1}, h_n}$$が入力文章からの重要な情報を含んだ文脈（上図の右上）となります。
RNNの処理の仕方として入力文章からの情報（トークンの埋め込みベクトル）を$${x_1}$$から順番に$${x_2}$$、$${x_3}$$、と$${x_n}$$まで処理していきます。このおかげで隠れ状態$${h_2}$$にはその前の隠れ状態$${h_1}$$からの情報がある程度受け継がれていきます。最終的には全てのトークンからの情報は$${h_n}$$へと集約されます。
上記の構造だと文章の解釈が順方向（前から後ろへ）だけで一方通行です。しかし、実際には後の情報がそれより前のトークンの解釈に役に立つこともあります。特に日本語では文章の最後に意味が逆転するようなケースは多々あります。このため、RNNベースの言語処理では順方向だけでなく逆方向も（$${x_n}$$から$${x_1}$$へと）処理して双方向（Bi-directional）の情報を組み合わせるなど処理を複雑化させる工夫の必要性も生じます。これらの複雑さはトランスフォーマーを使った言語処理では必要がなくなります。
さて、エンコーダからの出力された$${h_n}$$には、文章全体からの情報が反映されており、これをデコーダに渡すことでデコーダは文章の生成を始めることができます。ただし、RNNの手法の弱点としては、エンコーダの情報が最後の１つのベクトルに押し込められているため入力情報が多くなると対応しきれなくなることです。要するに入力文章のシーケンスが長いと前の方の離れたトークンの情報が弱まっていくという問題があります。例えば、$${h_1}$$からの情報は、$${h_n}$$ではあまり残っていないかもしれません。
そこで、一番最後の$${h_n}$$だけでなく、出力シーケンス$${h_1, h_2, \dots, h_{n-1}, h_n}$$全体からの文脈情報を利用しようというのがアテンションの考えになります。つまり、デコーダが$${h_n}$$の情報を使って出力したベクトルに、出力シーケンス$${h_1, h_2, \dots, h_{n-1}, h_n}$$から関連性の高い情報を抽出して合わせたものを最終的な出力とします。
上図では、右下のデコーダが開始のトークンとエンコーダからの最終の隠れ状態$${h_n}$$を受け取って最初の出力のために$${g_1}$$というベクトルを出力しています。これには出力文章の最初のトークンを生成するための情報が含まれています。しかし、前述したように$${h_n}$$に含まれている情報は、$${h_1, \dots, h_{n-1}}$$からの情報を十分に反映していないかもしれません。だったら、$${g_1}$$にある情報と関連の高い情報を文脈全体つまり$${h_1, \dots, h_n}$$から抽出して混ぜたもの（加重平均）を$${g_1}$$と組み合わせて使うことにしようということです。
下図はアテンションの部分だけを取り出したものです。
デコーダが出力したトークンの埋め込みベクトル$${g_j}$$とエンコーダからの文脈にある全ての埋め込みベクトル$${h_1, h_2, \dots, h_n}$$との関係の強さ（スコア）を計算します。スコアの計算には専用のニューラルネットワークとベクトル間の内積を合わせて使ったりします。ここではわかりやすさを重視して、単純に内積でベクトル間の関係の強さをスコアとして計算するものとしています。
このスコアをソフトマックス関数で重みに変換します。つまり、文脈シーケンスにある各ベクトルに合計すると１となる０から１までの重み（$${w_1, w_2}$$, $${\dots}$$, $${w_n}$$）を付与します。これによって全ての隠れ状態の値を加重平均したベクトルを計算できます。このベクトルには、エンコーダの文脈$${h_1, h_2, \dots, h_n}$$から抽出され、デコーダが出力した埋め込み$${g_j}$$に強く関連する情報が集約されています。
$$
\sum\limits_{i=1}^n w_i h_i
$$
このような計算を行う仕組みをアテンション機構（Attention Mechanism）と呼び、入力文章からの文脈をデコーダからの出力にうまく取り入れるために使われます。これによってデコーダから出力された値には、アテンションによって文脈全体から必要な情報が強められて含まれることになります。そして、その出力$${o_j}$$が次の$${j+1}$$のトークンの生成のためのインプットとなります。
下図は、RNNエンコーダ・デコーダの全体像をデコーダの観点から描いたものです。
アテンション機構は効果的なのですが、RNNには並列処理ができないという問題があります。なぜならば、時間軸に沿って順番に処理する必要があるからです。
しかし、よく考えてみるとアテンション機構そのものにはRNNによる回帰の処理が直接的には必要でないことがわかります。RNNはあくまでもシーケンスからの情報を抽出するためであり、アテンションはその情報間の関連性を計算するものだからです。
そこでトランスフォーマーではエンコーダとデコーダの構造やアテンション機構の精神を受け継ぎながら回帰を必要としないニューラルネットワークの構造を提案しています。これによって回帰による処理が複雑化していった経緯から逸脱してある意味より単純なモデルを構築しました。これまで見てきたように、トランスフォーマーの構造ではシーケンスから情報を抽出する部分をアテンション（とフィードフォワード）を中心とした処理に置き換えています。
こうしてみるとトランスフォーマーも以前の言語モデルの流れをある程度は汲んでいるのが見えてきます。なので、その論文自体にはこのような背景は詳しくは解説されていません。
それではトランスフォーマーの論文の続きを読みましょう。
3.2 アテンションのセクション
論文のセクション「3.2 Attention」はこう始まります。
アテンション関数は、クエリとキー・バリューのペアのセットを出力にマッピングするものとして記述できます。ここで、クエリ、キー、バリュー、そして出力はすべてベクトルです。出力はバリューの加重和として計算され、各バリューに割り当てられる重みは、クエリと対応するキーの互換性関数によって計算されます。
An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
そして前述の二つの図が登場します。
（左）スケールされた内積アテンション
（右）多頭アテンション
しかし、いきなりクエリ、キー、バリューが登場してなんのことだかわかりません。これらは全てベクトルだと言っています。また、キーとバリューはペア（対）になっているとも言及されています。さらに「互換性関数」という言葉も登場します。
この論文はこのあと次のようなサブセクションへと続きます。
3.2.1 Scaled Dot-Product Attention スケールされた内積アテンション
3.2.2 Multi-Head Attention 多頭アテンション
3.2.3 Applications of Attention in our Model 私たちのモデルにおけるアテンションの適用
つまり、上図の説明を知るには、このまま読み続けるしかありません。そこでは、クエリ、キー、バリュー、互換性関数の意味も解説されていると想定（期待）します。
スケールされた内積アテンション
セクション3.2.1はこう始まります。
私たちが使うアテンションを「スケールされた内積アテンション」と呼んでいます（図2参照）。入力は、次元が$${d_k}$$のクエリとキー、および次元が$${d_v}$$のバリューで構成されます。
We call our particular attention "Scaled Dot-Product Attention" (Figure 2). The input consists of queries and keys of dimension $${d_k}$$, and values of dimension $${d_v}$$.
そこで図２の左をよく見てみましょう。
Q、K、Vとあるのは、Query（クエリ）、Key（キー）、Value（バリュー）に関連していることは予想がつきます。また、これらがスケールされた内積アテンションの入力値だと理解できます。また、クエリとキーの次元数は$${d_k}$$で同じですが、バリューの次元数は$${d_v}$$と別扱いになっています。それでもキーとバリューはペアとなっている前述されているので、その点を心に留めておきます。
もう少し読み進めます。
クエリと全てのキーの内積（dot product）を計算し、それぞれを$${\sqrt{d_k}}$$で割り、ソフトマックス関数を適用してバリューに対する重みを得ます。
We compute the dot products of the query with all keys, divide each by $${\sqrt{d_k}}$$, and apply a softmax function to obtain the weights on the values.
論文で前述されていたようにクエリ、キー、バリューはベクトルです。また、クエリとキーは次元数が同じ$${d_k}$$なので内積を計算しその関係性の強さ（スコア）を計算することができます。「クエリと全てのキーの内積を計算し」とあるので、一つのクエリ（ベクトル）に対して、複数のキー（ベクトル）があり、その各々に対しての内積を計算していることになります。また、キーとバリューはペアになっているので、キーの数と同じだけバリューのベクトルも存在することになります。
つまり、ある一つのクエリ（ベクトル）と複数のキー（ベクトル）の内積を計算したものをソフトマックス関数によって変換したものが、各バリューへの重みにになるわけです。よってこのクエリ（ベクトル）に関しての重要度をキーとの内積で測ってバリューの値を加重平均でまとめています。これはRNNのエンコーダ・デコーダでも登場したアテンションの考え方と同じアプローチです。それが$${\sqrt{d_k}}$$で割られてスケールされているので、スケールされた内積アテンションとなっています。
ちなみに、内積のことを$${\cdot}$$で表現するのでドット積（dot product）とも呼びます。
内積をスケールする意図
では、なぜ内積の結果を$${\sqrt{d_k}}$$で割り、スケール（大きさ）を調整するのでしょうか。このことに関して論文の註釈４に次のように書かれています。
内積が大きくなる理由を説明するために、$${q}$$と$${k}$$の成分が平均0、分散1の独立した乱数変数であると仮定します。すると、その内積、$${q \cdot k = \sum_{i=1}^{d_k} q_i k_i }$$は、平均0、分散$${d_k}$$を持ちます。
To illustrate why the dot products get large, assume that the components of $${q}$$ and $${k}$$ are independent random variables with mean 0 and variance 1. Then their dot product, $${q \cdot k = \sum_{i=1}^{d_k} q_i k_i}$$, has mean 0 and variance $${d_k}$$.
よって、ベクトルの成分がそれぞれ平均０で分散１に従っているとすると、内積の分散は$${d_k}$$になります（分散が１の変数が$${d_k}$$個足されているから）。すると、ベクトルの成分の値を標準偏差$${\sqrt{d_k}}$$で割ることは内積の平均は０のままで分散を１にする、つまり正規化を行っているのと同じことになります。
実際にはベクトルの成分がそれぞれ完全に平均０で分散１に従っているわけではないですが、以前解説したレイヤー正規化によってある程度は似たようなスケールを持つようになっています。よって、このスケール調整はある種の正規化手法とみなすことができます。少なくとも、クエリとキーの内積の結果のスケールが調整されるので、大きな次元$${d_k}$$によって生じる分散の増加を補正していることになります。前回の記事で述べたように事前レイヤー正規化（Pre-LN）がより効果的なのは、アテンションの処理を行う前に正規化が行われるようになるからだと考えられます。
結果として得られる重みが適切な範囲内に収まり、ソフトマックス関数を通じても値が大きく偏ることなく勾配が安定し、効率的な学習が可能になります。オリジナルのトランスフォーマーでは大きな次元のベクトルが使われていますが、$${\sqrt{d_k}}$$でスケールすることで適切に機能するように調節していることになります。
内積後のマスクについて
なお、ソフトマックスの前にマスク（Mask）があります。これは前回の記事で触れた「マスクされた多頭アテンション」のためのものです。なので必要に応じて使われる機能なので、オプション（opt.）となっています。
マスクされるトークンの位置からの内積の値には大きな負の値（例えば、$${-10^9}$$など）が設定されます。するとソフトマックス関数がその重みをほぼゼロにするので実質的に無視されます。よってアテンションの計算から除外されます。つまり、マスクされることになります。
行列で内積の並列処理
では、読み進めます。
実際には、一連のクエリに対してアテンション関数を同時に計算します。これらのクエリは行列Qにまとめて格納されます。キーとバリューもそれぞれ行列KとVにまとめて格納されます。出力の行列は次のように計算されます：
In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix Q. The keys and values are also packed together into matrices K and V . We compute the matrix of outputs as:
そして次の式が提示されます。
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right)V
$$
これまでずっとクエリ、キー、バリューはベクトルとして解説していましたが、ここで実際にはクエリ、キー、バリューは複数ベクトルをまとめた行列だと知らされます。よって、Q、K、Vはそれぞれ行列です。
Qはクエリ（ベクトル）が行ベクトルとなって複数並んでいます。
$$
Q = \begin{bmatrix}
- \ \vec{q}_1 \ - \\
- \ \vec{q}_2 \ - \\
\vdots \\
- \ \vec{q}_m \ - \\
\end{bmatrix}
$$
Kはキー（ベクトル）がやはり行ベクトルとして複数並んでいます。
$$
K = \begin{bmatrix}
- \ \vec{k}_1 \ - \\
- \ \vec{k}_2 \ - \\
\vdots \\
- \ \vec{k}_m \ - \\
\end{bmatrix}
$$
$${QK^\top}$$の計算では、Kが転置されているので、各クエリ（ベクトル）とキー（ベクトル）の内積の組み合わせを全て並列処理によって計算することができます。
$$
\begin{aligned}
QK^\top &= \begin{bmatrix}
- \ \vec{q}_1 \ - \\
- \ \vec{q}_2 \ - \\
\vdots \\
- \ \vec{q}_m \ - \\
\end{bmatrix} \begin{bmatrix}
| & | & & | \\
\vec{k}_1 & \vec{k}_2 & \dots & \vec{k}_m \\
| & | & & | \\[1ex]
\end{bmatrix} \\[2ex]
&= \begin{bmatrix}
\vec{q}_1 \cdot \vec{k}_1 & \vec{q}_1 \cdot \vec{k}_2 & \dots & \vec{q}_1 \cdot \vec{k}_m \\
\vec{q}_2 \cdot \vec{k}_1 & \vec{q}_2 \cdot \vec{k}_2 & \dots & \vec{q}_2 \cdot \vec{k}_m \\
\vdots & \vdots & \vdots & \vdots \\
\vec{q}_m \cdot \vec{k}_1 & \vec{q}_m \cdot \vec{k}_2 & \dots & \vec{q}_m \cdot \vec{k}_m \\
\end{bmatrix}
\end{aligned}
$$
結果として、各行には、あるクエリ（ベクトル）に対する全てのキー（ベクトル）との内積が並ぶことになります。
さらに$${\sqrt{d_k}}$$でスケールされます。
$$
QK^\top = \frac{1}{\sqrt{d_k}}\begin{bmatrix}
\vec{q}_1 \cdot \vec{k}_1 & \vec{q}_1 \cdot \vec{k}_2 & \dots & \vec{q}_1 \cdot \vec{k}_m \\
\vec{q}_2 \cdot \vec{k}_1 & \vec{q}_2 \cdot \vec{k}_2 & \dots & \vec{q}_2 \cdot \vec{k}_m \\
\vdots & \vdots & \vdots & \vdots \\
\vec{q}_m \cdot \vec{k}_1 & \vec{q}_m \cdot \vec{k}_2 & \dots & \vec{q}_m \cdot \vec{k}_m \\
\end{bmatrix}
$$
そして、ソフトマックス関数を通ることで重みに変換されます。ここでのソフトマックスは各行ごとに行われるので、あるクエリに対してどのキーがより関連性が強いのかが計算されることになります。
$$
\text{softmax}\left(QK^\top\right) = \begin{bmatrix}
w_{11} & w_{12} & \dots & w_{1m} \\
w_{21} & w_{22} & \dots & w_{2m} \\
\vdots & \vdots & \vdots & \vdots \\
w_{m1} & w_{m2} & \dots & w_{mm} \\
\end{bmatrix}
$$
ここで各行$${r}$$ごとの重みの和は１になります。
$$
w_{r1} + w_{r2} + \dots + w_{rm} = \sum\limits_{c=1}^m w_{rc} = 1
$$
Vにはバリュー（ベクトル）が列ベクトルとして並んでいます。
$$
V = \begin{bmatrix}
- \ \vec{v}_1 \ - \\
- \ \vec{v}_2 \ - \\
\vdots \\
- \ \vec{v}_m \ - \\
\end{bmatrix}
$$
よって、クエリとキーからの重みで行列の積を計算すると加重平均されたバリューのベクトルが各クエリに対して計算されます。
$$
\begin{aligned}
\text{Attention}(Q, K, V) &= \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right)V \\
&= \begin{bmatrix}
w_{11} & w_{12} & \dots & w_{1m} \\
w_{21} & w_{22} & \dots & w_{2m} \\
\vdots & \vdots & \vdots & \vdots \\
w_{m1} & w_{m2} & \dots & w_{mm} \\
\end{bmatrix} \begin{bmatrix}
- \ \vec{v}_1 \ - \\
- \ \vec{v}_2 \ - \\
\vdots \\
- \ \vec{v}_m \ - \\
\end{bmatrix} \\[2ex]
&= \begin{bmatrix}
w_{11} \vec{v_1} + w_{12} \vec{v_2} + \dots + w_{1m} \vec{v}_m \\
w_{21} \vec{v_1} + w_{22} \vec{v_2} + \dots + w_{2m} \vec{v}_m \\
\qquad\quad \vdots \\
w_{31} \vec{v_1} + w_{32} \vec{v_2} + \dots + w_{3m} \vec{v}_m \\
\end{bmatrix}
\end{aligned}
$$
互換性関数とは
論文は、次のように続きます。
最も一般的に使用されるアテンション関数には、加算アテンションと内積（乗算）アテンションがあります。内積アテンションは、スケーリングファクター$${\frac{1}{\sqrt{d_k}}}$$を除いて、私たちのアルゴリズムと同一です。加算アテンションは、単一隠れ層を持つフィードフォワードネットワークを使用して互換性関数を計算します。二つは理論上の複雑さでは似ていますが、内積アテンションは、高度に最適化された行列乗算のコードを使用して実装できるため、はるかに高速で、よりメモリ使用の効率も良いです。
The two most commonly used attention functions are additive attention [2], and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $${\frac{1}{\sqrt{d_k}}}$$ . Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
こので互換性関数についての説明が登場しています。これによると、互換性関数とは、アテンション機構において、クエリとキーがどれだけ互いに適合しているか、つまりどれだけ「マッチ」しているかを数値化する関数のことを指します。
この関数によって計算された数値（互換性のスコア）は、ソフトマックス関数によってアテンションの重みを決定する際に使用されます。この重みは、あるクエリにとって、どのキー（とそれに対応するバリュー）がより重要かを示します。
また、互換性関数の例として、加算アテンションと内積アテンションの二つを例に挙げています。それぞれ計算方法が異なります。
加算アテンションでは、クエリとキーの組み合わせを入力として取り、単一の隠れ層を持つフィードフォワードニューラルネットワークを通して互換性スコアを計算します。この方法は、クエリとキーの関係を複雑な非線形関数（ニューラルネットワーク）でモデリングしています。
内積アテンションでは、クエリとキーの内積を計算し、その結果にスケーリングファクター$${\frac{1}{\sqrt{d_k}}}$$を掛けることで互換性スコアを得ます。この方法は計算がシンプルで、行列乗算を用いて効率的に実装できるため、大規模なデータセットやモデルに適しています。
よって、両方の方法は、クエリに対して最も関連性の高いキー（とそのバリュー）を識別するために使われますが、その計算の複雑さや実装の効率性において異なる特徴を持っています。トランスフォーマーではGPUをうまく利用して高速に大量のデータが処理できる内積アテンションが選ばれました。
なお、論文では、内積を$${\frac{1}{\sqrt{d_k}}}$$でスケールすることについて、さらに付け加えています。
$${d_k}$$の小さい値に対しては、二つの互換性関数は同様に機能しますが、$${d_k}$$の大きな値に対しては、スケーリングを行わない内積アテンションよりも加算アテンションの方が性能が良いです[3]。$${d_k}$$の大きな値に対しては、内積が大きな値になり、ソフトマックス関数を極端に小さい勾配の領域に押し込むと考えられます[4]。この効果を打ち消すために、内積を$${\frac{1}{\sqrt{d_k}}}$$でスケーリングします。
While for small values of $${d_k}$$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $${d_k}$$ [3]. We suspect that for large values of $${d_k}$$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients$${^4}$$. To counteract this effect, we scale the dot products by $${\frac{1}{\sqrt{d_k}}}$$.
つまり、論文の著者たちは、加算アテンションと内積アテンション（スケールなし）を比べて、ベクトルの次元数$${d_k}$$が大きい場合は、加算アテンションの方が性能が良いと分かったということです。
これは、ソフトマックス関数の入力値が非常に大きいまたは非常に小さい値になると、出力の確率分布において、ほとんどの確率が一つ（あるいは少数の）クラスに集中してしまう現象が生じるからです。この場合、入力値の小さな変化が出力にほとんど影響を与えなくなり、結果として勾配が非常に小さくなります。そして勾配消失の問題が生じます。
前述したように、キーの次元$${d_k}$$が大きくなると、クエリとキーの内積の値が大きくなりがちで、その結果、ソフトマックス関数の入力として非常に大きな値が与えられることになります。これがソフトマックス関数を極端に小さい勾配の領域に押し込む原因となります。
この問題に対処するために、スケールされた内積アテンションでは内積の結果を$${\frac{1}{\sqrt{d_k}}}$$でスケーリングすることで、入力値の範囲を制御し、ソフトマックス関数の勾配が適度な大きさを保つようにします。これにより、学習過程での勾配消失のリスクを減らし、より安定した学習が可能になります。
このように内積アテンションをどうにか高次元でも性能が落ちないように努力したのは、内積による計算はパラメータを必要とせず行列積として高速・効率的にGPU上で実行できるからでしょう。
次回予告
下図は、図２のスケールされた内積アテンションを日本語化したものです。
これを見てそれぞれ何をやっているのか分かるようになったでしょうか。この図を理解しておくと、スケールされた内積アテンションを利用する多頭アテンションの仕組みもわかりやすくなります。
次回は、3.2.2 Multi-Head Attention（多頭アテンション）から論文を読み続けます。
お楽しみに！
