Principle:Turboderp org Exllamav2 Layer Normalization
| Knowledge Sources | |
|---|---|
| Domains | Normalization, Transformer_Architecture, Deep_Learning |
| Last Updated | 2026-02-15 00:00 GMT |
Overview
Layer normalization stabilizes the hidden state activations within transformer layers by normalizing across the feature dimension, preventing gradient explosion and enabling deeper networks.
Description
Modern transformer architectures use normalization at various points within each layer. ExLlamaV2 implements several normalization variants:
- RMSNorm (Root Mean Square Normalization): A simplified normalization that divides by the root mean square of the input, without subtracting the mean. Uses fewer operations than standard LayerNorm and is the default in Llama-family models. Applied with a learned scale parameter (weight) but no bias.
- LayerNorm: The original normalization that subtracts the mean and divides by the standard deviation across features, then applies learned scale and bias parameters. Used by some architectures (e.g., certain GPT variants).
- Per-head QK normalization (HeadNorm): Normalizes the query and key projections independently for each attention head before computing attention scores. This prevents attention logit growth in deep models and is used by architectures like Cohere's Command-R.
All normalization modules support tensor parallelism, where the normalization is applied to partial tensors on each device before the results are gathered.
Usage
Layer normalization is applied at specific points in each transformer layer:
- Pre-attention norm: Applied to hidden states before the attention mechanism (pre-norm architecture)
- Post-attention norm: Applied after the attention output before the MLP
- Pre-MLP norm: Applied before the feed-forward network in some architectures
- Final norm: Applied to the final hidden state before the language model head
- QK normalization: Applied within the attention mechanism to Q and K projections per head
Theoretical Basis
RMSNorm
# Input: x of shape [batch, seq_len, hidden_dim]
# Weight: w of shape [hidden_dim]
# Epsilon: eps (small constant, typically 1e-6)
RMS(x) = sqrt(mean(x^2) + eps)
RMSNorm(x) = (x / RMS(x)) * w
LayerNorm
# Input: x of shape [batch, seq_len, hidden_dim]
# Weight: w, Bias: b of shape [hidden_dim]
mean = mean(x, dim=-1)
var = var(x, dim=-1)
LayerNorm(x) = ((x - mean) / sqrt(var + eps)) * w + b
Per-Head QK Normalization
# For each attention head h:
# Q_h, K_h of shape [batch, seq_len, head_dim]
# w_q_h, w_k_h of shape [head_dim]
Q_h_norm = RMSNorm(Q_h, w_q_h)
K_h_norm = RMSNorm(K_h, w_k_h)
# Use Q_h_norm, K_h_norm for attention score computation
Related Pages
Implemented By
- Implementation:Turboderp_org_Exllamav2_ExLlamaV2LayerNorm
- Implementation:Turboderp_org_Exllamav2_ExLlamaV2HeadNorm
- Implementation:Turboderp_org_Exllamav2_ExLlamaV2RMSNorm