Query-Key Normalization (QK-Norm): Mathematical Foundations, Attention Logit Drift, Entropy Collapse, and Large-Scale Transformer Training Stability

Large-scale Transformer pre-training at trillion-token scale frequently encounters sudden, irrecoverable training instabilities, often manifesting as gradient norm spikes and divergence in training loss. While gradient clipping, learning rate warmups, and pre-layer normalization (Pre-LN) mitigate optimization bottlenecks in shallow architectures, deeper and wider models exhibit an intrinsic failure mode in the standard scaled dot-product attention mechanism: unconstrained attention logit growth.

Query-Key Normalization (QK-Norm) introduces an architectural modification that normalizes query and key vectors prior to computing dot-product attention scores. Originally formulated by Henry et al. (2020) and subsequently adapted for massive foundation models by Dehghani et al. (2023) in ViT-22B, Wortsman et al. (2023), and the Gemma, Command R+, and Qwen2 model families, QK-Norm bounds the dynamic range of attention logits, preserves softmax entropy, and stabilizes gradient backpropagation across multi-trillion token training runs.

The Logit Drift and Entropy Collapse Problem

Standard multi-head self-attention computes query, key, and value representations from an input hidden sequence $X \in \mathbb{R}^{T \times d_{\text{model}}}$ using linear projections $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ for each attention head:

$$Q = X W_Q, \quad K = X W_K$$

The unnormalized attention logits matrix $S \in \mathbb{R}^{T \times T}$ is computed via scaled dot-product attention:

$$S_{ij} = \frac{q_i^\top k_j}{\sqrt{d_k}}$$

where $q_i, k_j \in \mathbb{R}^{d_k}$ represent the query and key vectors at sequence positions $i$ and $j$, and $d_k$ is the per-head projection dimension.

The Positive Feedback Loop

During unconstrained optimization, gradient updates systematically increase the spectral norm and Frobenius norm of projection weight matrices $W_Q$ and $W_K$. Because $q_i = X_i W_Q$ and $k_j = X_j W_K$, the $L_2$ norm of the resulting vectors grows multiplicatively:

$$\|q_i\|_2 \propto \|W_Q\|_F \|X_i\|_2, \quad \|k_j\|_2 \propto \|W_K\|_F \|X_j\|_2$$

This creates a destabilizing compounding dynamic:

1. Quadratic Logit Scaling: The logit magnitude scales as $S_{ij} \sim \mathcal{O}(\|W_Q\|_F \cdot \|W_K\|_F / \sqrt{d_k})$. In unnormalized models, maximum attention logits routinely drift from initial values of 4 to 8 up to values exceeding 50 to 100 as training progresses over hundreds of billions of tokens.
2. Softmax Saturation: The attention weights $A_{ij} = \frac{\exp(S_{ij})}{\sum_{m=1}^T \exp(S_{im})}$ become dominated by the maximum logit $\max_j S_{ij}$.
3. Entropy Collapse: The Shannon entropy of the attention distribution across sequence keys collapses toward zero:

$$H(A_i) = -\sum_{j=1}^T A_{ij} \ln A_{ij} \to 0$$ The attention distribution degenerates into a near one-hot indicator vector centered on a single token (often the initial sequence token or punctuation tokens acting as attention sinks).

1. Vanishing and Exploding Gradients: The Jacobian of the softmax operation with respect to pre-softmax logits is given by:

$$\frac{\partial A_{ij}}{\partial S_{ik}} = A_{ij}(\delta_{jk} - A_{ik})$$ When $A_{ij} \approx 1$ for $j = k$ and $A_{ij} \approx 0$ for $j \neq k$, the derivative $\frac{\partial A_{ij}}{\partial S_{ik}}$ vanishes for all non-dominant tokens, eliminating gradient flow across context positions. Conversely, when an outlier input shifts the dominant position, sudden regime changes induce massive gradient spikes that destabilize AdamW second-moment accumulator states ($v_t$).

Mathematical Formulation of QK-Norm

QK-Norm resolves unconstrained logit drift by inserting a normalization operator directly into the projection pipeline before the scaled dot product.

In modern foundation models such as Google Gemma 2 and Cohere Command R+, Root Mean Square Normalization (RMSNorm) is applied independently along the head dimension $d_k$ for each query and key vector:

$$\tilde{q}_i = \text{RMSNorm}(q_i) = \frac{q_i}{\text{RMS}(q_i)} \odot \gamma_Q$$

$$\tilde{k}_j = \text{RMSNorm}(k_j) = \frac{k_j}{\text{RMS}(k_j)} \odot \gamma_K$$

where the root mean square operator is defined as:

$$\text{RMS}(u) = \sqrt{\frac{1}{d_k} \sum_{m=1}^{d_k} u_m^2 + \epsilon}$$

and $\gamma_Q, \gamma_K \in \mathbb{R}^{d_k}$ are learnable gain parameters initialized to $1.0$.

The Cauchy-Schwarz Logit Bound

Under QK-Norm with unit gain ($\gamma_Q = \gamma_K = \mathbf{1}$), each normalized vector has an exact Euclidean norm:

$$\|\tilde{q}_i\|_2 = \sqrt{\sum_{m=1}^{d_k} \frac{q_{i,m}^2}{\frac{1}{d_k} \sum_{n=1}^{d_k} q_{i,n}^2 + \epsilon}} \approx \sqrt{d_k}$$

