Attention Mechanisms & Multi-Head Subspace Routing
How do query-key projections, softmax budgeting, and multi-head subspace synthesis route contextual information across tokens?
How this lesson unfolds.
Self-attention allows every word in a sequence to dynamically route information from every other word. By projecting token representations into query, key, and value subspaces, models compute directional alignment shadows that scale and synthesize into rich contextual representations.
How does the vector dot product act as a directional shadow measuring mutual alignment and compatibility?
Point a flashlight straight down at the ground: if you hold a ruler parallel to the ground, its shadow is long and full. If you tilt the ruler vertically, its shadow shrinks to a tiny dot. In attention mechanisms, the query vector acts like the beam and the key vector acts like the ruler. If both vectors point in the identical direction (0^∘), the shadow is maximal (1.0). If they are perpendicular (90^∘), the shadow vanishes completely (0.0). If they point in opposite directions (180^∘), the shadow turns negative (-1.0).
- $$Compute the scalar inner product q · k = |q||k| cos(theta) to quantify the extent to which two vector directions coincide.$$
- Learning goalDerive the geometric interpretation of the dot product and calculate pairwise alignment scores between query and key vectors.
- Helpful foundationsEuclidean vector spaces · Attention Is All You Need (Vaswani et al., 2017)
- Common trapBelieving dot product measures distance rather than alignment; two very long vectors can have a massive dot product even if their angle is wide, which is why normalization and scaling are mandatory.
- Next connectionRaw dot product scores can take arbitrary real values from -infinity to +infinity; how do we transform these unbounded scores into a shared, valid probability distribution? We look at budget allocation.
More context & mechanics
The algebraic inner product sums coordinate-wise products:
q · k = ∑_(i=1)^(d) qᵢ kᵢ
Geometrically, this is equivalent to:
q · k = \|q\| \|k\| cos θ
When vectors are normalized to unit length (\normq = \normk = 1), the dot product directly yields cos θ:
- θ = 0^∘ ⟹ q · k = 1 (perfect agreement)
- θ = 90^∘ ⟹ q · k = 0 (orthogonal independence)
- θ = 180^∘ ⟹ q · k = -1 (direct contradiction)
In the Scaled Dot-Product Attention of Transformers, every token's Query vector q takes dot products against all Key vectors k₁, ṡ, k_N. This raw compatibility score determines how much information from each surrounding word flows into the current token representation.
Josiah Willard Gibbs and Oliver Heaviside formalized vector dot products in the 1880s to compute physical work (W = F · d). Over a century later, Vaswani et al. adopted it as the fundamental scoring kernel of the Transformer architecture.
- Dot product is a spotlight casting a directional shadow.
- Parallel vectors scream; perpendicular vectors whisper; opposite vectors cancel.
How does the softmax function enforce a competitive zero-sum budget where all attention weights sum to exactly one?
Suppose three departments ask for project funding with raw priority scores of 2, 3, and 5. You have only one single 100% budget pie to divide. You cannot grant 200% total funding. To distribute the money fairly without negative payouts, you convert each score into a positive slice and divide by the total: 2/10 = 20%, 3/10 = 30%, and 5/10 = 50%. The departments now compete under a strict zero-sum constraint: if one department gains attention, another must sacrifice its share.
- Learning goalCompute softmax attention distributions across token logits and verify partition of unity on the probability simplex.
- Helpful foundationsExponential function rules · Probabilistic Interpretation of Feedforward Classification Networks (Bridle, 1989)
- Common trapConfusing sigmoid with softmax; sigmoid normalizes each element independently between 0 and 1 without enforcing that the sum across tokens equals 1.
- Next connectionWhat happens when attention scores grow extremely hot or freezing cold? We introduce the temperature scaling dial.
More context & mechanics
Given raw attention logits s = [s₁, s₂, ṡ, s_N]^T ∈ R^N, the softmax operator computes:
aᵢ = (e^(sᵢ)) / (∑_(j=1)^N e^(sⱼ))
This guarantees two vital mathematical properties:
- Strict positivity: aᵢ > 0 for all finite sᵢ because e^(sᵢ) > 0.
- Unit partition of unity: ∑_(i=1)^N aᵢ = 1.0.
The output vector a is strictly constrained to the standard (N-1)-simplex:
Δ^(N-1) = a ∈ R^N | aᵢ ≥ 0, ∑_(i=1)^N aᵢ = 1
Every Transformer attention head uses this budget allocation to synthesize a weighted context vector:
c = ∑_(i=1)^N aᵢ vᵢ
If the word 'bank' attends 85\% to 'river' and 10\% to 'water', the remaining 5\% is all that remains for the rest of the sentence.
Ludwig Boltzmann formulated this exponential distribution in 1868 to describe the statistical thermal equilibrium of gas molecules. In 1989, John Bridle introduced softmax into neural classification models.
- Softmax turns raw scores into slices of a 100% pie.
- Attention is a zero-sum economy: granting focus to one token deprives all others.
How do linear projection operators decouple a unified token embedding into specialized geometric attention subspaces?
Imagine looking at an intricate architectural sculpture. If you view it from only one fixed angle, the front pillar hides the courtyard behind it. Multi-head projection places 12 observers at 12 distinct vantage points around the building: Observer 1 sketches the roof structure, Observer 2 measures the foundation, Observer 3 checks the window alignment. By projecting the 3D sculpture into 12 different 2D sketches, no detail is hidden by perspective occlusions.
- $$Project high-dimensional token representations through learned matrices W_i^Q and W_i^K to isolate orthogonal syntactic and semantic aspects.$$
- Learning goalMap unified token representations into specialized attention head subspaces and quantify head specialization patterns.
- Helpful foundationsLinear transformations · In-context Learning and Induction Heads (Olsson et al., 2022)
- Common trapAssuming all attention heads learn identical features; training with random weight initialization and dropout causes heads to diverge into distinct functional roles.
- Next connectionHow do we recombine these decoupled perspectives into a single vector? We examine subspace synthesis.
More context & mechanics
Let X ∈ R^T × d_(model) be the sequence tensor.
For each head i ∈ 1, ṡ, h:
Qᵢ = X Wᵢ^Q, Kᵢ = X Wᵢ^K, Vᵢ = X Wᵢ^V
Where Wᵢ^Q, Wᵢ^K, Wᵢ^V ∈ R^d_(model) × dₖ.
The attention map for head i is:
Aᵢ = softmax( (Qᵢ Kᵢ^T) / (√(dₖ)) + M ) ∈ R^(T × T)
The projection matrices act as dimensionality-reducing linear maps that discard irrelevant global information, allowing Aᵢ to specialize:
- Head 1: Short-range local token interactions (A_(t, t-1) ≈ 1).
- Head 2: Long-range syntactic dependencies (attending to subject or root verb).
- Head 3: Coreference resolution (linking pronouns to proper nouns).
In transformer interpretability (e.g. Anthropic's circuits research), circuits of attention heads are mapped as induction heads, which perform in-context copying and pattern matching via specific subspace projections.
Vaswani et al. (2017) demonstrated that multi-head attention consistently outperformed single-head attention with the same computational budget, establishing subspace decoupling as deep learning doctrine.
- Different projections reveal different perspectives.
- Subspace decoupling lets 12 heads ask 12 independent questions.
How does the output projection matrix W_O perform linear information fusion across concatenated head representations?
Think of a symphony orchestra where 12 different instrument sections play simultaneously: violins play the melody, cellos provide the baseline, flutes add high trills, and trumpets deliver brass fanfare. If you simply recorded all 12 instruments on 12 separate audio tracks and played them raw, they would sound disconnected. The output projection matrix W_O is the master audio mixing board: it adjusts volume levels, blends harmonics, and masterfully mixes 12 separate audio tracks into a rich, unified stereo master track.
- $$Map the concatenated direct sum vector back into the residual stream via full-rank matrix multiplication y = Concat(head_1, ..., head_h) * W_O.$$
- Learning goalCalculate the linear synthesis of concatenated multi-head outputs via the projection matrix W_O.
- Helpful foundationsMatrix multiplication · Attention Is All You Need (Vaswani et al., 2017)
- Common trapBelieving W_O changes the vector dimension; in standard transformers, W_O is a square d_model x d_model matrix that transforms features while preserving tensor shape.
- Next connectionTo accelerate autoregressive decoding on edge devices, can we generate speculative candidate tokens in parallel trees? We explore speculative decoding tree projections.
More context & mechanics
Given h heads of dimension d_v, where h · d_v = d_(model):
H = [head₁ \| head₂ \| ṡ \| head_h] ∈ R^T × d_(model)
The linear synthesis step computes:
z = H W_O + b_O
Where W_O ∈ R^d_(model) × d_(model).
Each column of W_O forms a linear combination of all head outputs:
zⱼ = ∑_(k=1)^d_(model) H_(t, k) (W_O)_(k, j)
This allows features extracted by Head 1 (e.g. grammar) to directly influence features needed downstream for sentiment classification in the residual stream.
In transformer architectures, this output projection is critical for enabling the 'residual stream' concept, where layers read from and write to a shared communication bus.
The linear synthesis formulation was introduced in the seminal paper 'Attention Is All You Need' (2017) and remains unchanged across modern LLMs (LLaMA, GPT-4, Gemini).
- W_O is the mixing board blending 12 instruments into one symphony.
- Concatenation packs the vectors; W_O fuses their meaning.