Attention and Transformers

Read this: "Sam made soup and ate it." You know instantly that "it" is the soup, not Sam. You looked back at the other words, decided which ones mattered, and borrowed their meaning.

A language model has to do the same thing. After tokenization, every token becomes a vector, a list of numbers that stands for its meaning. On its own, the vector for "it" says almost nothing. To predict the next token well, "it" needs to pick up information from "soup". The mechanism that does this is called attention, and it is the core of the transformer, the design behind most of today's LLMs.

#The idea: ask, match, borrow

Picture each word as a person at a table, holding up three things:

• A query: what am I looking for? For "it", roughly "a thing that was mentioned earlier".
• A key: what do I offer? For "soup", roughly "I am a thing, a food".
• A value: the information I will hand over if someone picks me.

Each word compares its query with every key it can see. A good match earns a high score. The scores are turned into shares that add up to 1, and the word collects a blend of everyone's values, weighted by those shares. After this step, the vector for "it" contains mostly the meaning of "soup".

#Work it through with small numbers

Real vectors have hundreds or thousands of numbers. Here each one has just two, and we pretend the first number means "is a thing" and the second means "is a person". In a real model nobody chooses what the numbers mean; training finds whatever is useful.

Matching uses the dot product: multiply the numbers in the same positions and add. The query for "it" is [2, -1]: it wants a thing and is put off by a person.

for each visible word w:  score[w] = dot(query, key[w])weights = softmax(scores)output = sum of weights[w] * value[w]

#The code

This is one attention step for one word, exactly as traced. Real models also divide each score by a fixed number, the square root of the key vector's length, before the softmax. Long vectors give big dot products, and this keeps the scores in a range where softmax does not put almost everything on one word. With two numbers per vector it would not change the lesson, so it is left out.

const dot = (a, b) => a.reduce((sum, x, i) => sum + x * b[i], 0);
 
function softmax(scores) {
  const biggest = Math.max(...scores); // subtract the max so exp() cannot overflow
  const exps = scores.map((s) => Math.exp(s - biggest));
  const total = exps.reduce((a, b) => a + b, 0);
  return exps.map((e) => e / total);
}
 
function attend(query, keys, values) {
  const weights = softmax(keys.map((k) => dot(query, k))); // match, then share
  const out = new Array(values[0].length).fill(0);
  weights.forEach((w, j) => {
    values[j].forEach((v, i) => { out[i] += w * v; }); // blend the values
  });
  return { weights, out };
}
 
const keys = [[0, 2], [0.5, 0], [2, 0], [0, 0], [1, 0], [0.5, 0]];
const values = [[0, 1], [0, 0], [1, 0], [0, 0], [0.5, 0], [0, 0]];
attend([2, -1], keys, values).out; // about [0.85, 0.00]

import math
 
def dot(a, b):
    return sum(x * y for x, y in zip(a, b))
 
def softmax(scores):
    biggest = max(scores)  # subtract the max so exp() cannot overflow
    exps = [math.exp(s - biggest) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]
 
def attend(query, keys, values):
    weights = softmax([dot(query, k) for k in keys])  # match, then share
    out = [0.0] * len(values[0])
    for w, value in zip(weights, values):
        for i, v in enumerate(value):
            out[i] += w * v  # blend the values
    return weights, out
 
keys = [[0, 2], [0.5, 0], [2, 0], [0, 0], [1, 0], [0.5, 0]]
values = [[0, 1], [0, 0], [1, 0], [0, 0], [0.5, 0], [0, 0]]
attend([2, -1], keys, values)[1]  # about [0.85, 0.00]

In the trace, the queries, keys and values were typed in by hand. In a real model, each one is made from the word's current vector by multiplying it with a grid of numbers (a matrix) learned in training. Nobody writes rules like "pronouns look for things". Training finds whatever grids make the next-token chances better. And because a word's vector changes as it picks up context, the same word can ask different questions in different sentences.

#From one attention step to a transformer

A transformer repeats one block many times. Each block does two things to every token's vector:

1. Attention: gather information from the visible tokens, as above.
2. Feed-forward: a small neural network that works on each token's vector by itself, transforming what was gathered.

Each block adds its result to the vector instead of replacing it, so information from earlier blocks is kept. After the last block, the vector at the final position is turned into chances for the next token, and the loop from What Is an LLM picks one.

Three more details:

• Many heads. Each block runs several attention steps side by side, called heads, each with its own learned query, key and value numbers. A few heads end up with a job you can name, like linking a pronoun to its noun. Most do not.
• Looking back only. When generating, a token can attend to itself and the tokens before it, never after. Training uses the same rule, so the model never learns to rely on the future.
• Word order. Attention on its own treats the text like a bag of words: "dog bites man" and "man bites dog" would match the same way. So position information is mixed into the vectors, letting the model tell first from last.

#Cost: every token looks at every other

For a text of n tokens, each token compares itself with up to n others, so one attention step does about n²/2 comparisons. Double the text and the work quadruples. When generating, models keep the keys and values of earlier tokens in memory, so each new token costs about n comparisons rather than redoing everything.

This growth is one of the reasons models have a limit on how much text they can take in. The Context Window covers that limit.

#Common mistakes

#Try it yourself

#Check yourself

#What to practise next

Compute attention weights yourself in Pronoun Looks Back, which includes the square-root scaling. Then Training and Fine-Tuning shows how the learned numbers behind queries, keys and values get set in the first place.