QK-Norm（Query-Key Normalization）とは？（キューケーノーム）意味・特徴をわかりやすく解説

Q: QK-Norm（Query-Key Normalization）とは？

QK-Norm（Query-Key Normalization）は、Transformer の Self-Attention 機構において Query と Key のベクトルを正規化する手法であり、注意重みのスケールを安定させることで学習の安定性と推論の数値精度を向上させる。ViT-22B、Gemma 2、Chameleon など大規模モデルで採用されており、特に長系列・大次元のモデルで注意スコアのオーバーフローを防止する効果がある。

QK-Norm は Transformer の Attention 計算における Query ベクトルと Key ベクトルを L2 正規化（またはRMSNorm / LayerNorm）することで、Softmax に入力されるロジットのスケールを制御し、学習・推論の安定性を向上させる技術である。従来はスケーリング因子 1/√d_k のみでロジットのスケールを制御していたが、モデルの大規模化・系列長の増大に伴い、この制御が不十分になるケースが増加している。

概要

Transformer の Self-Attention では、Query（Q）と Key（K）の内積を計算してアテンションスコアを得る：

Attention(Q, K, V) = Softmax(QK^T / √d_k) × V

ここで d_k は Key の次元数であり、1/√d_k はスケーリング因子として内積の分散を正規化する役割を果たす。しかし、学習が進むにつれて Q と K のノルムが増大し、QK^T の値が非常に大きくなる「ロジット成長問題（logit growth）」が発生する。

ロジットが過大になると、以下の問題が起きる：

Softmax の飽和: 入力値が大きすぎると、Softmax 出力がほぼ 0 か 1 の二値に近づき（entropy collapse）、アテンションの「柔らかさ」が失われる
FP16 オーバーフロー: 半精度浮動小数点（FP16）では最大値が 65504 であり、大次元モデルでは QK^T が容易にこの範囲を超える
勾配消失: Softmax が飽和すると勾配がゼロに近づき、学習が停滞する

QK-Norm はこの問題に対するシンプルかつ効果的な解決策であり、Q と K を Softmax 計算の前に正規化してノルムを 1（または定数）に揃えることで、ロジットのスケールを確定的に制御する。

技術的詳細

QK-Norm の実装にはいくつかのバリエーションがある：

バリエーション 1: L2 正規化（コサイン類似度アテンション）

Q_norm = Q / ||Q||₂ K_norm = K / ||K||₂ Attention = Softmax(τ × Q_norm × K_norm^T) × V

ここで τ は学習可能な温度パラメータ（スカラー）である。Q と K を単位ベクトルに正規化することで、内積がコサイン類似度に等しくなり、値域が [-1, 1] に制限される。

バリエーション 2: RMSNorm/LayerNorm 適用

Q_norm = RMSNorm(Q) K_norm = RMSNorm(K) Attention = Softmax(Q_norm × K_norm^T / √d_k) × V

元のスケーリング因子 1/√d_k はそのままに、追加で Q と K に RMSNorm を適用する。この方式は Gemma 2 や Chameleon で採用されている。

バリエーション 3: ヘッド単位 LayerNorm

各アテンションヘッドの Q, K に独立して LayerNorm を適用する。ヘッド間の正規化は行わず、ヘッド内の特徴次元のみを正規化する。

バリエーション	正規化方法	追加パラメータ	採用モデル
L2 正規化			Q
RMSNorm	Q, K に RMSNorm	γ（d_k 次元/ヘッド）	Gemma 2, Chameleon
LayerNorm	Q, K に LayerNorm	γ, β（2×d_k 次元/ヘッド）	一部研究
CosNorm	ヘッド単位コサイン正規化	τ（温度、1スカラー）	初期研究

計算コストの影響

QK-Norm の追加計算コストは Attention 全体に対して非常に小さい。L2 正規化の場合、Q と K のノルム計算（各 d_k 次元の二乗和 + 平方根 + 除算）が追加されるが、QK^T の行列乗算（O(n² × d_k)）に比べて O(n × d_k) のコストは無視できる程度である。

操作	計算量	QK-Norm 追加	オーバーヘッド
Q, K 射影	O(n × d_model × d_k)	-	-
QK-Norm（L2）	-	O(n × d_k)	< 0.1%
QK^T 計算	O(n² × d_k)	-	-
Softmax	O(n²)	-	-
AV 計算	O(n² × d_v)	-	-
合計オーバーヘッド			< 0.5%

他のアテンション安定化手法との比較

