Transformers and Attention

Why sequences needed something new

Every chatbot you have used, most modern translation, the image models that read captions and the protein-structure model that won a Nobel Prize all share one building block. It is called attention, and the architecture built around it, the Transformer, was published in 2017 under a title that turned out to be accurate: attention is all you need.

What problem was it solving?

Language is a sequence, and the meaning of a word often depends on something far away. In "The keys to the cabinet are on the table", the verb must agree with "keys", not with the nearer "cabinet". Before 2017 the standard tool for sequences was the recurrent neural network (RNN), which reads one token at a time and carries a single hidden vector forward, like reading a book while only being allowed to remember one sticky note.

That design has two problems. First, everything the model knows about the past must squeeze through one fixed-size vector, a bottleneck. Second, training requires sending gradients backwards through every step. The gradient gets multiplied at each step, and products of many numbers either shrink towards zero or blow up, a problem analysed formally in the early 1990s. Gated cells such as the LSTM eased it but did not remove it, and they still had to process tokens one after another, which wastes most of a GPU.

How does the gradient vanish?

Consider the simplest possible recurrent network, where each step multiplies the gradient by the same factor. After k steps the signal has been multiplied by that factor k times. A factor of 0.9 sounds harmless, but 0.9³⁰ is about 0.04 and 0.9⁴⁰ is about 0.015. The word 40 tokens back barely influences learning at all. A factor of 1.1 does the opposite and explodes. Real RNNs multiply matrices rather than numbers, but the same exponential behaviour applies to their largest singular values.

Only a factor of exactly 1 keeps the signal intact, and a trained network cannot hold itself on that knife edge. Attention sidesteps the whole question: any token can read any other token in a single step, so the path between "keys" and "are" has length one no matter how far apart they are.

Why does it matter?

Removing the step-by-step chain did two things at once. It gave models a direct line to distant context, and it made training parallel: every position in a sequence can be processed simultaneously, which is exactly what GPUs are good at. That second property is what allowed models to grow from millions to hundreds of billions of parameters. If you want to see the recurrent approach at work first, the Recurrent Networks Lab lets you step an RNN through a sequence.

Key takeaways

• RNNs pass a single hidden vector along the sequence, which creates an information bottleneck and a long gradient path.
• Gradients through k recurrent steps scale roughly like a factor to the power k, so they vanish or explode over long distances.
• Attention connects every pair of tokens directly and lets all positions be computed in parallel.

Attention as a soft lookup

The cleanest way to understand attention is as a dictionary lookup where the match is allowed to be partial. You ask a question, every entry says how well it answers, and you get back a blend of the answers weighted by how well they matched.

What is it?

A Python dictionary has keys and values. You give it a query, it finds the key that equals the query exactly, and returns that one value. Attention keeps the same three roles but makes every step continuous:

• The query is a vector describing what this position is looking for.
• Each key is a vector describing what a position offers.
• Each value is the information that position hands over if it is chosen.

Instead of exact equality, similarity is the dot product of query and key. Instead of picking one winner, a softmax turns the similarities into weights that are positive and sum to one. The result is the weighted average of all the values.

Two things to notice. A longer query produces larger dot products and therefore a sharper, more decisive softmax, the same effect as the sharpness slider. And when the query sits between two keys, attention returns a genuine mixture of both, something a hard lookup can never do. That mixture is differentiable, so gradient descent can learn where queries and keys should point.

Where did the idea come from?

Attention predates the Transformer. In 2014 it was bolted onto an RNN translator so the decoder could look back at every source word while producing each output word, instead of relying on one summary vector. Translation quality on long sentences improved sharply. The 2017 insight was to throw the RNN away and build the whole network from attention.

Why does it matter?

Because the lookup is soft and learned, the model decides for itself what counts as relevant: the subject of a verb, the noun a pronoun refers to, the matching bracket in code, the object in an image patch next door. The same mechanism powers retrieval over documents, which is why the embeddings and retrieval lesson will feel familiar.

Key takeaways

• Attention is a dictionary lookup with partial matches: query against keys, then a weighted blend of values.
• Dot products measure the match and softmax turns matches into weights that sum to one.
• Everything is differentiable, so the model learns what to look for and what to advertise.

Scaled dot-product attention

Now the real formula, computed live on a ten-word sentence. Every number below comes from the same four operations a large language model runs billions of times per second.

Q, K and V stack the query, key and value vectors of every token as rows. d is the length of each query and key vector.

How does it work?

For one query token at a time, the recipe is:

