Lazy loaded image10.5 What is QK-Norm?

Motivation

This paper proposed a way to handle the attention mechanism, especially the part where we use softmax on dot products of queries and keys.
The key issue of the standard dot-product attention is that: the values going into softmax can get super large, which ends up making the attention too sharp or too “winner-take-all” like below:
Where both the 760 and 12 take most “attention”. That can actually hurt learning, especially when we want the model to spread its attention around a bit more. Also, the softmax only cares about the differences between values instead of their magnitudes.

Query-Key Normalization (QK-Norm)

Instead of using raw dot products between queries and keys, they normalize both vectors using L2 norm (so now we're working with cosine similarity).
However, the scale in the softmax still somehow matters since the range of cosine similarity is [-1, 1]. The expressive ability may be insufficient when cosine similarity is combined with Softmax.
Therefore, adding a learnable temperature parameter or something similar is important. The specific operation of the QK Norm is to
  • first normalize and and
  • then multiply them with a learnable scalar weight.
That way, the values going into softmax stay in a more useful range: not too big, not too small.
So normally:
With QK Norm:
where is that learnable scale factor and the denominator can be ignored; where and are and with -normalization applied along their head dimensions.

Difference with usual Normalizatiom

Note that in standard normalization techniques such as batch normalization, layer normalization, etc., we are actually doing standardization like below:
But in this QK Norm, we are doing real normalization, specifically the L2 norm.
The authors also mentioned that this QK Norm can be used with LayerNorm instead of replacing it.
Models like OLMoE, ViT-22B, and dots.llm1 use the QK Norm
Prev
10.4 Left Padding or Right Padding?
Next
10.6 Why LLM training rarely adopts Dropout?
Loading...
Article List
LLM Learning Roadmap
✨ Awesome-Anything
Lazy loaded image Awesome-AI-Tutorials
Lazy loaded image Awesome-AI-Researchs
Lazy loaded image Awesome-GitHub-Repos
Lazy loaded image Awesome-AI-Examples
🖼️ Digital Image Processing
Lazy loaded image DIP I Image Demosaicing and Histogram Manipulation
Lazy loaded image DIP II Image Denoising
Lazy loaded image DIP III Edge Detection
Lazy loaded image DIP IV Digital Half-toning
Lazy loaded image DIP V Geometric Image Modification
🍃 LLM Components
Lazy loaded image 1.1 Tokenization
Lazy loaded image 1.2 Embedding
Lazy loaded image 1.3 Attention
Lazy loaded image 1.4 FFN, Residual Addition, and LN
Lazy loaded image 1.5 Positional Encoding
Lazy loaded image 1.6 Architecture & Decoding Policy
🌱 LLM Pre-training
Lazy loaded image 2.1 Pre-training Definition
Lazy loaded image 2.2 Pre-training Data
Lazy loaded image 2.3 Pre-training Workflow
Lazy loaded image 2.4 Continued Pre-training (CPT)
Lazy loaded image 2.5 Pre-training Evaluation
☘️ LLM Post-Training
Lazy loaded image 3.1 Supervised Fine-Tuning (SFT)
Lazy loaded image 3.2 Reinforcement Learning
Lazy loaded image 3.3 RLHF: Reinforcement Learning from Human Feedback
Lazy loaded image 3.4 Optimizing PPO-Based Algorithms
Lazy loaded image 3.5 Direct Preference Optimization (DPO)
🍀 LLM Popular Models
Lazy loaded image 4.1 BERT & Variants
Lazy loaded image 4.2 PaLM Series
Lazy loaded image 4.3 GPT Series
Lazy loaded image 4.4 LLaMA Series
Lazy loaded image 4.5 GLM Series
Lazy loaded image 4.6 Qwen Series
Lazy loaded image 4.7 DeepSeek Series
Lazy loaded image 4.8 MOE Series
Lazy loaded image 4.9 Other
🪴 LLM Applications
Lazy loaded image 5.1 Prompt Engineering
Lazy loaded image 5.2 LLM-based Agents
Lazy loaded image 5.3 Retrieval-Augmented Generation (RAG)
Lazy loaded image 5.4 Enterprise-level RAG Practice (I)
Lazy loaded image 5.5 Enterprise-level RAG Practice (II)
🌿 LLM Optimization
Lazy loaded image 6.1 GPU Usage Analysis
Lazy loaded image 6.2 Latency Analysis for Inference
Lazy loaded image 6.3 FlashAttention
Lazy loaded image 6.4 PagedAttention
Lazy loaded image 6.5 Packing for Efficient Training
Lazy loaded image 6.6 PEFT
🌾 LLM Compression
Lazy loaded image 7.1 Quantization
Lazy loaded image 7.2 Pruning
Lazy loaded image 7.3 Knowledge Distillation
🌵 LLM Hands-on Practice
Lazy loaded image 8.1 nanoGPT: Training a GPT from Scratch
Lazy loaded image 8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
Lazy loaded image 8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
Lazy loaded image 8.4 OpenRLHF: Reward Model Training and Code Analysis
Lazy loaded image 8.5 OpenRLHF: PPO Tutorial and Code Analysis
🌴 LLM Must-read Papers
Lazy loaded image 9.1 Technical Reports Explained
Lazy loaded image 9.2 MedLLM
Lazy loaded image 9.3 o1 Techinical Paper
Lazy loaded image 9.4 Prompting
Lazy loaded image 9.5 LLM-as-Judge
🌳 LLM Q&A
Lazy loaded image 10.1 Why GRPO training loss starts at 0
Lazy loaded image 10.2 VLM-R1 Reproduction Results
Lazy loaded image 10.3 KL(P||Q) or KL(Q||P)?
Lazy loaded image 10.4 Left Padding or Right Padding?
Lazy loaded image 10.5 What is QK-Norm?
Lazy loaded image 10.6 Why LLM training rarely adopts Dropout?
Lazy loaded image 10.7 Dropout and DropPath
🐝 VLM Image Encoders
Lazy loaded image 11.0 Introduction
Lazy loaded image 11.1 Vision Transformer
Lazy loaded image 11.2 CLIP
Lazy loaded image 11.3 BLIP
Lazy loaded image 11.4 EVA
Lazy loaded image 11.6 BLIP-2
Lazy loaded image 11.7 SigLIP
📝 MISC.
🤔 Large Language Models Pass the Turing Test.
Lazy loaded image OpenAI ChatGPT Progression
⏳ Claude Sonnet 4 Test with Sandtris
🔍 What DeepSeek-R1's Surge Teaches Us About SEO
💰 CivitAI’s Payment Issue
🎓 arXivSub, an Academia Assistant