What attention does inside a language model

Follow one attention head

A Transformer turns each token’s current representation into several vectors. In one attention head, a query describes what the current position can look for, a key is used to score a position, and a value is the information mixed into the result.

1. Compare a query with each allowed key using a dot product. Scale the scores by the square root of the key size.
2. Use softmax to turn those scores into weights that sum to one.
3. Multiply the value vectors by the weights and add them to form the result for this head.

A causal mask limits a position to itself and earlier positions. Many decoder-only chat models use this kind of mask during training and generation.

Try two queries

This small program uses made-up two-number vectors. It keeps a third position out of view with a mask. Save and run the file with Python 3.

import math


def attend(query, keys, values, visible):
    scores = [
        sum(q * k for q, k in zip(query, key)) / math.sqrt(len(query))
        for key, allowed in zip(keys, visible)
        if allowed
    ]
    peak = max(scores)
    weights = [math.exp(score - peak) for score in scores]
    scale = sum(weights)
    allowed_weights = iter(weight / scale for weight in weights)
    weights = [next(allowed_weights) if allowed else 0.0 for allowed in visible]
    output = [sum(weight * value[index] for weight, value in zip(weights, values))
              for index in range(len(values[0]))]
    return weights, output


keys = [[1, 0], [0, 1], [1, 1]]
values = [[1, 0], [0, 2], [3, 3]]
visible = [True, True, False]

for query in ([2, 0], [0, 2]):
    weights, mixed = attend(query, keys, values, visible)
    print(f"Query {query}: weights {[round(w, 3) for w in weights]}, "
          f"mixed value {[round(value, 3) for value in mixed]}")

python compare_attention.py

Change visible to three true values and rerun. Notice how an allowed third key can change all the weights and the mixed value.

Connect attention to inference

Real models learn the transformations that produce queries, keys, and values. They run multiple attention heads and stack attention with feed-forward layers and other operations. The resulting representation helps predict the next token.

During generation, the runtime can keep keys and values for earlier tokens in a KV cache. The cache avoids recomputing those parts for each new token. It occupies memory and typically grows with the number of positions and active requests.