Transformer: Scaled Dot-Product Attention とスケーリングの重要性
TransformerのSelf-Attention機構では、Queryベクトル (Q) と Keyベクトル (K) の類似度を計算するために内積（Dot Product）がよく用いられます。しかし、単純な内積ではなく、Scaled Dot-Product Attention という、内積をスケーリング（調整）する手法が標準的に採用されています。
Scaled Dot-Product Attention の計算式
Queryベクトル \(\mathbf{q}\) と Keyベクトル \(\mathbf{k}\) のアテンションスコアは、以下のように計算されます。
なぜスケーリングが必要なのか？
スケーリングが必要となる主な理由は、内積の値の大きさと、後続のSoftmax関数の性質に関係しています。
- 内積の値と次元数: QueryベクトルとKeyベクトルの各要素が平均0、分散1の独立なランダム変数であると仮定すると、それらの内積 \(\mathbf{q} \cdot \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i\) の平均は0ですが、分散は \(d_k\) になります。つまり、ベクトルの次元数 \(d_k\) が大きくなるほど、内積の値（のばらつき）も大きくなる傾向があります。
- Softmax関数と勾配消失: アテンションスコアは、Softmax関数によって正規化され、アテンションウェイト（確率分布）に変換されます。Softmax関数 \( \text{softmax}(z_i) = e^{z_i} / \sum_j e^{z_j} \) は、全入力に同じ定数を足しても出力が変わらないため、飽和するかどうかは入力の絶対値ではなく入力どうしの差で決まります。あるスコアが他より大きく離れると出力がほぼone-hotになり、ヤコビアン \(\partial p_i/\partial z_j = p_i(\delta_{ij}-p_j)\) がほぼ0になります（一方、\((z_1,z_2)=(1000,1000)\) なら出力は \((0.5,0.5)\) で、対角成分は \(0.25\) のままです）。
- 問題点: もしスケーリングを行わず、次元数 \(d_k\) が大きい場合に内積 \(\mathbf{q} \cdot \mathbf{k}\) のばらつきが大きくなると、スコア間の差も大きくなってSoftmax関数が飽和領域に入り、勾配が消失します。これは、誤差逆伝播による学習がうまく進まなくなり、モデルの収束を妨げる原因となります。
スケーリングの効果
内積を \(\sqrt{d_k}\) で割ることにより、内積の値の分散が次元数 \(d_k\) に依存しないように（分散を1程度に）調整されます。これにより、Softmax関数への入力（スコア）の差が大きくなりすぎるのを防ぎ、勾配消失問題を緩和し、学習プロセス全体を安定化させることができます。(選択肢2が正しい)
他の選択肢の評価
- 選択肢1: スケーリングで内積の分散が次元数に左右されにくくなるのは事実ですが、それは手段です。問われている主な理由は、その結果として Softmax の飽和と勾配消失を防ぎ、学習を安定させることです。
- 選択肢3: このスケーリング操作は、計算量やメモリ使用量にはほとんど影響を与えません。
- 選択肢4: QueryとKeyの次元数を揃えることはAttention計算の前提ですが、スケーリングはそのための操作ではありません。
したがって、Scaled Dot-Product Attentionにおけるスケーリングは、特に高次元のベクトルを扱う場合に、学習を安定させるための重要なテクニックです。