Multi-Head Attention and Positional Encoding: Giving the Transformer a Sense of Direction

In Part 7, Lina saw how self-attention lets a model weigh the importance of different words — a spotlight on the most relevant parts of a sentence. The catch: if you scramble those words, basic attention doesn’t notice. To a computer, “the charge exceeded the refund” and “the refund exceeded the charge” look identical when we only look at the words themselves.

Today we’ll fix that. We’ll give our model a “GPS” so it knows where each word sits, and we’ll give it multiple “eyes” so it can track different patterns at once.

The Problem: A Model with No Sense of Time

Think of basic self-attention as a “bag-of-words” model. It treats a sentence like a pile of laundry on the floor rather than a neat line of clothes. Because the math behind attention sums weighted vectors, shuffling the order of the input words shuffles the order of the output vectors right along with them — nothing in the calculation knows what “first” or “third” means. In math terms, this is called permutation equivariance: permute the input, and the output permutes the exact same way.

What happens when we swap word order in a simple attention calculation?

import torch
import torch.nn.functional as F

def simple_attention_check(word_order):
    # Imagine these are 3-dimensional word embeddings
    # Word A, Word B, Word C -- always the same three vectors
    base_embeddings = torch.tensor([
        [1.0, 0.0, 0.0], # Word 1
        [0.0, 1.0, 0.0], # Word 2
        [0.0, 0.0, 1.0]  # Word 3
    ])

    # Arrange the sentence according to the order we're given
    embeddings = base_embeddings[word_order]

    scores = torch.matmul(embeddings, embeddings.T)
    probs = F.softmax(scores, dim=-1)
    output = torch.matmul(probs, embeddings)
    return output

sentence_1 = [0, 1, 2] # Word 1, 2, 3
sentence_2 = [2, 1, 0] # Word 3, 2, 1

print("Output 1:\n", simple_attention_check(sentence_1))
print("Output 2:\n", simple_attention_check(sentence_2))

• import torch / import torch.nn.functional as F — Imports PyTorch and its functional API (which provides softmax and other math ops without creating module objects).
• base_embeddings = torch.tensor([...]) — Three fixed, distinct one-hot vectors — Word 1, Word 2, Word 3 — that never change.
• embeddings = base_embeddings[word_order] — Actually uses the order we pass in: this reindexes the three base rows into whatever sequence word_order specifies, so sentence_1 and sentence_2 really do produce two differently-ordered inputs.
• torch.matmul(embeddings, embeddings.T) — Computes the full 3×3 matrix of pairwise dot products, in the current row order. Each entry [i, j] is the similarity between whatever word is in row i and whatever word is in row j.
• F.softmax(scores, dim=-1) — Applies softmax along the last dimension (across columns), turning each row of raw scores into a probability distribution that sums to 1.
• torch.matmul(probs, embeddings) — Multiplies the probability matrix by the embeddings, producing a weighted average of the embeddings for each word.
• Run both calls and compare: Output 2 is exactly Output 1 with its rows reversed, because sentence_2 reverses sentence_1’s order. Permute the input, and the output permutes the same way — that’s permutation equivariance, and it’s the property self-attention actually has, not invariance.

Look closely at what actually stayed the same, though: the value attached to each word’s identity didn’t change — only which row it landed in did. Word 1’s row in Output 1 and Word 1’s row in Output 2 are identical vectors; they just carry different row numbers depending on where Word 1 sits. The calculation itself has no internal sense that “Word 1” came first — it only knows “Word 1” is present, alongside the same two other words either way. To a Transformer, “Charged, not refunded” and “Refunded, not charged” hand each word the same contextual meaning regardless of which one came first, because both sentences contain exactly the same set of words attending to exactly the same set of words. We need to bake “position” into the data without breaking the math.

Positional Encoding: The ‘GPS’ for Your Words

We need to tell the model which word sits at index 0 and which sits at index 10. But we can’t just add the numbers 0, 1, 2… to the embeddings. If a sentence runs 1,000 words long, the number 1,000 dwarfs our embedding values. Those values usually hover between -1 and 1. The position signal would drown out the word’s actual meaning.

So we turn to Positional Encoding. Think of it as a timestamp or a GPS coordinate. We use sine and cosine waves at different frequencies to stamp each position with a unique signature.

Imagine a clock with many hands. One ticks every second, one every minute, one every hour. Read all the hands at once and you know exactly what time it is. The Transformer does something similar with waves.

import numpy as np
import matplotlib.pyplot as plt