Applying the Cauchy-Schwarz inequality to the attention dot product yields a strict mathematical upper bound:

$$|S_{ij}| = \left| \frac{\tilde{q}_i^\top \tilde{k}_j}{\sqrt{d_k}} \right| \le \frac{\|\tilde{q}_i\|_2 \cdot \|\tilde{k}_j\|_2}{\sqrt{d_k}} = \frac{\sqrt{d_k} \cdot \sqrt{d_k}}{\sqrt{d_k}} = \sqrt{d_k}$$

For a standard head dimension of $d_k = 128$ (or $d_k = 64$), the maximum possible attention logit before scaling is bounded by $\sqrt{128} \approx 11.31$ (or $\sqrt{64} = 8.0$), regardless of how large the underlying linear projection weights $W_Q$ and $W_K$ grow.

This strict bound guarantees:

• Attention logits cannot drift into numerical saturation regimes.
• Softmax entropy is bounded away from zero by construction.
• Gradient flow remains active across all contextual key tokens.
• Lower-precision FP8 and BF16 exponent ranges are protected against overflow during matrix multiplication.

Architectural Variants: LayerNorm vs. RMSNorm vs. Cosine Attention

Several variations of query-key normalization exist across literature and production architectures:

1. RMSNorm QK-Norm (Gemma, Command R+, Qwen2)

Applied per-head after RoPE (or prior to RoPE, depending on implementation). Because RMSNorm omits mean-centering offsets ($\beta$), it incurs minimal computational overhead and maps efficiently onto GPU SRAM kernels such as FlashAttention-2 and FlashAttention-3.

2. LayerNorm QK-Norm (ViT-22B, Wortsman et al.)

Dehghani et al. (2023) applied standard LayerNorm (with mean centering and variance scaling) to query and key projections in ViT-22B. This enabled stable training of a 22-billion parameter dense vision transformer on 4 billion images at a high learning rate ($1 \times 10^{-3}$) without divergence.

3. Scaled Cosine Attention (Swin Transformer v2)

Liu et al. (2022) replaced the dot product with scaled cosine similarity: $$S_{ij} = \frac{1}{\tau} \cdot \frac{q_i^\top k_j}{\|q_i\|_2 \|k_j\|_2}$$ where $\tau$ is a learnable per-head temperature parameter. While effective, full $L_2$ normalization per token can introduce numerical instability if vector norms approach zero, whereas RMSNorm with small $\epsilon$ provides smooth gradient dynamics.

PyTorch Reference Implementation

The following self-contained module demonstrates a multi-head attention block with per-head RMSNorm QK-Norm and Rotary Position Embeddings (RoPE):

import torch
import torch.nn as nn
import math

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight

class QKNormAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, eps: float = 1e-6):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Linear projections for Q, K, V, and Output
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

        # Per-head QK-Norm layers applied across head_dim
        self.q_norm = RMSNorm(self.head_dim, eps=eps)
        self.k_norm = RMSNorm(self.head_dim, eps=eps)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        # Linear projections: (B, T, d_model) -> (B, T, H, d_k) -> (B, H, T, d_k)
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Apply per-head QK Normalization
        q = self.q_norm(q)
        k = self.k_norm(k)

        # Scaled Dot-Product Attention: (B, H, T, d_k) @ (B, H, d_k, T) -> (B, H, T, T)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if mask is not None:
            scores = scores + mask

        attn_weights = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn_weights, v) # (B, H, T, d_k)

        # Recombine heads and project output: (B, T, d_model)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.o_proj(out)

Interaction with Multi-Head Latent Attention and Modern Optimizers

While QK-Norm has become standard across dense Multi-Head Attention (MHA) and Grouped-Query Attention (GQA) architectures, its integration with compressed attention mechanisms introduces specific trade-offs:

• Multi-Head Latent Attention (MLA): In DeepSeek-V2 and DeepSeek-V3, KV representations are compressed into a low-rank latent vector $c_t^{KV}$ to minimize KV-cache memory during inference. Applying QK-Norm requires materializing queries and keys per head before normalization, preventing complete pre-computation of compressed matrix products. Alternative stability mechanisms such as selective learning rate scaling for Q/K projections (e.g. QuacK) have been explored for MLA setups.
• Learning Rate Decoupling: Wortsman et al. (2023) demonstrated that QK-Norm allows pre-training runs to maintain higher peak learning rates without destabilizing loss trajectories, accelerating early convergence on massive corpora.
• Float8 (FP8) Mixed-Precision Serving: Bounded logits prevent exponent underflow and overflow in FP8 GEMM kernels, making QK-Norm a prerequisite for stable low-precision training and serving regimes.

Sources

• Query-Key Normalization for Transformers (Henry et al., 2020)
• Scaling Vision Transformers to 22 Billion Parameters (Dehghani et al., 2023)
• Small-Scale Proxies for Large-Scale Transformer Training Instabilities (Wortsman et al., ICLR 2024)
• Gemma: Open Models Based on Gemini Research and Technology (Google DeepMind, 2024)
• Swin Transformer V2: Scaling Up Capacity and Resolution (Liu et al., CVPR 2022)