AI企業の面接問題集

スケーリング付きドット積アテンションは、クエリとキーの内積をキー次元の平方根で割り、softmax で重みに変換して値を混合します。次元の増加によるスコア差の拡大を抑え、早すぎる飽和と学習信号の弱まりを防ぐのが目的で、重みを一様にする操作ではありません。

クエリ $Q$ は探す特徴、キー $K$ は照合する特徴、値 $V$ は集約する内容です。自己アテンションでは通常、入力 $X$ を別々の学習可能な行列で射影し、$Q=XW_Q$、$K=XW_K$、$V=XW_V$ を得ます。クエリ数が $n$、キー数が $m$、両者の次元が $d_k$、値の次元が $d_v$ なら、スコアは $n\times m$、出力は $n\times d_v$ です。$A=\operatorname{softmax}(QK^\top/\sqrt{d_k})$ を行ごとに計算し、$O=AV$ を得ます。各行の重みの和は一です。

平方根の根拠は、平均からの二乗偏差の期待値である分散です。$s=q^\top k=\sum_{i=1}^{d_k}q_i k_i$ とし、関係する全成分が独立、平均ゼロ、分散一と仮定します。各積の平均はゼロで、$\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$ より、平方根で割ると分散は一になります。これは分布の尺度の話で、個々の内積が必ず増えるという意味ではありません。学習後の相関は仮定を崩し得ます。

softmax は $p_i=e^{z_i}/\sum_j e^{z_j}$ であり、鋭さを決めるのはスコア差です。導関数は $\partial p_i/\partial z_j=p_i(\delta_{ij}-p_j)$。$\delta_{ij}$ は添字が等しければ一、それ以外はゼロです。一箇所に重みが集中すると、この局所的な導関数は小さくなります。独自の説明例で $d_k=4$、スコアが $[4,0]$ なら重みは約 $[0.982,0.018]$、二で割ると約 $[0.881,0.119]$ です。順位を保ちつつ、他の位置にも重みが残ります。

標準のスケーリングは、順位を変えずに次元の影響を補う簡単な既定策です。次元が小さい場合やスコアが別の仕組みで制御される場合、未スケーリングでも直ちに失敗するとは限りません。ただし次元やベクトルの大きさに鋭さが左右されます。逆に $d_k$ で割ると、仮定上の分散は $1/d_k$ となり、高次元では一様化して関連性の識別を弱め得ます。

追加の設計案に温度があります。$A=\operatorname{softmax}(QK^\top/(\sqrt{d_k}\tau))$、$\tau>0$ とすると、大きい温度は滑らかに、小さい温度は集中的にします。温度の学習は柔軟ですが、パラメータと検証事項が増えるため、既存モデルへ無検証で導入せず、学習と評価の根拠で選びます。

クエリとキーを正規化するコサイン類似度も選択肢です。ノルム依存を減らせますが、利用できる情報が変わり、尺度の再調整が必要です。一方、softmax 前に行の最大値を引く処理は指数のオーバーフロー対策です。確率もスコア差も変わらないため、スケーリングや飽和対策の代わりにはなりません。

以下は記事の配備仕様ではなく、推奨する実装手順です。バッチ、head、系列、特徴の軸を確認し、キーの最後の二軸だけを転置します。係数には埋め込み全体や値の次元ではなく、head ごとの $d_k$ を使います。マスクを加えてキー軸に softmax を適用し、$V$ を掛けます。因果マスクは未来、パディングマスクは無効な token を除外します。

手書きなら scores = (Q @ K.transpose(-2, -1)) / math.sqrt(Q.shape[-1])、加算マスク、weights = softmax(scores, dim=-1)、output = weights @ V の順です。除外箇所には負の無限大を使えますが、全箇所がマスクされた行にはゼロ出力などの明示的な処理が必要です。通常の softmax のままでは NaN になり得ます。ライブラリでは組み込みのスケーリングと真偽値マスクの意味を確認します。

形状だけでなく、先の二キー例で重みを照合します。値をそれぞれゼロと十にすれば、出力は約 $1.19$ です。有効行の重みの和が一に近いこと、除外位置の重みがゼロであること、順伝播と逆伝播が有限値であることを確認します。仮定を満たす独立乱数でも複数の次元を試し、前後の経験分散を比較します。有限標本では理論値との近似であり、完全一致は要求しません。

以下は追加の運用助言で、記事に本番の測定結果はありません。損失の不安定化、NaN、早期の極端な集中があれば、クエリとキーのノルム、スケーリング前後のスコア標準偏差、最大重み、勾配ノルムを標本調査します。集中度にはエントロピー $H=-\sum_i p_i\log p_i$ も使えますが、有効キー数と併せて解釈します。必要な位置への集中もあるため、低エントロピーだけで異常とは判断しません。

典型的な障害は、スケーリングの欠落や重複、head 次元の誤り、マスクの反転、全行マスク、低精度の積演算のオーバーフローです。安定な softmax でも、その前の内積が無限大になれば修復できません。小さな入力を高精度の参照計算と比較し、必要なら累積精度を調整します。学習時と推論時の係数、マスクの一致も確認します。

融合カーネルや head 数を変える際は、出力と勾配を許容誤差内で比較してから、品質、遅延、メモリを評価します。式だけから実性能は断言できません。面接では、$QK^\top$ で照合し、$1/\sqrt{d_k}$ で尺度を整え、softmax の重みで $V$ を集約するとまとめます。平方根は仮定付きの分散から導かれ、実装ではマスク、精度、分布の検証も必要です。