How transformers work: attention, with the numbers in it
Attention explained on numbers you can see: a robot's last twenty sensor readings, one query, twenty keys and values, the scores, the softmax and the weighted sum. The demos run in the simulator and print every number, and the page says plainly which parts a language model has that they do not.
A transformer is the kind of neural network behind large language models. Almost everything in it is ordinary: numbers multiplied by weights, added up and squashed, as in the neural network guide. The part that makes it a transformer is attention: before it works anything out, the network decides which of the things it has already seen are worth looking at now, and how much each one counts.
Attention is usually explained on words, where you cannot see the numbers. On this page it runs on something you can watch. A robot spins on the spot on an empty one metre mat, reading the distance in front of it four times a second, and the job is to guess the next reading from the last twenty. Each demo is a program you can change and run.
This is the same arithmetic a language model does, on numbers instead of words, with one head and nothing trained. The section near the end, "What this is and is not", says exactly which parts are the real thing and which are missing.
The problem: guess the next reading
The robot turns about 128 degrees a second, so one revolution takes 11 readings or so. Here are the first twenty-five readings of the spin in the demo below:
81 82 77 53 51 29 21 21 28 51 52 78 81 88 59 51 37 22 21 25 49 50 64 82 82
It is not random and it is not smooth. It repeats, roughly, every eleven readings, and in between it jumps by 20 or 30 centimetres in a quarter of a second as the sensor sweeps past a corner.
Guess "the next reading is the same as this one" and you are wrong by 12.1 cm on average. Take the average of the last three and you are wrong by 21.2 cm, because an average lags behind. Anything that only looks at the last few readings has the same trouble: the piece of the past that says what happens next is not the last few readings. It is the same place in the previous revolution, eleven steps back.
Attention is a way of letting the program find that for itself.
The idea in one line
guess = (how much moment 1 matches now) × (what happened after moment 1)
+ (how much moment 2 matches now) × (what happened after moment 2)
+ ...
with the "how much" numbers adding up to 1. The three words the papers use:
- The query is what is happening now, as numbers. Here it is the last five readings.
- A key is what was happening at some moment in the past, in the same form: five readings in a row.
- A value is what that past moment is offering: here, the reading that came next after it.
The query is compared with every key, which gives a score for each. The scores go through softmax, which turns them into weights that are all positive and add up to 1. The answer is the weighted sum of the values.
The scores
This demo spins the robot, keeps 25 readings, and works out the score for each of the twenty places it could look. Every reading is shifted and scaled first, (cm - 50) / 25, so the numbers are around -1 to 1 and the sums stay small.
The chart is a strip of the twenty weights in order, oldest on the left, with the value of each key beside it.
The program
from bugbot import *
import math
connect()
PATCH = 5 # a key is five readings in a row: the shape of one moment
KEYS = 20 # twenty moments in the past to choose between
DIVISOR = 0.5 # how sharp the softmax is
def scaled(cm):
return (cm - 50) / 25 # about -1 to +1, so the sums stay small
# spin on the spot, reading the distance four times a second
readings = []
for tick in range(49):
drive(0, 0, 100)
readings.append(distance())
wait(0.25)
stop()
history = readings[-26:-1] # the 25 readings before the last one
truth = readings[-1] # what actually came next
print("history, cm:", history)
query = [scaled(cm) for cm in history[-PATCH:]]
keys = [[scaled(cm) for cm in history[j:j + PATCH]] for j in range(KEYS)]
values = [history[j + PATCH] for j in range(KEYS)]
# the score: how much each key looks like the query
scores = [sum(q * k for q, k in zip(query, key)) / DIVISOR for key in keys]
# softmax: turn the scores into weights that are positive and add up to 1
biggest = max(scores)
weights = [math.exp(s - biggest) for s in scores]
total = sum(weights)
weights = [w / total for w in weights]
for j in range(KEYS):
plot("weight", weights[j])
plot("the reading after it, cm", values[j])
wait(0.1)
print("the query, cm:", history[-PATCH:])
print("the weights add up to:", round(sum(weights), 3))
for j in sorted(range(KEYS), key=lambda j: -weights[j])[:3]:
print("key", j, "is", KEYS - j, "steps back:", history[j:j + PATCH],
" score", round(scores[j], 1), " weight", round(weights[j], 3), " value", values[j])
print("the guess:", round(sum(w * v for w, v in zip(weights, values)), 1), "cm")
print("what came next:", truth, "cm")
Work through what it printed.
The query is [49, 53, 84, 81, 92]: the wall has just swung away from the sensor and the readings are climbing. The key that matches it best is eleven steps back, [50, 58, 88, 81, 78], which is the same climb one revolution earlier. Their dot product comes to 5.5, and dividing by 0.5 makes the score 11.0. Twelve steps back scores 8.9, one step back scores 7.9, and the keys from half a revolution back, where the sensor was pointing the other way, score below zero.
Softmax turns 11.0, 8.9 and 7.9 into 0.847, 0.094 and 0.037. It does two things at once: exp makes every number positive and pulls the big ones further ahead, and dividing by the total makes them add up to 1, so the result is a proper weighted average and cannot run off the scale.
The value attached to the best key is 54: the reading that followed that climb last time round. The weighted sum comes to 57.8, and the reading that actually came next was 57. The previous reading was 92.
Values, on every step of a run
Now the same three lines run inside the loop. Every quarter second the robot takes a reading, builds a query from the last five, scores the twenty keys, and guesses what the next reading will be. The chart shows the reading as it arrives and the guess that was made for it one step earlier.
To show that the window is what matters, and not the arithmetic, the program does it twice: once with all twenty keys, and once with only the three most recent, which is a fixed window of the kind attention replaced.
On the mat, the blue dots are the points the last 25 readings measured, and the green square is the one reading the guess mostly came from.
The program
from bugbot import *
import math
connect()
PATCH = 5
KEYS = 20
DIVISOR = 0.5
def scaled(cm):
return (cm - 50) / 25
def attention(history, how_many):
first = KEYS - how_many # use only the most recent how_many keys
query = [scaled(cm) for cm in history[-PATCH:]]
keys = [[scaled(cm) for cm in history[j:j + PATCH]] for j in range(first, KEYS)]
values = [history[j + PATCH] for j in range(first, KEYS)]
scores = [sum(q * k for q, k in zip(query, key)) / DIVISOR for key in keys]
biggest = max(scores)
weights = [math.exp(s - biggest) for s in scores]
total = sum(weights)
weights = [w / total for w in weights]
guess = sum(w * v for w, v in zip(weights, values))
best = first + max(range(how_many), key=lambda i: weights[i])
return guess, best
history = []
spots = []
wide = near = before = None
off_wide = []
off_near = []
off_same = []
for tick in range(70):
drive(0, 0, 100)
cm = distance()
x, y = position()
way = math.radians(heading())
spots.append((50 + x + cm * math.sin(way), 20 + y + cm * math.cos(way)))
if wide is not None: # score the guesses made last time round
plot("reading", cm)
plot("20 keys", wide)
plot("3 keys", near)
off_wide.append(abs(wide - cm))
off_near.append(abs(near - cm))
off_same.append(abs(before - cm))
history.append(cm)
if len(history) > KEYS + PATCH:
history.pop(0)
spots.pop(0)
before = cm
if len(history) == KEYS + PATCH:
wide, best = attention(history, KEYS)
near, _ = attention(history, 3)
draw("the last 25 readings", spots, "blue")
draw("the one it is using", [spots[best + PATCH]], "green", "squares", 8)
wait(0.25)
stop()
print("20 keys: off by", round(sum(off_wide) / len(off_wide), 1), "cm on average")
print("3 keys: off by", round(sum(off_near) / len(off_near), 1), "cm")
print("no change: off by", round(sum(off_same) / len(off_same), 1), "cm")
The three most recent keys are the last three quarter-seconds, and nothing in them says what is about to happen, so that version does no better than assuming nothing changes. The same arithmetic with twenty keys to choose from is wrong by less than half as much. Attention did not beat the window by being cleverer arithmetic. It beat it by being allowed to look further back and by working out for itself how far.
That is the point of the whole mechanism. A fixed window is a decision made by the programmer, once, for every situation. Attention makes it again on every step, from the numbers.
Softmax: how sharp is the looking
The divisor is the one number in the demos that was chosen by hand. It decides how much the best match wins by.
The program
from bugbot import *
import math
connect()
PATCH = 5
KEYS = 20
def scaled(cm):
return (cm - 50) / 25
def softmax(scores):
biggest = max(scores)
weights = [math.exp(s - biggest) for s in scores]
total = sum(weights)
return [w / total for w in weights]
readings = []
for tick in range(49):
drive(0, 0, 100)
readings.append(distance())
wait(0.25)
stop()
history = readings[-26:-1]
truth = readings[-1]
query = [scaled(cm) for cm in history[-PATCH:]]
keys = [[scaled(cm) for cm in history[j:j + PATCH]] for j in range(KEYS)]
values = [history[j + PATCH] for j in range(KEYS)]
dots = [sum(q * k for q, k in zip(query, key)) for key in keys]
print("the dot products:", [round(d, 1) for d in dots])
print("the average of all twenty values:", round(sum(values) / KEYS, 1))
for divisor in (2.24, 0.5, 0.12):
weights = softmax([d / divisor for d in dots])
guess = sum(w * v for w, v in zip(weights, values))
for j in range(KEYS):
plot("divide by " + str(divisor), weights[j])
wait(0.1)
print("divide by", divisor, ": biggest weight", round(max(weights), 3),
" guess", round(guess, 1), " truth", truth)
Both ends are useless in the same way: they stop the program from using the evidence. Divide by 2.24 and the weights are almost equal, so the guess slides towards the average of all twenty values, which is 51 and tells you nothing about now. Divide by 0.12 and one weight is 1 and the rest are 0, so the guess is whatever single value won, and one unlucky match takes the whole answer. In between, the best match leads and the runners-up still count.
A transformer divides by the square root of the length of the key vector, which here would be the square root of 5, or 2.24. That is the flat one. It is the right divisor in a real model because the query and the key have first been through learned matrices, and training is free to make those vectors as large as it needs for the softmax to be as sharp as it wants. Nothing on this page is trained, so the scale is ours to choose, and 0.5 is the choice.
Attention has no sense of order
Attention compares the query with each key separately and adds the results up. Adding is the same in any order, so the mechanism cannot tell which key was recent and which was long ago. Shuffle the keys and their values together and the answer does not move by one decimal place.
For a language model that is a serious problem, because "the robot pushed the ball" and "the ball pushed the robot" are the same set of words. The fix is to put the position into the numbers themselves: every key gets a number saying where it sits in the sequence, and the query gets one saying where "now" is. Positions that are close then add to the score.
This demo does both: it shuffles the keys, then adds a position number of its own.
The program
from bugbot import *
import math
connect()
PATCH = 5
KEYS = 20
DIVISOR = 0.5
def scaled(cm):
return (cm - 50) / 25
def guess_from(keys, values, query):
scores = [sum(q * k for q, k in zip(query, key)) / DIVISOR for key in keys]
biggest = max(scores)
weights = [math.exp(s - biggest) for s in scores]
total = sum(weights)
weights = [w / total for w in weights]
return sum(w * v for w, v in zip(weights, values)), weights
readings = []
for tick in range(49):
drive(0, 0, 100)
readings.append(distance())
wait(0.25)
stop()
history = readings[-26:-1]
truth = readings[-1]
query = [scaled(cm) for cm in history[-PATCH:]]
keys = [[scaled(cm) for cm in history[j:j + PATCH]] for j in range(KEYS)]
values = [history[j + PATCH] for j in range(KEYS)]
plain, plain_weights = guess_from(keys, values, query)
# the same twenty keys and values, in a jumbled order
order = [13, 2, 19, 7, 0, 11, 5, 9, 16, 3, 18, 8, 1, 14, 6, 12, 4, 17, 10, 15]
mixed, _ = guess_from([keys[j] for j in order], [values[j] for j in order], query)
print("in order:", round(plain, 4), " shuffled:", round(mixed, 4))
# now give every key a sixth number for where it sits, and the query one for now
for strength in (0.0, 0.6, 1.2, 1.8, 2.4):
placed = [keys[j] + [strength * j / (KEYS - 1)] for j in range(KEYS)]
with_place, weights = guess_from(placed, values, query + [strength])
near = KEYS - max(range(KEYS), key=lambda j: weights[j])
print("position worth", strength, ": guess", round(with_place, 1),
" the key it leans on is", near, "steps back")
if strength == 1.8:
for j in range(KEYS):
plot("no position", plain_weights[j])
plot("with position", weights[j])
wait(0.1)
print("the truth was", truth)
Position is not free. Turn it up far enough and recency drowns the evidence: the key one step back wins on being recent, the guess becomes "about the same as the last reading", and the model is back where the fixed window was. In a real transformer nobody picks that number by hand. Position goes in as a pattern added to each token's vector, and training decides how much of the score it is allowed to explain.
What this is and is not
Everything above is the real mechanism. A language model, given the words so far, works out a query from the position it is at, a key and a value from every earlier position, scores the query against every key with a dot product, softmaxes the scores into weights that add up to 1, and adds the values up in those proportions. That is what these demos do.
What a real model has that this page does not:
- Words, as long vectors. Each token, roughly a word or part of one, becomes a list of hundreds or thousands of numbers, learned during training. Here a token is one distance reading, and the vector is five of them.
- Learned queries, keys and values. A real model does not use the token's vector as its own query. It multiplies it by three matrices,
W_q,W_kandW_v, which are learned. Here those three matrices are missing, which is the same as using the identity: query, key and value are the raw numbers. - Many heads. A layer runs several attentions side by side with different learned matrices, so one head can follow grammar while another follows the subject of the sentence. Here there is one head.
- Many layers. The output of one attention block feeds a small ordinary network, and that feeds the next block, dozens of times over. Everything on this page happens once.
- Training. Every matrix in a real model is found by gradient descent over a very large amount of text, as in the neural network guide. Nothing here is trained at all. That is why the sharpness had to be chosen by hand.
- Words out, not numbers. A language model ends with a score for every word in its vocabulary, turned into probabilities by another softmax, and one is picked at random in proportion. This page ends with one number in centimetres.
Two things that are the same and are easy to miss. A language model may only look backwards when it is predicting the next word, which is called masking, and these demos only look backwards too. And the softmax weights are where "what is it paying attention to" actually lives: in a real model you can read them off in the same way this page prints them.
What this page cannot show you is why scale matters. The arithmetic is the same at every size. What changes with billions of weights, many heads and many layers is what the queries and keys come to mean, and that comes out of training, not out of the mechanism.
Where this is taught
- See as numbers: the robot's view is a list of numbers, which is where any of this has to start.
- Collecting data: recording readings and labelling them, the raw material for a key and a value.
- Nearest neighbour: find the most similar thing you have seen before. Attention is this, softened, and with every match counting a little.
- A tiny network: weights,
tanh, softmax, and training by gradient descent. - Models as text: what a trained model actually is once it is saved.
- Averaging: a weighted average, which is what the last line of attention is.
- Seeing the signal: plotting a sensor over time and reading its shape.
Questions
How do transformers work, in simple terms?
A transformer reads a sequence and, for every position, decides how much each earlier position matters to it right now. It does that by comparing a query taken from where it is with a key taken from every earlier place, turning those comparisons into weights that add up to 1, and mixing the earlier places' values in those proportions. Stack that with ordinary neural network layers, dozens of times, and train the lot on a great deal of text, and you have a large language model.
What are queries, keys and values?
Three vectors made from the same input, by three different learned matrices. The query says what this position is looking for, the key says what a position has to offer, and the value is what is actually handed over when it is chosen. On this page the query is the last five readings, a key is five readings from somewhere in the past, and a value is the reading that came next after them.
What is attention in a neural network?
A weighted average where the network works out the weights itself, from how well each thing matches what it is looking for. The weights are positive and add up to 1, so nothing can blow up, and any of them can be nearly zero, which is how the network ignores things.
Why is softmax used in attention?
Two reasons. exp makes every score positive and stretches the gaps, so a slightly better match gets a good deal more weight. Dividing by the total makes the weights add to 1, so the output is a proper weighted average of the values. On this page scores of 11.0, 8.9 and 7.9 become weights of 0.85, 0.09 and 0.04.
Why do transformers divide by the square root of d_k?
The dot product of two vectors of length d grows with d, so without that divisor the scores in a big model would be large, the softmax would be nearly one-hot and training would stall. Dividing by the square root of the key's length keeps the scores in a sensible range whatever the model's size. On this page the divisor is picked by hand instead, because there are no learned matrices to set the scale.
What is positional encoding and why is it needed?
Attention adds up its comparisons, and addition does not care about order, so a transformer cannot tell where anything was in the sequence. Positional encoding puts that back by adding a pattern for position into each token's vector before the attention happens. The demo above shows both halves: shuffling the keys leaves the answer identical to four decimal places, and adding a position number changes which key wins.
What is a head in multi-head attention?
One complete attention, with its own query, key and value matrices. A layer runs several at once and joins the results, so different heads can follow different relationships in the same sentence. The demos here have exactly one head.
Is attention the same as nearest neighbour search?
It is a soft version of it. Nearest neighbour picks the single most similar stored example; attention gives every stored item a share in proportion to how well it matches. The sharpness demo makes the link visible: divide the scores by a small enough number and attention turns into nearest neighbour, taking the single best value and nothing else.
Do transformers only work on text?
No. The mechanism only needs a sequence of vectors, so the same blocks are used on images cut into patches, on audio, and on sensor readings, which is what this page does. Text is where they were first used and where the largest models are.
Does this mean a language model is looking things up?
Partly, at the level of one head and one step: the softmax weights really are a lookup over what has already been read, and you can read them. What it produces is not copied from any one place. The value that arrives is mixed with all the others, pushed through several more layers, and turned into a score for every word in the vocabulary. The lookup is the mechanism, not the whole story.
Learn it step by step
These lessons build the same ideas one at a time, each with tasks the simulator marks.
- 8.1 See as numbers Learning, Robot club
- 8.2 Collecting data Learning, Robot club
- 8.3 Nearest neighbour Learning, Robot club
- 8.4 A tiny network Learning, Robot club
- 8.7 Models as text Learning, Robot club
- U1.5 Seeing the signal The robot as a system, University
- U4.2 Averaging Noise and filtering, University