Self-attention in plain words

The embedding of python is the same everywhere it appears. But in "run the python code" it means a language, and in "the python ate a mouse" it means an animal. The model needs a way for each word to look at the words around it and adjust. That mechanism is called attention.

The stand-up analogy

Picture a team stand-up meeting. Each person holds three things:

• a query: what I need to find out;
• a key: a label that says what I can help with;
• a value: the notes I would hand over if someone asks me.

Everyone compares their query with everyone's key. The better a key matches, the more attention that person receives, and the more of their value goes into your updated notes. After the round, each person's notes are a weighted mix of everyone's values, weighted by relevance.

In a real model the queries, keys and values are made from each token's vector by three learned tables of numbers (called weight matrices). Those are what training adjusts.

The code

def attention(q, k, v, causal=False):
    d = q.shape[-1]
    scores = q @ k.T / np.sqrt(d)               # query-key match, for every pair of tokens
    if causal:
        future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
        scores = np.where(future, -np.inf, scores)
    weights = softmax(scores, axis=-1)          # each row now sums to 1
    return weights @ v, weights                 # weighted mix of the values

Read it line by line. q @ k.T gives a score for every query against every key: a table with one row per token reading and one column per token being looked at. Dividing by the square root of the vector size keeps the scores from growing with the vector length, which would otherwise make the next step too extreme. (I measured it: the dot product of two random vectors of 1,000 numbers has a spread of about 31, which is 1,000 to the power one half.) The causal lines are explained in the next section. softmax turns each row of scores into positive weights that add up to 1. The last line mixes the values using those weights.

A two-token example you can do in your head

Take two tokens with queries and keys [1, 0] and [0, 1], and values [10, 0] and [0, 10]. Running attention gives these weights and outputs (real output):

weights   [[0.67 0.33]      output  [[6.7 3.3]
           [0.33 0.67]]              [3.3 6.7]]

The first token matches its own key best, so it takes 67% of its own value and 33% of the other's: 0.67 x [10, 0] + 0.33 x [0, 10] = [6.7, 3.3]. Every row of weights adds up to 1, and each output is a blend. A row such as [0.00, 0.93, 0.01, 0.01, 0.05] would read: "this token takes 93% of its new meaning from the second token".

Resolving "python"

The example code applies the idea to the toy table. It gives python some context words and lets it attend to them. Real results:

The word moved toward the meaning its neighbours support. Real models do this in every layer, for every token, with thousands of dimensions and many parallel heads. Each head looks for a different kind of relationship. A test checks that the shifted vector is closer to the right neighbour.