QK正規化（キューケーセイキカ）

QK正規化（QK-Norm）とは

QK正規化（Query-Key Normalization, QK-Norm）は、Transformerのスケールドドットプロダクトアテンション計算において、QueryベクトルとKeyベクトルにRMSNorm（またはLayerNorm）を適用してから内積を計算する技術です。

背景：アテンションエントロピー崩壊

大規模LLMの学習において、アテンションロジット（QとKの内積値）が異常に大きくなり、softmax後の確率分布が特定トークンに極端に集中する「アテンションエントロピー崩壊」（Attention Entropy Collapse）が問題となります。この現象が発生すると：

• 勾配が爆発または消失する
• 損失値が急激に増加して学習が発散する
• モデルが特定のアテンションパターンに固着する

QK-Normの動作原理

通常のAttention:
  scores = (Q @ K^T) / sqrt(d_k)

QK-Norm適用後:
  Q_norm = RMSNorm(Q)
  K_norm = RMSNorm(K)
  scores = (Q_norm @ K_norm^T) * scale  # 固定スケールパラメータ

QとKを正規化することで、内積値の絶対値が制御され、ロジット爆発を防止します。RMSNormを使用することで、計算オーバーヘッドは最小化されます。

採用モデルと効果

• Gemma 2（2B/9B/27B全サイズ）: 学習安定化の中核技術として採用
• LLaMA 3.1（8B/70B/405B）: 超大規模学習での安定化に必須
• Qwen2.5（7B以上）: 多言語学習での安定性向上

スケールパラメータの扱い

QK-Normを適用後はQとKのノルムが1に正規化されるため、通常のsqrt(d_k)でのスケーリングが意味を持たなくなります。代わりに学習可能なスケールパラメータを導入するか、固定値（ヘッド次元に比例した値）を使用します。

logit soft-cappingとの組み合わせ

Gemma 2ではQK-Normとlogit soft-cappingを両方組み合わせています：

• QK-Normが事前にQ・Kを正規化
• soft-cappingがアテンションスコアをtanhで制限

これにより二重の安定化機構が働きます。

実装コスト

RMSNormの計算量はO(d_model)と軽量で、全体のFLOPsへの影響は無視できる程度（0.1%未満）です。PyTorchではnn.RMSNormとして標準実装されており、既存のFlashAttentionカーネルとも組み合わせ可能です。

学習安定化への寄与

実験では、QK-Normを適用することで：

• 学習率を1.5〜2倍に引き上げ可能
• バッチサイズ増大時の安定性確保
• 長コンテキスト（128K+トークン）での学習安定化

が報告されており、超大規模モデルの学習では事実上の標準技術になりつつあります。