Transformers
Intermediate 16 min read Transformers

Self-Attention & Multi-Head Attention

Tokens attending to tokens, repeatedly, in parallel

Self-attention lets tokens attend to other tokens in the SAME sequence, and multi-head runs several attention maps in parallel so different subspaces learn different relationships.

self-attentionmulti-headQ K V projectionsattention scoresfreedom of position
Prerequisites: Attention
Before we start

What you will learn

  • Distinguish self-attention from cross-attention
  • Describe Q, K, V projections and heads
  • Explain the O(T²) cost and why length is the constraint
Introduction

What is this?

Self-attention is attention where queries, keys and values all come from the same sequence, letting each token condition on every other token — the mechanism that replaces recurrence. Multi-head attention runs several independent attention 'heads' in parallel over the same inputs (each with its own Q/K/V projections), concatenates, and projects: different heads can learn syntax, referents, or local patterns simultaneously.

Motivation

Why do you need it?

  • Recurrence reads one step at a time; self-attention reads all pairs in one parallel pass.
  • Position is only injected as an addition — the attend mechanism itself is position-agnostic, which lets it learn long-range dependencies recurrence can't.
  • Multi-head = many attention 'viewpoints'; heads specialize while the projection mixes them.
  • This pair of ideas IS the transformer block; everything else is framing.
Think about it

Mental model

Imagine

A committee, each member taking notes from every other member

Think of

Multi-head self-attention

At each meeting, each member (query) reads everyone's notes (keys/values). Because different members focus differently — one watches budgets, one watches deadlines — the combined notes (concatenated heads) capture far more than any single reader.

Mechanics

How it works

  • Project X → Q = XW_q, K = XW_k, V = XW_v (each (T, d) → (T, d') after matmul).
  • Scores = (Q·Kᵀ)/√d → softmax → weights (T, T).
  • Output = softmax(QKᵀ/√d)·V — each token a weighted mix of values.
  • Multi-head: split d' heads, run attention per head, concatenate, project W_o.
  • Cost: O(T²·d) from the T×T matrix — the famous quadratic bottleneck.

Self-attention with Q, K, V projections and the T×T score map.

X (T, d)
Q K V project
QKᵀ /√d
Softmax·V
H heads concat
Real world

In production

  • Encoder self-attention builds context-rich token embeddings for classification/retrieval.
  • Decoder self-attention with causal masking is every LLM's text-to-text engine.
  • Vision transformers flatten image patches and self-attend across them.
Guidelines

Best practices

  • Use head dimensions ~64 per head (8 heads × 64 = 512 common) and project at the end.
  • Keep causal masking on decoders: build the mask once, apply pre-softmax.
  • FlashAttention-style kernels avoid materializing the full T×T when you can.
  • Warm up LR: deep attention stacks are sensitive at initialization.
Watch out

Common mistakes

  • Using one head for everything — multi-head consistently wins for representation diversity.
  • Skipping the 1/√d scaling — saturated softmax destroys gradient flow.
  • Forgetting to apply the causal mask to the WHOLE matrix including padding.
  • Treating heatmaps as interpretable alignments without a mask/scale check.
Root cause analysis

Common issues

Error
All heads produce identical attention
Root cause

Shared projections (W_q identical across heads) or collapse in late training

Fix

Initialize heads independently and check projection weight rank.

Error
OOM on long sequences
Root cause

Materialized T×T attention matrix

Fix

Use flash/attention kernels, chunk, or reduce context length.

Error
Training unstable around warm-up
Root cause

Attention logits too large pre-scale

Fix

Enforce √d scaling and a smooth LR warm-up.

When things break

Debugging tips

  • Assert softmax rows sum to 1 and stay finite; NaN here poisons everything downstream.
  • Print the T×T scores on a 4-token input — you should be able to eyeball a sensible pattern.
  • If all heads converge identically, resize/init heads or add entropy/orthogonality pressure.
  • Measure peak memory: quadratic attention shows an O(T²) footprint clearly on long sequences.
Interviews

Interview checkpoint

Practice

Hands-on lab

Implement Self-Attention in NumPy
30 min · intermediate
In progress
Objectives
  • Implement Q/K/V projections and the attention math by hand
  • Verify softmax rows sum to 1 and symmetry behavior
  • Run a causal masking version and confirm no future leakage

Implement Self-Attention in NumPy

30 min · intermediate

0/3 steps
0%
sandbox · Linux sandbox
$
Try these commands
Test yourself

Quick quiz

0/2
1

Self-attention's Q, K, V all come from:

2

Self-attention cost grows quadratically with sequence length.

Go deeper

Official documentation

Transformer lecture (annotated transformer)
docs