AI企業の面接問題集

スケールド・ドット積注意機構は Query と Key を比較し、その重みで Value を混合します。式は次のとおりです。

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

$\sqrt{d_k}$ で割る目的は、次元増加に伴うスコアの尺度の拡大を抑え、softmax の早期飽和と勾配の弱まりを軽減することです。注意を一様にするためではありません。

自己注意では、入力表現 X を学習可能な行列で射影し、$Q=XW_Q$、$K=XW_K$、$V=XW_V$ を作ります。Query は探す特徴、Key は照合する特徴、Value は取り出す内容です。n 個の token に対し、単一 head の Q、K は $n\times d_k$、V は $n\times d_v$ にできます。$QK^\top$ の各行は一つの Query と全 Key の比較です。行ごとの softmax で総和一の重みに変換し、V を掛けると $n\times d_v$ の文脈表現になります。

$s=\sum_{i=1}^{d_k}q_i k_i$ とし、関連する各成分が互いに独立、平均ゼロ、分散一と仮定します。分散は平均の周りのばらつきです。$\mathbb{E}[q_i k_i]=0$、$\mathbb{E}[q_i^2k_i^2]=\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]=1$ で、異なる積の共分散はゼロなので $\operatorname{Var}(s)=d_k$、標準偏差は $\sqrt{d_k}$ です。$\operatorname{Var}(s/c)=\operatorname{Var}(s)/c^2$ より、$c=\sqrt{d_k}$ なら分散は一になります。この結果は仮定の下では厳密ですが、学習後の相関した Q、K に対する保証ではありません。

softmax の重みは次の式で定義されます。

$$p_i=\frac{\exp(z_i)}{\sum_j\exp(z_j)}$$

この重みは、共通のずれではなくスコア差に依存します。差が大きいと一つの位置に重みが集中します。導関数は次のとおりです。

$$\frac{\partial p_i}{\partial z_j}=p_i(\delta_{ij}-p_j)$$

$\delta_{ij}$ は $i=j$ なら一、それ以外はゼロです。飽和付近では多くの導関数が小さくなります。独自の説明例として $d_k=4$、スコア [4,0] なら重みは約 [0.982,0.018]、二で割ると約 [0.881,0.119]。Value が [10,0] なら出力は約 9.82 から 8.81 になります。順位は同じでも弱い一致の寄与が増えます。実測値ではありません。

標準の係数は、次元に応じた統計的根拠のある既定値です。低次元や小さなスコアなら無スケーリングでも直ちに問題になるとは限りませんが、head が広いほど初期化時から過度に鋭い分布になる危険が増えます。一つの位置を選ぶこと自体は正常な場合もあり、問題は学習した関連性ではなく尺度が集中を生むことです。

$d_k$ で割ると、同じ仮定の下で分散は $1/d_k$ となり、高次元ほどスコア差を小さくしすぎます。固定定数では特定の尺度にしか対応できません。小さいスコアほどよいわけではなく、平坦すぎる重みも有用な内容を薄めます。

追加の設計案として、学習可能な温度で softmax の鋭さを調整できますが、検証資源と監視が必要です。Q、K を単位長にする cosine attention は、長さの情報を捨てて類似度の定義を変えます。温度も別途設計すべきで、元の分散導出はそのまま使えません。既存 checkpoint の係数を変える場合も検証が必要です。

以下は推奨実装です。まず検査しやすい参照版を作り、射影後に head を分け、Q の最終次元を $d_k$ とします。scores=(Q @ K.transpose(-2,-1))/sqrt(d_k) を計算し、マスク、Key 軸の softmax、V との積を順に適用します。全 embedding 幅や系列長で割らず、Q と K の特徴次元、K と V の位置数の整合性も確認します。

マスクは padding や未来位置を除外し、通常は禁止スコアを負の無限大にします。各行に有効な Key が必要です。全位置が禁止なら、処理を飛ばす、ゼロを返す、入力を拒否するなどの方針を定義します。安定版 softmax は行の最大値を引いて指数のオーバーフローを防ぎますが、差は変わらないので平方根スケーリングの代わりにはなりません。

二位置の例を再現し、出力形状、有効行の重みの総和が一、禁止位置の重みがゼロ、逆伝播の勾配が有限であることを確認します。独立な標準正規乱数の Q、K で次元を変え、分散が縮放前は $d_k$、縮放後は一に近いかも調べます。有限標本では誤差があります。融合 attention API は係数を内蔵する場合があるため、仕様を確認し、出力と勾配を参照版と比較して二重適用を防ぎます。

記事には本番環境の測定結果はなく、以下は推奨運用です。layer、head ごとにスコアの縮放前後の標準偏差、Q/K のノルム、最大注意重み、勾配ノルムを抽出して監視します。重みの分散具合を表すエントロピー $H=-\sum_i p_i\log p_i$ も使えますが、比較時は有効 Key 数をそろえ、系列長による違いを障害と誤認しないようにします。

学習初期から多くの head で低エントロピーと弱い勾配が見られたら、係数の欠落、embedding 幅の誤使用、射影重みの増大、入力分布の変化を調べます。不自然に一様なら二重適用や $d_k$ による除算を疑います。ただし診断の手掛かりにすぎず、損失とタスク品質も確認します。低エントロピーだけで故障とは言えず、係数だけで全勾配問題を解決できるわけでもありません。

NaN や無限値では、全遮蔽行、非有限の入力、softmax 前の低精度内積のオーバーフロー、カーネルの累積精度を確認します。精度、融合カーネル、head 構成を変更したら、制御した入力で出力と勾配を比較し、妥当な浮動小数点誤差を許容しつつ、因果マスクから未来情報が漏れないことを検証します。面接での要点は、$QK^\top$ が一致度、softmax が混合重み、V が内容を与え、$\frac{1}{\sqrt{d_k}}$ が明示的な仮定の下で尺度を安定化することです。実際の挙動は実装テストと観測で確認します。