def get_positional_encoding(max_seq_len, d_model):
    # Create a matrix of zeros
    pe = np.zeros((max_seq_len, d_model))
    for pos in range(max_seq_len):
        for i in range(0, d_model, 2):
            # Use sine for even indices, cosine for odd
            pe[pos, i] = np.sin(pos / (10000 ** (i / d_model)))
            pe[pos, i + 1] = np.cos(pos / (10000 ** (i / d_model)))
    return pe

# Let's visualize a small encoding for a 50-word sentence
pe_matrix = get_positional_encoding(50, 128)
plt.imshow(pe_matrix, cmap='RdBu')
plt.xlabel("Embedding Dimension")
plt.ylabel("Word Position")
plt.title("Positional Encoding 'GPS' Map")
plt.show()

• pe = np.zeros((max_seq_len, d_model)) — Initializes a matrix where each row is a position (0 to max_seq_len - 1) and each column is one dimension of the positional encoding vector.
• for pos in range(max_seq_len) — Iterates over every position in the sequence.
• for i in range(0, d_model, 2) — Steps through the embedding dimensions two at a time (because each pair gets one sine and one cosine).
• np.sin(pos / (10000 ** (i / d_model))) — Computes the sine for an even-indexed dimension. The divisor 10000 ** (i / d_model) controls the frequency: lower indices produce high-frequency waves (fast-changing positions), higher indices produce low-frequency waves (slow-changing positions). This mimics a clock with hands of varying speeds.
• np.cos(pos / (10000 ** (i / d_model))) — Computes the cosine for the next odd-indexed dimension, paired with the sine above.
• plt.imshow(pe_matrix, cmap='RdBu') — Renders the encoding matrix as a heatmap where red and blue show positive and negative values, revealing the wave patterns visually.

Each row in that heatmap is a unique “barcode.” Add that barcode to our word embeddings and the model can “feel” where it sits in the sentence. The reason it also picks up on relative position isn’t periodicity — it’s a trig identity: for any fixed offset $k$, the encoding at position $pos + k$ can always be written as a fixed linear combination of the encoding at position $pos$, for every frequency pair, because $\sin$ and $\cos$ of a shifted angle expand into a fixed combination of $\sin(pos)$ and $\cos(pos)$. That fixed relationship is the same no matter where in the sentence the offset happens, which is what lets the model learn “3 spots behind” as one reusable pattern instead of re-learning it at every position.

Positional Encoding Formulas

For a word at position $pos$ in the sequence, the positional encoding for embedding dimension $2i$ (even) and $2i+1$ (odd) is:

$PE_{(pos,\;2i)} = \sin\!\left(\frac{pos}{10000^{\,2i/d_{model}}}\right)$

$PE_{(pos,\;2i+1)} = \cos\!\left(\frac{pos}{10000^{\,2i/d_{model}}}\right)$

Plain English	Statistical symbol	Python equivalent
Word position in the sequence (0, 1, 2, …)	$pos$	pos (loop variable)
Embedding dimension index (0, 1, 2, …)	$i$	i (loop variable)
Total embedding vector size	$d_{model}$	d_model
Frequency denominator (controls wave speed)	$10000^{2i/d_{model}}$	10000 ** (i / d_model)
Encoding at even dimension	$PE_{(pos,\;2i)}$	pe[pos, i] = np.sin(...)
Encoding at odd dimension	$PE_{(pos,\;2i+1)}$	pe[pos, i + 1] = np.cos(...)

The constant $10000$ doesn’t actually touch the fastest wave: at dimension index $2i = 0$, the exponent is 0, so that dimension is always $\sin(pos)$ and $\cos(pos)$ no matter what constant you pick. What $10000$ actually controls is the slowest wave, at the highest dimension index — a bigger constant stretches that wave’s period further out, letting the lowest-frequency dimensions keep changing smoothly (instead of wrapping around and repeating) across a longer sequence. Each dimension acts like a different “hand” on the multi-handed clock analogy from the prose; $10000$ sets how slowly the slowest hand ticks.

Why One Head Isn’t Enough: The ‘Blind Men and the Elephant’ Problem

Now that the model knows where words are, it needs to understand what they’re doing. A word can have several jobs in a single sentence. Take “The battery charge is low, so I called support about the billing charge” — the word “charge” carries two different meanings.

A single attention head is like one person looking at an elephant in the dark. One feels the trunk and says “It’s a snake!” Another feels a leg and says “It’s a tree!”

Multi-Head Attention is like having a team of experts. One head might focus on grammar (subject-verb agreement). Another could track rhyming patterns. A third might pick up on factual relationships.