手法	アプローチ	メリット	デメリット
標準スケーリング (1/√d_k)	固定スケール	シンプル	大モデルで不十分
QK-Norm (L2)	Q,K のノルム正規化	確定的スケール制御	温度パラメータが必要
QK-Norm (RMSNorm)	Q,K に RMSNorm	追加学習パラメータで柔軟	やや実装複雑
ログスケール Attention	log-space で計算	オーバーフロー完全防止	実装が複雑
クランピング	ロジットの上限設定	シンプル	勾配が不連続
Sigma-Reparam	スペクトルノルム制御	理論的に厳密	計算コスト高

QK-Norm のメリット

ロジットオーバーフローの防止: QK^T の値域を確定的に制限でき、FP16/BF16 でも安全
学習の安定化: entropy collapse を防止し、アテンションの多様性を維持できる
実装の簡便さ: 既存の Attention 実装への変更は Q, K に1行の正規化を追加するのみ
FlashAttention との互換性: Q, K の前処理として適用するため、FlashAttention のカーネルをそのまま使用できる
計算コストがほぼゼロ: Attention 全体に対するオーバーヘッドは 0.5% 未満

QK-Norm のデメリット

表現力の制約: Q, K のノルム情報を捨てるため、理論的には表現力が低下する。ただし実験的には問題にならないことが多い
温度パラメータの学習: L2 正規化版ではτの初期値と学習率の設定がやや敏感
既存モデルとの非互換: QK-Norm を適用したモデルと非適用モデルは重みの互換性がない

採用モデルと実装例

モデル	公開年	QK-Norm 方式	適用理由
ViT-22B（Google）	2023	L2 正規化 + 学習可能温度	22B ViT の学習安定化
Gemma 2（Google）	2024	RMSNorm	Soft-capping と併用
Chameleon（Meta）	2024	LayerNorm	マルチモーダルの安定化
nGPT（NVIDIA）	2024	L2 正規化（全パラメータ）	超球面上の最適化
Cohere Command R+	2024	QK-Norm	長系列対応
InternLM 2（Shanghai AI Lab）	2024	GroupNorm	グループ単位正規化

PyTorch での実装例（L2 正規化版）

import torch
import torch.nn as nn
import torch.nn.functional as F

class QKNormAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        # 学習可能な温度（ヘッドごと）
        self.tau = nn.Parameter(torch.ones(n_heads, 1, 1) * (self.d_k ** 0.5))

    def forward(self, x):
        B, L, _ = x.shape
        q = self.W_q(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)

        # QK-Norm: L2 正規化
        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)

        attn = torch.matmul(q, k.transpose(-2, -1)) * self.tau
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, L, -1)
        return self.W_o(out)

FAQ

Q: QK-Norm はどのような場合に特に有効か？

A: QK-Norm は以下のケースで特に効果が大きい：（1）モデルの隠れ次元が大きい場合（d_model ≥ 4096）— Q, K の内積のスケールが大きくなりやすい、（2）系列長が長い場合（8K トークン以上）— アテンションパターンの集中が起きやすい、（3）混合精度学習（FP16）を使用する場合 — ロジットのオーバーフローリスクが高い、（4）マルチモーダルモデル — テキストと画像で特徴のスケールが大きく異なる。逆に、小規模モデル（d_model ≤ 1024）や短系列（< 2048トークン）では効果が限定的であり、追加する必要性は低い。

Q: QK-Norm と Softmax Capping は併用すべきか？

A: 併用可能であり、Gemma 2 は実際に QK-Norm と Soft-capping（Softmax のロジットを tanh で制限）を併用している。QK-Norm がロジットの「入力」を制御し、Soft-capping がロジットの「出力」を制限するため、二重の安全策として機能する。ただし、両方を適用するとアテンションの表現力が過度に制約される可能性があるため、一方で十分な場合はどちらかを選択するのが一般的である。FlashAttention 2/3 はカーネル内で Soft-capping をネイティブサポートしているため、実装の追加コストは小さい。

Q: QK-Norm と Rotary Position Embedding（RoPE）の適用順序は？

A: QK-Norm を先に適用し、その後に RoPE を適用するのが一般的である。RoPE は Q, K に回転行列を掛ける操作であり、L2 正規化後に回転行列を掛けてもノルムは保存される（回転行列は直交行列のため）。しかし、RMSNorm 版の QK-Norm を使用する場合は、RoPE の後に適用する方が安定するケースもある。Gemma 2 の実装では、RoPE 適用後に RMSNorm を掛ける順序が採用されている。実装時はモデルの公式リポジトリのコードを参照して順序を確認することが重要である。

QK-Norm（Query-Key Normalization）（キューケーノーム）

概要