LLM + RoboticsUnit 3 · The Transformer era · Day 13

Attention Is All You Need

Drop the recurrence, keep the attention. Every token looks at every other token in one parallel step, and that turned out to scale better than anything before it.

AuthorsVaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin
Year / venue2017 · NIPS (NeurIPS)
Core ideaSelf-attention instead of recurrence
Study time4–6 h

Original: arXiv:1706.03762 · Code, notebook and exercises: 30u30 Day 13

Start here (no maths needed)

A Transformer lets every word in a sentence look at every other word at the same time, and decide which ones matter most for understanding it.

Think of it like this

Imagine you scroll back in a group chat and see the message "She said yes!" To understand it, you look back at the other messages to find out who "she" is and what the question was. You don't reread every message equally. You pay most attention to the few that help. In a Transformer, every word does this for itself, and all the words do it at once instead of waiting in a queue.

How it works

  1. The sentence is split into small pieces called tokens, usually words or parts of words.
  2. Each token makes a "question" (what am I looking for?) and a "label" (what do I have to offer?).
  3. Each token compares its question with every other token's label. Good matches get a bigger share of its attention. The shares always add up to 100%.
  4. Each token then mixes in information from the tokens it paid attention to. All tokens do this at the same time, which makes it very fast on modern chips.
  5. This is repeated in several layers, so the understanding gets deeper each time.

Why robots care

A robot looking at a table can treat each thing it sees as a token: the gripper, a cup, a plate, a toy. Attention lets it link "gripper" with "cup" when the task is "pick up the cup", without anyone writing special rules. Many modern robot brains are built from this exact block.

Try it

Words to know

Token
One small piece of the input, such as a word.
Attention
Deciding how much to focus on each other token.
Query, key, value
A token's question, its label, and the information it passes on.
Transformer
A neural network built from attention layers, with no step-by-step memory loop.
Parallel
Doing many things at the same time instead of one after another.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Summary

TL;DR

Before 2017, the best translation systems were recurrent encoder–decoders with attention bolted on (Day 15). This paper kept only the attention, "dispensing with recurrence and convolutions entirely". The result, the Transformer, trained faster because every position is computed in parallel, and it set new translation records. Almost every large model since, from GPT and BERT to ViT and today's robot policies, is a variant of this block.

2. The core idea

Pack the queries, keys and values of all tokens into matrices Q, K, V (one row per token):

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

head_i      = Attention(X·W_i^Q, X·W_i^K, X·W_i^V)
MultiHead(X) = Concat(head_1, …, head_h) · W^O

FFN(x) = max(0, x·W_1 + b_1)·W_2 + b_2
x ← LayerNorm(x + Sublayer(x))      # after every sublayer

Positions are injected once, at the bottom of each stack:

PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

Each pair of dimensions is a sine/cosine at a different frequency, from fast to very slow. The authors chose this because a fixed offset k becomes a linear function of the encoding, which should make relative positions easy to learn. Learned position embeddings gave nearly identical results.

3. Deep dive

Why throw away recurrence

An RNN (Days 01–02) must process token 1 before token 2. For a sequence of length n that is n steps you can't parallelise, and information from the start must survive n state updates to reach the end. The paper's Table 1 compares layer types:

When n is smaller than d (true for sentences), attention is cheaper and fully parallel. The cost is the n² term, which is why long-context models later invented linear and hybrid attention. See Module 05 for how Kimi K3 mixes linear attention with full attention to reach a million tokens.

Three jobs for one mechanism

Multiple heads matter: in the paper's ablation, a single head was 0.9 BLEU worse than the best setting, and too many heads also hurt. At inference the decoder caches each layer's keys and values so old tokens aren't recomputed. That cache is the memory bottleneck that DeepSeek's MLA attacks in Module 05.

Why divide by √d_k

If query and key entries are independent with mean 0 and variance 1, their dot product has variance d_k. With d_k = 64 raw scores of ±8 or more are normal, and the softmax then puts nearly all weight on one key. Gradients through a saturated softmax are tiny, so learning stalls. Dividing by √d_k brings the variance back to 1 regardless of width.

Interactive: scaled vs unscaled attention

One query against 8 random keys, entries drawn from N(0, 1). Left (red): softmax(q·k). Right (green): softmax(q·k / √d_k). Raise d_k and watch the unscaled weights collapse onto one key.

Max weight, unscaled
–
Max weight, scaled
–
Entropy unscaled (bits)
–
Entropy scaled (bits)
–

Training recipe, in one paragraph

Adam with a learning rate that rises linearly for 4,000 warmup steps and then decays with the inverse square root of the step. Dropout 0.1 on each sublayer output. Label smoothing 0.1, which the authors note hurts perplexity but improves BLEU. Beam search of size 4 at test time. Day 14 turns each of these into running code. The same architecture also reached 91.3 F1 on English constituency parsing trained only on the WSJ data, a sign it wasn't tuned only for translation.

4. Why it matters for robotics

5. How to study it

  1. Read the paper sections 3 (model) and 4 (why self-attention) carefully, and skim 5–6. Redraw Figure 1 from memory.
  2. In the 30u30 Day 13 folder, read implementation.py top to bottom, then run train_minimal.py. Open notebook.ipynb to look at attention maps.
  3. Do the exercises in order: scaled dot-product, multi-head attention, positional encoding, encoder block, full Transformer. Try each before opening solutions/.
  4. Check your attention by hand: every row of the weight matrix must sum to 1, and a causal mask must zero the upper triangle.
  5. Stretch: practise the attention problems in the TensorTonic Transformers track (Module 02, Stage D).

6. Self-check

What do Q, K and V each do?

The query says what a token is looking for. Keys say what each token offers to be matched against. Values are what actually gets mixed into the output once the weights are set.

Why scale the scores by 1/√d_k?

Dot products of random d_k-dimensional vectors have variance d_k. Large scores saturate the softmax, which makes gradients vanishingly small. Scaling restores unit variance.

Why does the Transformer need positional encodings at all?

Attention is a weighted sum over a set. Shuffle the input tokens and each output is just shuffled the same way. Without position signals, "dog bites man" and "man bites dog" look identical.

What does the decoder mask do, and why is it needed in training?

It sets scores for future positions to −∞, so token i can only see tokens up to i. Training feeds the whole target sentence at once, so without the mask the model could copy the answer.

When is self-attention more expensive than a recurrent layer?

When sequence length n exceeds the model width d, since attention costs O(n²·d) per layer and recurrence O(n·d²). Long contexts are where this bites.

7. Further reading