1. Score: take the dot product of this token's query with every key.
2. Scale: divide every score by √d.
3. Normalise: apply softmax across the scores so they become weights summing to 1.
4. Mix: add up the value vectors, each multiplied by its weight.

In a trained model the query, key and value of token i come from its embedding x_i multiplied by three learned matrices: q = x W_Q, k = x W_K, v = x W_V. Their dimensions have no human names. For this demo we hand-set the vectors with four named dimensions (noun, animate, action, object) so you can read every score. The rules are toy rules, but the arithmetic is exactly the real thing.

Look at it: its query asks for a living noun, cat offers exactly that, and about half the attention lands there. The output vector for it now carries the "animal" feature, so later layers can treat the pronoun as the cat. Notice also that tokens with nothing to offer still receive some weight, because softmax never gives exactly zero. Trained models often park that leftover weight on the first token or punctuation.

Why divide by √d?

If the components of q and k are roughly independent with variance 1, their dot product is a sum of d such terms and has variance d. With d = 128, typical scores would be around ±11, and softmax of numbers that far apart is essentially one-hot. One-hot softmax has a near-zero gradient, so learning stalls. Dividing by √d brings the variance back to about 1. In the demo d = 4, so the divisor is 2: switch scaling off and every row gets visibly peakier.

The whole matrix at once

Doing this for every query gives the attention matrix: one row per query, one column per key, every row summing to one. On a GPU the whole thing is two matrix multiplications and a softmax, computed for all tokens in parallel. That is the practical reason transformers train so much faster than RNNs on the same hardware.

Why does it matter?

These three lines of linear algebra are the only place in a transformer where tokens communicate. Everything a model knows about context, from agreement in grammar to which function a variable came from, arrives through weights like the ones in that heatmap. The Attention Visualizer lab lets you explore the pattern on your own sentences.

Key takeaways

• Score with q·k, scale by √d, softmax each row, then take the weighted sum of values.
• Scaling keeps scores in a range where softmax still has useful gradients.
• The attention matrix is computed for all tokens in parallel with two matrix multiplications.

Multi-head attention

One softmax can only express one notion of relevance at a time. But a word needs several kinds of context at once: its subject, its neighbour, the bracket it closes. So transformers run several attentions side by side.

What is it?

Multi-head attention splits the model width into h smaller heads. Each head has its own W_Q, W_K, W_V, computes its own attention matrix, and returns its own output. The outputs are concatenated and multiplied by one more learned matrix W_O. The original Transformer used 8 heads of size 64 in a 512-wide model, so the total cost is about the same as one full-width head.

Do real heads specialise like this?

Often, yes, though rarely so tidily. Interpretability research on trained models finds previous-token heads very much like head 2, and "induction heads" that build on them to copy patterns: if the text earlier contained "Mr Dursley", an induction head looking at "Mr" attends to what followed the previous "Mr" and predicts "Dursley". Other heads track syntax, and many do things nobody has yet named. Heads are also redundant: pruning a large fraction of them after training often costs little accuracy.

Why does it matter?

Multiple heads give each layer several independent channels of communication for roughly the price of one. Stack dozens of layers and the model can compose them: one head finds the previous token, a later head uses that to find repeats, a later one still uses repeats to complete a pattern. The Transformer Explorer shows heads across layers.

Key takeaways

• Each head has its own query, key and value projections and computes its own attention pattern.
• Head outputs are concatenated and mixed by W_O, so later layers see several kinds of context together.
• Trained heads often specialise, for example in previous-token and induction patterns, and later layers compose them.

Positional encoding

Here is a surprising fact: attention on its own has no idea about word order. Shuffle the input and each token gets exactly the same output as before, just in a different place. Order has to be added deliberately.

Why is order missing?

The score between two tokens depends only on their vectors, not on where they sit. So without extra information "dog bites man" and "man bites dog" produce the same set of representations. Mathematicians call this permutation equivariance. It is great for sets of objects and terrible for language.

How does the original fix work?

The 2017 paper added a fixed vector to each token embedding based on its position. Dimension pairs hold a sine and cosine at geometrically spaced frequencies, from a fast wave that changes every token to slow waves spanning thousands of tokens:

pos is the position, i indexes dimension pairs, d is the model width.

Think of a clock with many hands moving at different speeds. The seconds hand distinguishes neighbours, the hour hand distinguishes distant positions, and together they give every position a unique fingerprint.