Here’s the tricky part: we don’t actually run 8 separate models, and we don’t literally chop the raw embedding into 8 disjoint pieces either. Each head first gets its own learned projection of the entire embedding — a dense matrix multiply that mixes all 512 input dimensions together — and only after that projection do we reshape the result into 8 chunks of 64. So head 3’s 64 numbers aren’t “dimensions 192 through 255 of the original embedding.” They’re a learned combination of all 512 original dimensions, shaped by weights the network trains specifically for head 3. That’s what actually lets heads specialize: the specialization is learned in the projection, not an accident of which raw slice each head happened to get.

# Conceptual split of a vector into 4 heads -- projection first, then split
d_model = 12 # Total size of our word vector
num_heads = 4
head_dim = d_model // num_heads

# A single word vector
word_vector = torch.randn(1, d_model)

# Each head gets a LEARNED projection of the WHOLE embedding first --
# heads never see raw, un-projected slices
projection = torch.nn.Linear(d_model, d_model)
projected_vector = projection(word_vector)

# Only THEN do we reshape the projected output into 4 'expert' views
heads = projected_vector.view(num_heads, head_dim)

print(f"Original vector size: {word_vector.shape}")
print(f"Split into {num_heads} heads, each with size: {heads.shape[1]}")

• d_model = 12 — The total embedding dimension. In real Transformers this is typically 512 or 768; here it’s kept small for readability.
• num_heads = 4 — The number of attention heads. Each head will operate on a learned slice of the projected embedding, not a raw one.
• head_dim = d_model // num_heads — Integer division: each head gets 12 // 4 = 3 dimensions. This must divide evenly.
• torch.randn(1, d_model) — Generates a random 1×12 vector representing a single word’s embedding.
• projection = torch.nn.Linear(d_model, d_model) / projected_vector = projection(word_vector) — The step the naive “just slice it” picture skips: a learned dense layer touches every one of the 12 input dimensions before anything gets split.
• projected_vector.view(num_heads, head_dim) — Now reshapes into a 4×3 matrix. Each row is one head’s view, but every number in that row is a learned mixture of all 12 original dimensions, not a raw fragment of the embedding. The .view() method reinterprets the same memory without copying data.
• heads.shape[1] — Returns the second dimension of the reshaped tensor (3), confirming each head sees 3 features.

What this means in practice: the model looks at the data in parallel. Instead of one big messy calculation, you get 4 (or 8, or 16) focused ones running at once.

Single-Head vs. Multi-Head Attention

Approach	What it does	Best for	Tradeoff
Single-Head Attention	Runs one set of Q, K, V projections over the full d_model vector. Every dimension of the embedding participates in one unified attention calculation.	Simple relationships; toy models; debugging; low-compute settings.	Can only learn one pattern of relevance at a time. If the sentence needs simultaneous grammatical and semantic disambiguation (e.g., “charge” as both subject and financial term), a single head must compress both into the same weighted sum — it can’t specialize.
Multi-Head Attention	Projects d_model through learned Q/K/V matrices, splits the result into num_heads chunks, runs independent attention per chunk, then concatenates and re-projects. Each head is free to specialize in a different relationship.	Real-world language; any task where multiple simultaneous patterns matter (syntax, semantics, coreference, rhyme, etc.).	Total parameters and compute stay essentially fixed as num_heads changes — the Q/K/V and output projections are always d_model-by-d_model, whatever the head count. What actually changes is how that fixed budget is split: more heads means a smaller head_dim = d_model / num_heads per head, so each head attends within a lower-rank subspace. Too many heads with too small a head_dim can underperform — the head doesn’t have enough dimensions to learn a meaningful pattern.

What heads learn to specialize in (examples):

• Head 1 — Syntax: Focuses on subject–verb agreement (“charge” ↔ “exceeded”).
• Head 2 — Semantics: Focuses on word meaning (“charge” ↔ “billing,” disambiguating financial vs. electrical).
• Head 3 — Coreference: Focuses on pronoun resolution (“it” ↔ the most recent noun).
• Head 4 — Positional patterns: Focuses on adjacent or nearby tokens, leveraging the positional encoding to find local context.

Rule of thumb: Lina would reach for multi-head attention whenever the same word participates in multiple simultaneous relationships that can’t be captured by a single weighted sum. The cost isn’t more parameters — head count barely moves the total — it’s that a fixed embedding budget gets divided into narrower subspaces, plus the constraint that d_model must be divisible by num_heads.

The Concatenation Trick: Putting the Experts Back Together

Once each head has done its work and chosen which words matter, we’re left with 8 different answers. The next layer of the network expects a single vector per word.

