1 背景
1.1 注意機構の概要
注意機構(Attention Mechanism)は、深層学習モデルが入力系列中の要素間の依存関係を動的に重み付けする手法である。従来の系列モデルでは固定長のコンテキストベクトルに情報を圧縮する必要があったが、注意機構は各出力ステップで入力全体を参照することを可能にし、長距離依存関係の捕捉を改善した。初期の応用例として機械翻訳におけるBahdanau注意が有名であり、その後様々な変種が提案されてきた。
1.2 点積注意の登場
点積注意(Dot-Product Attention)は、クエリとキーのドット積を類似度尺度として直接利用する方式である。加算型注意に比べて計算が高速で、行列演算による並列化が容易である。しかし、キーの次元が大きくなると内積の値が大きくなり、ソフトマックス関数の勾配が極端に小さくなる問題が生じた。この問題に対処するため、スケーリング因子を導入したスケーリングドット積注意が提案され、Transformerモデルの中核演算として定着した。
2 数学的定義
2.1 基本公式
2.1.1 クエリ、キー、バリューの定義
スケーリングドット積注意は、三つの行列で表される入力を扱う。クエリ行列Q、キー行列K、バリュー行列Vはそれぞれ、系列中の各要素を線形変換して得られる。通常、Q∈R^(n×d_k)、K∈R^(m×d_k)、V∈R^(m×d_v)であり、nはクエリ数、mはキー・バリュー数、d_kはキーの次元、d_vはバリューの次元を表す。
2.1.2 スケーリング因子の導入
スケーリング因子として、キーの次元d_kの平方根√d_kが用いられる。この値でドット積の結果を割ることで、内積の分散が1程度に調整される。これにより、d_kが大きい場合でもソフトマックス関数の入力が適切な範囲に収まり、勾配消失を防ぐ。
2.2 ソフトマックス関数による重み付け
スケーリング後の値に対して、ソフトマックス関数を各行に適用する。ソフトマックスは各クエリについて、すべてのキーとの類似度を正の値に変換し、合計が1となる確率分布を出力する。この重み行列A∈R^(n×m)は、各クエリがどのキーにどの程度注目すべきかを表す。
2.3 加重和の計算
得られた重み行列Aとバリュー行列Vの積を計算することで、出力行列を得る。各クエリに対応する出力は、重みで加重平均されたバリューのベクトルとなる。この操作により、注目すべき情報が強調され、不要な情報が抑制される。
3 動作原理
3.1 計算フロー
3.1.1 ドット積の算出
まず、クエリ行列Qとキー行列Kの転置の行列積を計算する。これにより、各クエリと各キーの間の内積からなる類似度行列S=QK^Tが得られる。内積は二つのベクトルの方向と長さの情報を含むため、類似度の指標として自然である。
3.1.2 スケーリングの役割
内積の値はd_kの平方根に比例して増加する傾向がある。スケーリングを行わない場合、d_kが大きいとソフトマックス関数の入力が極端に大きな値となり、その勾配がほぼゼロになる。スケーリングにより入力の分散を1程度に抑えることで、ソフトマックス関数が適切な勾配を維持できる。
3.2 数値的安定性と勾配の制御
スケーリングは数値計算上の安定性も向上させる。大きな値の指数計算はオーバーフローの原因となるが、スケーリングによってこれを回避する。また、勾配の流れが改善されるため、深いネットワークでも学習が安定する。特にTransformerのように多層の注意機構を積み重ねる場合、この効果は重要である。
4 応用
4.1 Transformerアーキテクチャ
4.1.1 エンコーダー内の自己注意
Transformerのエンコーダーでは、各位置のクエリ、キー、バリューがすべて同じ系列から生成される。これにより、系列内の全てのペア間の依存関係を捉えることができる。例えば、文中の単語が他のどの単語に注目すべきかを学習する。
4.1.2 デコーダー内のマスク注意
デコーダーでは、未来の情報を参照しないようにマスクが適用される。具体的には、ソフトマックス演算の前に、未来に対応する内積値を非常に小さな値(-∞)に置き換える。これにより、自己回帰的な生成が可能となる。
4.2 派生モデル(BERT, GPT等)
BERTはエンコーダーのみを用い、双方向の文脈を利用する。GPTはデコーダーのみを用い、左から右への言語モデルを学習する。いずれもスケーリングドット積注意を基本構成要素としており、その効率的な計算が大規模モデルの学習を支えている。
5 拡張と変種
5.1 マルチヘッド注意
マルチヘッド注意は、複数の異なる線形変換により複数の注意ヘッドを並列に計算し、それらを連結して再度線形変換する。各ヘッドは異なる表現部分空間に注目できるため、モデルの表現力が向上する。スケーリングドット積注意は各ヘッド内で独立に適用される。
5.2 相対位置エンコーディング
元のTransformerは絶対位置エンコーディングを用いるが、相対位置エンコーディングではクエリとキーの相対的な位置関係をスケーリングドット積注意の計算に組み込む。これにより、系列長に対してロバストな注意重みを学習できる。
5.3 線形注意と効率化手法
スケーリングドット積注意の計算量は系列長の二乗に比例するため、長い系列では非効率である。線形注意はカーネル近似を用いて計算量を線形に削減する手法であり、スケーリング因子の扱いを工夫することで効率と精度のバランスを取る。他にも、スパース注意やストリーミング注意などの効率化手法が提案されている。
6 利点と限界
6.1 計算効率の高さ
スケーリングドット積注意は、行列演算を用いてGPU上で効率的に実装できる。特にバッチ処理やマルチヘッドの並列化が容易であり、Transformerの学習速度を大幅に向上させた。加算型注意と比較して、専用ハードウェアによる高速化の恩恵を受けやすい。
6.2 長い系列に対する課題
系列長が長くなると、計算量とメモリ使用量が二乗で増加する。また、すべての位置間の相互作用を考慮するため、極端に長い系列では疎な注意が必要となる。近年のLongformerやBigBirdなどの研究は、この限界を克服しようとする試みである。