Why multi-head self attention works: math, intuitions and 10+1 hidden insights | AI Summer

Abstraction: Mathematical and intuitive deep-dive into multi-head self-attention mechanisms

Key points:

• Self-attention is defined as two matrix multiplications: QK^T (similarity) scaled and softmaxed, then multiplied by V; critically, self-attention is NOT symmetric because Q and K use different projection matrices
• Multi-head attention decomposes the embedding dimension into h parallel lower-dimensional heads, enabling independent parallel computations per head — conceptually multiple "linear views" of the same sequence
• Voita et al. pruning study: 17 out of 48 encoder heads retained with nearly the same BLEU score; important heads classified as positional, syntactic, or rare-word-attending; encoder-decoder (cross) attention heads are most critical and hardest to prune
• After softmax, self-attention is low-rank (Linformer insight): most information recoverable from the top-128 singular values, motivating linear complexity approximations
• Without skip connections and MLPs, attention converges exponentially to rank-1 (token uniformity collapse); layer norm plays no role in preventing rank collapse
• Fine-tuning only layer norm parameters (0.1% of weights) achieves near full-fine-tune performance on low-data regimes; Q/K/V matrices pretrained on large NLP corpora encode transferable computational primitives

Connections: Bert · Google Brain · Transformers · Attention Mechanism · Self Attention

Source: https://theaisummer.com/self-attention/