We solve this with the Concatenation Trick. We glue the heads’ outputs back together, side by side. Then that long vector passes through one final linear layer — a weight matrix — letting the heads mix their findings into one coherent representation.

class MultiHeadAttentionSimple(torch.nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Linear layers for Q, K, V
        self.q_linear = torch.nn.Linear(d_model, d_model)
        self.k_linear = torch.nn.Linear(d_model, d_model)
        self.v_linear = torch.nn.Linear(d_model, d_model)

        # Final output layer
        self.out_proj = torch.nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape

        # 1. Project, then split into heads: (batch, heads, seq, head_dim)
        q = self.q_linear(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_linear(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_linear(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 2. Scaled dot-product attention, run independently inside each head
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = torch.softmax(scores, dim=-1)
        attended = torch.matmul(attn_weights, v)  # (batch, heads, seq, head_dim)

        # 3. Concatenate heads back together
        combined = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)

        # 4. Final summary report
        return self.out_proj(combined)

• class MultiHeadAttentionSimple(torch.nn.Module) — Defines a custom PyTorch module inheriting from nn.Module, which gives it __call__, parameter tracking, and .to(device) support for free.
• super().__init__() — Calls the parent constructor to initialize internal PyTorch bookkeeping (parameter registry, gradient hooks, etc.).
• self.q_linear = torch.nn.Linear(d_model, d_model) — A learnable linear layer that projects the input into the Query space. The same is done for Key (k_linear) and Value (v_linear). Each has its own separate weights.
• self.out_proj = torch.nn.Linear(d_model, d_model) — The final output projection that lets the concatenated heads “talk” to each other and merge their findings.
• batch_size, seq_len, d_model = x.shape — Unpacks the input tensor’s three dimensions: number of sequences in the batch, length of each sequence, and embedding dimension.
• self.q_linear(x).view(...).transpose(1, 2) — Projects the input through the Query linear layer, reshapes from (batch, seq, d_model) to (batch, seq, num_heads, head_dim), then swaps the seq and num_heads axes so each head’s slice can run attention independently. Same for K and V.
• scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) — The scaled dot-product attention from Part 7, computed inside every head at once: each head gets its own similarity scores between its queries and its keys, scaled by √head_dim.
• attn_weights = torch.softmax(scores, dim=-1) — Turns each head’s raw scores into attention weights that sum to 1, exactly like Part 7’s softmax step.
• attended = torch.matmul(attn_weights, v) — The weighted sum of values inside each head — this is the actual attention calculation the original snippet only commented as “(Simplified attention calculation here…).”
• combined = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model) — Swaps the axes back and reshapes the per-head attention outputs (not raw v) into the full d_model width — the real Concatenation Trick.
• return self.out_proj(combined) — Passes the concatenated vector through the output linear layer, producing the final multi-head attention output for the next layer.

We take the notes from all our experts, tape them into one long scroll, and summarize that scroll into a single report for the next layer.

Let’s Check the Data: Does it Actually Work?

Does adding these features actually help? Rather than assert it, let’s build real (if tiny) embeddings, actually add the positional encoding computed earlier, and check whether the model can tell “the charge exceeded the refund” apart from “the refund exceeded the charge” — two sentences built from the exact same three content words in reversed order.

# Let's compare attention paid FROM 'exceeded' to its two neighbors
# in 'the charge exceeded the refund' vs 'the refund exceeded the charge'

torch.manual_seed(0)
d_model = 8

# Real (random but fixed) embeddings for the three content words
word_embeddings = {
    'charge': torch.randn(d_model),
    'exceeded': torch.randn(d_model),
    'refund': torch.randn(d_model),
}

sentence_a = ['charge', 'exceeded', 'refund']  # "the charge exceeded the refund"
sentence_b = ['refund', 'exceeded', 'charge']  # "the refund exceeded the charge"

# Reuse the exact get_positional_encoding function defined earlier
pe = torch.tensor(get_positional_encoding(max_seq_len=3, d_model=d_model), dtype=torch.float32)

def attention_from_middle_word(words, add_pe):
    x = torch.stack([word_embeddings[w] for w in words])
    if add_pe:
        x = x + pe
    scores = torch.matmul(x, x.T) / (d_model ** 0.5)
    weights = F.softmax(scores, dim=-1)
    return weights[1]  # attention distribution FROM "exceeded" -- always the middle word

def weight_by_identity(words, weights):
    return {word: round(weights[i].item(), 3) for i, word in enumerate(words)}

no_pe_a = attention_from_middle_word(sentence_a, add_pe=False)
no_pe_b = attention_from_middle_word(sentence_b, add_pe=False)
pe_a = attention_from_middle_word(sentence_a, add_pe=True)
pe_b = attention_from_middle_word(sentence_b, add_pe=True)

print("Without PE, 'exceeded' attends to (sentence A):", weight_by_identity(sentence_a, no_pe_a))
print("Without PE, 'exceeded' attends to (sentence B):", weight_by_identity(sentence_b, no_pe_b))
print("With PE,    'exceeded' attends to (sentence A):", weight_by_identity(sentence_a, pe_a))
print("With PE,    'exceeded' attends to (sentence B):", weight_by_identity(sentence_b, pe_b))

• word_embeddings = {...} — Real (if small and random) embeddings for the three content words — not hand-typed relevance scores.
• sentence_a / sentence_b — The same three words, in reversed order: “charge, exceeded, refund” vs. “refund, exceeded, charge.”
• pe = torch.tensor(get_positional_encoding(...)) — Reuses the exact function defined earlier in this article, instead of generating a matrix that’s plotted once and never touched again.
• attention_from_middle_word(words, add_pe) — Builds the sentence, optionally adds the positional encoding, runs real scaled dot-product attention (scores, then softmax), and returns the row for “exceeded” — the word sitting in the middle of both sentences.
• weight_by_identity(...) — Relabels the returned weights by word, not by column position, so we can compare “how much attention did ‘charge’ get” directly across the two sentences.
• The four print calls show the actual result: without PE, the weights are bit-for-bit identical between sentence A and B — “charge” and “refund” get the exact same attention regardless of which one came first. With PE added, the weights differ between the two sentences.

The numbers make the point concretely. Without PE, “exceeded” pays 0.126 attention to “charge” and 0.208 to “refund” — in both sentences, no matter which one actually came first. Swap their positions and the model doesn’t notice; it can’t, because nothing in the calculation reads position. Add PE, and that symmetry breaks: “charge” gets 0.147 attention when it precedes “exceeded” but only 0.136 when it follows it; “refund” gets 0.185 when it follows “exceeded” but 0.2 when it precedes it. The differences here are modest — this is a 3-word toy example with random, untrained embeddings — but they’re real, measured numbers, not typed-in ones, and they only show up once position joins the calculation. Pair this with Multi-Head Attention, and a trained model has room to let one head lean harder on exactly this kind of positional signal while another head tracks word content.

Recap: What we learned

1. Self-Attention is order-blind: Without help, it treats sentences like a pile of words.
2. Positional Encoding is a GPS: It uses sine waves to give every word a unique address.
3. Multi-Head Attention is a team of experts: It splits the work so the model can see grammar, meaning, and context at once.
4. Concatenation merges the views: We glue the experts’ findings back together into a single representation.

Lina’s model now has a sense of direction and multiple eyes. She has all the pieces — attention, multi-head, and positional encoding — and she’s ready to assemble them into a working model. In Part 9, we’ll build a miniature Transformer from scratch and see how these pieces fit together.

Check Your Understanding

Questions below move from simple recall up to open-ended design, roughly following Bloom’s Taxonomy.

Remember What does “permutation equivariance” mean, and why is it a problem for basic self-attention?

Understand In your own words, explain why Positional Encoding uses sine and cosine waves instead of just adding the raw position number (0, 1, 2, …) to each word’s embedding.

Apply Using the article’s head-splitting formula (head_dim = d_model // num_heads), calculate head_dim for a model with d_model = 64 and num_heads = 8.

Analyze The article says one head might focus on “who is doing the action” while another focuses on “the mood of the sentence.” Walk through why splitting the embedding into separate chunks per head—rather than just running the same full-size attention calculation multiple times—is what actually lets each head specialize in something different.

Evaluate The Concatenation Trick glues all the heads’ outputs back together and passes them through one final linear layer. Critique this: what would likely go wrong if you skipped that final linear layer and just used the concatenated output directly as the next layer’s input?

Create Design a scenario with a short sentence (5-6 words) where you’d want at least 3 distinct attention heads, and describe what each of the 3 heads should specialize in to fully understand that sentence.

---

Related articles

• The Attention Mechanism, Finally Explained (Without the Matrix Algebra))
• Building a Miniature Transformer from Scratch: The Full Picture)

References & Further reading

• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems (NeurIPS 2017). — The paper that introduced the Transformer architecture, including both multi-head attention (splitting embeddings into parallel attention subspaces) and sinusoidal positional encoding (the sine/cosine GPS scheme built in this article).

---