The similarity curve peaks at the chosen position, nearby positions score higher than distant ones, and the curve simply slides along as you move the slider. That is the useful property: the dot product between two encodings depends on how far apart they are, not on where they are. Because a shift in position is a rotation of each sine and cosine pair, a linear layer can learn to express "three tokens back".

Rotary embeddings (RoPE)

Most current open models, including the Llama family, use rotary position embedding instead of adding a vector. RoPE takes each pair of dimensions in the query and the key and rotates it by an angle proportional to the token's position. When you then take q·k, the two rotations partly cancel and the score depends only on the relative offset between the tokens. Relative position ends up exactly where it is used, inside the attention score. Extending context windows often works by adjusting RoPE's frequencies, which is why the method matters for long-context models.

Key takeaways

• Self-attention is blind to order; without positions a sentence is a bag of tokens.
• Sinusoidal encodings give each position a multi-frequency fingerprint whose similarity depends on relative distance.
• RoPE rotates queries and keys by position so attention scores depend directly on relative offset.

The transformer block

Attention is the headline, but a transformer is a stack of identical blocks, and each block has four parts. Understand one block and you understand the whole model.

What are the parts?

How does information flow?

A useful mental picture is the residual stream: a vector per token that runs straight up through the network. Each attention layer reads the stream, gathers information from other tokens and writes a correction back. Each MLP reads the stream at one token, transforms it, and writes back. Nothing overwrites; everything adds. After the last block, a final LayerNorm and a linear layer turn each token's vector into scores over the vocabulary.

The original paper put LayerNorm after each residual addition (Post-LN). Almost every large model since places it before each sublayer (Pre-LN, as drawn), because Pre-LN trains stably without a delicate learning-rate warm-up. Many models also swap in RMSNorm, a cheaper variant, and gated activations such as SwiGLU in the MLP. The skeleton is unchanged.

Why does it matter?

The block is the same whether the input is text, image patches, audio frames or protein residues. Vision Transformers cut images into 16 by 16 patches and treat each as a token, and the rest of the block is unchanged. That generality is why one architecture now spans language, vision and multimodal models.

Key takeaways

• Each block is attention then an MLP, each wrapped in LayerNorm and a residual addition.
• The residual stream carries each token upward; sublayers only add corrections to it.
• About two thirds of a standard transformer’s parameters are in the per-token MLPs.

Encoders, decoders and masks

The same block builds three families of model. What separates them is not new machinery but which tokens are allowed to see which: a mask on the attention matrix.

What distinguishes them?

Encoder-only models let every token see every other token, in both directions. BERT was trained by hiding 15% of the words and predicting them from both sides. These models produce excellent representations for classification, search and embeddings, but they do not naturally generate text.

Decoder-only models apply a causal mask: token i may attend only to tokens 0 to i. Masked scores are set to minus infinity before the softmax, so their weights become exactly zero. Trained to predict the next token, they can generate text one token at a time, feeding each output back as input. GPT, Llama, Claude and Gemini are all in this family.

Encoder-decoder models, the original 2017 design, encode the input bidirectionally and then decode with causal self attention plus cross-attention, where queries come from the output being written and keys and values come from the encoded input. This suits tasks with a clear input and output, such as translation and speech recognition.

How does the mask make training efficient?

With the causal mask, a single forward pass over a 1,000-token document gives 1,000 next-token predictions at once, each one only using the tokens before it. Every position is a training example and all are computed in parallel. Go back to the step-by-step demo and switch on the causal mask: the upper triangle of the matrix disappears, and it can still find cat because the cat came first.

Why does it matter?

Decoder-only models won the scale race because next-token prediction needs no labels, any text works as training data, and a single model can then be prompted into translation, summarising or coding. The Large Language Models lesson picks up the story from here.

Key takeaways

• Encoder-only models see in both directions; decoder-only models use a causal mask; encoder-decoders add cross-attention.
• A causal mask sets future scores to minus infinity so their softmax weights are exactly zero.
• Causal masking turns every position of a document into a parallel next-token training example.

Scale and efficiency

Transformers won because they turned hardware into capability more efficiently than anything before them. The price is that attention cost grows with the square of the context length, and much of the engineering since 2020 is about paying less of that price.

Why did transformers win?

An RNN needs n sequential steps for a sequence of length n; a transformer layer needs a constant number, because all positions are processed together. In 2017 the base Transformer reached state-of-the-art English to German translation after 12 hours of training on 8 GPUs, a fraction of the cost of the recurrent systems it beat. Once training is parallel, bigger models and bigger datasets just need more chips. Researchers then found that loss falls as a smooth power law in parameters, data and compute, which made scaling a predictable investment rather than a gamble.

