QK-norm
In plain wordsA small change to how a language model pays attention — rescaling two sets of internal numbers before comparing them — that prevents a common way large models' training goes haywire.
Normalization of the query and key vectors in attentionattention, n.: A way for an AI model to choose what to focus on. For each word it processes, the model scores how relevant every other word is and draws information mostly from the relevant ones. — typically LayerNormlayer normalization, n.: A layer that rescales the numbers inside a network, separately for each example, so they keep a steady average and spread. It, or a simpler variant, is a standard part of the models behind today's chatbots. or RMSNormRMSNorm, n.: A simpler, faster version of the layer that keeps a network's internal numbers at a steady scale. Most recent large language models use it. applied per head, sometimes L2 normalization with a learned scale — before their dot productscaled dot-product attention, n.: The standard formula for attention in transformers: score how well each question matches each label, scale the scores down, turn them into percentages, and use those to blend the stored information. is taken. Without it, query and key norms can grow during training until attention logits become very large, the softmax collapses onto single tokens and the loss diverges; bounding the norms keeps attention logits in a stable range and tolerates higher learning rates. Applied at scale in Google's 22-billion-parameter vision transformer (2023), it has since been adopted in many language models, including AI2's OLMo 2, Alibaba's Qwen3 and Google's Gemma 3, and is a standard training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it. mitigation alongside z-lossz-loss, n.: A small extra penalty added during training that stops the raw scores a language model produces before choosing a word from growing huge, which helps keep very large training runs from blowing up..
Origin
Alex Henry and colleagues proposed query-key normalization for transformers in 2020 (Findings of EMNLP); Mostafa Dehghani and colleagues at Google popularized LayerNorm on queries and keys as a stability fix in 'Scaling Vision Transformers to 22 Billion Parameters' (2023).
Related
- training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it.
- z-lossz-loss, n.: A small extra penalty added during training that stops the raw scores a language model produces before choosing a word from growing huge, which helps keep very large training runs from blowing up.
- attentionattention, n.: A way for an AI model to choose what to focus on. For each word it processes, the model scores how relevant every other word is and draws information mostly from the relevant ones.
- query, key and valuequery, key and value, n.: The three roles each word plays inside attention: the query is what it is looking for, the key is what it can be found by, and the value is the information it hands over when chosen.
- layer normalizationlayer normalization, n.: A layer that rescales the numbers inside a network, separately for each example, so they keep a steady average and spread. It, or a simpler variant, is a standard part of the models behind today's chatbots.
- loss spikeloss spike, n.: A sudden jump in a model's error partway through training. Sometimes the model recovers on its own; sometimes the run falls apart and engineers must rewind to an earlier saved copy.
Sources
- Henry, A., Dachapally, P. R., Pawar, S. S. & Chen, Y. (2020). Query-Key Normalization for Transformers. Findings of EMNLP 2020, 4246–4253. aclanthology.org/2020.findings-emnlp.379
- Dehghani, M. et al. (2023). Scaling Vision Transformers to 22 Billion Parameters. ICML 2023.
Mentioned in
training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it., loss spikeloss spike, n.: A sudden jump in a model's error partway through training. Sometimes the model recovers on its own; sometimes the run falls apart and engineers must rewind to an earlier saved copy., RMSNormRMSNorm, n.: A simpler, faster version of the layer that keeps a network's internal numbers at a steady scale. Most recent large language models use it., z-lossz-loss, n.: A small extra penalty added during training that stops the raw scores a language model produces before choosing a word from growing huge, which helps keep very large training runs from blowing up.