The Scribe Reads the Room — Part 2
Self-Attention: Or, How to Listen to Everyone at Once Without Going Completely Meshuggeneh
[This is Part 2. Part 1 — The Forgetful Scribe covers autoencoders.]
Mathityahu’s aunt Rivka was a very good listener. At any family gathering, she could track fourteen conversations simultaneously — who said what to whom, who owed whom an apology, which cousin still hadn’t called his mother. She did not need to wait until the end of dinner to understand the beginning of dinner. Everything informed everything else, in real time.
“This,” said Mathityahu, “is what we need the machine to do.”
The autoencoder, as we saw in Part 1, forgets the order. It compresses everything into a single small representation and reconstructs from that. For images, this works. For language — for anything where sequence matters — it doesn’t. A sentence is not a bag of words. It is a structure.
The field tried to solve this with sequential models. Read left to right, carry the memory forward. But the memory fades. By the time you reach the end of a long sentence, the beginning is a blur — like trying to remember the first course at a Pesach seder after you’ve already reached the afikomen.
The solution was something else entirely. Instead of reading left to right and forgetting, look at everything at once — and learn what to look at.
This is self-attention.
The Bottleneck Problem
Before we solve it, let us be precise about what is broken.
Imagine encoding the sentence “The matzah, which the bubbe made from scratch and which took her three hours and two arguments with Uncle Shimon, was excellent” into a single fixed-size vector.
By the time you encode “was excellent”, the model must somehow still remember “matzah” — the subject — despite everything that came in between. In practice, it doesn’t. It remembers the most recent things well, and the earlier things poorly.
This is the bottleneck. Not size this time — time. The information has to travel through a long chain, and it degrades.
The Fix: Look at Everything at Once
Self-attention abandons the sequential constraint entirely. Instead of passing information down a chain, every token looks at every other token directly and decides for itself what to pay attention to.
But before any of that, the model needs to convert words into numbers. A neural network cannot operate on the word “matzah” — it needs a vector. So each word in the vocabulary is mapped to a fixed-length list of numbers, say 512 of them, called an embedding. These embeddings are learned during training. Words that appear in similar contexts end up with similar embeddings — “matzah” and “bread” will be closer together than “matzah” and “accountant.”
So the input to self-attention is not words. It is a sequence of vectors — one per token, each 512 numbers long. For the sentence “The bubbe made matzah ball soup again,” that is seven vectors sitting side by side.
Now the attention mechanism has something to work with.
Take the word “ball.” It needs context. “Ball” alone could mean anything — basketball, formal dance, a good time. What resolves the ambiguity is “matzah,” two positions back. Self-attention lets “ball” reach directly to “matzah” and borrow meaning from it. No chain. No fading memory. A direct connection.
Here is how. Each embedding is projected into three new vectors through learned weight matrices:
Query — what this token is looking for. “Ball” generates a Query that is, roughly, asking: “Is there a food-type modifier nearby that would tell me what kind of ball I am?”
Key — what this token is advertising. “Matzah” generates a Key that says: “I am a type of food, specifically unleavened bread, relevant to compound nouns.”
Value — the actual information to hand over if selected. “Matzah”’s Value is its full embedding — the meaning it carries.
The model computes a score between “ball”’s Query and every other token’s Key. Where Query and Key align, the score is high. “Ball” scores high against “matzah” and low against “again” and “The.” Those scores are turned into weights with softmax — they add up to 1.0, like a probability distribution over the sentence. Then the model takes a weighted sum of all the Values, with “matzah”’s Value contributing most.
The output for “ball” is now a blend of the whole sentence, leaning heavily toward “matzah.” It knows what kind of ball it is.
Mathematically, for a single head:
The √d scaling keeps the dot products from getting too large and pushing the softmax into a regime where one score dominates everything and the gradients vanish. A technical nuisance, not a deep idea.
One attention head produces one perspective on the sentence — one particular notion of what is relevant to what. That is already powerful. But language carries multiple kinds of relationships simultaneously. “The bubbe made matzah ball soup again” — one head might notice that “made” attends to “bubbe” (verb to subject). Another might notice that “ball” attends to “matzah” (noun to modifier). A third might notice that “again” attends to “made” (adverb to the verb it modifies). No single head can learn all of this at once without the relationships interfering with each other.
“So run it multiple times,” said Devorah.
Exactly. Multi-head attention runs h independent attention operations in parallel, each with its own W_Q, W_K, W_V matrices. Each head attends to the sentence through a different learned lens. Then the h outputs are concatenated and passed through one final linear layer W_O to produce the result.
The key insight: the weights are not fixed. They are computed fresh for every input. Each head learns to look for different kinds of relationships — one head might learn syntax (verbs attending to their subjects), another semantics (nouns attending to modifiers). The network learns what relevance means, independently, in parallel, eight times over.
The Code
In PyTorch, self-attention is a clean computation:
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, d_model):
super().__init__()
self.d_model = d_model
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
def forward(self, x):
# x: (batch, seq_len, d_model)
Q = self.W_q(x)
K = self.W_k(x)
V = self.W_v(x)
# Scaled dot-product attention
scores = torch.bmm(Q, K.transpose(1, 2)) / (self.d_model ** 0.5)
weights = F.softmax(scores, dim=-1) # (batch, seq_len, seq_len)
output = torch.bmm(weights, V) # (batch, seq_len, d_model)
return output, weights
Let us run it on a toy example — five tokens, embedding size 16:
batch_size, seq_len, d_model = 1, 5, 16
x = torch.randn(batch_size, seq_len, d_model)
attn = SelfAttention(d_model)
output, weights = attn(x)
print(f"Input shape: {x.shape}") # (1, 5, 16)
print(f"Output shape: {output.shape}") # (1, 5, 16)
print(f"Weights shape: {weights.shape}") # (1, 5, 5) — each token attends to all tokens
Each token in the output is a blend of all tokens in the input. The sequence length is preserved. The information is not bottlenecked into a single vector.
What the Weights Look Like
Below is an illustrative attention map — what a trained model might learn for the sentence “The bubbe made matzah ball soup again”:
Read it row by row. Each row is one query token — the word that is “looking.” Each column is a key token — the word being “looked at.” Darker means stronger attention.
Notice: “ball” attends strongly to “matzah” — because “ball” alone means nothing; its meaning depends on what preceded it. “made” attends to “bubbe” — the verb looks for its subject.
This is not programmed. It is learned.
How Did It Learn That?
“But wait,” said Devorah, who had been watching over Mathityahu’s shoulder. “In that chart — ‘ball’ is already looking at ‘matzah.’ How does it know to do that? Someone told it?”
Nobody told it. It learned. But to understand how, you need to understand what it was learning from — and what it was trying to do.
The data
Raw text. No labels. Books, articles, Wikipedia, code, the entire digitized output of people writing things down. The labels come from the text itself: every word in a sentence is the correct answer to the question “what comes after the previous words?”
“So someone had to label all of that?” asked Devorah.
“No,” said Mathityahu. “The text labels itself. If I give you ‘The bubbe made matzah ball’ and ask what comes next — the answer is already in the next line of the book.”
This is called self-supervised learning. The model never needs a human to annotate anything. It just needs text, and there is no shortage of text.
The task
At each position in a sequence, the model predicts the next token. Feed in “The bubbe made matzah ball” — five tokens. The model runs every one of them through the attention layers and produces, for each position, a probability distribution over the entire vocabulary. At position 5 (“ball”), it outputs something like: P(“soup”) = 0.34, P(“game”) = 0.12, P(“lamp”) = 0.001.
The correct answer is “soup.” The model assigned it 0.34. That is wrong enough to learn from.
The loss
The mistake is measured with cross-entropy loss:
If the model assigned 0.34 to “soup”, the loss is −log(0.34) ≈ 1.08. If it had assigned 0.01 — confidently wrong — the loss is −log(0.01) ≈ 4.6. The loss is high when the model is wrong, higher when it is confidently wrong.
criterion = nn.CrossEntropyLoss()
# logits: (batch, seq_len, vocab_size)
# targets: (batch, seq_len) — the actual next tokens
logits = model(input_tokens)
loss = criterion(
logits.view(-1, vocab_size),
targets.view(-1)
)
Backpropagation
Once you have the loss, you differentiate it with respect to every parameter in the network using the chain rule. The gradient flows backward through each operation in reverse order:
- Through the final linear projection to vocabulary size
- Through W_O, the output projection of multi-head attention
- Through each head’s attention computation — through the softmax, through the QKᵀ dot product, into W_Q, W_K, and W_V
- Through the feedforward layers
- All the way back into the embedding matrix
Every parameter receives a number saying: “move in this direction by this much to reduce the loss.” The optimizer — usually Adam — applies those updates.
loss.backward() # compute all gradients via chain rule
optimizer.step() # update every parameter
optimizer.zero_grad()
At the start of training, all matrices are random. “ball” attends to “The” just as much as it attends to “matzah.” The attention map looks like television static.
After millions of updates, the gradient has consistently pushed in one direction: make “ball” look at “matzah,” because that pattern reliably predicts “soup.” Make “made” look at “bubbe,” because verbs need their subjects. The relationships are not programmed. The task demanded them. The loss enforced them.
The scale
“How many sentences?” asked Devorah.
“GPT-2 trained on forty gigabytes of text,” said Mathityahu. “GPT-3, three hundred billion tokens. Weeks of compute on thousands of processors.”
Devorah was quiet for a moment. “And all of that just to predict the next word?”
“All of that just to predict the next word.”
The architecture is not complicated. The scale is what produces the behavior.
“So the attention map is the network’s notes,” said Devorah finally. “What it had to learn to get the answers right.”
“Exactly,” said Mathityahu.
The Remaining Problem
Rivka could track fourteen conversations, but she knew who said what when. The words arrived in order. She knew which story came first.
Self-attention, as written above, has no such knowledge. If you shuffle the tokens — “soup again bubbe matzah made ball The” — the attention computation gives exactly the same result. Position has no meaning.
For text, this is catastrophic. “The dog bit the man” must be different from “The man bit the dog.” If the model cannot tell position 1 from position 5, it cannot tell subject from object.
We need a way to tell the model: “this token is first, that one is sixth.” We need to bake position into the representation itself.
And once we solve that — once every token knows both what it is and where it sits — we have everything we need to build something much larger.
“Call everyone in,” said Mathityahu. “Multi-head. All of them.”
[Continue to Part 3 — The Whole Room Is Listening]
All code in this post runs on CPU with no training required. Full code: Jewpyter notebook repository.