Where does it hurt?

Every query scores every key, so one head in one layer computes n² numbers. During generation the model also stores the keys and values of every past token in every layer, the KV cache, so that it does not recompute them. The calculator below uses the real shape of Llama 3 70B: 80 layers, 64 query heads and 8 shared key/value heads of size 128.

How modern models cope

Why does it matter?

Almost every practical limit you meet with an LLM, from price per token to maximum document length to how many users one GPU can serve, comes back to these costs. Knowing that attention is quadratic in context and the KV cache is linear explains why long prompts cost more, why providers cache prompts, and why architecture papers obsess over memory rather than arithmetic.

Key takeaways

• Parallel training plus predictable scaling laws is why transformers displaced RNNs.
• Attention scores grow with n squared; the KV cache grows linearly with n and dominates memory at inference.
• FlashAttention, grouped-query attention, mixture of experts and hybrid layers attack these costs without changing the core idea.

Check your understanding

Six scenarios. Each one asks you to apply an idea from the lesson rather than recall it.

A model reads "The trophy did not fit in the suitcase because it was too big." In a trained transformer, which mechanism most directly lets the representation of "it" absorb information from "trophy"?

References

The papers behind this lesson. The 2017 paper is short and readable; section 3 is the core of everything above.

References

1. [1]

   Learning long-term dependencies with gradient descent is difficult(opens in a new tab)

   Bengio, Simard, Frasconi, 1994

   Showed why gradients through many recurrent steps tend to vanish or explode.

2. [2]

   Long Short-Term Memory(opens in a new tab)

   Hochreiter, Schmidhuber, 1997

   The gated recurrent cell designed to carry information across long gaps.

3. [3]

   Neural Machine Translation by Jointly Learning to Align and Translate(opens in a new tab)

   Bahdanau, Cho, Bengio, 2014

   Added an attention mechanism to an RNN translator so the decoder could look back at every source word instead of one fixed vector.

4. [4]

   Attention Is All You Need(opens in a new tab)

   Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin, 2017

   Introduced the Transformer: an encoder-decoder built entirely from attention and feed-forward layers, trained in parallel.

5. [5]

   RoFormer: Enhanced Transformer with Rotary Position Embedding(opens in a new tab)

   Su, Lu, Pan, Murtadha, Wen, Liu, 2021

   Rotary position embedding: rotate queries and keys by an angle proportional to position so their dot product depends on relative offset.

6. [6]

   On Layer Normalization in the Transformer Architecture(opens in a new tab)

   Xiong, Yang, He, Zheng, Zheng, Xing, Zhang, Lan, Wang, Liu, 2020

   Explains why placing LayerNorm before each sublayer (Pre-LN) trains more stably than the original Post-LN arrangement.

7. [7]

   An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale(opens in a new tab)

   Dosovitskiy et al., 2020

8. [8]

   BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding(opens in a new tab)

   Devlin, Chang, Lee, Toutanova, 2018

9. [9]

   Scaling Laws for Neural Language Models(opens in a new tab)

   Kaplan, McCandlish, Henighan, Brown, Chess, Child, Gray, Radford, Wu, Amodei, 2020

   Language-model loss falls as a smooth power law in parameters, data and compute.

10. [10]

    The Llama 3 Herd of Models(opens in a new tab)

    Llama Team, AI @ Meta, 2024

    Architecture details for Llama 3, including 80 layers and 8 key/value heads for the 70B model.

11. [11]

    FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness(opens in a new tab)

    Dao, Fu, Ermon, Rudra, Re, 2022

    Computes exact attention in tiles held in fast on-chip memory, never writing the full n x n matrix to GPU memory.

12. [12]

    GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints(opens in a new tab)

    Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebron, Sanghai, 2023

13. [13]

    Mixtral of Experts(opens in a new tab)

    Jiang et al., 2024

    A sparse mixture-of-experts model: each token has access to 47B parameters but uses 13B per token.

14. [14]

    Mamba: Linear-Time Sequence Modeling with Selective State Spaces(opens in a new tab)

    Gu, Dao, 2023

15. [15]

    Lost in the Middle: How Language Models Use Long Contexts(opens in a new tab)

    Liu, Lin, Hewitt, Paranjape, Bevilacqua, Petroni, Liang, 2023

    Models retrieved information best from the start or end of a long input and worst from the middle.