QK norm

Query-key normalization that rescales attention queries and keys before their dot products, changing the attention score path rather than the whole residual stream.

QK norm stabilizes attention by normalizing the query and key vectors right before their dot products, which narrows the change to the score path instead of renormalizing the entire hidden stream around every transformer sublayer.

At a glance

Released

October 2020

Authors

Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, et al.

Optimizes

  • Attention Logit Stability

What It Is

QK norm applies normalization directly to the query and key vectors used in attention scoring. Instead of normalizing the whole residual stream, it targets the part of the block that decides how sharply one token attends to another.

Why It Exists

QK norm mainly targets attention-logit stability. It reduces the chance that raw query or key magnitude will make the softmax too sharp, especially in deeper models or longer-context settings.

How It Works

The block projects hidden states into queries and keys, normalizes those vectors, then forms attention scores from the normalized versions before the usual softmax and value mixing. The attention pattern depends more on direction and less on raw vector length.

Math Or Compute Schema

The formula below shows the attention score path after normalized queries and keys are substituted into the standard dot product. The hats mark the normalized forms.
QK normalization inside attention
Attention(Q^,K^,V)=softmax ⁣(Q^K^⊤dk)V\mathrm{Attention}(\hat{Q}, \hat{K}, V) = \mathrm{softmax}\!\left(\frac{\hat{Q}\hat{K}^{\top}}{\sqrt{d_k}}\right) V
Q^\hat{Q}
Normalized query vectors.
K^\hat{K}
Normalized key vectors.
VV
Value vectors used after the attention weights are formed.
dkd_k
Key dimension used in the score scaling factor.

Compared To Nearby Modules

Layer norm and RMSNorm usually normalize the hidden state that flows through a whole transformer block. QK norm is narrower: it only changes the query-key score path inside attention. That makes it more of a targeted attention tweak than a general residual-stream norm.
Comparison dimensionQK NormLayer NormRMSNorm
Normalization targetQuery and key vectors onlyFull token or hidden vectorFull token or hidden vector
PlacementInside attention before the dot productAround transformer sublayersAround transformer sublayers
Main benefitKeeps attention logits from becoming too sharpStabilizes the broader residual streamStabilizes scale with less work than layer norm

Example Architectures

QK norm appears in some transformer variants that want extra control over attention sharpness, especially when long context or deep stacks make logits harder to keep in a comfortable range.

Limitations And Tradeoffs

QK norm only solves one narrow problem. It does not replace the broader stability role of layer norm or RMSNorm across the residual stream, so many architectures still need a separate hidden-state normalization strategy.

Why It Still Matters

QK norm is a good reminder that normalization is not only about whole-layer placement. Sometimes the most useful intervention is inside one sensitive operation, such as attention scoring.

Tags

References

  1. Henry, Alex, et al. "Query-Key Normalization for Transformers." arXiv, 2020, https://arxiv.org/abs/2010.04245.

Layer norm

Previous Page

RMSNorm

Next Page

On this page

What It IsWhy It ExistsHow It WorksMath Or Compute SchemaCompared To Nearby ModulesExample ArchitecturesLimitations And TradeoffsWhy It Still MattersRelatedTagsReferences